TL;DR: A subtle bug in the fine-tuning batch sampler caused resumed training jobs to repeatedly skip the beginning of the dataset after checkpoint recovery. Training never crashed, but the model quietly learned from an incomplete view of the data. Making the sampler epoch-aware fixed the issue and restored correct, deterministic data coverage.
“The checkpoint loaded successfully. GPUs were fully utilized. Training resumed without a single error.”
Yet something was wrong.
Every resumed run showed unusual loss behavior. Nothing crashed, no assertions failed, and every monitoring dashboard suggested the training job was perfectly healthy.
The model was training.
It just wasn’t training on the right data.
Those are my favorite kinds of bugs — not because they’re easy, but because they’re deceptive. They’re the bugs that quietly produce incorrect results while every system insists everything is working.
Recently, while contributing to Megatron Bridge, I investigated one such issue. What initially looked like a training instability turned out to be a subtle state-management bug inside the fine-tuning data sampler.
This is the story of how I tracked it down — PR #4601 (merged July 6, 2026), fixing issue #4565.
The Symptom
The original issue was straightforward:
After resuming fine-tuning from a checkpoint, users observed abnormal loss fluctuations.
At first glance, this looked like a familiar optimization problem.
My first instinct was to suspect the usual suspects:
- learning rate scheduling,
- optimizer state restoration,
- randomness introduced by distributed training.
Resume bugs often originate from these components.
But one observation completely changed the direction of the investigation.
The behavior only appeared after resuming from a checkpoint.
Fresh training runs behaved exactly as expected.
That suggested something important: the model itself was probably fine.
The problem likely lived somewhere in the state restored during checkpoint recovery.
The Resume Contract We Expect
Checkpointing exists so that long-running training jobs can continue exactly where they stopped.
When we resume training, we expect continuity across multiple pieces of state:
- model weights,
- optimizer state,
- learning rate scheduler,
- random number generators,
- data iteration progress.
That last item is surprisingly easy to overlook.
Imagine an epoch consisting of ten batches. Suppose training stops after Batch 4.
A correct checkpoint resume should continue from Batch 5, finish the current epoch, and then naturally begin the next epoch from the beginning.
That’s the behavior every user expects.
But that wasn’t what was happening.
Following the Data Pipeline
Since fresh training behaved correctly, I shifted my attention away from the optimizer and started tracing the data pipeline.
Eventually, the investigation led to the fine-tuning batch sampler.
Specifically, I focused on how it tracked one small piece of state:
consumed_samples.
This variable records how much of the dataset has already been processed before a checkpoint is written.
Restoring it during resume is exactly the right thing to do.
The problem wasn’t that it was restored.
The problem was that it never truly stopped influencing future epochs.
The Hidden Bug
The sampler correctly restored the number of consumed samples when training resumed.
So far, so good.
However, that resume offset was effectively treated as permanent instead of temporary.
Instead of applying the offset once to finish the interrupted epoch, it kept influencing every subsequent epoch.
The beginning of the dataset was repeatedly skipped. The tail was repeatedly replayed.
Nothing crashed. Nothing logged an error. The sampler continued producing perfectly valid batches.
But over time, the model was learning from a biased subset of the training data — some samples seen multiple times, others effectively disappearing after the checkpoint.
That imbalance changed the optimization trajectory, which eventually showed up as unstable loss curves. The loss wasn’t the root cause. It was simply the first visible symptom.
Why This Was Such a Difficult Bug
From an infrastructure perspective, everything appeared healthy.
- Checkpoint restoration succeeded.
- GPUs stayed fully utilized.
- Training loops completed normally.
- No assertions failed.
- No exceptions were raised.
The failure wasn’t operational.
It was semantic.
Every subsystem behaved correctly in isolation.
The bug only emerged from how one piece of state persisted across epoch boundaries.
These are often the hardest bugs to diagnose because nothing obviously looks broken.
Instead, the system quietly drifts away from the behavior users expect.
Designing the Fix
The solution wasn’t simply resetting a counter.
The sampler still needed to resume exactly where training had stopped.
The challenge was ensuring that the resume offset only affected the interrupted epoch.
After that, epoch boundaries needed to behave like true boundaries again.
The fix made the sampler epoch-aware in data/samplers.py:
consumed_samplesnow advances as batches are produced,- the resume offset is computed relative to the current epoch,
- every new epoch starts from the correct boundary.
Shuffle and seed parameters were threaded through build_pretraining_data_loader, with shuffle enabled for fine-tuning train loading in data/loaders.py.
In other words:
Resume once.
Finish the interrupted epoch.
Then continue exactly as if training had never been interrupted.
Making Training More Reproducible
While working on the sampler, I realized there was another opportunity to improve the training pipeline.
Previously, fine-tuning data was processed sequentially.
The updated sampler now supports deterministic per-epoch reshuffling.
Each epoch derives its ordering from:
- a user-provided seed,
- the current epoch number.
This provides two important benefits.
First, every epoch sees the data in a different order, which is generally preferable for model training.
Second, the ordering remains completely reproducible.
If two engineers resume from the same checkpoint using the same seed, they’ll observe exactly the same sequence of batches.
That level of determinism makes debugging training regressions dramatically easier.
Guarding Against Future Regressions
Data pipeline bugs are particularly dangerous because they rarely fail loudly.
Instead, they silently change what the model learns.
To prevent this issue from returning, I added regression tests in tests/.../data/test_samplers.py covering:
- mid-epoch checkpoint resume,
- epoch-boundary resume,
- deterministic per-epoch reshuffling,
- checkpoint resume parity,
- seed reproducibility.
These tests don’t simply verify that training continues.
They verify that training continues on the correct data.
That’s an important distinction.
What This Taught Me
One lesson from this bug has stayed with me.
In machine learning systems, correctness isn’t just about tensors, kernels, or gradients.
It’s equally about data.
A perfectly implemented model can still produce degraded results if the data pipeline quietly behaves differently than intended.
In this case, every individual component looked reasonable.
Checkpoint loading worked.
The sampler generated batches.
Training progressed normally.
The bug only appeared because one small piece of state - consumed_samples - persisted longer than it should have.
Those are often the most interesting engineering problems.
Not because they’re complicated.
But because they’re almost correct.
Final Thoughts
I enjoy working on bugs like this because they sit at the intersection of software systems and machine learning.
A few lines of sampler state management determined whether the model saw its complete training distribution - or repeatedly learned from only the tail of the dataset.
When training results drift after a checkpoint, it’s tempting to blame optimizers, hyperparameters, or model architecture.
Sometimes, though, the answer is hiding somewhere much simpler.
Sometimes it’s hiding inside the data pipeline.
And those are exactly the bugs worth hunting.