Conversation
This was referenced May 22, 2026
Connor-XY
force-pushed
the
pr4798-2-hybrid-stack-grouped
branch
7 times, most recently
from
June 3, 2026 19:21
5e5dee3 to
91e86bb
Compare
44 of 72 tasks
Connor-XY
force-pushed
the
pr4798-2-hybrid-stack-grouped
branch
from
June 29, 2026 16:29
91e86bb to
7a0106e
Compare
Connor-XY
force-pushed
the
pr4798-2-hybrid-stack-grouped
branch
2 times, most recently
from
July 16, 2026 00:36
360afee to
b3467e6
Compare
Contributor
Author
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Connor-XY
force-pushed
the
pr4798-2-hybrid-stack-grouped
branch
from
July 20, 2026 16:33
eb80207 to
3242e57
Compare
Contributor
Author
|
/claude review |
Connor-XY
force-pushed
the
pr4798-2-hybrid-stack-grouped
branch
from
July 20, 2026 16:38
3242e57 to
b80f2c9
Compare
Contributor
There was a problem hiding this comment.
Light review complete. The hybrid grouping / bracketed-pattern feature is well-covered by new unit tests (layer allocation, sharded state dict, forward equivalence, EP-overlap callables, postprocess). One potential bug flagged inline in router.py: the aux_loss MTP-slot indexing fix (line 595) was not applied to the identical z_loss path (line 698), which can still index past the metrics tracker when MTP wraps a HybridStack and z_loss is enabled.
Contributor
Author
|
/claude review |
Contributor
Author
|
/ok to test 4aa1be7 |
Contributor
Author
|
/ok to test 12407a4 |
guihong-nv
approved these changes
Sep 10, 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>
Connor-XY
force-pushed
the
pr4798-2-hybrid-stack-grouped
branch
from
September 29, 2026 22:19
12407a4 to
9f08c2f
Compare
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Sep 29, 2026
Map logical grouped checkpoint keys and nested FSDP paths to physical GPT attention and MLP layers. Detect homogeneous versus indexed source keys and retarget model and optimizer state consistently. Carry the grouped checkpoint regression tests with this implementation, including parallel-layout roundtrips and model, optimizer, and FSDP loads. This is the checkpoint slice of e8db9dd, extracted from PR NVIDIA#4942. Signed-off-by: Yan Xu <yxu1@nvidia.com>
This was referenced Sep 29, 2026
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Connor-XY
added a commit
to Connor-XY/Megatron-LM
that referenced
this pull request
Sep 29, 2026
Map logical grouped checkpoint keys and nested FSDP paths to physical GPT attention and MLP layers. Detect homogeneous versus indexed source keys and retarget model and optimizer state consistently. Carry the grouped checkpoint regression tests with this implementation, including parallel-layout roundtrips and model, optimizer, and FSDP loads. This is the checkpoint slice of e8db9dd, extracted from PR 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>
Contributor
Author
|
/ok to test 3b24fca |
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>
Contributor
Author
|
/ok to test ad603d2 |
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
6 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.
What does this PR do?
Adds bracketed HybridStack groups and EP communication overlap to HybridModel. Patterns such as
[*-]and[*E]make several physical layers one logical layer for scheduling and checkpoint storage. Groups reject nesting, nonterminal MoE layers, and colliding checkpoint namespaces. This is part 2 of the original #4798 split by @Wohox, @Connor-XY, and @guihong-nv; the common scheduler prerequisite #4941 has merged.Grouping is what makes EP overlap possible for hybrid models. The overlap scheduler overlaps one microbatch's MoE all-to-all with another microbatch's compute at logical-layer granularity, so an attention (or Mamba/GDN) layer and the MoE layer after it must form one logical layer, e.g.
[*E].The overlap callables support attention, MLA, Mamba, GDN, and MoE layers, with per-microbatch norm-offload state and correct delayed-weight-gradient hook ownership. Hybrid MTP is not supported with EP overlap: as on main,
HybridModelraises for an MTP pattern withoverlap_moe_expert_parallel_comm. Existing ungrouped Hybrid checkpoint names and direct-stack offsets remain compatible.HybridModel.sharded_state_dictomits the output layer's (empty) TE extra state, asGPTModeldoes. GPT checkpoint import into grouped models is isolated in #7192.Change to the common MTP callable (affects GPTModel).
build_mtp_layer_callablesinmodels/common/fine_grained_callables.pynow detaches the per-depth hidden-state chunks it stores for the MTP post-process slot (node.detach), and the schedule node merges their gradients. Before, the post-process slot backpropagated through the pre-dispatch slot's graph, so the final norm's backward graph was traversed twice whenever the pre-dispatch slot applies the final norm. The code path runs for GPTModel withmtp_num_layers >= 1andoverlap_moe_expert_parallel_comm=True. The existing GPT chunk-level 1F1B tests with MTP (tests/unit_tests/a2a_overlap/test_schedule_chunk_1f1b.py) pass unchanged.Current-main integration and boundaries
Rebased on
b4d72b79c6b7e30010894ea0827a06c3477387fd. The refresh preserves eager hash routing, sequence-parallel token alignment, precomputed MTP embeddings, CP layout inputs, inference caches, and tensor observation hooks. Group parsing accounts for the current CSA/HCA/window attention symbols. MTP metrics retain main's explicit depth mapping. Ungrouped shortcut checkpoints keep physical storage indices, and static inference accepts caller-provided embeddings without token IDs.With EP overlap,
HybridModelrejects CUDA graphs, mixed CP-layout conversion, mHC, and MoE shortcuts at construction;TransformerConfigalready rejects hash-routed MoE and wide residuals with overlap. These features require dedicated schedule integration before they can be enabled together.Review and dependent PRs
The dependent PRs target
mainand contain this PR's commits until it merges, so this PR merges first.Validation
c331df5efplus Support GPT checkpoint loading into grouped HybridModel #7192d2385590f). Those heads contain this PR's commits unchanged, and the dependent PRs don't modify any file this PR touches. Each rank passed 18 FSDP, 19 regression, 31 checkpoint, and two un-packed Mamba replay cases (70 total). The four packed replay cases failed in that container before compute because itsselective_scan_cudabuild didn't match the platform. A rerun of all six ordinary/deferred Mamba replay cases on unchanged source in a matching container passed on every rank, including both packed layouts. No test expectations or dependency code were changed.14a16f198,305c0f9a7,ad603d21c), the unit tests for the files this PR changes andtests/unit_tests/a2a_overlapwere rerun atad603d21con an internal GB200 node. Results match the previous head's run, with no new failures. CI at the previous head3b24fcadcpassed apart from the code-owner approval gate.