[feat] HybridStack grouped syntax + checkpoint compat + EP-overlap (2/4 of #4798) - #4942
[feat] HybridStack grouped syntax + checkpoint compat + EP-overlap (2/4 of #4798)#4942Connor-XY wants to merge 9 commits into
Conversation
5e5dee3 to
91e86bb
Compare
91e86bb to
7a0106e
Compare
360afee to
b3467e6
Compare
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if in_inference_mode or is_spec_decode: | ||
| # Cache decoder hidden states for serial MTP computation after | ||
| # speculative token verification. | ||
| self._decoder_hidden_states_cache = hidden_states |
There was a problem hiding this comment.
This looks like a regression for hybrid MTP inference / speculative decoding. The cached hidden states are stored on self._decoder_hidden_states_cache, but nothing in the codebase reads that attribute. The pre-refactor hybrid code (and GPTModel._postprocess) write to inference_context.mtp_decoder_hidden_states, which is what TextGenerationController reads back (context.mtp_decoder_hidden_states, e.g. line 904). With this change has_mtp will always be False for hybrid models, so the serial MTP step after speculative-token verification never runs.
The GPT path also preserves the block-scope CUDA-graph copy_() into the pre-allocated buffer, which is dropped here. Consider mirroring GPT:
if is_spec_decode:
assert inference_context is not None
if self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block:
assert inference_context.mtp_decoder_hidden_states is not None
inference_context.mtp_decoder_hidden_states[: hidden_states.shape[0]].copy_(hidden_states)
else:
inference_context.mtp_decoder_hidden_states = hidden_states
elif not in_inference_mode:
hidden_states = process_mtp_loss(...)Note the condition also changed from spec-decode-only to in_inference_mode or is_spec_decode; the original only cached during is_spec_decode.
eb80207 to
3242e57
Compare
|
/claude review |
3242e57 to
b80f2c9
Compare
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.
|
/claude review |
|
/ok to test 4aa1be7 |
|
/claude strict-review |
| output_processor=None, | ||
| output_processor_context=None, |
There was a problem hiding this comment.
[IMPORTANT Compatibility] _postprocess accepts output_processor and output_processor_context but never uses them in the method body (the only two occurrences of output_processor in this file are these declarations).
Why it matters: GPTModel._postprocess implements a real branch — when output_processor is not None it short-circuits into the custom hook and returns early. The docstring says this method "mirrors GPTModel._postprocess", and PostProcessNode.forward_impl (megatron/core/models/common/utils.py) forwards chunk_state.output_processor/output_processor_context into _postprocess unconditionally. For a HybridModel these arrive and are silently discarded — the model runs the default logits/loss path instead of the caller's hook, producing wrong output with no error.
In this PR the hybrid build_schedule_plan never populates these (they default to None), so the divergence is latent today, but silently ignoring an argument that alters semantics elsewhere is a foot-gun that surfaces the moment a caller wires an output processor through the hybrid path.
Suggestion: either implement the same if output_processor is not None: branch that GPTModel._postprocess has, or drop the two parameters and let PostProcessNode fail loudly. If deferring, add assert output_processor is None, "output_processor not yet supported for HybridModel" plus a TODO.
There was a problem hiding this comment.
Fixed in 08e141e. _postprocess now has the same if output_processor is not None: early return as GPTModel._postprocess, forwarding the same kwargs (context=output_processor_context, compute_language_model_loss, scale_logits, config), placed at the same point — after the MTP block, before the materialize_only_last_token_logits slicing.
Left HybridModel.forward alone; it doesn't take an output processor today, and PostProcessNode is the path that actually forwards one.
New test: test_hybrid_postprocess_uses_output_processor_hook.
| if name == 'final_norm': | ||
| replace_prefix_for_sharding( | ||
| module_sharded_state_dict, module_prefix, f'{prefix}final_layernorm.' |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This renames the sharded checkpoint key for the final norm from decoder.final_norm.* to decoder.final_layernorm.* for all HybridStack instances — the branch is unconditional on name == 'final_norm', not restricted to bracketed-group models.
Why it matters: the registered submodule is still final_norm (line 238), and prior hybrid dist-checkpoints were saved with the final_norm. sharded prefix. After this change, loading an existing (non-grouped) hybrid checkpoint will request decoder.final_layernorm.* and miss the old decoder.final_norm.* entries — a load failure under raise_unexpected strictness, or silently-uninitialized final-norm weights otherwise. I don't see a load-side alias mapping old final_norm keys, so pre-existing hybrid checkpoints appear to become incompatible.
Suggestion: confirm this rename is intended for existing hybrid checkpoints (the goal of GPT key-compat is clear, but the break for old hybrid ckpts may be unintended). If the break is acceptable, note it in the PR description; otherwise add a load-time key alias/migration (final_norm → final_layernorm), or gate the rename so it only affects the top-level model that needs GPT compatibility.
There was a problem hiding this comment.
You're right that this was too broad, and it breaks more than old checkpoints. docs/user-guide/hybrid-model-migration.md documents decoder.final_norm as the HybridModel key, and both tools/checkpoint/gpt_hybrid_conversion.py and tools/checkpoint/remap_gpt_dsa_to_mamba.py rewrite GPT's decoder.final_layernorm.* into decoder.final_norm.* when converting. With the rename unconditional, every checkpoint those tools produce would have failed to load.
Now gated, in 08e141e: HybridStack takes a transformer_sharded_keys flag, and HybridModel sets it from whether the full layer pattern contains a bracketed group — full pattern rather than this rank's PP segment, so every stage agrees on the naming. Grouped patterns get final_layernorm (their logical layers are 1:1 with transformer layers, which is the whole point of the GPT cross-load); everything else keeps final_norm, so pre-existing hybrid checkpoints and converter output still load. The converter itself only accepts flat patterns — [ and ] aren't in its VALID_LAYER_SYMBOLS — so it can't produce a grouped target and stays on the final_norm side.
Tests: test_ungrouped_sharded_state_dict_keeps_hybrid_final_norm_key and test_sharded_state_dict_keeps_final_norm_key_without_transformer_keys. The existing grouped assertions are unchanged; test_group_sharded_state_dict_uses_logical_layer_keys builds a HybridStack directly so it now passes transformer_sharded_keys=True itself.
|
Strict Review Summary Findings by severity: CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 0 This is a large, careful refactor (bracketed HybridStack groups, EP-overlap schedule plans, checkpoint-key compat, Mamba backward_dw, MTP aux-loss modulo fix). The code is well-commented, the tricky correctness rationale (viewless tensors, the _run_moe_combine recompute-hook fix, shared-expert wgrad slot ordering) is documented inline, and unit coverage is strong (grouped forward-equivalence, sharded-key, aux-loss slot). Most new identifiers were verified to have real use paths — logical_layer_offset, is_layer_group_stack, the physical/logical count helpers, _resolve_free_input, and the new backward_dw methods all check out. The router MTP modulo fix and its guard on mtp_num_layers is not None are correct. Most impactful findings:
Overall risk: Low-to-moderate. No correctness bugs found in the parallelism, dtype, or scheduling logic. The two flagged items are backward-compatibility concerns — the checkpoint-key rename is the one worth an explicit confirmation before merge. The dependency noted in the PR description means the diff will look larger until the prerequisite part merges (expected). |
| output_processor=None, | ||
| output_processor_context=None, |
There was a problem hiding this comment.
[IMPORTANT Compatibility] _postprocess accepts output_processor and output_processor_context but never uses them anywhere in the method body (confirmed: the only two occurrences of output_processor in this file are these two parameter declarations).
Why it matters: GPTModel._postprocess implements a real branch on these — when output_processor is not None it short-circuits into the custom hook (used by RL / custom output paths) and returns early. The docstring here states this method "mirrors GPTModel._postprocess", and PostProcessNode.forward_impl (megatron/core/models/common/utils.py) forwards chunk_state.output_processor/output_processor_context into _postprocess unconditionally. For a HybridModel these currently arrive and are silently discarded — the model runs the default logits/loss path instead of the caller's hook, producing wrong output with no error.
In this PR the hybrid build_schedule_plan never populates these (they default to None), so the discrepancy is latent today. But accepting an argument that changes semantics elsewhere and silently ignoring it is a foot-gun that will surface the moment a caller wires an output processor through the hybrid path.
Suggestion: either implement the same if output_processor is not None: branch that GPTModel._postprocess has, or drop the two parameters entirely and let PostProcessNode fail loudly (TypeError) rather than silently no-op. If deferring, add an explicit assert output_processor is None, "output_processor not yet supported for HybridModel" plus a TODO.
There was a problem hiding this comment.
Fixed in 08e141e. _postprocess now has the same if output_processor is not None: early return as GPTModel._postprocess, forwarding the same kwargs (context=output_processor_context, compute_language_model_loss, scale_logits, config), placed at the same point — after the MTP block, before the materialize_only_last_token_logits slicing.
Left HybridModel.forward alone; it doesn't take an output processor today, and PostProcessNode is the path that actually forwards one.
New test: test_hybrid_postprocess_uses_output_processor_hook.
| if name == 'final_norm': | ||
| replace_prefix_for_sharding( | ||
| module_sharded_state_dict, module_prefix, f'{prefix}final_layernorm.' |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This renames the sharded checkpoint key for the final norm from decoder.final_norm.* to decoder.final_layernorm.* for all HybridStack instances (the branch is unconditional on name == 'final_norm'), not just bracketed-group models.
Why it matters: the registered submodule is still final_norm (line 238), and prior hybrid checkpoints were saved with the final_norm. sharded prefix. After this change, loading an existing (non-grouped) hybrid dist-checkpoint will look for decoder.final_layernorm.* and miss the old decoder.final_norm.* entries — either a load failure under raise_unexpected strictness or silently-uninitialized final-norm weights. The PR validated Transformer↔HybridModel cross-load, but I don't see a load-side fallback that maps old final_norm keys, so pre-existing hybrid checkpoints appear to become incompatible.
Suggestion: confirm this is intended for existing hybrid checkpoints (not just the new grouped ones), and if so add a load-time key alias / migration for final_norm → final_layernorm (or note it as a breaking change in the PR description). If the rename should only affect grouped stacks, gate it on self.is_layer_group_stack / the top-level hybrid model.
There was a problem hiding this comment.
You're right that this was too broad, and it breaks more than old checkpoints. docs/user-guide/hybrid-model-migration.md documents decoder.final_norm as the HybridModel key, and both tools/checkpoint/gpt_hybrid_conversion.py and tools/checkpoint/remap_gpt_dsa_to_mamba.py rewrite GPT's decoder.final_layernorm.* into decoder.final_norm.* when converting. With the rename unconditional, every checkpoint those tools produce would have failed to load.
Now gated, in 08e141e: HybridStack takes a transformer_sharded_keys flag, and HybridModel sets it from whether the full layer pattern contains a bracketed group — full pattern rather than this rank's PP segment, so every stage agrees on the naming. Grouped patterns get final_layernorm (their logical layers are 1:1 with transformer layers, which is the whole point of the GPT cross-load); everything else keeps final_norm, so pre-existing hybrid checkpoints and converter output still load. The converter itself only accepts flat patterns — [ and ] aren't in its VALID_LAYER_SYMBOLS — so it can't produce a grouped target and stays on the final_norm side.
Tests: test_ungrouped_sharded_state_dict_keeps_hybrid_final_norm_key and test_sharded_state_dict_keeps_final_norm_key_without_transformer_keys. The existing grouped assertions are unchanged; test_group_sharded_state_dict_uses_logical_layer_keys builds a HybridStack directly so it now passes transformer_sharded_keys=True itself.
…essor Two compatibility fixes from the strict review on NVIDIA#4942: - HybridModel._postprocess accepted output_processor / output_processor_context but silently discarded them, so a caller that wired an output hook through PostProcessNode would get the default logits/loss path with no error. Implement the same early-return branch GPTModel._postprocess has. - The final-norm sharded key rename (final_norm -> final_layernorm) was unconditional, so it also changed the keys of non-grouped hybrid models whose existing dist checkpoints were saved under final_norm. Gate it on a new transformer_sharded_keys flag that HybridModel derives from the full layer pattern, so only bracketed-group models (whose logical layers map one-to-one onto transformer layers, which is what the GPT cross-load compatibility is for) get the transformer-style key. Signed-off-by: Yan Xu <yxu1@nvidia.com>
|
Both strict-review findings addressed in ffcae6d + 08e141e, with replies inline.
|
Add bracketed HybridStack group syntax (e.g. ``[*-]``, ``M[M*]-``) with nested HybridStack instances, rejecting invalid recursion. Migrate grouped HybridStack checkpoints to Transformer-compatible logical layer keys and make ``HybridModel.sharded_state_dict()`` drop the empty ``output_layer._extra_state`` to match GPT behavior. Extend EP-overlap scheduling to HybridStack: add the hybrid fine-grained callables and ``HybridStackModelChunkSchedulePlan``, expose ``HybridModel.build_schedule_plan`` and add the ``return_schedule_plan`` path in ``pretrain_hybrid.py``. Add Mamba ``backward_dw`` so the hybrid schedule node can register Mamba pre-layer weight grads alongside attention and GDN pre-layers. Fix the MoE TopKRouter MTP layer-number indexing when the MTP block wraps a HybridStack so the aux-loss tracker is not indexed past its size. Part 2/4 of splitting NVIDIA#4798 (original changes by @Wohox). Depends on the common combined-1F1B refactor in part 1/4 (#TBD). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
Carries over upstream commit ce6e229 from NVIDIA#4798: in HybridStack's ``_run_moe_combine`` (A2A overlap path), ``layer._forward_post_mlp`` registers a second ``discard_output_and_register_recompute`` hook on ``mlp_output_with_bias[0]``. The hook fires during combine_bwd's autograd backward and triggers the LN recompute ahead of mlp_bwd / pre_dispatch_bwd. In bracketed-hybrid logical layers (``[*E]``), this corrupts gradients in attention's autograd chain (grad_norm explodes from iter 2). Fix: stop calling ``_forward_post_mlp`` from ``_run_moe_combine``; inline the ``bda + offload_mlp_norm + make_viewless_tensor`` steps directly, mirroring GPT's ``submodule_combine_forward``. The first recompute hook on ``expert_output`` (registered in ``_run_moe_experts``) already fires the LN recompute in mlp_bwd, so the second hook is redundant. Part 2/4 of splitting NVIDIA#4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>
The recent merge of origin/main introduced `name=(name + f".layers.{i}")`
into every layer-type branch of HybridStack's build loop, but didn't change
the local loop header `for layer_type in self.layer_type_list:` to surface
`i`. Result: `NameError: name 'i' is not defined` at HybridStack init for
all hybrid runs (GPT path unaffected).
Trigger: any hybrid_stack_spec model crashes on init, including the 16-node
Bug 2a repro and the 8-node GPT-vs-Hybrid perf comparison runs.
Fix: convert the loop to `for i, layer_type in enumerate(...)`. Keep the
existing `physical_layer_offset` counter (used for FP8/FP4 contexts and
`layer_number`) because bracket groups count >1 physical layer per logical
entry — these are separate from the logical index `i` used for module names.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Co-authored-by: Pingtian Li <pingtianl@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Publish speculative-decoding hidden states through the inference context, including the fixed buffer used by block-scope CUDA graphs. Preserve the canonical inference-mode check and the inputs required to derive RL MTP labels. Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
Signed-off-by: Yan Xu <yxu1@nvidia.com>
…essor Two compatibility fixes from the strict review on NVIDIA#4942: - HybridModel._postprocess accepted output_processor / output_processor_context but silently discarded them, so a caller that wired an output hook through PostProcessNode would get the default logits/loss path with no error. Implement the same early-return branch GPTModel._postprocess has. - The final-norm sharded key rename (final_norm -> final_layernorm) was unconditional, so it also changed the keys of non-grouped hybrid models whose existing dist checkpoints were saved under final_norm. Gate it on a new transformer_sharded_keys flag that HybridModel derives from the full layer pattern, so only bracketed-group models (whose logical layers map one-to-one onto transformer layers, which is what the GPT cross-load compatibility is for) get the transformer-style key. Signed-off-by: Yan Xu <yxu1@nvidia.com>
test_group_sharded_state_dict_uses_logical_layer_keys builds a HybridStack directly rather than through HybridModel, so it has to set transformer_sharded_keys itself now that the final-norm key rename is gated. Add the mirror-image case asserting the default keeps final_norm. Signed-off-by: Yan Xu <yxu1@nvidia.com>
08e141e to
6b3772f
Compare
|
Rebased onto current main (076f61f). Conflicts were with the MLA-in-HybridModel port (#4452) and the mamba prefix-caching kernels (#5866):
Verified on 1 GB200 (hsg): 169 passed, 3 failed. The three are |
4 participants
What does this PR do?
Part 2 of 4 splitting #4798 by @Wohox, @Connor-XY, and @guihong-nv.
Summary
Add the core HybridStack feature work from #4798:
[*-],[*E],M[M*]-): build bracketed groups as nestedHybridStackinstances; reject invalid recursion (nested brackets inside a group); keep MoE constrained to the last symbol inside a group for EP-overlap scheduling.HybridModel.sharded_state_dict()drops the emptyoutput_layer._extra_stateto match GPT behavior, and grouped stacks publish the final norm under GPT'sdecoder.final_layernorm.*instead ofdecoder.final_norm.*. The key rename is gated ontransformer_sharded_keys, whichHybridModelsets only when the full layer pattern contains a bracketed group, so non-grouped hybrid models keepdecoder.final_norm.*— the key existing hybrid checkpoints were saved with and the onetools/checkpoint/gpt_hybrid_conversion.pywrites. Verified cross-load (Transformer ↔ HybridModel) in [feat] Hybrid model ep overlapping main #4798.HybridStackModelChunkSchedulePlan,hybrid/fine_grained_callables.py, exposeHybridModel.build_schedule_plan, and wire thereturn_schedule_planpath inpretrain_hybrid.py.backward_dw: register Mamba pre-layer wgrad alongside attention/GDN pre-layers so the schedule node iterates a uniform set of callables.self.layer_numbervia modulo so the aux-loss tracker is not indexed past its size when MTP wraps a HybridStack (e.g.*Efor one depth).pre_mlp_layernormrecompute hook in moe_combine that corrupted attention gradients in bracketed-hybrid logical layers ([*E]).Why this slice
Touches 6 reviewer groups:
core-adlr,core-nemo,hybrid-model,hybrid-mamba,mixture-of-experts-adlr,mixture-of-experts-devtech.Dependencies
_BackwardDWWrapperlocation.GitHub will show #4941's diff in this PR until #4941 merges; expected.
Validation
The full integrated change was validated in #4798:
test_hybrid_layer_allocation,test_hybrid_block,test_hybrid_model::test_grouped_sharded_state_dict_uses_transformer_checkpoint_keys— 86 passed.--dist-ckpt-strictness raise_unexpected: STATUS 0 for save / transformer→hybrid / hybrid→transformer.[*-][*-]|[*-][*-]|[*E][*E]|...pattern + VPP=2 + A2A + recompute layernorm: 100/100 iters identical between baseline and A2A overlap.Issue tracking
Linked issue: part of #4798.
Pre-checks
🤖 Generated with Claude Code