[Connor-XY]

Why. save_checkpoint and load_checkpoint read the global parallel grid for every group they were not given:

• legacy shard names and the ranks that write a legacy checkpoint;
• the RNG shard key and gather;
• replica ids and the fully parallel save and load groups;
• the local checkpoint group, the layer-wise optimizer file index and the log lines;
• the fsdp_dtensor optimizer metadata, which also replaced the caller's DP group with the global one.

A model built on its own grid was checkpointed under the global ranks. The training loop copied seven groups out of the model's collection into save_checkpoint and five into load_checkpoint, and any group it did not find fell back to the global grid.

What.

• save_checkpoint and load_checkpoint take pg_collection, defaulting to the model's own collection. Only a model built without one uses the global grid.
• The groups match what the training loop already passed: dp and expt_dp for RNG state and writers, dp_cp_gtp_remat for replica ids, dp_gtp_remat for dataloader state. A collection without the GTP-remat fields falls back to the plain groups.
• Explicit group arguments still override the matching group. Legacy shard names, and the layer-wise optimizer file index on load, always follow the collection. Every in-repo caller that passes groups takes them from that same collection.
• get_checkpoint_name reads the ranks it is not given from the collection.
• The training loop passes the model's collection instead of copying groups out of it.
• The checkpoint converter's loader sets the rank of its stand-in pipeline group, which now names legacy shards.

Tests. New tests/unit_tests/training/test_checkpointing_pg_collection.py: legacy and torch_dist round trips on a TP2×PP2 collection over a TP1×PP1 global grid, with checkpointing's global accessors replaced by a stand-in that fails on any read, plus CPU tests for shard naming and collection resolution.

On 4 GPUs against main: test_checkpointing_pg_collection.py passes all 5 of its tests, which are new, including the 3 round trips. test_checkpointing.py: 102 passed here vs 102 on main. dist_checkpointing/test_layer_wise_optimizer.py: 75 passed here vs 75 on main, with 31 skipped on both sides; 44 of its tests fail on both sides because they need 8 ranks. No test regresses. Before the sync with main, on 4 CPU gloo ranks, the 2 CPU tests pass and the 3 round trips need CUDA. In the integration run (main with all #6307 PRs merged), this PR's test files show no regressions against main, and the 44 layer-wise optimizer tests pass on 8 GPUs (two nodes). That run covers the full unit test suite; test_checkpointing.py, test_attention.py, test_multi_token_prediction.py, test_moe_single_grouped_weight_numerics.py, test_training.py, tests/unit_tests/tools/checkpoint/ and the Megatron-FSDP v1 fused-MLA checkpoint test show no regressions there. End-to-end training smokes on that tree, including train-and-resume runs, give bit-identical losses to main, and resumed runs match their uninterrupted runs. The converter loader change has no unit test in this PR. The ckpt_converter functional test exercises it, but its recipe sits in the mr-broken scope.

Compatibility. Only new keyword arguments. On the standard grid the groups, names and contents are unchanged. For callers that pass no groups:

• With GTP remat, RNG state and legacy writers use the replicate dp and expt_dp, like the training loop, instead of the GTP-inclusive defaults. With expert GTP remat, the fully parallel ep_dp save and load group moves from the GTP-inclusive expert-DP group to expt_dp.
• fsdp_dtensor loads ask DCP for rng_state.(pp, tp) instead of (0, 0), so each rank restores its own RNG state. A strict load into a different TP/PP layout fails on the missing key.
• Legacy shard names follow the model's collection, not the _MPU_* setters. An external converter that emulates stages with those setters must also set its stand-in pipeline group's rank, as tools/checkpoint/loader_base.py does here. Otherwise it reads stage 0's shard for every stage.
• Passing a MultiModuleProcessGroupCollection raises ValueError.

Not covered here:

• fsdp_dtensor expert keys still come from the global EP grid through get_ep_layer_offset; Name FSDP-DTensor experts by the model's expert-parallel layout #8045 fixes that.
• The checkpoint helpers (get_checkpoint_name, get_rng_state, _build_sharded_state_dict_metadata, maybe_save_dataloader_state, ...) still fall back to the global grid silently for callers without a collection and for hand-built collections that lack a field. Warn when training checkpoint helpers fall back to the global grid #8758 adds warnings to them.

Landing: #8033 adds the same set_rank call in the converter loader's per-stage loop, on the other side of set_pipeline_model_parallel_rank. Git merges both without a conflict, so whichever PR lands second drops one copy. #8744 rewrites that loop and conflicts with this PR in tools/checkpoint/loader_base.py. Whichever lands second takes that PR's loop. #8037 merges cleanly with this PR.

Synced with main 77391e7 on 2026-11-01 by a merge commit that lists each conflict and its resolution, so existing commits and review links stay valid.

Part of #6307.

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Connor-XY]

Superseded by #8891, which carries this change as commit 73be752 under the same title, with "Originally #8740." in its message. Closing so review happens in one place. The branch stays on the fork.