[Connor-XY]

Why. SequentialMLP builds each local expert MLP with only tp_group=expt_tp and no collection. The expert fc1 and fc2 then:

• resolve their GTP-remat group through resolve_gtp_remat_group(None, is_expert=True), which reads the global expt_gtp_remat;
• stamp no replica group on their GTP shards, so checkpointing elects the shard writer from parallel_state.get_expert_data_parallel_rank().

A model on its own grid gets its experts sharded over the global group, or left unsharded when the globals do not exist, and reads the global grid to save them. TEGroupedMLP already passes its collection to its grouped linears.

What. Pass SequentialMLP's pg_collection to each expert MLP, which hands it to fc1 and fc2. They take expt_gtp_remat as the shard group and stamp expt_dp as the replica group of their GTP shards.

Tests. New tests/unit_tests/transformer/moe/test_sequential_mlp_process_groups.py, with local and TE linears:

• builds SequentialMLP from a full collection while every parallel_state accessor and use_mpu_process_groups() raise;
• with expert GTP-remat on distinct communicators over the global ranks, checks each expert weight's shard and replica group, and that sharded_state_dict reads no global group.

In the integration run on 4 GPUs, test_sequential_mlp_process_groups.py gives 2 passed and 2 skipped; main has no such file, so all 4 tests are new. In the integration run, this PR's test files show no regressions against main. The 2 passing tests are the raising-accessor cases; a 4-GPU run against main before the 2026-11-01 sync showed them failing on main and passing here. The two GTP cases skip on the test image because GTP needs Transformer Engine 2.19 or later, so the GTP-sharded expert path has not run on GPU yet. Do you want a run on a Transformer Engine 2.19 image before review?

Compatibility.

• On the standard grid the collection holds the global groups, so sharding, initialization and checkpoint replica ids are unchanged.
• A collection without expt_gtp_remat keeps the fallback to the globals.
• build_inference_pg_collection sets expt_gtp_remat=None on purpose. SequentialMLP experts on such a collection now stay unsharded, matching TEGroupedMLP. On main they took the global expt_gtp_remat.
• Out-of-tree linear builders such as Kitchen now receive SequentialMLP's real collection instead of None. One that takes its TP group from pg_collection.tp would switch the experts from expert TP to dense TP. That can't be checked in this repo.
• SequentialMLP.sharded_state_dict with metadata=None still reads the global DP x CP group. Use the owner's DP x CP group for replica ids in SSM, MoE and TE layers #8732 fixes that and merges cleanly with this change.
• The test helper uses the checker's private _is_deprecated_accessor. It can switch to Add test helpers to forbid global grid reads and build collections #8017's forbid_global_process_groups once Add test helpers to forbid global grid reads and build collections #8017 lands.
• Needs determinism-exempt: experts.py is a registered kernel file, and only process groups change.

The branch was synced with main 77391e7 on 2026-11-01 by a merge commit, with no conflicts, 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 #8899, which carries this change as commit 18f3713 under the same title, with "Originally #8724." in its message. Closing so review happens in one place. The branch stays on the fork.