Skip to content

Name FSDP-DTensor experts by the model's expert-parallel layout - #8045

Closed
Connor-XY wants to merge 6 commits into
NVIDIA:mainfrom
Connor-XY:yx/pstate-b2-expert-offsets
Closed

Name FSDP-DTensor experts by the model's expert-parallel layout#8045
Connor-XY wants to merge 6 commits into
NVIDIA:mainfrom
Connor-XY:yx/pstate-b2-expert-offsets

Conversation

@Connor-XY

@Connor-XY Connor-XY commented Oct 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Why. The FSDP-DTensor checkpoint helpers turn local expert indices into global ones with the global expert-parallel rank and size. A model with its own EP layout then saves and loads optimizer state under other experts' names. The distributed optimizer's parameter names and the param-group dump have the same problem.

What.

  • get_ep_layer_offset, handle_experts_in_state_dict, expert_param_local_key and get_global_unique_param_name take keyword-only integer ep_rank and ep_size. Integers keep offline tools without process groups working. Passing only one raises ValueError; passing neither keeps the global fallback.
  • get_ep_rank_and_size(pg_collection) reads them from a collection. An unset ep keeps the fallback, and ep=None means expert parallelism is off.
  • DistributedOptimizer's parameter names, the get_megatron_optimizer param-group dump, handle_swiglu_in_state_dict and the training loop's FSDP-DTensor preprocessing pass their model's layout.

Tests. New tests check expert offsets, state-dict keys and parameter names on an EP grid that differs from the global one, for the helpers, the distributed optimizer and the param-group dump. Measured before the 2026-11-01 sync, they fail on main and pass here on 4 GPUs, and test_fsdp_dtensor_checkpoint.py matches main. In the integration run on 4 GPUs, all 14 tests in this PR's test files pass: 12 in test_fsdp_dtensor_expert_offsets.py (8 added here, 4 added by a later PR in the series), 1 in test_distrib_optimizer_param_name.py and 1 in test_param_group_map_expert_names.py; main has none of these files. In the integration run, this PR's test files show no regressions against main.

Compatibility. Names on the standard grid do not change. MoE models must carry a pg_collection at some wrapper level for these paths, as all Megatron models do.

Synced with main 77391e7 on 2026-11-01 by a merge commit with no conflicts, so existing commits and review links stay valid. The merge also updates the process-group allowlist to record the reads main added since this PR was written.

Part of #6307.

🤖 Generated with Claude Code

@Connor-XY

Connor-XY commented Oct 9, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f8ea050

@copy-pr-bot

copy-pr-bot Bot commented Oct 9, 2026

Copy link
Copy Markdown

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 and others added 5 commits October 26, 2026 13:59
@Connor-XY @claude
Take expert offsets for FSDP-DTensor names as integers
d69ce7d 
get_ep_layer_offset read the global expert-parallel rank and size, so
handle_experts_in_state_dict, expert_param_local_key and
get_global_unique_param_name named a model's local experts by the
global grid even when the model has its own EP layout. With a custom
EP=4 grid next to a global EP=2 grid, FSDP-DTensor state was saved
under other experts' keys.

Add keyword-only ep_rank and ep_size integers, so offline tools without
process groups keep working, and get_ep_rank_and_size to read them from
a collection (vars() semantics: an unset ep keeps the global fallback,
ep=None means EP is off). handle_swiglu_in_state_dict passes its
model's layout. Omitting both keeps the global read; passing one of
them raises ValueError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
@Connor-XY @claude
Name the distributed optimizer's experts by the model's EP layout
85cda6e 
DistributedOptimizer._param_name, which keys FSDP-DTensor optimizer
state, offset local experts by the global expert-parallel rank. Pass
the EP rank and size of the model chunk's own collection.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
@Connor-XY @claude
Name experts in the param-group dump by the collection's EP layout
e09f5de 
get_megatron_optimizer's dump_param_to_param_group_map named expert
parameters by the global expert-parallel layout even when the caller
passed a collection. Pass the collection's EP rank and size; without a
collection the names are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
@Connor-XY @claude
Pass the model's expert-parallel layout to FSDP-DTensor checkpoints
3b00d72 
preprocess_fsdp_dtensor_state_dict renamed expert keys by the global
expert-parallel layout. Pass the EP rank and size of the model's own
collection, which match the global ones on the standard grid.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
@Connor-XY @claude
Refresh the process-group usage allowlist
772b148 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
@Connor-XY
Connor-XY force-pushed the yx/pstate-b2-expert-offsets branch from f8ea050 to 772b148 Compare October 26, 2026 22:26
@Connor-XY

