Warn on the FSDP process-group fallbacks to the global grid by Connor-XY · Pull Request #8749 · NVIDIA/Megatron-LM

Connor-XY · 2026-10-30T03:01:44Z

Depends on #8017, #8045, #8049 (which contains #6293) and #8051. Please review the diff from the merge of their heads: scoped diff. The #8049 merge regenerates the process-group allowlist, whose total conflicted.

Land this after #8725, which rejects GTP weight remat in Megatron-FSDP v1. The commits apply without it, but the FSDP v1 fallback selects the same groups as before only because v1 runs without GTP weight remat. Land it after #8731 too. Until then FullyShardedOptimizer names experts without an expert layout, so Megatron-FSDP v2 MoE runs reach the new get_ep_layer_offset warning once, with the same names as before.

Why. When the caller passes no process groups, three FSDP paths read the global grid in parallel_state with no warning and no removal date:

Megatron-FSDP v1 has 11 group accessor calls in _init_dist_index.
Torch FSDP reads its shard group.
The FSDP-DTensor expert offsets (get_ep_layer_offset) read the global EP rank and size. They did this for dense models too, although a dense model's offset is always 0. get_global_unique_param_name names every parameter through this helper.

Separately, MFSDP v2 checked dp_cp with hasattr, which is always true for a declared collection field, so an unset dp_cp reached DeviceMesh.from_group as None.

What.

Each fallback goes through resolve_process_groups. It emits a ProcessGroupFallbackWarning once per owner (deprecated in 0.21, removed in 0.23), then returns the same groups as before.
- FSDP v1 builds a collection with every field its collection path reads, then takes that path. The collection holds the replicate DP groups, while the old fallback read the GTP-inclusive accessors. With GTP weight remat off, these are the same group objects.
- Torch FSDP resolves dp_cp_gtp_remat, which is the group it read before: get_data_parallel_group(with_context_parallel=True).
- get_ep_layer_offset returns 0 for a model without experts and reads nothing. With experts, it warns and takes the rank and size of the global EP group. If parallel_state has no EP group, it raises RuntimeError (before, this was an assertion or a division by zero).
MFSDP v2 checks dp_cp with vars(), so an unset or None dp_cp raises ValueError.
The checker goes from 173 reads to 162: 14 accessor reads become 3 guarded resolve_process_groups calls. The allowlist also records the reads that main added since this PR was written.

Megatron-LM training already passes its groups to FSDP v1, Torch FSDP and the checkpoint helpers, so stock training does not warn once #8731 is in.

Tests.

mfsdp_v1/test_mcore_fsdp_dist_index.py: in eight layouts, the fallback warns, reads the grid only through the resolver, and selects the same groups as an explicit collection. This replaces the old test of the fallback's ranks.
test_torch_fully_sharded_parallel.py: the fallback warns and picks get_data_parallel_group(with_context_parallel=True). An explicit group reads nothing global.
transformer/test_fsdp_dtensor_expert_offsets.py: a dense model reads nothing. With experts at EP=2, the fallback warns and matches the global layout. Without an EP group, it raises.
mfsdp_v2/test_mcore_dense.py: a collection with dp_cp unset or None is rejected.
mfsdp_v2/test_checkpoint.py passes its collection's expert layout.

On 4 GPUs against its base, mfsdp_v1/test_mcore_fsdp_dist_index.py passes 20 tests here and 20 on the base; the 8 layout tests that replace the old fallback-rank test pass. mfsdp_v2/test_checkpoint.py passes 11 tests on both sides. mfsdp_v2/test_mcore_dense.py passes 14 here vs 12 on the base, and its 2 new tests pass. transformer/test_fsdp_dtensor_expert_offsets.py passes 12 here vs 8 on the base, and its 4 new tests pass. No test regresses in these four files. test_torch_fully_sharded_parallel.py did not finish on either side on 4 GPUs; its test_fsdp2_swiglu_sharded_tensor_factory fails on both sides there because it hard-codes an 8-rank TP2xDP4 layout. On 8 GPUs (two nodes) in the integration run, the file's 3 tests pass, as on main. In the integration run, this PR's test files show no regressions against main. The dist-index and expert-offset test files are not on main; they come from the PRs below this one, and their 20 and 12 tests pass in that run.

