[Connor-XY]

Depends on #8757 (which contains #8741, #8037 and #8049) and #8740. Please review the diff from their merge: scoped diff.

• Save and load checkpoints on the model's process groups #8740 moves save_checkpoint and load_checkpoint onto the model's groups, and this PR's tests extend its test file.
• train_step resolves the model's collection with Run the training loop's remaining reads on its process-group collection #8741's _get_model_pg_collection and the owner argument that Finish moving the training loop onto its process-group collection #8757 adds.
• Both merges are clean.

Why. save_checkpoint and load_checkpoint pass the model's groups to every helper they call. The helpers still read the global grid, silently, for any group a direct caller omits. save_grads always reads it, so the wgrad, param, dgrad and activation dumps of a model built on its own grid are named and written by the global ranks.

What.

• These fallbacks now emit ProcessGroupFallbackWarning (deprecated in 0.21, removed in 0.23), naming the function and the omitted argument, and then read the same global groups as before:
  • get_checkpoint_name, for a rank or flag that is neither passed nor in a collection;
  • get_rng_state, for dp_group and dp_cp_group;
  • _build_sharded_state_dict_metadata and the default metadata of generate_state_dict;
  • maybe_save_dataloader_state, once per omitted group, before it branches on parallel ranks;
  • the restore group of a local non-persistent checkpoint with sharded state;
  • the fully parallel load group;
  • the collection of a model built without one, owned by save_checkpoint or load_checkpoint.
• Paths that read nothing global stay silent: explicit ranks, a base-directory name and a disabled dataloader save.
• save_grads takes a keyword-only pg_collection, which must set tp, pp, ep and expt_dp. The writers are the first ranks of expt_dp_gtp_remat, or of expt_dp for a collection without it. Without a collection it warns and reads the global grid as before.
• DataGradLogger, ActivationLogger and their module-level wrappers forward the collection. train_step passes the model's collection to all three dumps, with owner="train_step".

Tests.

• New tests/unit_tests/training/test_checkpointing_fallback_warnings.py:
  • on CPU, each fallback warns and calls the same accessor, and each explicit path neither warns nor reads the global grid;
  • on 4 GPUs, dumps on a TP2×PP2 grid unlike the global one;
  • on 4 GPUs, on the standard TP2×PP2×EP2 grid, checkpoint names, RNG shard keys, dump files and dataloader files match the fallback's.
• In test_checkpointing_pg_collection.py, the model-grid save/load round trips now also fail on any fallback warning, and the collection-resolution test checks the warned path.
• Under -W error::megatron.core.process_groups_config.ProcessGroupFallbackWarning, 7 existing maybe_save_dataloader_state tests in tests/unit_tests/test_checkpointing.py fail. They omit cp_group, which now warns. CI sets no such filter.

On 4 GPUs against its base, test_checkpointing_fallback_warnings.py passes 35 tests, all of them new, and test_checkpointing_pg_collection.py passes 5 here vs 5 on the base. In the integration run, which merges all #6307 PRs onto main and runs the full unit-test suite against main, this PR's test files show no regressions: 35 and 5 tests pass on 4 GPUs. End-to-end smokes on the same code, 22 training configurations on 4 GPUs including resumes, give bit-identical losses to main; resumed runs match their uninterrupted runs, and no process-group fallback warning is recorded.

Before the 2026-11-01 sync, the branch ran both test files on 4 CPU ranks with Gloo: 35 passed, the same per test as on the stacked version, and 33 of them are new. Two new 4-GPU tests and three existing ones fail there for lack of CUDA.

Compatibility.

• Every new argument is keyword-only and defaults to the old behavior, which now warns. Stock training and the checkpoint converter pass their groups and do not warn.
• On the standard grid, checkpoint names, writers and contents are unchanged.
• generate_state_dict builds its default torch_dist metadata from the dp_cp_group it is given, instead of the global DP×CP group. That is the same group on the standard grid, and save_checkpoint and load_checkpoint do not use the default.
• save_grads with a collection raises ValueError if tp, pp, ep or expt_dp is missing. That only matters when a dump is enabled, and stock collections set all four. A collection with expt_dp=None and no expt_dp_gtp_remat makes every data-parallel replica write the same file.
• The multimodal example's dataloader restore calls get_checkpoint_name without tensor_rank or pipeline_parallel (examples/multimodal/dataloader_provider.py). Until Use the model's process groups in ModelOpt, multimodal and MIMO examples #8761 passes them, resuming its dataloader warns once, and raises under --strict-process-group-fallbacks.
• Needs determinism-exempt only while Finish moving the training loop onto its process-group collection #8757 is unlanded. CI diffs against main, and Finish moving the training loop onto its process-group collection #8757 changes utils/common_utils.py, a registered kernel file. This PR's own diff passes the coverage check.

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

[Connor-XY]

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