[tylerpoon]

Summary

Fix multimodal checkpoint resume with --train-full-dataset and strict dataloader-state restoration.

The training-horizon probe runs before model/optimizer checkpoint setup. Previously it reset args.iteration to zero and invoked the real dataloader provider with restoration enabled. A run resuming a valid nonzero checkpoint consequently looked for dataloader state at iteration zero and failed before training.

• Extract the existing full-dataset sample-count probe into a focused helper.
• Skip restoration only for that temporary, length-only loader; preserve args.iteration.
• Keep restoration enabled by default for the actual training loader after checkpoint setup, including strict missing-state errors.
• Add regression coverage for probe behavior, resolved checkpoint iteration, strict restoration after probing, argument validation, and CP-deduplicated sample counts.

Only providers opting into supports_train_full_dataset use this path. Opt-in providers must accept the new restore_dataloader_state keyword. No model math, optimizer, checkpoint format, or ordinary non-opt-in training path is changed.

Validation

• The fix on the earlier main base passed 30 focused tests on each of four GPU ranks, including the distributed CP sample-count check.
• A 256-GPU multimodal SFT restore test resumed checkpoint 100 through iteration 110 successfully; subsequent segments reached iteration 270 and saved model and dataloader state. The longer comparison was split by scheduler preemption, not uninterrupted.
• Replayed loss windows differed by at most 0.055% from the uninterrupted reference; this is not a claim of bitwise equivalence or original-production parity.
• Rebased onto main 7de072aa9e; all 30 focused tests passed again on each of four GPU ranks on PR commit 8a6f9d674d.
• A separate four-GPU check verified that the private diagnostic logger reports both averaging conventions correctly without altering gradients. Two sequential 256-GPU, 10-update diagnostic probes are queued; their results are not yet available.

The private launchers, datasets, production-specific scripts, and diagnostic-only metric hooks are deliberately excluded. Historical MTP/norm reporting differences are being investigated separately; no speculative model/optimizer changes are included here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tylerpoon]

/ok to test 982d863

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/36495942402

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/36508505601