Before the sync, a fix to the no-experts test fixture was checked separately: the 2 fixture tests fail without it. Before the sync, the branch also ran on 4 CPU gloo ranks: no regressions, the 4 new offset tests pass, and every other failure also occurs on the base, mostly for lack of CUDA. The MFSDP v2 tests cannot be collected on CPU.

Compatibility.

On the standard grid, groups, expert offsets and checkpoint names are unchanged.
Callers that omit groups get one warning per class or function. In this repo, the MFSDP v1 end-to-end tests, test_checkpointing.py and test_save_run_config.py still build FSDP v1 without a collection and now warn. They keep covering the default path.
The old EP accessors honored the overrides that offline tools set with set_expert_model_parallel_world_size and set_expert_model_parallel_rank. The fallback reads the group instead, so such callers must pass ep_rank and ep_size. No in-repo tool calls these helpers without them. Megatron-Bridge calls handle_experts_in_state_dict without an expert layout, so it warns for MoE models until it passes one.
Suppose the global grid runs GTP weight remat while the FSDP v1 config does not. The old fallback silently used the GTP-inclusive groups; the new one fails at mesh construction. Every consistent configuration selects the same groups.
Not changed here: MFSDP v2 with MoE at EP=1 and no expt_dp still passes None to DeviceMesh.from_group, because validation checks expt_dp only for EP>1.
Needs determinism-exempt only while the PRs below this one are unlanded. CI diffs against main, and the base carries registered kernel files. This PR's own diff passes the coverage check.

The branch was 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.

Later on 2026-11-01 the branch also merged main bfab98f. Its #8471 moves FullyShardedDataParallelV1 and FullyShardedDataParallelV2 from mcore_fsdp_adapter.py into mcore_adapter_v1.py and mcore_adapter_v2.py, and keeps mcore_fsdp_adapter.py as the factory that re-exports them. This PR's edits to those classes are applied unchanged where the code moved, and the process-group allowlist follows the move. On 4 GPUs against its new base, the Megatron-FSDP adapter tests and this PR's own tests (10 files) show no regressions: 103 passed here, 97 on the base. The two CUDA-graph end-to-end cases of test_mcore_fully_sharded_data_parallel.py fail on both sides, as they do on main. test_torch_fully_sharded_parallel.py stops on both sides at a case that needs 8 GPUs.

Part of #6307.

🤖 Generated with Claude Code

Carry the optional dense and expert optimizer Gloo groups in explicit collections and resolve the legacy global collection at the public optimizer factory boundary. Preserve the disabled-Gloo path and permit collection construction when Gloo groups were not created. Refresh against current main, preserve the nested optimizer regressions, and cover optional accessor checks plus Gloo presence and enablement. Signed-off-by: Yan Xu <yxu1@nvidia.com>

…o groups With use_gloo_process_groups=True, a process-group collection without Gloo groups used to fall back silently to running the distributed optimizer without Gloo. The global path raised in that case before this change. Raise ValueError instead, except under GTP_remat, which never builds Gloo groups. Signed-off-by: Yan Xu <yxu1@nvidia.com>

get_megatron_optimizer resolves the collection before calling its helpers, so a missing collection there is a caller error. Report it with ValueError instead of an assert, which is skipped under python -O. The previous hasattr(tp) check could never fail, because ProcessGroupCollection returns None for unset fields, so it is dropped. Signed-off-by: Yan Xu <yxu1@nvidia.com>

