[Silas-Zeng]

Problem

With cyclic dataloading and data sharding enabled, resuming after a data-parallel size change can silently repeat and skip samples because the sampler's buckets and resume offset depend on the effective DP size.

Fix

Persist explicit cyclic dataloader metadata in checkpoint arguments alongside the consumed sample count: the effective dataloader DP size and the sharding mode. On a normal resume, validate that metadata before adopting the checkpoint's training progress and fail clearly when the sharded layout cannot be preserved. Older checkpoints fall back to their existing DP arguments for compatibility; finetune and release loads keep their existing behavior.

Validation

• python -m py_compile megatron/training/checkpointing.py tests/unit_tests/test_checkpointing.py
• git diff --check
• Parameterized tests cover mismatch rejection, matching layouts, missing legacy metadata, sharding-mode changes, non-target modes, finetune/release, and equivalent effective DP sizes.
• pytest -q tests/unit_tests/test_checkpointing.py — 47 passed on NVIDIA RTX 5090 32GB, CUDA 12.8, PyTorch 2.8.0+cu128, Triton 3.4.0, Python 3.12.3.

Fixes #7645

[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.

[guihong-nv]

@Silas-Zeng The branch includes ~9 unrelated upstream commits (biggest structural issue). Of the 14 commits in this PR, only 5 are from you. The rest are merged mainline commits: MXFP8 via MFSDP adapter, HybridStack residuals, MIMO fixes, DCP V2, SSM fusion tests, GTP teardown fix, Docker parallelization. The total diff is +3,537/−279 lines, but the actual fix is almost certainly <100 lines touching checkpointing.py and test_checkpointing.py. This inflates the review surface dramatically and makes it hard to audit what the PR actually changes. The branch should be rebased onto current main with only the fix commits, or at minimum the maintainers should review only the your commits directly.

[svcnvidia-nemo-ci]

🛠️ Claude fix commit 7483a59d3bd6 (attempt 2)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Reformatted two SimpleNamespace(...) call sites in pr-head／tests／unit＿tests／test＿checkpointing.py to match black's expected output: collapsed the multi-line argument lists in test＿validate＿cyclic＿dataloader＿resume＿uses＿effective＿dp＿size (checkpoint＿args) and test＿validate＿cyclic＿dataloader＿resume＿rejects＿sharding＿mode＿change (args) onto single lines, dropping trailing commas. Whitespace only; no logic change, no new／deleted／renamed files, nothing committed.

Files changed by Claude

• tests/unit_tests/test_checkpointing.py

Why
Only the linting job failed: black --check --diff said "would reformat tests／unit＿tests／test＿checkpointing.py", exit 1. With skip＿magic＿trailing＿comma=true, black joins those arg lists onto one line as both fit in line＿length=100. I applied black's two printed hunks verbatim; my line deltas match the log hunk headers. black and pytest are unavailable (no network), so I verified the AST is identical pre／post, confirming a cosmetic edit. No unit failures in the log.

DCO
@Silas-Zeng, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test 7483a59

[svcnvidia-nemo-ci]

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.

[Silas-Zeng]

Thanks for the fixes. The Claude workflow timed out waiting for CI on 7483a59. Run 36857019494 is still waiting at cicd-wait-in-queue, with a pending deployment for the test environment; lint has passed. @dimapihtar, could you check what is needed for the CI run to proceed?

[Silas-Zeng]

Hi @dimapihtar @guihong-nv @wujingyue, thanks for enabling CI for these fixes. The CICD runs for #7648 and #7649 have now passed.

Could you help route these two PRs to the appropriate code owner for review? They address repeated or skipped samples when resuming cyclic dataloaders after a DP-size change. I'm happy to provide any additional validation needed.

Thanks for your time!