Connor-XY commented Oct 26, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 772b148

This was referenced Oct 30, 2026
Closed
Closed
Closed
Closed
Closed
Closed
Closed
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the model-based expert layout from NVIDIA#8045
bda5278 
The MFSDP v2 optimizer change that follows names its experts with
NVIDIA#8045's get_model_ep_rank_and_size and the ep_rank and ep_size
arguments of get_global_unique_param_name, and its test reuses NVIDIA#8045's
expert-offsets test helpers, so merge NVIDIA#8045 here. NVIDIA#8045's head already
contains main, so the merge is clean and its tree is NVIDIA#8045's.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the integer FSDP-DTensor expert offsets from NVIDIA#8045
f415981 
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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the fallback warning helpers from NVIDIA#8049
fc2355c 
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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the model-based expert layout from NVIDIA#8045
6e03b72 
NVIDIA#8045 names experts by the model's expert layout with
get_model_ep_rank_and_size. The config-logger change that follows adds
an import next to NVIDIA#8045's import of that helper in
distrib_optimizer.py, so merge NVIDIA#8045 here. The merge is clean.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the optimizer's process-group fallback warnings from NVIDIA#8752
926787e 
NVIDIA#8752 gives Float16OptimizerWithFloat16Params a keyword-only
pg_collection and makes get_megatron_optimizer resolve its collection
once, through resolve_process_groups, before it builds the optimizers.
NVIDIA#8752 also carries NVIDIA#6303, whose imports its tests use. The config-logger
change that follows forwards that collection to the optimizers' config
log, so merge it here.

Two files conflict with NVIDIA#8042, which this branch carries:
- optimizer/__init__.py: keep both import changes, NVIDIA#8752's
  resolve_process_groups and NVIDIA#8045's get_ep_rank_and_size; in
  _get_megatron_optimizer_based_on_param_groups keep NVIDIA#8752's ValueError
  for a missing collection and NVIDIA#8042's set_duplicate_filter_groups call,
  which replaces the tp_group and expert_tp_group setattrs that NVIDIA#8752's
  side still read.
- optimizer.py: import the union of both sides from
  process_groups_config (ProcessGroupCollection,
  require_process_groups, resolve_gtp_remat_group,
  resolve_process_groups).

The process-group usage allowlist conflicts on NVIDIA#8042's entries; it is
regenerated on the merged tree and equals NVIDIA#8752's file (170 reads: 149
accessor, 18 shim, 3 guarded).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the FSDP process-group fallback warnings from NVIDIA#8749
9d54837 
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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the optimizer's process-group fallback warnings from NVIDIA#8752
f05c731 
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>
Connor-XY added a commit to Connor-XY/Megatron-LM that referenced this pull request Nov 1, 2026
@Connor-XY @claude
Merge the config logger's process groups from NVIDIA#8763
eb2a679 
This base holds every other PR of the parallel_state deprecation series
(NVIDIA#6307): the end-to-end test that follows asserts that no process-group
fallback or deprecation warning fires anywhere on a collection that is
not the global grid, and the strict-fallback functional tests turn
every remaining fallback into an error, so both need the whole series
below them. NVIDIA#8763 brings NVIDIA#8752, NVIDIA#8755 (with NVIDIA#6303), NVIDIA#8042 and NVIDIA#8045.

Four files conflict:
- extensions/transformer_engine.py: the process_groups_config import
  adds NVIDIA#8755's resolve_process_groups (eac0b1c).
- models/T5/t5_model.py: drop the self.tp_group line that NVIDIA#6293 kept;
  NVIDIA#6303's b732c67, which NVIDIA#8763 brings, removes it because
  LanguageModule already sets self.tp_group from the collection.
- optimizer/fully_sharded_optimizer.py: keep NVIDIA#8763's
  ProcessGroupCollection import (c17120b) and NVIDIA#8731's
  get_model_ep_rank_and_size import (3877940).
- parallel_state.py, get_data_parallel_group_gloo and
  get_expert_data_parallel_group_gloo: take NVIDIA#6259's check_initialized
  argument and Args section, which NVIDIA#8763 brings, with NVIDIA#7991's
  deprecation paragraph before the Args section.
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>
@Connor-XY @claude
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>
This was referenced Nov 3, 2026
Draft
Draft
Draft
@Connor-XY

Connor-XY commented Nov 3, 2026

Copy link
Copy Markdown
Contributor Author

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

@Connor-XY Connor-XY closed this Nov 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

determinism-exempt

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Connor-XY