get_megatron_optimizer now resolves a missing pg_collection with ProcessGroupCollection.use_mpu_process_groups() up front, so the default path goes through the explicit-collection branch of setup_process_groups_for_optimizer. That branch returned no intra_dist_opt group unless ddp_config.use_distributed_optimizer was set. Megatron-FSDP leaves that flag False while OptimizerConfig enables the distributed optimizer, so its gradient-statistics group became None (TestMegatronFsdpGradStatsGroup in mfsdp_v1). Pass intra_dist_opt through from the collection whenever it is present, as the parallel_state branch always did. The initializer-cleanup factory tests patch every optimizer helper and never initialize model parallelism; stub the MPU fallback there instead of reaching the global grid. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

With Muon and the distributed optimizer, the emerging-optimizer path builds a separate DistributedOptimizer for the parameters Muon does not manage, and it set up that optimizer's process groups with Gloo hard-coded off. The collection now carries the Gloo groups, so honor the caller's use_gloo_process_groups instead. Without a Gloo group, the first memory-efficient fully reshardable checkpoint save failed the data_parallel_group assert in get_parameter_state_dp_zero. Tests: the layer-wise optimizer's DistributedOptimizer gets the collection's Gloo group (and none when Gloo is off), and the Muon + Lion checkpoint round trip also covers the memory-efficient fully reshardable format. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Without the distributed optimizer, the emerging-optimizer path reduced the gradient norm and zero count of expert parameters over tp_ep_pp, which leaves out the EGTP_remat axis. EGTP_remat peers hold distinct shards of each expert weight, so each peer computed the norm of its own shards only, and with gradient clipping the peers applied different clip factors and their replicas drifted apart. Use tp_ep_pp_with_egtp_remat, as the standard optimizer path does, falling back to tp_ep_pp for collections without it (the same group when EGTP_remat is off). Test: Muon with EP=2 and EGTP_remat=2 reports the norm over every expert shard on every rank. Before this change the two EGTP_remat peers reported 53.1 and 73.3 on 4 ranks, against a true norm of 89.1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The distributed optimizer shards its state over the collection's data-parallel group (dp_cp, or intra_dp_cp with more than one optimizer instance) and gathers that state over the matching Gloo group for checkpoint I/O. Taking the Gloo groups from the collection means they can now disagree: for example, a caller who starts from use_mpu_process_groups() and replaces dp_cp with their own group keeps parallel_state's Gloo group, and the save would then gather over other ranks, failing on a size mismatch or silently scrambling the optimizer checkpoint. Compare the rank lists of each Gloo group and the group it mirrors (dense and expert) and raise ValueError when they differ. Only expert parameters use the expert Gloo group, so require intra_expt_dp_gloo only when the collection has an expert data-parallel group. A dense model on its own grid can then enable Gloo by supplying intra_dp_cp_gloo alone. Tests: groups built by parallel_state pass with one and two optimizer instances; a Gloo group over other ranks raises, dense and expert; replacing dp_cp in an MPU collection with a same-size group over other ranks raises; a dense collection with only intra_dp_cp_gloo is accepted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The Gloo-collection test built its collection with use_mpu_process_groups(), so its Gloo groups were the parallel_state globals and the identity checks also passed for code that read the globals directly. Swap freshly created Gloo groups into the collection instead, and add an expert parameter so the expert distributed optimizer's Gloo group is checked too. The case where Gloo is requested in a job built without Gloo groups is now its own test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The CI ratchet from NVIDIA#6258 landed on main after this branch was written, so its allowlist still lists the two optimizer-helper fallbacks this change removes and lacks the one boundary read it consolidates them into. Refresh it: 210 -> 209 grandfathered reads. The addition was reviewed by hand: get_megatron_optimizer resolves a missing pg_collection with use_mpu_process_groups() once, at the entry point, replacing the fallbacks in _get_megatron_emerging_optimizer and _get_megatron_optimizer_based_on_param_groups that were already grandfathered. After adding it, `python tools/check_process_group_usage.py --update` dropped the two stale entries. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Signed-off-by: Yan Xu <yxu1@nvidia.com>

