[Connor-XY]

This PR makes the optimizer, the tensor-parallel helpers and the process-group setup helpers warn when a caller leaves out a process group. Before, they used WORLD or the global parallel_state grid without any signal. Each fallback now emits a ProcessGroupFallbackWarning (deprecated in 0.21, removed in 0.23) and returns the same groups as before. Callers that build optimizers or DDP, or seed the RNG, without a ProcessGroupCollection are affected, as are maintainers of the optimizer, GTP, tensor-parallel and DDP setup code. It consolidates 2 draft PRs unchanged, one commit per member, so please review commit by commit. A final commit regenerates tools/process_group_usage_allowlist.json. It builds on #6303, #8890 (core process-group APIs: #8017, #8049), #8889 (optimizer plumbing: #8042, #8043) and #8895 (training loop and strict mode: #8031). It lands after #8898 and #8895, which make stock training pass its collection to DistributedDataParallel, get_megatron_optimizer and the RNG seed, so stock runs reach none of the new warnings.

Review

Commits of this unit only: Connor-XY/Megatron-LM@ea3b434...ba30732

The full PR diff also shows the earlier units until they merge.

Members

• Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734: Warn on optimizer, TP and helper fallbacks to WORLD or the global grid. The grad-stats getters of DistributedOptimizer, MegatronOptimizer and LayerWiseDistributedOptimizer warn when the norm reduction would run over WORLD, and the duplicate filters, model_parallel_cuda_manual_seed, gtp_replica_rank and the DDP, optimizer and GTP-remat group helpers warn before they read the global grid.
• Warn when the optimizer falls back to the global process groups #8752: Warn when the optimizer falls back to the global process groups. get_megatron_optimizer without pg_collection and clip_qk without dp_cp_group warn and resolve the same global groups, and Float16OptimizerWithFloat16Params takes a keyword-only pg_collection that its GTP-remat checkpoint rebuild reads instead of the global TP and DP groups.

Behavior changes for callers

• On the standard grid, groups, ranks, numerics, seeds, duplicate filtering and checkpoints do not change.
• resolve_gtp_remat_group gains a keyword-only owner. Existing calls keep working.
• A distributed optimizer with intra_dist_opt=None, a non-distributed optimizer with mp=None and a layer-wise optimizer without intra_dist_opt warn once per process that the grad-norm reduction uses WORLD. They keep working.
• Hand-built optimizers, DistributedDataParallel without pg_collection, and direct calls without groups to the duplicate filters, model_parallel_cuda_manual_seed, setup_process_groups_for_optimizer, get_megatron_optimizer or clip_qk warn once per owner, and raise under an error filter.
• About 700 test calls of model_parallel_cuda_manual_seed(123) without a collection now warn once per process. They are not converted here.
• The fallbacks read group ranks. A rank forced with set_tensor_model_parallel_rank or the expert rank setters no longer replaces the group rank there. With lazy_mpu_init, the seed uses the TP group rank.
• The GTP-remat checkpoint rebuild raises ValueError at the first save that needs it if the collection lacks tp or both dp_cp and dp_cp_gtp_remat. It only runs with GTP, distributed Muon and Mamba or GDP in_proj weights.
• Needs determinism-exempt: the registered kernel files change only by warnings, group plumbing and owner= keywords.

Teams

core-adlr, core-nemo, data-parallelism, dist-optimizer, gtp, mcore-emerging-optimizers, mcore-optimizer, mixture-of-experts-adlr, mixture-of-experts-devtech, reinforcement-learning, training, transformer

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.

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.