[Connor-XY]

This PR moves the optimizer's duplicate filters and expert naming off the global parallel state, and makes the tensor-parallel and GTP helpers warn when they still fall back to it. It's part of the parallel_state deprecation (#6307). Gradient-norm and zero-count filtering now treat a TP or GTP group that the model's collection sets to None as an axis that's off, instead of reading the global rank, so a model on its own grid gets the same norm as on the standard grid. The distributed optimizer, the param-group dump and the MFSDP v2 optimizer name experts by the model's expert-parallel layout. The duplicate filters, model_parallel_cuda_manual_seed, resolve_gtp_remat_group, gtp_replica_rank and the DDP and optimizer group setup helpers emit a ProcessGroupFallbackWarning when a caller leaves out a group. Callers that build optimizers, DDP or linears, or seed the RNG, without a ProcessGroupCollection are affected, as are maintainers of the optimizer, distributed optimizer and GTP code. It carries #8042, three parts of #8045, #8731, and the core and GTP parts of #8734. Each commit is one original change, so please review commit by commit. A final commit regenerates tools/process_group_usage_allowlist.json. It builds on #8890 (#8049's resolve_process_groups and warn_global_process_group_fallback, which the new warnings use) and #8892 (#8045's get_ep_rank_and_size and get_model_ep_rank_and_size, which the optimizers call). It lands after #8891, which passes stock training's collection to DDP and the optimizer, #8895, which seeds the RNG from it, and #8898, which passes it to the optimizer with Gloo groups on, so stock runs reach none of the new warnings.

Review

This PR's own commits: Connor-XY/Megatron-LM@25ecaf8...41edc4b

The full PR diff also shows the earlier PRs of the stack until they merge.

Commits

• Filter optimizer grad-norm duplicates by the collection's TP and GTP groups #8042: Filter optimizer grad-norm duplicates by the collection's TP and GTP groups. A new set_duplicate_filter_groups(optimizer, pg_collection) in optimizer/optimizer.py sets the TP, expert-TP and GTP groups for the grad-norm and zero-count filters, and replaces the group assignments in _get_megatron_optimizer_based_on_param_groups, _get_megatron_emerging_optimizer and LayerWiseDistributedOptimizer. param_is_not_tensor_parallel_duplicate, param_is_not_gtp_duplicate and count_zeros_fp32 take a keyword-only use_global_fallback; optimizers set up from a collection pass False, so a None group means the axis is off instead of "read the global rank". This fixes an under-counted gradient norm, and so too-weak clipping, when the model's collection turns TP or GTP off next to a global grid that has them on. New cases in test_grad_norm_gtp_invariance.py cover axes off in the collection, the collection's expert TP and a collection without tp.
• Name FSDP-DTensor experts by the model's expert-parallel layout #8045 (mcore-optimizer part): get_megatron_optimizer reads the EP rank and size from its pg_collection with get_ep_rank_and_size and passes them to get_global_unique_param_name when it dumps the param-to-param-group map. The other parts land in Use modules' own process groups in BERT, GPT, MLA, attention and checkpoints #8892 (the FSDP-DTensor helpers), Save and load checkpoints on the model's process groups #8981 (training) and in this PR (dist-optimizer, tests).
• Name FSDP-DTensor experts by the model's expert-parallel layout #8045 (dist-optimizer part): DistributedOptimizer names its parameters with the EP rank and size of its first model chunk (get_model_ep_rank_and_size) when the model has experts, and passes them to handle_experts_in_state_dict. The other parts land in Use modules' own process groups in BERT, GPT, MLA, attention and checkpoints #8892, Save and load checkpoints on the model's process groups #8981 and in this PR (mcore-optimizer, tests).
• Name FSDP-DTensor experts by the model's expert-parallel layout #8045 (tests part): Adds test_fsdp_dtensor_expert_offsets.py, test_distrib_optimizer_param_name.py and test_param_group_map_expert_names.py, which check expert offsets from integers, collections and wrapped models, and that the distributed optimizer and the param-group dump follow a custom expert layout. The other parts land in Use modules' own process groups in BERT, GPT, MLA, attention and checkpoints #8892, Save and load checkpoints on the model's process groups #8981 and in this PR (mcore-optimizer, dist-optimizer).
• Name the MFSDP v2 optimizer's experts by the model's EP layout #8731: Name the MFSDP v2 optimizer's experts by the model's EP layout. FullyShardedOptimizer passes the EP rank and size from the model's collection, found by get_model_ep_rank_and_size, when the model has experts. Dense models skip the lookup. Tests in test_fully_sharded_optimizer_param_name.py and the MFSDP v2 MoE test.
• Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734 (core part): The duplicate filters, model_parallel_cuda_manual_seed, setup_process_groups_for_optimizer and setup_process_groups_for_ddp warn before they read the global grid. The TP filter and the seed read group ranks through resolve_process_groups, and the GTP filter and the seed resolve GTP groups with resolve_gtp_remat_group, instead of the parallel_state rank getters. resolve_gtp_remat_group gains a keyword-only owner and warns when it falls back to a global GTP group of more than one rank; the TE, inference and tensor-parallel linears and VocabParallelEmbedding pass their class name. The intra_dist_opt error message now says what the group is for. The other parts land in Warn on optimizer process-group fallbacks and pass groups to the config logger #8903 (optimizer grad-stats warnings and their tests), Warn on finalize_model_grads, DSGQA and FSDP fallbacks to global groups #8905 (data-parallelism), Deprecate the parallel_state grid accessors and quiet training's fallback reads #8907 (training) and Enforce the process-group checks and warn on the last core grid reads #8909 (MoE), plus the GTP part here.
• Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734 (gtp part): gtp_replica_rank resolves its global fallback through resolve_process_groups (dp_cp, or expt_dp for routed experts), which warns, and raises RuntimeError when no such group exists. The other parts land as listed above.
• Tests moved in from Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734: test_grad_norm_gtp_invariance.py stops patching the rank getters that tensor_parallel.layers no longer imports. They land here with the core part, which they exercise.

Behavior changes for callers

• A collection passed to the optimizer builders must set tp (None when TP is off). Otherwise set_duplicate_filter_groups raises ValueError; before, an unset tp read as None and the filters used the global TP rank.
• In that collection, tp, expt_tp, gtp_remat or expt_gtp_remat set to None now means the axis is off for the duplicate filters, not "use the global group". An unset expt_tp defaults to tp.
• An unset gtp_remat or expt_gtp_remat resolves through resolve_gtp_remat_group, as the model's layers resolve it, instead of getattr(..., "gtp_remat", None). When the field is unset, these two can return different groups.
• Expert names on the standard grid don't change, so checkpoints from the standard grid load unchanged. Optimizer state saved on a custom EP grid before this change uses global-layout names and won't match.
• On the standard grid, groups, ranks, seeds, duplicate filtering and checkpoints don't change.
• DistributedDataParallel without pg_collection, setup_process_groups_for_optimizer without one, and direct calls without groups to the duplicate filters, model_parallel_cuda_manual_seed or gtp_replica_rank warn once per owner and argument, and raise under an error filter. resolve_gtp_remat_group warns only when the global GTP group has more than one rank.
• About 520 test calls of model_parallel_cuda_manual_seed 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.
• Needs determinism-exempt: the registered kernel files change only by warnings, group plumbing and owner= keywords.

Teams

core-adlr, core-nemo, dist-optimizer, gtp, mcore-optimizer

Testing

I ran the 7 unit-test files that cover this change at this PR's commit on one 4-GPU GB300 node. I compared each failure with the same test at the end of the stack (#8989, same code) and at the commit the stack starts from.

• All 50 tests pass. Tests that skip themselves on this setup aren't counted.
• The process-group allowlist check passes at this commit.

Test files

• tests/unit_tests/distributed/mfsdp_v2/test_mcore_moe.py
• tests/unit_tests/optimizer/test_distrib_optimizer_param_name.py
• tests/unit_tests/optimizer/test_fully_sharded_optimizer_param_name.py
• tests/unit_tests/optimizer/test_grad_norm_gtp_invariance.py
• tests/unit_tests/optimizer/test_param_group_map_expert_names.py
• tests/unit_tests/test_process_groups_config.py
• tests/unit_tests/transformer/test_fsdp_dtensor_expert_offsets.py

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