A collection without an embd field, or tied pipeline embedding/output stages without a usable embedding group, is a caller error. Report it with ValueError instead of an assert, which is skipped under python -O. Signed-off-by: Yan Xu <yxu1@nvidia.com>

Six sites in megatron/core called parallel_state accessors even though the correct group was already stored on the object (self.tp_group / self.cp_group). Pure substitution: no signature changes, no new arguments, no caller impact. SelfAttention.run_realtime_tests 3 sites -> self.tp_group SelfAttention.get_query_key_value_tensors 1 site -> self.tp_group MultiTokenPredictionLayer 1 site -> self.tp_group LLaVAModel 1 site -> self.cp_group Global process-group reads in megatron/core: 163 -> 157. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…scope Two more sites where the correct group was reachable without any API change. transformer_engine.py: in the cp_comm_type == 'a2a+p2p' branch the code read the global hierarchical CP groups, even though the caller-supplied branch directly above asserts hasattr(pg_collection, 'hcp') for exactly that case. The assert was already guaranteeing the group the code then declined to use. language_module.py: the word-embedding replica id derived its DP-CP rank from the global grid while the sharded tensor nine lines later is built against metadata['dp_cp_group']. Now both use the same group, so replica id and sharding agree under a caller-supplied grid. Global process-group reads in megatron/core: 157 -> 155. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

megatron/core/models/multimodal/context_parallel.py hard-coded the global CP grid in 11 places, across module-level helpers with no group parameter. This is the vision-encoder path, so it is exactly where an independent encoder grid silently used the LLM's groups instead. cp_group is now a required argument on the public helpers and is threaded through the private ones and the autograd Function. get_pg_size/get_pg_rank preserve the previous behaviour when CP is disabled (size 1, rank 0). The call graph is closed inside this repo: LLaVAModel already holds self.cp_group, and the only other callers are unit tests. The module is not re-exported and has no users in Megatron-Bridge or nemo-rl. Global process-group reads in megatron/core: 155 -> 144. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

_gather_from_tensor_parallel_region read the global TP grid for its size, rank and both collectives. Both callers are subclasses that already hold the group (self.tp_group on ColumnParallelLinear, self._tp_group on the TE variant), and the mapping helpers it calls already accept an explicit group. Global process-group reads in megatron/core: 144 -> 141. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

… in tests Caught by a distributed run: changing the autograd Function signature missed one direct .apply() call site in the unit tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

hybrid_context_parallel_forward_backward read the global TP and DP-CP grids for its broadcast source rank, broadcast group, rank checks and barrier. Its only caller is forward_backward_no_pipelining, which already holds pg_collection. dp_cp is asserted in the hybrid-CP branch rather than unconditionally, matching how TEDotProductAttention asserts hcp only for cp_comm_type == 'a2a+p2p'; the legacy schedule collection already requests dp_cp. get_tensor_model_parallel_src_rank() becomes torch.distributed.get_global_rank(tp_group, 0), which is its definition. hybrid_cp_schedule.py no longer imports parallel_state at all. Global process-group reads in megatron/core: 141 -> 136. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

run_realtime_tests compared Q/K layernorm parameters across the global DP grid. The DP group is required only by this debug check, so it is asserted there rather than in __init__, which asks use_mpu_process_groups for tp/cp only -- the same idiom TEDotProductAttention uses for hcp. megatron/core/transformer/attention.py no longer references parallel_state. Global process-group reads in megatron/core: 136 -> 133. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

mtp_on_this_rank read the global PP rank to decide whether MTP layers live on this stage. Both model callers (GPTModel, HybridModel) already hold self.pp_group from LanguageModule.__init__, which runs before the call. Global process-group reads in megatron/core: 125 -> 124. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

