Fix Megatron-FSDP GDN optimizer checkpoint metadata for FSDP DTensor - #2
Open
conver334 wants to merge 5 commits into
Open
Conversation
Reviewer's GuideFixes Megatron-FSDP DTensor checkpoint metadata for GDN optimizer and adds MoE sequence-packing / HybridEP support and example recipes, by attaching explicit DTensor chunk metadata where the generic uneven-DTensor logic fails, adding TP-aware optimizer offsets, padding THD/HybridEP token flows, and wiring new large-scale DeepSeek/Qwen3 recipes and configs. Sequence diagram for MoE HybridEP token dispatch with sequence packingsequenceDiagram
participant TD as TokenDispatcher
participant HybridEP as hybrid_ep_dispatch
participant HybridEPComb as hybrid_ep_combine
participant Dist as torch.distributed
TD->>TD: setup_metadata(routing_map, probs)
TD->>TD: _original_num_tokens = routing_map.shape[0]
TD->>Dist: all_reduce(max_num_tokens_across_ep)
TD->>TD: _padded_num_tokens computed
TD->>TD: pad routing_map, probs to _padded_num_tokens
TD->>TD: dispatch(hidden_states)
TD->>TD: pad hidden_states to _padded_num_tokens
TD->>HybridEP: hybrid_ep_dispatch(x=hidden_states, ...)
HybridEP-->>TD: dispatched_hidden, handle
TD->>HybridEPComb: hybrid_ep_combine(..., handle)
HybridEPComb-->>TD: hidden_states
TD->>TD: trim to _original_num_tokens
TD-->>TD: reset _original_num_tokens, _padded_num_tokens
Flow diagram for FSDP DTensor optimizer checkpoint metadata handlingflowchart LR
A[make_fsdp_dtensor in split_gdn_fused] --> B[_validate_explicit_dtensor_chunk_metadata]
B --> C[_set_explicit_dtensor_chunk_metadata]
C --> D[preprocess_state_dict_for_uneven_dtensor]
D --> E{_has_explicit_chunk_metadata?}
E -- Yes --> F[skip update_uneven_dtensor_chunk_metadata]
E -- No --> G[update_uneven_dtensor_chunk_metadata]
subgraph Flat_optimizer_state
H[_pack_hybrid_optimizer_fsdp_state]
H --> I[_get_flat_fsdp_dtensor_tp_offset]
I --> J[_set_flat_fsdp_dtensor_chunk_metadata]
J --> D
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- In
_get_thd_padded_seqlens(test_token_dispatcher), you assumeseqlensis non-empty and mutate the last element in-place; consider asserting non-emptiness or handling the empty-sequence case explicitly to avoid a silent IndexError in edge-case tests. - For the explicit DTensor chunk metadata helpers (
_set_explicit_dtensor_chunk_metadata,_validate_explicit_dtensor_chunk_metadata, and_set_flat_fsdp_dtensor_chunk_metadata), it may be worth centralizing the marker/hook-setting logic to avoid divergence between the different code paths that tag tensors with_megatron_fsdp_explicit_chunk_metadata.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `_get_thd_padded_seqlens` (test_token_dispatcher), you assume `seqlens` is non-empty and mutate the last element in-place; consider asserting non-emptiness or handling the empty-sequence case explicitly to avoid a silent IndexError in edge-case tests.
- For the explicit DTensor chunk metadata helpers (`_set_explicit_dtensor_chunk_metadata`, `_validate_explicit_dtensor_chunk_metadata`, and `_set_flat_fsdp_dtensor_chunk_metadata`), it may be worth centralizing the marker/hook-setting logic to avoid divergence between the different code paths that tag tensors with `_megatron_fsdp_explicit_chunk_metadata`.
## Individual Comments
### Comment 1
<location path="megatron/core/transformer/fsdp_dtensor_checkpoint.py" line_range="59" />
<code_context>
from megatron.core.utils import get_attr_wrapped_model
+def _set_explicit_dtensor_chunk_metadata(dtensor: "DTensor", offsets, sizes):
+ """Attach DCP chunk metadata for a DTensor whose local shard is not rank-contiguous."""
+ chunk_meta = ChunkStorageMetadata(offsets=tuple(offsets), sizes=tuple(sizes))
</code_context>
<issue_to_address>
**issue (complexity):** Consider extracting the DTensor chunk arithmetic and metadata wiring into small reusable helpers to keep split_gdn_fused focused on high-level orchestration.
You can keep the new behavior but reduce the local complexity by pulling the DTensor–chunk wiring into small helpers.
### 1. Encapsulate GDN chunk arithmetic + DTensor construction
The block in `split_gdn_fused` that computes `trailing_numel`, `component_start`, `chunk_offsets`, `chunk_sizes`, and then calls `_validate_explicit_dtensor_chunk_metadata` / `_set_explicit_dtensor_chunk_metadata` can be turned into a focused helper.
For example:
```python
def _build_gdn_component_dtensor(
comp_data: torch.Tensor,
meta: torch.Tensor,
dist_index,
split_dim: int,
shard,
comp_slice,
per_tp_rank_shape,
tp_mesh,
) -> "DTensor":
meta_shape = list(per_tp_rank_shape)
meta_shape[split_dim] = comp_data.shape[split_dim]
trailing_numel = math.prod(meta_shape[split_dim + 1 :])
assert trailing_numel > 0, f"Invalid GDN component shape: {meta_shape}"
component_start = 0
if shard.start != shard.stop:
component_start = shard.start - comp_slice.start
assert component_start % trailing_numel == 0, (
f"GDN component shard is not aligned with tensor rows: "
f"component_start={component_start}, trailing_numel={trailing_numel}, "
f"meta_shape={meta_shape}, shard={shard}, comp_slice={comp_slice}"
)
assert comp_data.numel() % trailing_numel == 0, (
f"GDN component shard size is not aligned with tensor rows: "
f"numel={comp_data.numel()}, trailing_numel={trailing_numel}, "
f"meta_shape={meta_shape}, shard={shard}, comp_slice={comp_slice}"
)
chunk_offsets = [0] * len(meta_shape)
chunk_offsets[split_dim] = component_start // trailing_numel
chunk_sizes = list(comp_data.shape)
tp_partition_dim = get_mcore_tensor_parallel_partition_dim(meta)
if tp_partition_dim is not None:
tp_rank = dist.get_rank(tp_mesh.get_group())
chunk_offsets[tp_partition_dim] += tp_rank * meta_shape[tp_partition_dim]
dtensor = make_fsdp_dtensor(
comp_data.data,
meta,
dist_index=dist_index,
is_expert_param=False,
run_check=False,
update_uneven_dtensor_chunk_meta=False,
)
_validate_explicit_dtensor_chunk_metadata(dtensor, chunk_offsets, chunk_sizes)
_set_explicit_dtensor_chunk_metadata(dtensor, chunk_offsets, chunk_sizes)
return dtensor
```
Then the body in `split_gdn_fused` is reduced to the higher‑level orchestration:
```python
meta_shape = list(per_tp_rank_shape)
meta_shape[split_dim] = s
meta = torch.empty(*meta_shape, device="meta")
copy_tensor_model_parallel_attributes(meta, dist_param)
if (
tp_mesh.mesh.numel() > 1
and split_dim == 0
and global_shape[split_dim] == total_split
and not is_mcore_tensor_model_parallel(meta)
):
meta._tensor_parallel_mode = "column"
dtensor = _build_gdn_component_dtensor(
comp_data=comp_data,
meta=meta,
dist_index=dist_index,
split_dim=split_dim,
shard=shard,
comp_slice=comp_slice,
per_tp_rank_shape=per_tp_rank_shape,
tp_mesh=tp_mesh,
)
results.append(dtensor)
```
This keeps all the same checks/behavior but pushes the low‑level arithmetic and validations into a single well‑named helper.
### 2. Deduplicate DTensor + local‑tensor metadata wiring
`_set_explicit_dtensor_chunk_metadata` mirrors the pattern used in your optimizer’s `_set_flat_fsdp_dtensor_chunk_metadata`. You can share the mechanics of “attach closures + flag to DTensor and its local tensor” while allowing different metadata/content.
For example, introduce a generic attachment helper:
```python
def _attach_chunk_metadata_to_dtensor(dtensor: "DTensor", chunk_list_fn, write_items_fn):
dtensor.__create_chunk_list__ = chunk_list_fn
dtensor.__create_write_items__ = write_items_fn
dtensor._megatron_fsdp_explicit_chunk_metadata = True
local = dtensor._local_tensor
local.__create_chunk_list__ = chunk_list_fn
local.__create_write_items__ = write_items_fn
local._megatron_fsdp_explicit_chunk_metadata = True
```
Then `_set_explicit_dtensor_chunk_metadata` becomes:
```python
def _set_explicit_dtensor_chunk_metadata(dtensor: "DTensor", offsets, sizes):
chunk_meta = ChunkStorageMetadata(offsets=tuple(offsets), sizes=tuple(sizes))
def _chunk_list():
return [chunk_meta]
def _write_items(fqn: str, tensor: "DTensor"):
if tensor.to_local().numel() == 0:
return []
return [
WriteItem(
type=WriteItemType.SHARD,
index=MetadataIndex(fqn, chunk_meta.offsets),
tensor_data=TensorWriteData(
chunk=chunk_meta,
properties=TensorProperties.create_from_tensor(tensor.to_local()),
size=tensor.size(),
),
)
]
_attach_chunk_metadata_to_dtensor(dtensor, _chunk_list, _write_items)
```
The optimizer helper can call `_attach_chunk_metadata_to_dtensor` with its own closures as well, so the low‑level DTensor/local‑tensor wiring is defined once and reused.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
wplf
force-pushed
the
jinliang/fix-mfsdp-optimizer-offload
branch
2 times, most recently
from
May 27, 2026 03:54
2f2442b to
21ba0b9
Compare
Owner
|
review |
wplf
added a commit
that referenced
this pull request
May 29, 2026
… split Mirror PR #2's GDN fix in split_gdn_fused over to split_swiglu_linear_fc1: both helpers wrap fused-parameter component shards as DTensors during checkpoint save, and both previously called make_fsdp_dtensor with run_check=True, update_uneven_dtensor_chunk_meta=True. That path issues a blocking all_gather_object inside the checkpoint preprocessing, which deadlocks under MFSDP + HybridDeviceOptimizer (issue NVIDIA#4910). PR #2 fixed split_gdn_fused but left split_swiglu_linear_fc1 untouched. Qwen3.5-35B-A3B has both GDN (1/4 of layers via linear_attention_freq=4) and SWiGLU MLP (the other 3/4), so save still hangs in the SWiGLU branch. Apply the same explicit chunk metadata pattern there.
wplf
added a commit
that referenced
this pull request
May 29, 2026
The split_gdn_fused() fast-path for already-TP-local DTensors (added by NVIDIA#4799) calls split_dtensor(..., update_uneven_dtensor_chunk_meta=True). split_dtensor() itself unconditionally calls gather_and_compute_chunk_metadata (line 457 in uneven_dtensor.py), which issues an all_gather_object on the DTensor's shard groups. Inside checkpoint save preprocessing, this triggers the same deadlock as issue NVIDIA#4910: ranks enter collectives on different mesh groups in different orders and the NCCL watchdog aborts after 10 min. PR #2's explicit-chunk-metadata fix only covered the slow rebuild path; the fast-path was untouched. Disable it so GDN model state falls through to the slow path.
Signed-off-by: conver334 <conver334@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
#4623 fixes the small/dense Megatron-LM optimizer CPU offload repro from #4910, but it does not fix the full Qwen3.5-35B-A3B failure reported in the latest NVIDIA#4910 comment.
The remaining failure happens when saving a Megatron-FSDP
fsdp_dtensoroptimizer checkpoint with optimizer CPUoffload enabled. The run reaches checkpoint save, then fails while splitting GDN optimizer state in
handle_gdn_in_state_dict()because the generated DTensor chunk metadata does not describe the actual FSDP/TP- local shard layout.This PR adds explicit DCP chunk metadata for those cases instead of letting the generic uneven-DTensor preprocessing recompute it incorrectly.
Problem
The full-model repro in NVIDIA#4910 still fails after merging NVIDIA#4623 into latest
dev:Some ranks also have non-empty local shards but still fail the same boundary check. This indicates that the
checkpoint metadata does not correctly describe where each GDN component shard belongs in the global tensor.
The problematic path is specific to FSDP DTensor checkpoint preprocessing for GDN fused parameters and optimizer
state.
Changes
Why
For TP-sharded parameters, fsdp_slice.start alone is not enough to describe the global checkpoint offset. Each TP rank owns a different TP-local portion of the logical global tensor, so DCP metadata needs the TP-rank offset as well.
For GDN fused parameters, optimizer state is split into logical components such as query/key/value/z/beta/alpha. After this split, rank-local component shards are not always representable by the generic rank-order uneven-DTensor metadata computation. The checkpoint writer needs explicit offsets and sizes for the component shard that this rank actually owns.
Related Issues / PRs
Summary by Sourcery
Fix FSDP DTensor optimizer and GDN fused-parameter checkpointing metadata, improve MoE HybridEP compatibility with sequence packing and transport backends, and add end-to-end MoE recipes for DeepSeek-V3 and Qwen3-235B across multiple GPU platforms.
Bug Fixes:
Enhancements:
Build:
Deployment:
Documentation:
Tests: