Validation Loss Plateaus only after 2000 steps Pretraining Regardless of LR Schedule (Follow-up to #3868)

[jefcoder]

Hi everyone,

Following up on my earlier post (#3868) about multi-dataset sampling and learning rate selection .

Thanks so much to this community for all the help — special thanks to @tianyu-l, @felipemello1, @slimfrkha, @rakkit, and @haydn-jones . With your input, I've made significant progress, and validation loss across all datasets is now decreasing consistently .

However, I'm now running into a related but more fundamental problem: the validation loss plateaus hard only after 2000 steps of continued pretraining, regardless of the learning rate schedule I use .

Setup
Hardware: 8×B200s
Run length: 12 hours
Total loss reduction: ~0.0008 over the full run
LR schedules tried: Cosine annealing decay: 5e-6 → 1e-6

Observed behavior
Regardless of schedule, the loss curve is heavily front-loaded:

~70% of total reduction happens in the first third of training
~20% in the second third
~10% or less in the final third — essentially flat

I initially assumed the main culprit was the low LR combined with cosine annealing — especially since LR values below 3e-6 seemed to produce a negligible learning signal, essentially "useless" for training.

So I switched to a constant LR of 5e-6, expecting this to resolve the issue . But I observed the same symptom regardless .

I'm now genuinely puzzled as to why the loss only reduces meaningfully at the very beginning of training, then flattens out almost entirely only after 2000 steps.

I'd really appreciate hearing from anyone who's hit a similar wall in large-scale continued pretraining. As before, happy to discuss further via DM, and open to consulting/collaboration if useful.

Thanks in advance!

[haydn-jones]

My first guess is LR is too low if loss is moving by ~0.0008. I would try something higher like 3e-5 (keep a warmup period) and decay down to 3e-6.

[jefcoder]

Hi @haydn-jones
Thanks for your reply.
The model has trained on provided dataset quite a long time.
I also tried various LR but if I set higher LR, then the loss of each dataset just increase.

[haydn-jones]

Have you tried higher LR with a WSD schedule? Also, what model are you training?

[jefcoder]

It's Qwen3.5 9B model and the team has already trained the model for provide dataset for a month in 8xB200.
I just tried warmup + constant/cosine, but when constant/cosine phase, the loss just boosts when lr is higher than it.

This is the log when I used 1e-5.

Same setup, but when lr is 5e-6, it improved a lot.

[rakkit]

Are you sure your training data and Val data is properly pre-shuffled?

And I suggest you do 2d sweep for bs and learning rate for at least 1K steps. I think lr scheduler is not the main reason for what you observed.

[slimfrkha]

Can you describe how you are splitting each dataset into train / val ? Maybe there is a shift between train and val if you re not globally shuffling the whole dataset ?
Also can you run some benchmarks (linked to to each dataset domain) to check their trend ? Val loss can plateauing but benchmarks are improving.
For batch size, make sure you're in the sweet spot for global batch size (4M to 60M tokens) as described in the ultra scale playbook. if you re already in that interval, sweeping the bs would only yield very little improvement so not worth it.

[jefcoder]

Hi @slimfrkha

Thanks for your valuable feedback.
I evaluate the checkpoint with randomly and globally sampled 5000 sequences for each dataset and find out that evaluation result is quite different from what it shows in the wandb log.

This is my another training. Here I trained the only mix of 4 datasets.

I evaluate the loss of those 4 datasets with 5000 randomly sampled sequences.

Then only the loss of automathtext has reduced around 0.00092.
Then loss of finewebedu increased 0.00014,
big-1 increased 0.00005
small-1 incrased 0.00050

This means the initial sharp drop across all DBs isn't contributing to any real loss reduction. And if the loss is fluctuating afterward, that suggests it hasn't actually decreased in a meaningful way.

For the training data, I used random shuffling. For evaluation, I randomly sampled 512 sequences from each DB. I'm currently investigating why this mismatch is occurring — could you give me some tips on how to debug this?

Additionally, could you explain why the loss for all DBs drops so quickly at the beginning, even though this drop doesn't seem to correspond to actual model improvement?

Thanks.

[slimfrkha]

let me clarify some points to be sure we are on the same page

Monitoring during training

• Train loss is the primary real-time signal: watch for spikes/divergence, steady decrease, eventual plateau.
• A plateau under high LR (WSD stable phase, cosine high regime) is often fine. but can also signal a bad mix, stale data, or LR too high.
• Held-out val loss detects overfitting/repetition
• Downstream benchmarks (MATH500, AIME if you train on math data; coding benchmarks for code data etc) track capability
• Val loss and benchmarks aren't exclusive. Use both.

If doing val loss, the correct order

• Dedup first (remove near-duplicates. web corpora are full of them).
• Stratify by domain/source when a source orders the data by domain otherwise split would yield a shifted distribution.
• Split within each stratum into train/val (document/source level), so every domain appears in both in the same proportions.
• Shuffle after splitting
• Train only on train splits, validate only on val splits. NEVER mix.

---

Finally, here is the expected behavior of both val loss and benchs for CPT with WSD narrated by claude:

Here's the expected behavior, phase by phase. Two caveats up front that shape everything: (1) in continual pretraining you start from an already-trained model, so there's no "loss falling from random-init" — you start low and the dynamics are about adapting to the new mixture; (2) which val set matters enormously — new-domain val vs. original-distribution val behave differently, and this is where the interesting signal lives.

Warmup phase
Val loss (new-domain / new mixture): often rises or stays flat briefly before falling. Two things fight: the LR is ramping up from near-zero (weak updates early) and the optimizer states / the model are adjusting to the new distribution. A small bump at the very start is normal in continual pretraining, not alarming.
Val loss (original distribution): frequently ticks up here — this is the onset of forgetting as you pull the model toward the new mixture. Expected, not a bug.
Benchmarks: essentially flat / noise. Too few tokens, and benchmarks are laggy anyway. Don't read anything into warmup-phase evals.

Constant (stable) LR phase
Val loss (new-domain): steady decrease, then a plateau. This is the bulk of learning. The plateau at high LR is expected — high LR keeps the model bouncing in a wide basin, so loss can't settle to its minimum. A plateau here does not mean learning has stopped; the decay phase will extract further gains. This is the single most important thing to internalize about WSD: a flat stable-phase loss is not convergence.
Val loss (original distribution): continues drifting up or stabilizes at an elevated level, depending on mixture ratio and replay. If you're mixing in replay data from the original distribution, this stays more controlled.
Benchmarks: slow, noisy improvement on domains matched to your injected data (math benchmarks if you injected math, etc.). Often still mostly flat and jumpy — benchmarks tend to lag loss and may not move meaningfully until the decay phase. Don't over-trust small mid-plateau swings.

Decay phase
Val loss (new-domain): a characteristic sharp additional drop as LR anneals to (near) zero. This is the WSD signature — the model finally settles into the basin it was hovering over, and loss falls noticeably below the stable-phase plateau. This drop is the payoff and the reason WSD works; the steepness/size of it is a good health signal.
Val loss (original distribution): can partially recover (dip back down) if there's replay in the mix, or at least stabilize. With no replay it may just settle at its elevated level.
Benchmarks: this is where they finally move. Capability gains that were latent during the plateau tend to materialize in the decay phase — benchmark scores jump here more than anywhere else. If your data injection worked, this is where you see it on MATH500/AIME/coding evals. If benchmarks don't move in decay despite a clean loss drop, that's a signal the loss improvement is compression/formatting rather than transferable capability.