RotaryEmbedding and MultimodalRotaryEmbedding fell back to the global CP grid via a ternary when cp_group was not passed. 10 of the 22 constructor calls in this repo already passed one; the remaining 12 (3 in core, 9 in tests) now do too, so the fallback is gone rather than merely unused. megatron/core/models/common/embeddings/rotary_pos_embedding.py no longer references parallel_state. Global process-group reads in megatron/core: 124 -> 122. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…nd T5Model Both classes fell back to ProcessGroupCollection.use_mpu_process_groups() when no collection was passed. The fallback is now an assertion naming the missing argument, which is the 'fails loudly rather than silently' half of the goal. DotProductAttention: already dead in-repo -- Attention always passes pg_collection when building core_attention, and the one direct test construction passes it positionally. T5Model: 4 of 5 constructors did not pass one (3 tests, 1 example); they now do. This is a deliberate breaking change for external callers that construct these directly without a collection, which is the intent of the migration. Global process-group reads in megatron/core: 122 -> 120 (shims 20 -> 18). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…lock The use_mpu_process_groups fallback becomes an assertion naming the missing argument. All 7 constructors without one live in a single test module and now pass a collection. Global process-group reads in megatron/core: 120 -> 119 (shims 18 -> 17). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

LLaVAModel is the case the migration exists for: a vision encoder and an LLM may run on independent parallel grids, so falling back to the global grid silently picks the wrong one. The fallback becomes an assertion that says exactly that. All 8 constructors without a collection (7 tests, 1 example) now pass one. Global process-group reads in megatron/core: 119 -> 118 (shims 17 -> 16). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Attention, MultiLatentAttention (x2) and AbsorbedMLASelfAttention each fell back to use_mpu_process_groups when no collection was passed. All are dead in-repo: TransformerLayer sets attention_optional_kwargs['pg_collection'] unconditionally before build_module, and every direct constructor already passes one (verified by AST across megatron/, tests/ and examples/). Global process-group reads in megatron/core: 118 -> 114 (shims 16 -> 12). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Global process-group reads in megatron/core: 114 -> 113 (shims 12 -> 11). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

These were based on a faulty inventory. My constructor scan filtered on len(node.args) < 2, which silently skipped every call site passing arguments positionally, so 'no caller is missing pg_collection' was wrong. A paired run showed 85 real failures: 63x AssertionError: MultiLatentAttention requires an explicit pg_collection 22x AssertionError: Attention requires an explicit pg_collection 10x TypeError: DSAIndexer.__init__() got multiple values for 'pg_collection' (my inserted keyword duplicated an existing positional argument) Reverting to the last validated state. The earlier fallback removals (DotProductAttention, T5Model, MultiTokenPredictionBlock, LLaVAModel) were each validated against a pristine baseline and are retained. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Redo of the reverted attempt, this time with a signature-aware inventory: both DSAIndexer constructors already pass pg_collection positionally, so the fallback is dead and the tests need no change. The previous attempt added a redundant keyword and produced 'got multiple values for argument pg_collection'. Global process-group reads in megatron/core: 118 -> 117 (shims 16 -> 15). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Second attempt, this time driven by the signature-aware scanner: 15 of 16 SelfAttention constructors were missing pg_collection (all in test modules) and now pass one. The earlier attempt used a scan that skipped positional call sites and had to be reverted. Global process-group reads in megatron/core: 117 -> 116 (shims 15 -> 14). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…sses Second attempt, driven by the signature-aware scanner. MultiLatentAttention (x2) and AbsorbedMLASelfAttention no longer fall back to use_mpu_process_groups. 18 of 19 MLASelfAttention constructors were missing pg_collection, all in one test module, and now pass one; AbsorbedMLASelfAttention's single constructor already did. Global process-group reads in megatron/core: 116 -> 113 (shims 14 -> 11). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Paired run: 24 failed / 43 passed against a pristine 20 / 57, with 110 'requires an explicit pg_collection' assertions. The failures are concentrated in TestFusedMLA*, i.e. a MultiLatentAttention subclass my scan never covered -- I scanned MLASelfAttention and AbsorbedMLASelfAttention by name and assumed those were the only constructors reaching the base __init__. The Attention removal in the preceding commit validated clean (5F/121P on both trees) and is retained. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

