[Connor-XY]

This PR makes Megatron DDP, finalize_model_grads and Megatron-FSDP use only the process groups of the model they serve. Today, NCCL user-buffer barriers, the DCP optimizer-key gather and several gradient reductions run on WORLD or on the global grid. They hang or pair the wrong ranks when one job holds several models or grids, or when ranks own different buffers. Callers that pass a pg_collection and maintainers of the data-parallel wrappers are affected. The PR consolidates 5 draft PRs from #6307 without changes. Each member is one commit, so review it commit by commit. It has no code dependency on other units or open PRs. The last commit regenerates tools/process_group_usage_allowlist.json. Members 4 and 5 already target the mcore_adapter_v1.py / mcore_adapter_v2.py split from #8471.

Review

The unit's own commits: Connor-XY/Megatron-LM@e82e858...0d09ab7

The branch is a linear chain on main. The full PR diff also shows the PRs and units below this one until they merge.

Members

• Scope DDP's NCCL user-buffer warmup barriers to the DP group #8815: Scope DDP's NCCL user-buffer warmup barriers to the DP group. With nccl_ub=True, _ParamAndGradBuffer runs its two warmup barriers on the buffer's data-parallel group instead of WORLD, so pipeline stages with different expert buffers no longer hang at setup.
• Use Megatron-FSDP's own groups for UBR barriers and DCP key gather #8816: Use Megatron-FSDP's own groups for UBR barriers and DCP key gather. It drops the WORLD barrier in the buffer constructor, runs the manual_buffer_registration() barriers on the user-buffer groups, and gathers DCP optimizer-state keys over the DP group, plus the expert DP group under EP.
• Read only the caller's process groups in finalize_model_grads #8027: Read only the caller's process groups in finalize_model_grads. With a pg_collection, it uses only that collection's groups, raises ValueError for a missing one before any communication, and no longer reduces num_tokens over WORLD, which made per-token-loss gradients 4x too small on TP=2 x PP=2.
• Create Megatron-FSDP v1 TP groups collectively and validate pg_collection #8051: Create Megatron-FSDP v1 TP groups collectively and validate pg_collection. It creates single-rank TP groups once with new_subgroups_by_enumeration on every rank, only for the axes the model uses, and raises ValueError naming any group the collection lacks.
• Reject GTP weight remat in Megatron-FSDP v1 #8725: Reject GTP weight remat in Megatron-FSDP v1. FullyShardedDataParallelV1 raises ValueError when dense or expert GTP weight remat is above 1 in the config or in the collection, because v1 cannot shard or reduce GTP shards correctly.

Behavior changes for callers

• Without nccl_ub, without preproc_state_dict_for_dcp_ckpt and without a pg_collection, nothing changes.
• With nccl_ub, DDP and Megatron-FSDP buffer setup no longer synchronizes ranks outside the buffer's groups. Megatron-FSDP runs one fewer barrier at construction.
• Under EP, the DCP key gather adds one small gather over the expert DP group.
• finalize_model_grads with a pg_collection requires tp and pp, embd when PP > 1, and the DP x CP and TP x DP x CP groups its options use. A missing group raises ValueError. A missing tp_dp_cp used to raise AssertionError.
• Megatron-FSDP v1 raises ValueError instead of AttributeError for a collection missing a needed group.
• Megatron-FSDP v1 with GTP weight remat now fails at construction. Expert GTP > 1 is rejected even for a dense model.

The members touch files registered for determinism coverage but not their reductions, so this PR needs the determinism-exempt label.

Teams

core-adlr, core-nemo, data-parallelism, megatron-fsdp

Testing

Each member was GPU-tested against its base before consolidation. The consolidated chain (every unit in landing order, on current main) was run against main on 4 GPUs: 258 test files, and no test fails on the chain that passes on main. Tests in 28 files fail on both, mostly tests written for 8 GPUs, and one file times out on both.

🤖 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.