[Connor-XY]

Depends on #8044. Please review the diff from the merge of #8044 into main: scoped diff.

Why. Called without metadata, sharded_state_dict() of these modules filled in the global DP x CP group:

• the Mamba, GDN and GDP mixers and HybridStack;
• the MoE expert classes TEGroupedMLP and SequentialMLP;
• TELayerNormColumnParallelLinear.

A module on its own grid then computed replica ids over another grid's DP ranks, so several ranks, or none, wrote a shard. A model built from explicit groups could not save these modules without a global grid. #8044 lets modules pass their own group to ensure_metadata_has_dp_cp_group and uses it for the language model; these modules still used the default.

What.

• Each module passes its collection's dp_cp_gtp_remat, or dp_cp when that is unset, to ensure_metadata_has_dp_cp_group. A group already in the metadata still wins.
• HybridStack fills the metadata before it visits its layers. The GDP mixer hands its group to its gated norm, as it does the TP group.
• The experts fill in the dense group, the one a model passes down. Their replica ids still take the DP entry from the expert data-parallel group.
• TEGroupedMLP, SequentialMLP and TELayerNormColumnParallelLinear keep the collection they are built with.

Tests. With the global DP x CP accessor patched to raise, new tests compare (key, offset, replica_id) against the module's own group for a Mamba mixer, a hybrid stack, both expert classes and the TE layer:

• tests/unit_tests/ssm/test_ssm_sharded_state_dict_dp_cp_group.py and tests/unit_tests/transformer/moe/test_experts_sharded_state_dict_dp_cp_group.py are new.
• test_te_layernorm_column_parallel_linear.py also checks the TE layer on a TP=1 grid that differs from the global TP=2 grid, and without a collection. The no-collection test guards the global default and passes on Compute checkpoint replica ids over the module's own DP x CP group #8044 too.
• test_gdp_tp_checkpoint.py covers the GDP mixer and its norm on CPU, including the dp_cp fallback. Its stand-in collection is now a ProcessGroupCollection, because the mixer reads both DP x CP fields.

The _GDNBase change has no test of its own; it is the same one-line change as the Mamba and GDP mixers.

On 4 GPUs against its base (main with #8044 merged): test_te_layernorm_column_parallel_linear.py: 8 passed here vs 5 on the base, and the 3 new tests pass. #8044's test_sharded_state_dict_dp_cp_group.py passes 6 tests on both sides. An earlier version of the experts test failed on both sides, because its config used a top-1 softmax router, which TransformerConfig rejects. It now routes each token to two experts and passes. Before the sync, the branch also ran on 4 CPU gloo ranks with no regressions; the 3 new GDP cases pass, the SSM tests need CUDA, and the TE and MoE tests skip without Transformer Engine. In the integration run (all #6307 PRs merged onto main), this PR's test files show no regressions against main, and every test in the four files passes on 4 GPUs: test_gdp_tp_checkpoint.py 8 passed vs 5 on main, test_te_layernorm_column_parallel_linear.py 8 vs 5, and the new SSM and experts files 2 each.

Compatibility.

• Replica ids on the standard grid do not change. Modules built without a collection keep the global default. No new global reads.
• The experts fill the dense DP x CP group, not expt_dp. This keeps standard-grid replica ids as they are and matches what a model passes down. The DP entry of each expert replica id already comes from expt_dp.
• dp_cp_gtp_remat or dp_cp is evaluated even when the metadata already carries a group, and the GDP mixer reads both fields at construction. A duck-typed collection without these attributes, such as a SimpleNamespace, now raises AttributeError. No in-repo or Megatron-Bridge caller builds one.
• Pass the owner's collection to FP8/FP4 quantization contexts #8035 adds the same self._pg_collection = pg_collection line to TELayerNormColumnParallelLinear.__init__, a few lines away. Git merges both without a conflict, so whichever PR lands second should drop its copy.
• Merges cleanly with Give SequentialMLP experts the expert groups of their collection #8724, which edits the expert MLP construction next to this change in experts.py.
• Not covered here: the MambaLayer, MoELayer and Router entry points without a collection, the MLA fused q/kv-down split and the global TP read in the Mamba mixer's ExtendedRMSNorm. Warn when checkpoint helpers fall back to the global process groups #8753 handles them. SharedExpertMLP.sharded_state_dict with metadata=None and the shared-expert gate still raises TypeError. No PR in [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 fixes that yet.
• Needs determinism-exempt: experts.py, transformer_engine.py, mamba_mixer.py, gated_delta_product.py and gated_delta_net/common.py are registered kernel files, and only checkpoint metadata groups change. While Compute checkpoint replica ids over the module's own DP x CP group #8044 is unmerged, its edit to the registered transformer/utils.py needs the label too.

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