copy-pr-bot · 2026-10-30T03:01:49Z

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.

The FSDP fallback tests build their collections with NVIDIA#8017's build_test_pg_collection and new_group_with_same_ranks, and check with forbid_global_process_groups that the explicit paths do not read the global grid, so merge NVIDIA#8017's head (which contains NVIDIA#6259's Gloo fields of the collection). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

NVIDIA#8045 takes expert offsets for FSDP-DTensor names as integers and adds get_model_ep_rank_and_size. NVIDIA#8749's offset warning edits get_ep_layer_offset and the expert-offset tests that NVIDIA#8045 adds, so merge NVIDIA#8045's head. The merge is clean. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

NVIDIA#8049 adds resolve_process_groups and ProcessGroupFallbackWarning, which all three of NVIDIA#8749's code commits and three of its test files use. Merge NVIDIA#8049's head (which contains NVIDIA#6293). The only conflict is the allowlist's total (NVIDIA#8049's 173 against the 195 of NVIDIA#8017 and NVIDIA#8045). Regenerate it with refresh_allowlist.py on the merged tree: 172 reads, which is NVIDIA#8049's file with NVIDIA#6259's single get_megatron_optimizer shim read in place of the two optimizer shim reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

NVIDIA#8051 creates the Megatron-FSDP v1 TP groups collectively and validates the pg_collection it is given. The Megatron-FSDP v1 fallback warning changes the same adapter code and the dist-index tests that NVIDIA#8051 creates, so merge NVIDIA#8051's head. The merge is clean. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

FullyShardedDataParallelV1._init_dist_index read eleven groups from parallel_state when its caller passed no pg_collection: TP, expert TP, EP, the data-parallel groups and, with several distributed-optimizer instances, the partial and inter-instance groups. Nothing told the caller, and the fallback had no removal date. Without a collection, the adapter now builds one through resolve_process_groups, which warns once per class (deprecated in 0.21, removed in 0.23), with every field the collection path reads, and then takes the collection path. The fallback read the GTP-inclusive data-parallel groups, while the collection holds the replicate groups. With GTP weight remat off, which the adapter requires, they are the same groups, so the standard grid selects the same groups as before. MFSDP v2 checked dp_cp with hasattr, which is always true for a declared collection field, so a collection without dp_cp reached DeviceMesh.from_group with None. It now checks vars() and raises a ValueError for an unset or None dp_cp. Megatron-LM training passes a collection, so it does not warn. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

TorchFullyShardedDataParallel read get_data_parallel_group( with_context_parallel=True) from parallel_state when its caller passed no process_group, without a warning. It now resolves the same group, dp_cp_gtp_remat, through resolve_process_groups, which warns once (deprecated in 0.21, removed in 0.23). Megatron-LM training already passes the group of its collection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

