Repository navigation
Conversation
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
8 times, most recently
from
June 3, 2026 19:21
6fc6e32 to
2e09a54
Compare
44 of 72 tasks
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
from
June 29, 2026 16:29
2e09a54 to
1f3b873
Compare
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
5 times, most recently
from
July 20, 2026 17:14
36b00db to
f06507c
Compare
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
2 times, most recently
from
July 27, 2026 23:33
e378675 to
0b44daf
Compare
Contributor
Author
|
Rebased onto the updated #4942 branch to pick up the two review fixes there (ffcae6d, 08e141e): the |
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
from
July 28, 2026 00:22
0b44daf to
bc5a0f7
Compare
Contributor
Author
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
from
September 8, 2026 20:30
bc5a0f7 to
9a766fc
Compare
Connor-XY
force-pushed
the
pr4798-3-fsdp-hybrid
branch
2 times, most recently
from
September 29, 2026 22:23
9a766fc to
e7ad92c
Compare
This was referenced Sep 29, 2026
Refresh PR NVIDIA#4942 against main while preserving hash routing, MTP, CP layouts, and ungrouped checkpoint compatibility. Integrate the cross-slot MTP gradient fix and reject unsupported direct-overlap execution modes. Original grouped HybridStack work by @Wohox, @Connor-XY, and @guihong-nv in NVIDIA#4798/NVIDIA#4942. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Rewrite comments that only made sense relative to the diff or to a past incident: drop a line-number reference in favor of the function name, state the recompute-hook invariant instead of the symptom it caused, and describe the checkpoint keys each pattern publishes instead of calling them historical or pre-existing. Signed-off-by: Yan Xu <yxu1@nvidia.com>
The schedule plan checked static config on every build_schedule_plan() call, i.e. every microbatch, walked all modules each time, and mixed assert with ValueError. Move the checks to HybridModel.__init__ and raise ValueError: - CUDA graphs (the message no longer claims the check is specific to grouped patterns; it applies to any HybridModel with EP overlap), - mHC connections and MoE shortcut connections, - mixed context-parallel layouts, checked on every HybridStack including nested bracketed groups. Drop the hash-routed MoE and wide-residual checks: TransformerConfig already rejects both with overlap_moe_expert_parallel_comm. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Read known attributes directly and pass arguments unconditionally: - HybridStack sets final_norm to None when it has no final norm (like TransformerBlock.final_layernorm), so the final_layernorm alias and the overlap callables read it directly. The flextron hook checks for a None final norm instead of a missing attribute. - The schedule plan dispatches on type: HybridStack for the layer-type symbol, and TransformerLayer / MultiTokenPredictionLayer for the layer-level quantization context. Its imports move to module level. - The MoE preprocess slot always sets shared_expert_output; the combine slot reads it and mlp_norm_manager directly. The Mamba pre-layer call no longer probes the chunk state for an inference context the training schedule never sets. - HybridStack.forward drops rotary_pos_cos / rotary_pos_sin / rotary_pos_cos_sin, which nothing passes in, and HybridModel._postprocess drops is_spec_decode, which no caller passes. - A bracketed group stack never applies full recomputation itself; gate on the static is_layer_group_stack instead of threading the private _checkpointed_forward_in_parent flag through forward. checkpointed_forward identifies group layers with isinstance(layer, HybridStack) and passes input_ids, packed_seq_params_by_layout and cp_layout_plan unconditionally. - MambaLayer.backward_dw and MambaMixer.backward_dw call their delayed wgrad unconditionally. Building the overlap callables with delay_wgrad_compute now raises for a Mamba layer whose mixer cannot defer its wgrad (e.g. GatedDeltaProductMixer) instead of silently skipping those weight gradients. Signed-off-by: Yan Xu <yxu1@nvidia.com>
MoELayer.route takes the batch-first [b, s] padding mask and transposes it for the router, and MoELayer.preprocess takes the same mask so that dropless HybridEP drops padded rows before dispatch. The hybrid overlap callable transposed the mask before calling route, so the router flattened a [b, s] mask against sequence-first tokens, which misaligns the mask whenever the micro-batch size is greater than 1. It also never passed the mask to preprocess. Pass node.chunk_state.padding_mask to route and preprocess unchanged, as MoELayer.forward and the GPT overlap callables do. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The EP-overlap schedule plan calls the layer callables directly and bypasses HybridStack.forward, which builds the packed-sequence CP metadata that chunkwise linear CP needs and rejects padding masks for it. With context_parallel_size > 1 and linear_cp_mode='chunkwise', the model built and the first packed step failed inside the Gated Delta Product mixer. Reject the combination at construction for every HybridStack, including nested groups and MTP stacks, like the other features the overlap schedule cannot run. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Import HybridStack and the hybrid callable builders at module level in common/fine_grained_callables.py; there is no import cycle. - Document bracketed layer groups in docs/user-guide/hybrid-model-migration.md, and point Symbols.GROUP_START, the --hybrid-layer-pattern help and the HybridModel docstring to it. - Add a module docstring to hybrid/fine_grained_callables.py and name GPTModel explicitly in the hybrid schedule-plan and callable docstrings. - Move the layer-symbol to checkpoint-namespace map and the duplicate- namespace check for layer groups into hybrid/layers/utils.py. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A bracketed group stored all of its layers under one logical layer index, and grouped patterns renamed the final norm to final_layernorm, so a checkpoint saved with [*-][*-] could not be loaded into a *-*- model or the reverse. Following the canonicalization approach of NVIDIA#4463, every layer now keeps the key it has in the pattern without brackets, its global layer number, as in ungrouped stacks. Grouped and ungrouped distributed checkpoints now load into each other. This removes the logical-offset plumbing (logical_layer_offset and select_pipeline_segment_with_logical_offset), the TransformerBlock-style key switch, and the rule that a group's layers need distinct checkpoint namespaces, which existed only because they shared one index. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
HybridStack.layer_type_list reports a bracketed group as one tuple of symbols, but MambaInferenceStateConfig.from_model compared its entries with single symbols. A pattern whose Mamba or GDN layers all sit inside groups, such as [M*E], returned no recurrent state config, and [M*E][G*E] did not raise the Mamba/GDN mixing error. Flatten the symbols before inspecting them and pass the flat list to ssm_chunking and DynamicInferenceContext, which index by physical layer. HybridStack.physical_layers() now also unpacks group stacks so it lines up with the flat list, and mamba_state_shapes_per_request zips it with the flattened configs. Annotate layer_type_list as returning LayerPatternItem entries. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With inter-document masking and micro-batch size above 1, pretrain_hybrid flattens a batch to [mbs*s, 1, h] and sets tokens_per_sample. The eager TransformerLayer._forward_mlp path restores [s, mbs, h] before the MoE layer so the router computes its sequence-level aux loss per sample. The EP-overlap MoE slot skipped that step, so the router treated the whole micro-batch as one sequence and its weight gradient differed from eager. Call _maybe_unflatten_for_moe after the pre-MLP norm in _run_moe_preprocess and _maybe_reflatten_from_moe after postprocess in _run_moe_combine, before the residual add. The micro-batch size travels in layer_state because another microbatch can run between the two slots. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On main, NVIDIA#6429 replaced TransformerConfig's ban on full recompute with EP overlap by layer-segment replay in the common schedule plan. The hybrid schedule plan inherits that replay once this branch is rebased, but it has only been validated with GPTModel layers, not with Mamba or grouped HybridStack layers. Reject recompute_granularity='full' in _validate_ep_overlap_support, next to the other features the hybrid overlap schedule does not run yet, until a hybrid port of NVIDIA#6429's parity tests passes. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PagedStashRunner collected MoE layers by reading .mlp on each entry of decoder.layers. A bracketed group is a nested HybridStack without .mlp, so with a pattern such as [M*E][M*E] the runner tracked no MoE layer. With --moe-expert-rank-capacity-factor or --moe-paged-stash, HybridEP and NCCL EP over-budget token drops in grouped layers then went undetected and triggered no dropless rerun. Descend into grouped HybridStack layers when collecting MoE layers for the decoder and MTP. This is the paged-stash part of NVIDIA#6958's commit 5d26416, moved here because this PR introduces the groups. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Main's GPTModel overlap callables now pass the chunk's packed_seq_params to MoELayer.route, as MoELayer.forward does, so dynamic-CP aware aux losses see the sequence layout (NVIDIA#7046). They also treat the DeepEP v2 flex backend ("deepepv2") like DeepEP (NVIDIA#5153). Do the same in the hybrid MoE slots. Without the backend check, a "deepepv2" run skipped handing the detached probs to the comm manager, reconnecting the backward graph across schedule nodes. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Main's hybrid packed-schedule test (NVIDIA#7047) and SP-embedding test (NVIDIA#8274) call HybridModel.forward on stubs. Bind the _preprocess/_postprocess hooks it now delegates to, and assert on the unwrapped decoder output, since the stub decoder unwraps its WrappedTensor like HybridStack does. - Main's TransformerLayerSchedulePlan.run reads recompute_segment (NVIDIA#6429), so the schedule-run test sets it on its __new__-built plan. - Main's DSA indexer FLOPs count (NVIDIA#7043) numbered every pattern character as a layer, so [ and ] shifted the top-k skip positions. Count the flattened symbols, which keep the ungrouped numbering. - Drop an unused test import and a comment about a chunkwise-CP padding-mask check that NVIDIA#7047 removed. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Main's NVIDIA#7054 added KDA ('K') hybrid layers. Like GDN, a KDA layer is a TransformerLayer whose self_attention is a GDN-family mixer with its own backward_dw, so the EP-overlap pre-dispatch slot runs it through _forward_attention and registers its delayed wgrad through the layer's backward_dw wrapper. Without this, a grouped pattern such as [K*E] raised "does not support layer type 'K'" on the first overlapped forward. Also count 'K' in the grouped layer-count test, since get_hybrid_layer_counts now reports every symbol in the table. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adjust the mcore-FSDP adapter and the megatron-FSDP core so HybridStack (including nested grouped HybridStack instances) is a valid FSDP unit and participates in the EP-overlap schedule plan. Add the ``test_fsdp_hybrid_overlap`` integration test exercising the FSDP + grouped HybridModel forward/backward path. Part 3/4 of splitting NVIDIA#4798 (original changes by @Wohox). Depends on the HybridStack changes in part 2/4 (#TBD). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Address review feedback on the FSDP HybridStack support: - Import HybridStack at module level in the MCore FSDP adapter instead of inside __init__. - Make fsdp_unit_filter the one mechanism that selects FSDP units. MegatronFSDP and BucketingPolicy now consult only the filter; the public fsdp_unit_modules argument is kept as shorthand for a class-based filter (FSDPUnitTypeFilter), and passing both raises a ValueError. - The MCore adapter passes only fsdp_unit_filter, built from its unit types, and leaves out the outer HybridStack of a grouped hybrid pattern. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Main's NVIDIA#7751 makes Megatron-FSDP reject gradient_reduce_div_fusion when one coalesced bucket group mixes gradient dtypes. The test builds its model without Float16Module, so MambaMixer's A_log and D stay FP32 next to the BF16 weights of the same [M*E] group FSDP unit, and both [M*E][M*E] cases failed at wrap time. Turn the fusion off in the test's DDP config; the overlap/no-overlap parity it checks does not depend on it. Signed-off-by: Yan Xu <yxu1@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch was successfully deployed
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
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.
Depends on #4942 — merge #4942 first. Scoped diff of this PR's own changes. The full diff also shows #4942's commits until it merges.
What does this PR do?
Add Megatron FSDP v1 support for bracketed HybridStack groups in the 1F1B EP-overlap schedule. Each bracket group becomes one FSDP unit; the outer decoder stack is excluded.
Apply the same instance filter to parameter buckets and lifecycle hooks, so separate groups never share the outer stack's bucket identity. Preserve the existing constructor's positional arguments. Regression coverage includes grouped, mixed, and ungrouped unit selection, plus three-step loss/parameter parity against non-overlapped FSDP for grouped attention and Mamba patterns.
Dependencies and review scope
Validation
Restacked on #4942 at
980af2d63(main through5052ce8c4); this PR's own commits applied without conflicts. Black, isort, repository-pinned pylint, Ruff, and diff checks pass for this slice. A limited CPU harness using PyTorch 2.14.0 passes all six unit-selection cases and fails the filtered case when the bucketing fix is removed. It executes the source grouping function with ordinary CPU parameters; this does not validate CUDA/FSDP collectives.An internal GB200 run tested the integrated stack at #6960
c331df5ef3f2d6e057ad11cee28cf6341d9cdd8bplus #7192d2385590f5ed4e0189e19e434d262ed5c7ac02e7on four GB200 GPUs (PyTorch2.12.0a0+0291f960b6.nv26.04.48445190, CUDA 13.2, Transformer Engine2.14.0+f031cf87). All 18 FSDP tests passed on every rank: 12 three-step loss/final-parameter parity cases across attention/Mamba/MLA, alltoall/flex dispatchers, shared experts off/on, plus six unit/bucket tests. That run validated the integrated source combination before the restack.After the restack,
tests/unit_tests/a2a_overlap/test_fsdp_hybrid_overlap.pyat the top of the stack (#6958) passed all six alltoall cases ([*E][*E],[M*E][M*E], and the MLA[+E][+E]case that #6958 adds, shared experts off/on) on every rank across eight GB300 GPUs; the flex-dispatcher cases were deselected because DeepEP/HybridEP cannot initialize in that container. On four GPUs,tests/unit_tests/distributed/mfsdp_v1/test_mfsdp_fully_shard.pypassed 26 (main: 19), with the same 315 skipped eight-rank cases and 6 xfailed as main.