[Connor-XY]

Depends on #6259, #6303 (which contains #6293), #8017, #8043 and #8049. Please review the diff from the merge of their heads: scoped diff. The base takes #8017's update that merges #6259 and sets the Gloo fields of build_test_pg_collection. #6303 is in the base because the test-caller commit uses ProcessGroupCollection in tests/unit_tests/dist_checkpointing/utils.py and tests/unit_tests/training/test_param_norm.py, which import it only with #6303. The base's process-group allowlist records the 172 reads of its tree, without the T5Model shim read that #6303 removes.

Land this after #8738. With Gloo groups on, the default, setup_model_and_optimizer passes no collection to get_megatron_optimizer. Without #8738, stock training would warn once there, and so would test_training.py, test_fp8_param.py, test_fp4_param.py and test_cuda_graphs.py.

Why. Three optimizer paths read the global parallel_state groups without a warning when the caller omits them:

• get_megatron_optimizer without pg_collection builds one with use_mpu_process_groups();
• clip_qk without dp_cp_group reduces the attention logits over the global DP×CP group;
• Float16OptimizerWithFloat16Params.sharded_state_dict rebuilds the checkpoint entry of a GTP-remat shard that has no model entry, such as the Mamba and GDP in_proj weights under distributed Muon. The rebuild read the global TP and DP groups, even for an optimizer built from a collection.

A model on its own grid, such as a MIMO module, then sharded its optimizer state, clipped its weights or described its shards with another model's ranks. Without a global grid, the rebuild hit an assertion.

What.

• get_megatron_optimizer and clip_qk fall back through resolve_process_groups. It warns once with a ProcessGroupFallbackWarning (deprecated in 0.21, removed in 0.23) and returns the same groups as before:
  • get_megatron_optimizer builds every collection field, as use_mpu_process_groups() did. Several of the optimizers' readers treat an unset field differently from one set to None.
  • clip_qk resolves dp_cp_gtp_remat, and only when a layer has logits to reduce.
• Float16OptimizerWithFloat16Params takes a keyword-only pg_collection, which get_megatron_optimizer and LayerWiseDistributedOptimizer pass. The GTP rebuild takes tp from it, plus dp_cp_gtp_remat or, when that is unset, dp_cp, as the model's own entries do. A collection without tp, or with neither DP×CP group, raises ValueError. Without a collection, the rebuild warns and uses the global groups. It still resolves groups only when a shard needs a rebuild.
• 46 unit-test calls to get_megatron_optimizer and get_megatron_muon_optimizer pass ProcessGroupCollection.use_mpu_process_groups(). That is the collection the fallback builds, so these tests run on the same groups without the warning. Two of these calls are in main's new test_shortcut_moe_training_and_clipping, which the sync merge adapts so that it does not hit the warned fallback. The initializer-cleanup tests pass an empty collection instead of stubbing use_mpu_process_groups().
• The checker goes from 172 reads to 171. The three fallbacks become guarded reads, and the rebuild reads one collection instead of two accessors.

Tests.

• New tests/unit_tests/optimizer/test_optimizer_process_groups.py:
  • the factory builds on the collection it is given, with global reads forbidden and the warning made an error;
  • the layer-wise optimizer hands its collection to the optimizers it wraps;
  • without a collection, the factory warns and builds the same collection as use_mpu_process_groups();
  • the GTP rebuild uses the collection's groups for four field combinations and rejects a collection without a DP×CP group;
  • the rebuild warns without a collection, and needs no groups when every shard has a model entry.
• test_qk_clip.py: a passed group is used with no global read. An omitted group warns and reduces over the same global group. A model without logits needs no group.

On 4 GPUs against its base, 22 of the 23 test files that exercise this change finish. In those 22, no test that passes on the base fails here. Across the 19 files with a pytest summary, 642 tests pass here vs 628 on the base. The 12 tests of the new test_optimizer_process_groups.py pass. test_qk_clip.py renames one test and adds one: 3 pass here vs 2 on the base. One test in test_layer_wise_optimizer.py fails on the base and passes here. The three generalized_tensor_parallel test files run no tests on either side; they skip because the container's Transformer Engine is older than GTP needs. distributed/mfsdp_v2/test_mcore_hybrid.py timed out on both sides. 78 tests in dist_checkpointing/test_optimizer.py fail on both sides because they need 8 ranks; they pass on 8 GPUs (two nodes) in the integration run. The two CUDA-graph end-to-end cases in distributed/mfsdp_v1/test_mcore_fully_sharded_data_parallel.py fail on both sides, the same way as on main.

In the integration run, which merges all #6307 PRs onto main 77391e7, this PR's test files show no regressions against main on 4 GPUs and 8 GPUs (two nodes). test_mcore_hybrid.py timed out there on both sides as well. End-to-end training smokes on that tree (22 configurations on 4 GPUs) give bit-identical losses to main and record no process-group fallback warning. Before the sync, the branch also ran on 4 CPU gloo ranks, one test file at a time, with no regressions. On CPU, 9 tests in the new file pass and 3 need CUDA. The checker tests pass.

Compatibility.

• On the standard grid the groups are the same, so checkpoints do not change. The new constructor argument is optional and keyword-only. GraphedGradCopyFloat16Optimizer forwards it.
• The rebuild checks its groups at the first checkpoint save that needs a rebuild, not at construction. Such a save needs GTP, distributed Muon and Mamba or GDP in_proj. setup_process_groups_for_optimizer accepts a collection with neither dp_cp nor dp_cp_gtp_remat when CP is 1, and the rebuild then raises. The DP×CP group does not change the rebuilt entry, because a GTP shard elects its writer over the group stamped on the parameter.
• Without a collection, the rebuild resolves dp_cp along with tp and dp_cp_gtp_remat. That is one global group more than before, and the standard grid never uses it.
• The rebuild warns under the name Float16OptimizerWithFloat16Params, also for GraphedGradCopyFloat16Optimizer.
• Not changed here: LayerWiseDistributedOptimizer reads expt_dp through getattr(..., None). A collection without it treats the expert data-parallel axis as size 1 and passes None, which torch.distributed reads as WORLD, to the expert all-gather. That predates this PR, and every in-repo collection sets expt_dp.
• The test-caller commit touches 19 test files across several CODEOWNERS groups. Should it move into its own PR?
• Needs determinism-exempt: optimizer/__init__.py and optimizer/optimizer.py are registered kernel files. The changes there are group plumbing and warnings; numerics do not change.

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 #8903, which carries this change as commit c388533 under the same title, with "Originally #8752." in its message. Closing so review happens in one place. The branch stays on the fork.