get_ep_layer_offset read the global expert-parallel rank and size when its caller passed neither ep_rank nor ep_size, without a warning. It also read them for models without experts, whose offset is 0 whatever the layout; get_global_unique_param_name names every parameter through it. - Without experts, the offset is 0 and nothing is read. - With experts, the fallback resolves the global expert-parallel group through resolve_process_groups, which warns once (deprecated in 0.21, removed in 0.23), and takes its rank and size. On the standard grid these are the values the accessors returned. If parallel_state has not created the group, it raises a RuntimeError instead of an assertion or a division by zero. - The MFSDP v2 checkpoint test passes the expert layout of its collection. The size and rank accessors honored the overrides that offline tools set with set_expert_model_parallel_world_size and _rank; the fallback reads the group instead. No in-repo tool calls these helpers without an expert layout. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The Megatron-FSDP v1 adapter, Torch FSDP and the FSDP-DTensor expert offsets read the global grid only through resolve_process_groups: 14 accessor reads become 3 guarded reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The checker classification that follows requires every allowlisted read in megatron/core to be warned or guarded. NVIDIA#8749 warns the global fallbacks of Megatron-FSDP v1, Torch FSDP and the FSDP-DTensor expert offsets, and brings NVIDIA#8017, NVIDIA#8045 and NVIDIA#8051. The guard tests that follow edit the mfsdp_v1 test_mcore_fsdp_dist_index.py that NVIDIA#8749 extends. The merge conflicts only in the process-group usage allowlist, which is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The checker classification that follows requires every allowlisted read in megatron/core to be warned or guarded, and its test lists the only two reads allowed to stay unwarned. NVIDIA#8752 warns before get_megatron_optimizer's use_mpu_process_groups() fallback, before clip_qk's global data-parallel group (with NVIDIA#8043, which NVIDIA#8752 brings) and describes the GTP shards in _backfill_gtp_sharded_param_map with the optimizer's own groups. Without it those four reads are unwarned and the classification test fails. megatron/core/optimizer/__init__.py conflicts on its imports: NVIDIA#8045's e09f5de (which NVIDIA#8749 brings) imports get_ep_rank_and_size next to get_global_unique_param_name, and NVIDIA#8752's 2d4d3cc adds resolve_process_groups to the process_groups_config import. The import takes both. NVIDIA#8752 also carries NVIDIA#6303, which this branch does not have yet, so tests/unit_tests/export/trtllm/test_trtllm_distributed_gpu_converter.py conflicts on its imports: NVIDIA#6303 imports ProcessGroupCollection and NVIDIA#8059 adds ProcessGroupFallbackWarning to the same import. The import takes both. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Conflicts: - examples/t5/pretrain_t5.py: Main adds `cfg = get_run_config()` to model_provider and drops the `build_tokenizer` import in favour of `get_tokenizer()`. The PR moves that import for isort and passes a ProcessGroupCollection into T5Model. The result keeps main's `cfg` line and the PR's pg_collection fallback and argument. It drops the `build_tokenizer` import, which would otherwise be unused. - megatron/core/models/common/embeddings/rotary_pos_embedding.py: Main imports `register_capture_constant` for its shared rotary cache. The PR imports `warn_global_process_group_fallback`. Both imports are kept, along with main's cache refactor and the PR's warning cp_group fallback. - pretrain_gpt.py, pretrain_hybrid.py: Main (NVIDIA#8341, NVIDIA#8592) moves `is_dataset_built_on_rank` into megatron/training/datasets/providers.py and the dataset config builder into argument_utils. The PR adds `pp_group` to the `mtp_on_this_rank` call in the local predicate. The result takes main's layout and keeps the PR's `pp_group` in get_batch. Adapted to main outside the conflicts: - megatron/training/datasets/providers.py: The PR's predicate edit is re-applied in main's shared `is_dataset_built_on_rank_explicit`. It passes `pp_group=mpu.get_pipeline_model_parallel_group()` to `mtp_on_this_rank`, so middle pipeline stages do not hit the fallback warning. The function's signature is unchanged. The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…ir update to main Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Conflicts: - megatron/core/distributed/fsdp/mcore_fsdp_adapter.py: Main adds the ShortcutMoEBlock import (NVIDIA#8312) and the PR adds the resolve_process_groups import. The result keeps both in isort order, along with the PR's _init_dist_index rewrite, which no longer reads parallel_state. - megatron/core/models/common/embeddings/rotary_pos_embedding.py: Main adds the register_capture_constant import for its embedding cache (NVIDIA#8358) and the PR adds the warn_global_process_group_fallback import. The result keeps both, plus the PR's warned cp_group fallback. - pretrain_gpt.py, pretrain_hybrid.py: Main moves is_dataset_built_on_rank and core_gpt_dataset_config_from_args into megatron/training/datasets/providers.py (NVIDIA#8341, NVIDIA#8592). The result takes main's side and keeps the PR's pp_group in get_batch. - examples/t5/pretrain_t5.py: The result keeps main's tokenizer refactor (NVIDIA#8391) and the PR's pg_collection default passed to T5Model. The unused build_tokenizer import that main dropped stays dropped. Adapted to main outside the conflicts: - megatron/training/datasets/providers.py: The PR's pp_group argument to mtp_on_this_rank is re-applied in is_dataset_built_on_rank_explicit, where main moved the code. The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Main's NVIDIA#8471 moves FullyShardedDataParallelV1 and FullyShardedDataParallelV2 out of mcore_fsdp_adapter.py into mcore_adapter_v1.py and mcore_adapter_v2.py, and keeps mcore_fsdp_adapter.py as the factory that re-exports them. This branch's edits to the two classes and their helpers are applied unchanged where the code moved, and the process-group allowlist follows the move. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Main's NVIDIA#8471 moves FullyShardedDataParallelV1 and FullyShardedDataParallelV2 out of mcore_fsdp_adapter.py into mcore_adapter_v1.py and mcore_adapter_v2.py, and keeps mcore_fsdp_adapter.py as the factory that re-exports them. This branch's edits to the two classes and their helpers are applied unchanged where the code moved, and the process-group allowlist follows the move. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Connor-XY · 2026-11-03T14:48:26Z

Superseded by #8905, which carries this change as commit 858b819 under the same title, with "Originally #8749." in its message. Closing so review happens in one place. The branch stays on the fork.

Connor-XY and others added 30 commits October 7, 2026 16:02

Preserve caller process groups in T5Model initialization

cf5f71e

Signed-off-by: Yan Xu <yxu1@nvidia.com>

Require an explicit process-group collection in DSAIndexer

aa5f2e7

Global process-group reads in megatron/core: 114 -> 113 (shims 12 -> 11). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

This was referenced Oct 30, 2026

Deprecate the parallel_state accessors of the global grid #8765

Closed

Make the process-group checks enforce what the deprecation notes claim #8770

Closed

Connor-XY and others added 8 commits October 30, 2026 13:22

Connor-XY force-pushed the yx/pstate-c2-fsdp-fallbacks branch from ce1cc89 to 74e6882 Compare November 1, 2026 01:38

Connor-XY mentioned this pull request Nov 1, 2026

Name the MFSDP v2 optimizer's experts by the model's EP layout #8731

Closed

Connor-XY and others added 11 commits October 31, 2026 23:00

Merge main into yx/pstate-f3-fsdp-v1-dist-index

8286160

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Merge main into yx/pstate-gloo-in-collection

7a951c3

The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Merge main into yx/pstate-b2-expert-offsets

2038339

The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Merge NVIDIA#6259 after its update to main

c2d65b2

The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Merge NVIDIA#6293 after its update to main

d89ecce

The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Merge NVIDIA#8049, NVIDIA#8017, NVIDIA#8045 and NVIDIA#8051 after the…

3115978

…ir update to main Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Merge main into the bases after the MFSDP adapter split

c9cdc74

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

This was referenced Nov 3, 2026

Warn on global process-group fallbacks in finalize_model_grads and FSDP #8905

Draft

Warn on the last core grid reads, enforce the checks, and document them #8909

Draft

Connor-XY closed this Nov 3, 2026

Navigation Menu

Repository navigation

Warn on the FSDP process-group fallbacks to the global grid - #8749

Warn on the FSDP process-group fallbacks to the global grid#8749
Connor-XY wants to merge 92 commits into
NVIDIA:mainfrom
Connor-XY:yx/pstate-c2-fsdp-fallbacks

Connor-XY commented Oct 30, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Oct 30, 2026

Uh oh!

Connor-XY commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Connor-XY commented Oct 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Oct 30, 2026

Uh oh!

Connor-XY commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Connor-XY commented Oct 30, 2026 •

edited

Loading