From ba1e5c0a93f84bee5c3a08ac9019a9f1666899ef Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Thu, 21 May 2026 15:12:28 -0700 Subject: [PATCH 01/11] [feat] HybridStack grouped syntax + checkpoint compat + EP-overlap 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 #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) Signed-off-by: Yan Xu --- .../models/common/fine_grained_callables.py | 11 + .../models/hybrid/fine_grained_callables.py | 374 ++++++++++++++ megatron/core/models/hybrid/hybrid_block.py | 154 +++++- .../models/hybrid/hybrid_layer_allocation.py | 281 ++++++++-- megatron/core/models/hybrid/hybrid_model.py | 485 ++++++++++++------ .../hybrid/model_chunk_schedule_plan.py | 138 +++++ megatron/core/ssm/mamba_layer.py | 11 + megatron/core/ssm/mamba_mixer.py | 15 + megatron/core/transformer/moe/router.py | 15 +- pretrain_hybrid.py | 33 +- tests/unit_tests/models/test_hybrid_model.py | 23 + tests/unit_tests/ssm/test_hybrid_block.py | 163 +++++- .../ssm/test_hybrid_layer_allocation.py | 57 ++ 13 files changed, 1518 insertions(+), 242 deletions(-) create mode 100644 megatron/core/models/hybrid/fine_grained_callables.py create mode 100644 megatron/core/models/hybrid/model_chunk_schedule_plan.py diff --git a/megatron/core/models/common/fine_grained_callables.py b/megatron/core/models/common/fine_grained_callables.py index 8f46711d553..184e853d0c4 100644 --- a/megatron/core/models/common/fine_grained_callables.py +++ b/megatron/core/models/common/fine_grained_callables.py @@ -144,9 +144,14 @@ def rng_context_wrapper(func, *args, **kwargs): def get_layer_moe_metadata(layer): """Return ``(is_moe, num_local_experts)`` for schedule-node construction.""" + from megatron.core.models.hybrid.hybrid_block import HybridStack if isinstance(layer, MultiTokenPredictionLayer): return get_layer_moe_metadata(layer.mtp_model_layer) + if isinstance(layer, HybridStack): + from megatron.core.models.hybrid.fine_grained_callables import get_hybrid_stack_moe_metadata + + return get_hybrid_stack_moe_metadata(layer) if isinstance(layer, TransformerLayer): is_moe = isinstance(layer.mlp, MoELayer) num_local_experts = layer.mlp.num_local_experts if is_moe else None @@ -160,9 +165,15 @@ def build_layer_callables(layer): Returns ``(forward_funcs, backward_dw)``. """ + from megatron.core.models.hybrid.hybrid_block import HybridStack if isinstance(layer, MultiTokenPredictionLayer): return build_mtp_layer_callables(layer) + if isinstance(layer, HybridStack): + from megatron.core.models.hybrid.fine_grained_callables import build_hybrid_stack_callables + + forward_funcs, backward_dw, _, _ = build_hybrid_stack_callables(layer) + return forward_funcs, backward_dw if isinstance(layer, TransformerLayer): return build_transformer_layer_callables(layer) diff --git a/megatron/core/models/hybrid/fine_grained_callables.py b/megatron/core/models/hybrid/fine_grained_callables.py new file mode 100644 index 00000000000..7e45881518c --- /dev/null +++ b/megatron/core/models/hybrid/fine_grained_callables.py @@ -0,0 +1,374 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from contextlib import nullcontext +from functools import partial +from typing import Optional + +import torch +from torch import Tensor + +from megatron.core.enums import Fp8Recipe +from megatron.core.fp4_utils import get_fp4_context +from megatron.core.fp8_utils import get_fp8_context +from megatron.core.models.common.utils import TransformerLayerNode, should_free_input +from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_layer_allocation import LayerPatternItem +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols +from megatron.core.models.hybrid.hybrid_layer_allocation import is_layer_group +from megatron.core.pipeline_parallel.utils import ScheduleNode +from megatron.core.transformer.transformer_layer import make_viewless_tensor + + +class _SharedExpertBackwardDWWrapper: + """Run MoE shared-experts wgrad as part of the ``pre_dispatch_computation`` slot. + + Why: shared-experts forward is part of ``_run_moe_preprocess`` (which runs in the + pre_dispatch slot), and TE's delay-wgrad model only ``put``s to the wgrad queue + inside the autograd backward (dgrad). So shared-experts' ``backward_dw`` must + fire *after* the pre_dispatch slot's autograd backward — registering it in the + ``mlp`` slot would call it before that dgrad and trigger + ``RuntimeError: Pop empty queue`` from TE. + """ + + def __init__(self, layer): + self.layer = layer + self.shared_expert_dw_callable = None + if layer.mlp.use_shared_expert and not layer.mlp.shared_expert_overlap: + self.shared_expert_dw_callable = partial( + layer.mlp.backward_dw, routed_experts=False, shared_experts=True + ) + + def backward_dw(self): + """Run shared-expert backward wgrad after pre-dispatch autograd backward.""" + if self.shared_expert_dw_callable is not None: + self.shared_expert_dw_callable() + self.layer = None + self.shared_expert_dw_callable = None + + def parameters(self): + """Expose shared-expert params so post_wgrad_grad_acc_hook discovery works. + + The schedule node's backward_dw iterates module.parameters() looking for + post_wgrad_grad_acc_hook on each param. _SharedExpertBackwardDWWrapper is + a plain class (not nn.Module), so we forward to layer.mlp.shared_experts. + Returns an empty iterator after backward_dw has cleared the layer ref. + """ + if self.layer is None or self.layer.mlp.shared_experts is None: + return iter([]) + return self.layer.mlp.shared_experts.parameters() + + +class HybridStackNode(TransformerLayerNode): + """Schedule node for HybridStack-built fine-grained callables. + + Subclassed from ``TransformerLayerNode`` so the runtime backbone (forward / + backward / backward_dw plumbing, detach bookkeeping, output-grad release) + is shared. The hybrid path keeps a separate node class so its free-input + policy can diverge from the GPT defaults — for example, the + ``pre_dispatch_computation`` slot here covers the whole pre-dispatch loop + (mamba + attention + …) rather than a single attention block, and + group-level decisions about whether the input is needed in backward may + differ from ``should_free_input`` in ``gpt/fine_grained_callables.py``. + Keep this override thin until a hybrid counter-example forces it to + diverge; the explicit subclass exists so the divergence can be made + surgically without touching the GPT class. + """ + + @staticmethod + def _resolve_free_input(name, is_moe, config, num_local_experts): + """Hybrid free-input policy. + + Currently mirrors the GPT default: dense layers always retain their + input for backward; MoE-only "moe_dispatch", "mlp", and "moe_combine" + slots can free, subject to the dispatcher / cuda-graph constraints + encoded in ``should_free_input``. Hybrid groups have a + "pre_dispatch_computation" slot whose semantics differ (it covers a + loop over Mamba/attention/GDN sub-layers, not a single attention + block), but its policy resolves to ``False`` in + ``should_free_input``, which is correct: pre-layer outputs are needed + for backward through the loop. Override here when a hybrid-specific + rule is needed. + """ + return should_free_input(name, is_moe, config, num_local_experts) + + +def _get_inner_quant_context(layer): + config = layer.config + if config.fp8 and config.fp8_recipe != Fp8Recipe.delayed: + return get_fp8_context(config, layer.layer_number - 1) + if config.fp4: + return get_fp4_context(config, layer.layer_number - 1) + return nullcontext() + + +def _as_hybrid_layers(layer, layer_type: Optional[LayerPatternItem]): + """Return ``(layer_type, layer)`` pairs for a hybrid logical layer.""" + if isinstance(layer, HybridStack): + return list(zip(layer.layer_type_list, layer.layers)) + assert layer_type is not None, "Hybrid layer scheduling requires the layer type symbol." + return [(layer_type, layer)] + + +def _split_hybrid_layers_for_overlap(layer, layer_type: Optional[LayerPatternItem]): + layer_items = _as_hybrid_layers(layer, layer_type) + if any(is_layer_group(item_type) for item_type, _ in layer_items): + raise ValueError("Nested HybridStack groups are not supported in overlap scheduling.") + + terminal_idx = None + for idx, (item_type, _) in enumerate(layer_items): + if item_type in (LayerSymbols.MLP, LayerSymbols.MOE): + terminal_idx = idx + break + + if terminal_idx is not None and terminal_idx != len(layer_items) - 1: + raise ValueError("HybridStack overlap requires MLP/MoE to be the last layer in a group.") + + terminal_type = layer_items[terminal_idx][0] if terminal_idx is not None else None + terminal_layer = layer_items[terminal_idx][1] if terminal_idx is not None else None + pre_layers = layer_items[:terminal_idx] if terminal_idx is not None else layer_items + is_moe = terminal_type == LayerSymbols.MOE + num_local_experts = terminal_layer.mlp.num_local_experts if is_moe else None + return pre_layers, terminal_type, terminal_layer, is_moe, num_local_experts + + +def get_hybrid_stack_moe_metadata(layer, layer_type: Optional[LayerPatternItem] = None): + """Return ``(is_moe, num_local_experts)`` for one HybridStack schedule layer.""" + _, _, _, is_moe, num_local_experts = _split_hybrid_layers_for_overlap(layer, layer_type) + return is_moe, num_local_experts + + +def _maybe_apply_final_norm(node: ScheduleNode, hidden_states: Tensor): + final_norm = getattr(node.chunk_state.model.decoder, "final_norm", None) + final_norm = final_norm or getattr(node.chunk_state.model.decoder, "final_layernorm", None) + if not node.is_mtp and final_norm is not None and node.is_last_layer: + hidden_states = final_norm(hidden_states) + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + return hidden_states + + +def _get_moe_padding_mask(node: ScheduleNode): + padding_mask = node.chunk_state.padding_mask + if padding_mask is not None: + # MoELayer.forward receives [batch, seq] and transposes before routing. + padding_mask = padding_mask.transpose(0, 1).bool() + return padding_mask + + +def _run_moe_preprocess(layer, node: ScheduleNode, hidden_states: Tensor): + pre_mlp_layernorm_output = layer._forward_pre_mlp_layernorm(hidden_states) + if isinstance(pre_mlp_layernorm_output, tuple): + if len(pre_mlp_layernorm_output) != 2: + raise ValueError( + f"When the output of pre_mlp_layernorm is a tuple, it is expected to have " + f"2 elements (output, residual), but got {len(pre_mlp_layernorm_output)}" + ) + pre_mlp_layernorm_output, residual = pre_mlp_layernorm_output + else: + residual = hidden_states + + if layer.config.fp32_residual_connection: + residual = residual.float() + + shared_expert_output = layer.mlp.shared_experts_compute(pre_mlp_layernorm_output) + probs, routing_map = layer.mlp.route(pre_mlp_layernorm_output, _get_moe_padding_mask(node)) + local_tokens, probs = layer.mlp.preprocess(pre_mlp_layernorm_output, probs, routing_map) + + node.layer_state.residual = node.detach(residual) + if layer.mlp.use_shared_expert and not layer.mlp.shared_expert_overlap: + node.layer_state.shared_expert_output = node.detach(shared_expert_output) + + return local_tokens, probs + + +def _run_moe_experts(layer, node: ScheduleNode, dispatched_tokens: Tensor): + dispatched_probs = node.layer_state.dispatched_probs + enable_hybridep = ( + layer.config.moe_token_dispatcher_type == "flex" + and layer.config.moe_flex_dispatcher_backend == "hybridep" + ) + enable_deepep = ( + layer.config.moe_token_dispatcher_type == "flex" + and layer.config.moe_flex_dispatcher_backend == "deepep" + ) + token_dispatcher = layer.mlp.token_dispatcher + if enable_deepep or enable_hybridep: + token_dispatcher._comm_manager.dispatched_probs = dispatched_probs + + expert_output, _ = layer.mlp.routed_experts_compute(dispatched_tokens, dispatched_probs) + + if enable_hybridep: + tokens_per_expert = token_dispatcher._comm_manager.get_number_of_tokens_per_expert() + node.layer_state.tokens_per_expert = tokens_per_expert + + if layer.recompute_pre_mlp_layernorm: + layer.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(expert_output) + + return expert_output + + +def _run_moe_combine(layer, node: ScheduleNode, output: Tensor): + residual = node.layer_state.residual + shared_expert_output = getattr(node.layer_state, 'shared_expert_output', None) + output = layer.mlp.combine(output) + output = layer.mlp.postprocess(output, shared_expert_output) + output = layer._forward_post_mlp((output, None), residual) + + node.layer_state.residual.record_stream(torch.cuda.current_stream()) + if shared_expert_output is not None: + shared_expert_output.record_stream(torch.cuda.current_stream()) + + node.layer_state.residual = None + node.layer_state.shared_expert_output = None + + return _maybe_apply_final_norm(node, output) + + +def build_hybrid_stack_callables(layer, layer_type: Optional[LayerPatternItem] = None): + """Create fine-grained callables for one logical HybridStack layer. + + A logical layer may be a bracketed nested ``HybridStack`` (for example ``[M*E]``) + or a single legacy hybrid layer symbol. The split is: + pre-dispatch compute -> dispatch -> MLP/experts -> combine. + """ + pre_layers, terminal_type, terminal_layer, is_moe, num_local_experts = ( + _split_hybrid_layers_for_overlap(layer, layer_type) + ) + + def pre_dispatch_computation(node: ScheduleNode, hidden_states: Tensor): + for item_type, item_layer in pre_layers: + with _get_inner_quant_context(item_layer): + if item_type == LayerSymbols.MAMBA: + hidden_states = item_layer( + hidden_states=hidden_states, + attention_mask=node.chunk_state.attention_mask, + inference_context=getattr(node.chunk_state, "inference_context", None), + packed_seq_params=node.chunk_state.packed_seq_params, + ) + elif item_type in ( + LayerSymbols.ATTENTION, + LayerSymbols.DS_ATTENTION, + LayerSymbols.GDN, + ): + # Use _forward_attention rather than __call__: an attention half-layer has + # mlp=IdentityOp / mlp_bda=IdentityFuncOp by default, and TransformerLayer's + # __call__ would route through _forward_mlp + mlp_bda, double-applying the + # post-attention residual. + hidden_states, _ = item_layer._forward_attention( + hidden_states=hidden_states, + attention_mask=node.chunk_state.attention_mask, + rotary_pos_emb=node.chunk_state.rotary_pos_emb, + rotary_pos_cos=node.chunk_state.rotary_pos_cos, + rotary_pos_sin=node.chunk_state.rotary_pos_sin, + packed_seq_params=node.chunk_state.packed_seq_params, + sequence_len_offset=node.chunk_state.sequence_len_offset, + ) + # _forward_attention returns the bias_dropout_add output which can be a + # view tensor (the mlp_bda's add into the post-attention residual produces + # a view from a fused/JIT kernel). Downstream cuBLAS matmuls — including + # the terminal MLP/MoE's pre_mlp_layernorm and the next attention's QKV + # projection in a multi-pre-layer group — pick algorithms based on input + # strides; a view's non-canonical strides can lead to different algo + # selection across processes and produce ~1e-5 bit drift on the forward + # output. TransformerLayer's full forward() inserts this exact call at the + # MLP exit (transformer_layer.py:895) for the same reason; the + # _forward_attention shortcut here doesn't get that cleanup, so we add it + # explicitly. Same idea as the make_viewless_tensor in _maybe_apply_final_norm. + hidden_states = make_viewless_tensor( + inp=hidden_states, + requires_grad=hidden_states.requires_grad, + keep_graph=True, + ) + else: + raise ValueError( + f"HybridStack overlap does not support layer type '{item_type}' before " + "the terminal MLP/MoE layer." + ) + + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + if terminal_type == LayerSymbols.MOE: + with _get_inner_quant_context(terminal_layer): + return _run_moe_preprocess(terminal_layer, node, hidden_states) + + if terminal_type is None: + return _maybe_apply_final_norm(node, hidden_states) + + return hidden_states + + def dispatch(node: ScheduleNode, local_tokens: Tensor, probs: Tensor): + enable_hybridep = ( + terminal_layer.config.moe_token_dispatcher_type == "flex" + and terminal_layer.config.moe_flex_dispatcher_backend == "hybridep" + ) + enable_deepep = ( + terminal_layer.config.moe_token_dispatcher_type == "flex" + and terminal_layer.config.moe_flex_dispatcher_backend == "deepep" + ) + token_dispatcher = terminal_layer.mlp.token_dispatcher + if enable_deepep or enable_hybridep: + token_dispatcher._comm_manager.token_probs = probs + with _get_inner_quant_context(terminal_layer): + dispatched_tokens, dispatched_probs = terminal_layer.mlp.dispatch(local_tokens, probs) + node.layer_state.dispatched_probs = node.detach(dispatched_probs) + return dispatched_tokens + + def mlp(node: ScheduleNode, hidden_states: Tensor): + if terminal_type == LayerSymbols.MLP: + with _get_inner_quant_context(terminal_layer): + hidden_states = terminal_layer._forward_mlp( + hidden_states, padding_mask=node.chunk_state.padding_mask + ) + return _maybe_apply_final_norm(node, hidden_states) + if terminal_type == LayerSymbols.MOE: + with _get_inner_quant_context(terminal_layer): + return _run_moe_experts(terminal_layer, node, hidden_states) + return hidden_states + + def combine(node: ScheduleNode, output: Tensor): + with _get_inner_quant_context(terminal_layer): + return _run_moe_combine(terminal_layer, node, output) + + def raise_not_implemented(*args): + raise NotImplementedError("This callable is not implemented for non-MoE hybrid layers.") + + backward_dw = {} + pre_bwd_dw = [] + for item_type, item_layer in pre_layers: + if item_type in (LayerSymbols.ATTENTION, LayerSymbols.DS_ATTENTION, LayerSymbols.GDN): + # TransformerLayer-backed pre-layers go through the standard + # _BackwardDWWrapper which coordinates attn / shared-expert wgrad + # with cuda-graph replay scopes. + item_layer.init_backward_dw_wrapper() + pre_bwd_dw.append(item_layer.backward_dw_wrapper) + elif item_type == LayerSymbols.MAMBA: + # MambaLayer is not a TransformerLayer, so init_backward_dw_wrapper + # would assert. MambaLayer.backward_dw delegates to its mixer, which + # in turn calls backward_dw on the in_proj / out_proj linears. The + # schedule node iterates this list and calls .backward_dw() on each; + # registering the layer directly is sufficient. + pre_bwd_dw.append(item_layer) + if is_moe: + # MoELayer.backward_dw default kwargs (routed_experts=True, shared_experts=False) + # handle the routed-experts wgrad in the mlp slot. The shared-experts wgrad goes + # into the pre_dispatch_computation slot so it runs after that slot's autograd + # backward (where TE's wgrad_store.put fires); the wrapper is a no-op when + # shared_expert_overlap is enabled. + shared_expert_dw = _SharedExpertBackwardDWWrapper(terminal_layer) + if shared_expert_dw.shared_expert_dw_callable is not None: + pre_bwd_dw.append(shared_expert_dw) + backward_dw["mlp"] = terminal_layer.mlp + elif terminal_type == LayerSymbols.MLP: + backward_dw["mlp"] = terminal_layer.mlp + + if pre_bwd_dw: + backward_dw["pre_dispatch_computation"] = pre_bwd_dw + + forward_funcs = [ + pre_dispatch_computation, + dispatch if is_moe else raise_not_implemented, + mlp, + combine if is_moe else raise_not_implemented, + None, + ] + return forward_funcs, backward_dw, is_moe, num_local_experts diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 0042cbea010..073fd0385ce 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -21,7 +21,12 @@ from megatron.core.fp8_utils import get_fp8_context from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.inference.utils import InferenceMode +from megatron.core.models.hybrid.hybrid_layer_allocation import LayerPatternItem from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols +from megatron.core.models.hybrid.hybrid_layer_allocation import ( + get_layer_type_physical_count, + is_layer_group, +) from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.recompute import checkpointed_forward @@ -82,8 +87,10 @@ def __init__( config: TransformerConfig, submodules: HybridStackSubmodules, pre_process: bool = True, - layer_type_list: Optional[list[str]] = None, + layer_type_list: Optional[list[LayerPatternItem]] = None, pp_layer_offset: int = 0, + logical_layer_offset: int = 0, + is_layer_group_stack: bool = False, post_layer_norm: bool = True, post_process: bool = True, device=None, @@ -101,6 +108,8 @@ def __init__( self.post_layer_norm = post_layer_norm self.post_process = post_process self.is_mtp_layer = is_mtp_layer + self.logical_layer_offset = logical_layer_offset + self.is_layer_group_stack = is_layer_group_stack assert pg_collection is not None, "pg_collection must be provided for HybridStack" @@ -122,16 +131,39 @@ def __init__( # Build layers from the pre-selected segment self.layers = nn.ModuleList() - for i, layer_type in enumerate(self.layer_type_list): - layer_number = i + 1 + pp_layer_offset - if self.config.fp8: - quant_init_context = get_fp8_context(self.config, i + pp_layer_offset, is_init=True) + physical_layer_offset = pp_layer_offset + for layer_type in self.layer_type_list: + layer_number = physical_layer_offset + 1 + if is_layer_group(layer_type): + quant_init_context = nullcontext() + elif self.config.fp8: + quant_init_context = get_fp8_context( + self.config, physical_layer_offset, is_init=True + ) elif self.config.fp4: - quant_init_context = get_fp4_context(self.config, i + pp_layer_offset, is_init=True) + quant_init_context = get_fp4_context( + self.config, physical_layer_offset, is_init=True + ) else: quant_init_context = nullcontext() with quant_init_context: - if layer_type == LayerSymbols.MAMBA: + if is_layer_group(layer_type): + layer = HybridStack( + config=self.config, + submodules=submodules, + pre_process=True, + layer_type_list=list(layer_type), + pp_layer_offset=physical_layer_offset, + logical_layer_offset=logical_layer_offset + len(self.layers), + is_layer_group_stack=True, + post_layer_norm=False, + post_process=False, + device=device, + dtype=dtype, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + ) + elif layer_type == LayerSymbols.MAMBA: layer = build_module( submodules.mamba_layer, config=self.config, @@ -188,6 +220,7 @@ def __init__( layer_number=layer_number, pg_collection=pg_collection, add_layer_offset=False, + is_mtp_layer=is_mtp_layer, name=(name + f".layers.{i}") if name is not None else None, ) elif layer_type == LayerSymbols.GDN: @@ -203,6 +236,7 @@ def __init__( else: raise ValueError("unexpected layer_type") self.layers.append(layer) + physical_layer_offset += get_layer_type_physical_count(layer_type) if self.config.cuda_graph_impl == "local": annotate_first_last_layer(self.layers) @@ -238,6 +272,17 @@ def _fuse_mla_down_proj(self, submodules: HybridStackSubmodules) -> HybridStackS } return submodules + @property + def final_layernorm(self): + """Alias for ``final_norm`` matching the attribute name on TransformerBlock. + + Lets generic decoder consumers (e.g. ``GPTModel.PostProcessNode``) discover the + final norm via the same attribute name they use for non-hybrid decoders, while + keeping ``final_norm`` as the registered submodule so existing hybrid checkpoint + keys are unchanged. + """ + return getattr(self, "final_norm", None) + def set_input_tensor(self, input_tensor: Tensor): """Set input tensor to be used instead of forward()'s input. @@ -254,7 +299,11 @@ def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int if this block contains Mamba layers (this may not be the case with PP > 1). """ for layer_type, layer in zip(self.layer_type_list, self.layers): - if layer_type == LayerSymbols.MAMBA: + if is_layer_group(layer_type): + shapes = layer.mamba_state_shapes_per_request() + if shapes is not None: + return shapes + elif layer_type == LayerSymbols.MAMBA: return layer.mamba_state_shapes_per_request() return None @@ -264,6 +313,10 @@ def forward( attention_mask: Tensor, inference_context: Optional[BaseInferenceContext] = None, rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + sequence_len_offset: Optional[Tensor] = None, *, inference_params: Optional[BaseInferenceContext] = None, packed_seq_params: Optional[PackedSeqParams] = None, @@ -304,7 +357,9 @@ def forward( inference_context.max_seqlen = inference_context.max_sequence_length inference_context.seqlen_offset = inference_context.sequence_len_offset - if ( + if sequence_len_offset is not None: + pass + elif ( (self.config.cuda_graph_impl == "local" or self.config.flash_decode) and inference_context and inference_context.is_static_batching() @@ -361,16 +416,35 @@ def get_inner_quant_context(config, layer_number): else: for layer in self.layers: # Layers have 1-indexed layer numbers attribute. - inner_quant_context = get_inner_quant_context( - self.config, layer.layer_number - 1 - ) + if isinstance(layer, HybridStack): + inner_quant_context = nullcontext() + else: + inner_quant_context = get_inner_quant_context( + self.config, layer.layer_number - 1 + ) with inner_quant_context: - if isinstance(layer, TransformerLayer): + if isinstance(layer, HybridStack): + hidden_states = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + ) + elif isinstance(layer, TransformerLayer): hidden_states, _ = layer( hidden_states=hidden_states, attention_mask=attention_mask, inference_context=inference_context, rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, sequence_len_offset=sequence_len_offset, packed_seq_params=packed_seq_params, padding_mask=padding_mask, @@ -423,17 +497,48 @@ def sharded_state_dict( dict: The sharded state dictionary for the current object. """ + return self._sharded_state_dict( + prefix=prefix, + sharded_offsets=sharded_offsets, + metadata=metadata, + sharded_layer_prefix=None, + ) + + def _sharded_state_dict( + self, + prefix: str = '', + sharded_offsets: Optional[tuple] = None, + metadata: Optional[dict] = None, + sharded_layer_prefix: Optional[str] = None, + ) -> ShardedStateDict: + sharded_offsets = sharded_offsets or () sharded_state_dict = {} layer_prefix = f'{prefix}layers.' + if sharded_layer_prefix is None: + sharded_layer_prefix = layer_prefix - for local_layer_idx, layer in enumerate(self.layers): - - global_layer_offset = layer.layer_number - 1 # self.layer_number starts at 1 - state_dict_prefix = ( - f'{layer_prefix}{local_layer_idx}.' # module list index in HybridStack + for local_layer_idx, (layer_type, layer) in enumerate( + zip(self.layer_type_list, self.layers) + ): + state_dict_prefix = f'{layer_prefix}{local_layer_idx}.' + logical_layer_idx = ( + self.logical_layer_offset + if self.is_layer_group_stack + else self.logical_layer_offset + local_layer_idx ) - sharded_prefix = f'{layer_prefix}{global_layer_offset}.' + if is_layer_group(layer_type): + sharded_state_dict.update( + layer._sharded_state_dict( + state_dict_prefix, + sharded_offsets, + metadata, + sharded_layer_prefix=sharded_layer_prefix, + ) + ) + continue + + sharded_prefix = f'{sharded_layer_prefix}{logical_layer_idx}.' sharded_pp_offset = [] layer_sharded_state_dict = layer.sharded_state_dict( @@ -447,15 +552,10 @@ def sharded_state_dict( # Add modules other than self.layers for name, module in self.named_children(): if not module is self.layers: - sharded_state_dict.update( - sharded_state_dict_default( - module, - f'{prefix}{name}.', - sharded_offsets, - metadata, - tp_group=self.tp_group, - ) + module_sharded_state_dict = sharded_state_dict_default( + module, f'{prefix}{name}.', sharded_offsets, metadata, tp_group=self.tp_group ) + sharded_state_dict.update(module_sharded_state_dict) return sharded_state_dict diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index 83a6163b88d..86ccd7dd1d6 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -2,7 +2,7 @@ import logging from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, Union import torch @@ -23,6 +23,8 @@ class Symbols: MOE = 'E' PIPE = '|' MTP_SEPARATOR = "/" + GROUP_START = "[" + GROUP_END = "]" VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLA, MLP, MOE} @classmethod @@ -38,6 +40,57 @@ def name_sorted_valid_layer_symbols(cls) -> list[str]: return [value for (_, value) in valid_layer_attrs] +LayerPatternItem = Union[str, Tuple[str, ...]] + + +def is_layer_group(layer_type: LayerPatternItem) -> bool: + """Return whether a parsed layer item is a bracketed group.""" + return isinstance(layer_type, tuple) + + +def flatten_layer_type_list(layer_type_list: List[LayerPatternItem]) -> List[str]: + """Flatten bracketed layer groups into their physical layer symbols.""" + flattened = [] + for layer_type in layer_type_list: + if is_layer_group(layer_type): + flattened.extend(layer_type) + else: + flattened.append(layer_type) + return flattened + + +def get_layer_type_physical_count(layer_type: LayerPatternItem) -> int: + """Return the number of physical layers represented by a parsed layer item.""" + return len(layer_type) if is_layer_group(layer_type) else 1 + + +def get_layer_type_logical_count(layer_type: LayerPatternItem) -> int: + """Return the number of logical layers represented by a parsed layer item.""" + return 1 + + +def get_layer_type_list_physical_count(layer_type_list: List[LayerPatternItem]) -> int: + """Return the number of physical layers represented by a parsed layer list.""" + return sum(get_layer_type_physical_count(layer_type) for layer_type in layer_type_list) + + +def get_layer_type_list_logical_count(layer_type_list: List[LayerPatternItem]) -> int: + """Return the number of logical layers represented by a parsed layer list.""" + return sum(get_layer_type_logical_count(layer_type) for layer_type in layer_type_list) + + +def layer_type_item_to_str(layer_type: LayerPatternItem) -> str: + """Render one parsed layer item back to pattern syntax.""" + if is_layer_group(layer_type): + return f"{Symbols.GROUP_START}{''.join(layer_type)}{Symbols.GROUP_END}" + return layer_type + + +def layer_type_list_to_str(layer_type_list: List[LayerPatternItem]) -> str: + """Render a parsed layer list back to pattern syntax.""" + return ''.join(layer_type_item_to_str(layer_type) for layer_type in layer_type_list) + + @dataclass class ParsedHybridPattern: """Result of parsing a unified hybrid pattern string. @@ -139,7 +192,10 @@ def get_hybrid_total_layer_count(pattern: str) -> int: """ main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] _validate_pattern(main_pattern, "main", allow_pipe=True) - return len(main_pattern.replace(Symbols.PIPE, '')) + return sum( + get_layer_type_list_physical_count(validate_segment_layers(segment)) + for segment in main_pattern.split(Symbols.PIPE) + ) def get_hybrid_total_pipeline_segment_count(pattern: str) -> int: @@ -184,15 +240,14 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: # Count main decoder layers (skip '|' pipe separators) if parsed.main_pattern: - for char in parsed.main_pattern: - if char in counts: + for segment in parsed.main_pattern.split(Symbols.PIPE): + for char in flatten_layer_type_list(validate_segment_layers(segment)): counts[char] += 1 # Count MTP layers (pattern repeated mtp_num_depths times) if parsed.mtp_pattern and parsed.mtp_num_depths > 0: - for char in parsed.mtp_pattern: - if char in counts: - counts[char] += parsed.mtp_num_depths + for char in flatten_layer_type_list(validate_segment_layers(parsed.mtp_pattern)): + counts[char] += parsed.mtp_num_depths return counts @@ -267,6 +322,16 @@ def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern: _validate_pattern(mtp_pattern, "MTP", allow_pipe=False) + # MTP layers are themselves a fused unit (each MTP depth contains its own attention + # + MLP), so it does not make sense to wrap them in a HybridStack group. Reject + # bracketed groups inside MTP patterns to keep downstream construction simple. + if Symbols.GROUP_START in mtp_pattern or Symbols.GROUP_END in mtp_pattern: + raise ValueError( + f"In MTP pattern, layer groups '{Symbols.GROUP_START}...{Symbols.GROUP_END}' " + f"are not supported because each MTP depth is already a fused unit. " + f"Got MTP pattern: '{mtp_pattern}'." + ) + return ParsedHybridPattern( main_pattern=main_pattern if main_pattern else None, mtp_pattern=mtp_pattern, @@ -285,20 +350,94 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) Raises: ValueError: If pattern contains invalid symbols """ - valid_chars = Symbols.VALID_LAYERS | {Symbols.PIPE} if allow_pipe else Symbols.VALID_LAYERS - for char in pattern: - if char not in valid_chars: + valid_chars = ( + Symbols.VALID_LAYERS + | {Symbols.GROUP_START, Symbols.GROUP_END} + | ({Symbols.PIPE} if allow_pipe else set()) + ) + if not allow_pipe and Symbols.PIPE in pattern: + raise ValueError( + f"In {pattern_name} pattern, '{Symbols.PIPE}' is not a valid layer symbol. " + f"Valid symbols are: {valid_chars}" + ) + flat_layers = [] + for segment in pattern.split(Symbols.PIPE): + flat_layers.extend( + flatten_layer_type_list( + _parse_segment_layers(segment, pattern_name, valid_chars=valid_chars) + ) + ) + + # Disallow Attention + MLA/DSA hybridity. + if Symbols.ATTENTION in flat_layers and ( + Symbols.DS_ATTENTION in flat_layers or Symbols.MLA in flat_layers + ): + raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + + +def _parse_segment_layers( + segment: str, pattern_name: str, valid_chars: Optional[set[str]] = None +) -> List[LayerPatternItem]: + """Parse a pipe-free pattern segment into symbols and bracketed groups.""" + if valid_chars is None: + valid_chars = Symbols.VALID_LAYERS | {Symbols.GROUP_START, Symbols.GROUP_END} + + layer_type_list: List[LayerPatternItem] = [] + flat_layers = [] + i = 0 + while i < len(segment): + layer_char = segment[i] + if layer_char == Symbols.GROUP_START: + group_end = segment.find(Symbols.GROUP_END, i + 1) + if group_end == -1: + raise ValueError( + f"In {pattern_name} pattern, '[' starts a layer group without a matching ']'." + ) + group = segment[i + 1 : group_end] + if group == "": + raise ValueError(f"In {pattern_name} pattern, layer groups cannot be empty.") + if Symbols.GROUP_START in group or Symbols.GROUP_END in group: + raise ValueError( + f"In {pattern_name} pattern, nested layer groups are not supported." + ) + for group_char in group: + if group_char not in Symbols.VALID_LAYERS: + raise ValueError( + f"In {pattern_name} pattern, '{group_char}' is not a valid layer symbol. " + f"Valid symbols are: {valid_chars}" + ) + if Symbols.MOE in group[:-1]: + raise ValueError( + f"In {pattern_name} pattern, MoE layer '{Symbols.MOE}' must be the last " + f"symbol inside a layer group." + ) + group_tuple = tuple(group) + layer_type_list.append(group_tuple) + flat_layers.extend(group_tuple) + i = group_end + 1 + continue + if layer_char == Symbols.GROUP_END: + raise ValueError( + f"In {pattern_name} pattern, ']' closes a layer group that was not opened." + ) + if layer_char not in Symbols.VALID_LAYERS: raise ValueError( - f"In {pattern_name} pattern, '{char}' is not a valid layer symbol. " + f"In {pattern_name} pattern, '{layer_char}' is not a valid layer symbol. " f"Valid symbols are: {valid_chars}" ) + layer_type_list.append(layer_char) + flat_layers.append(layer_char) + i += 1 - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in pattern and (Symbols.DS_ATTENTION in pattern or Symbols.MLA in pattern): + if Symbols.ATTENTION in flat_layers and ( + Symbols.DS_ATTENTION in flat_layers or Symbols.MLA in flat_layers + ): raise ValueError("Not supported to have both Attention and MLA/DSA in one model") + return layer_type_list + -def validate_segment_layers(segment: str) -> List[str]: +def validate_segment_layers(segment: str) -> List[LayerPatternItem]: """Validate and convert a single pipeline segment pattern to a layer type list. This is used after the main pattern has been split by '|' into segments. @@ -313,19 +452,95 @@ def validate_segment_layers(segment: str) -> List[str]: Raises: ValueError: If segment contains invalid layer symbols. """ - layer_type_list = list(segment) - for layer_char in layer_type_list: - if layer_char not in Symbols.VALID_LAYERS: + return _parse_segment_layers(segment, "hybrid layer pattern segment") + + +def _slice_layer_type_list_by_physical_range( + layer_type_list: List[LayerPatternItem], offset: int, count: int +) -> List[LayerPatternItem]: + """Slice parsed layer items by physical layer range without splitting groups.""" + selected = [] + cursor = 0 + end = offset + count + for layer_type in layer_type_list: + item_count = get_layer_type_physical_count(layer_type) + item_end = cursor + item_count + if item_end <= offset: + cursor = item_end + continue + if cursor >= end: + break + if cursor < offset or item_end > end: raise ValueError( - f"In hybrid layer pattern segment, '{layer_char}' is not " - f"one of {Symbols.VALID_LAYERS}" + "Pipeline splitting would split a bracketed hybrid layer group. " + "Add pipe ('|') separators around bracketed groups to define valid boundaries." ) + selected.append(layer_type) + cursor = item_end + return selected + + +def _get_logical_offset_from_physical_offset( + layer_type_list: List[LayerPatternItem], offset: int +) -> int: + """Return the logical item count before a physical-layer offset.""" + logical_offset = 0 + cursor = 0 + for layer_type in layer_type_list: + item_count = get_layer_type_physical_count(layer_type) + item_end = cursor + item_count + if item_end <= offset: + logical_offset += get_layer_type_logical_count(layer_type) + cursor = item_end + continue + if cursor == offset: + return logical_offset + raise ValueError( + "Pipeline splitting would split a bracketed hybrid layer group. " + "Add pipe ('|') separators around bracketed groups to define valid boundaries." + ) + if cursor == offset: + return logical_offset + raise ValueError(f"Physical layer offset {offset} is out of range for hybrid layer pattern.") - # Disallow Attention + MLA/DSA hybridity. - if Symbols.ATTENTION in segment and (Symbols.DS_ATTENTION in segment or Symbols.MLA in segment): - raise ValueError("Not supported to have both Attention and MLA/DSA in one model") - return layer_type_list +def select_pipeline_segment_with_logical_offset( + main_pattern: str, + pp_group: Optional[torch.distributed.ProcessGroup], + vp_stage: Optional[int], + first_stage_layers: Optional[int] = None, + last_stage_layers: Optional[int] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, +) -> Tuple[List[LayerPatternItem], int, int]: + """Select a pipeline segment and return physical and logical offsets.""" + layer_type_list, layer_offset = select_pipeline_segment( + main_pattern, + pp_group, + vp_stage, + first_stage_layers=first_stage_layers, + last_stage_layers=last_stage_layers, + tp_group=tp_group, + dp_cp_group=dp_cp_group, + ) + + segments = main_pattern.split(Symbols.PIPE) if main_pattern else [''] + if len(segments) == 1: + full_layer_type_list = validate_segment_layers(segments[0]) + logical_layer_offset = _get_logical_offset_from_physical_offset( + full_layer_type_list, layer_offset + ) + else: + pp_rank = torch.distributed.get_rank(pp_group) if pp_group is not None else 0 + pp_size = torch.distributed.get_world_size(pp_group) if pp_group is not None else 1 + vp_rel = vp_stage if vp_stage is not None else 0 + segment_index = vp_rel * pp_size + pp_rank + logical_layer_offset = sum( + get_layer_type_list_logical_count(validate_segment_layers(segments[i])) + for i in range(segment_index) + ) + + return layer_type_list, layer_offset, logical_layer_offset def select_pipeline_segment( @@ -400,7 +615,7 @@ def select_pipeline_segment( ) full_pattern = segments[0] layer_type_list = validate_segment_layers(full_pattern) - num_layers = len(layer_type_list) + num_layers = get_layer_type_list_physical_count(layer_type_list) if first_stage_layers is not None or last_stage_layers is not None: first = first_stage_layers or 0 @@ -443,12 +658,13 @@ def select_pipeline_segment( offset = pp_rank * layers_per_rank count = layers_per_rank - selected = layer_type_list[offset : offset + count] + selected = _slice_layer_type_list_by_physical_range(layer_type_list, offset, count) log_on_each_pipeline_stage( logger, logging.INFO, f"HybridModel: pp_rank={pp_rank}/{pp_size}, vp_stage={vp_stage}, " - f"layers='{''.join(selected)}' ({len(selected)} layers), " + f"layers='{layer_type_list_to_str(selected)}' " + f"({get_layer_type_list_physical_count(selected)} layers), " f"layer_offset={offset} (auto-split)", tp_group=tp_group, dp_cp_group=dp_cp_group, @@ -474,7 +690,10 @@ def select_pipeline_segment( f"the current PP/VPP configuration." ) - layer_offset = sum(len(segments[i]) for i in range(segment_index)) + layer_offset = sum( + get_layer_type_list_physical_count(validate_segment_layers(segments[i])) + for i in range(segment_index) + ) my_segment = segments[segment_index] layer_type_list = validate_segment_layers(my_segment) @@ -484,7 +703,7 @@ def select_pipeline_segment( logging.INFO, f"HybridModel: pp_rank={pp_rank}/{pp_size}, vp_stage={vp_rel}, " f"segment_index={segment_index}/{len(segments)}, " - f"layers='{my_segment}' ({len(layer_type_list)} layers), " + f"layers='{my_segment}' ({get_layer_type_list_physical_count(layer_type_list)} layers), " f"layer_offset={layer_offset}", tp_group=tp_group, dp_cp_group=dp_cp_group, @@ -493,14 +712,16 @@ def select_pipeline_segment( return layer_type_list, layer_offset -def get_layer_maps_from_layer_type_list(layer_type_list: list[str]) -> dict[str, dict[int, int]]: +def get_layer_maps_from_layer_type_list( + layer_type_list: list[LayerPatternItem], +) -> dict[str, dict[int, int]]: """ Returns maps from global layer index to the corresponding layer index for each valid layer type (those in Symbols.VALID_LAYERS) given a layer type list. """ layer_types = [symbol for symbol in Symbols.name_sorted_valid_layer_symbols()] layer_maps = {layer_type: {} for layer_type in layer_types} - for global_layer_idx, layer_type in enumerate(layer_type_list): + for global_layer_idx, layer_type in enumerate(flatten_layer_type_list(layer_type_list)): layer_map = layer_maps[layer_type] local_layer_idx = len(layer_map) layer_map[global_layer_idx] = local_layer_idx diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f750c77e05b..7d0feb44ec8 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -1,12 +1,13 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging -from typing import Literal, Optional +from typing import Dict, Literal, Optional from torch import Tensor from megatron.core import tensor_parallel from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk +from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.inference.utils import InferenceMode from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding @@ -193,7 +194,7 @@ def __init__( # determine the pipeline segment for this model instance. from megatron.core.models.hybrid.hybrid_layer_allocation import ( parse_hybrid_pattern, - select_pipeline_segment, + select_pipeline_segment_with_logical_offset, ) parsed = parse_hybrid_pattern(self.hybrid_layer_pattern) @@ -202,13 +203,15 @@ def __init__( logging_pg_kwargs = _hybrid_logging_pg_kwargs(self.pg_collection) - layer_type_list, layer_offset = select_pipeline_segment( - parsed.main_pattern or '', - self.pg_collection.pp, - vp_stage, - first_stage_layers=self.config.num_layers_in_first_pipeline_stage, - last_stage_layers=self.config.num_layers_in_last_pipeline_stage, - **logging_pg_kwargs, + layer_type_list, layer_offset, logical_layer_offset = ( + select_pipeline_segment_with_logical_offset( + parsed.main_pattern or '', + self.pg_collection.pp, + vp_stage, + first_stage_layers=self.config.num_layers_in_first_pipeline_stage, + last_stage_layers=self.config.num_layers_in_last_pipeline_stage, + **logging_pg_kwargs, + ) ) # Determine if MTP is needed (based on pattern parsing) @@ -278,6 +281,7 @@ def __init__( pre_process=self.pre_process, layer_type_list=layer_type_list, pp_layer_offset=layer_offset, + logical_layer_offset=logical_layer_offset, post_process=self.post_process, dtype=config.params_dtype, pg_collection=self.pg_collection, @@ -348,115 +352,36 @@ def set_input_tensor(self, input_tensor: Tensor) -> None: assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert' self.decoder.set_input_tensor(input_tensor[0]) - def preprocess_for_fine_grained_offloading(self): - """Preprocess for fine-grained activation offloading.""" - off_interface.init_chunk_handler( - pp_rank=self.pg_collection.pp.rank(), - vp_size=self.config.virtual_pipeline_model_parallel_size, - vp_stage=self.vp_stage, - min_offloaded_tensor_size=self.config.min_offloaded_tensor_size, - delta_offload_bytes_across_pp_ranks=self.config.delta_offload_bytes_across_pp_ranks, - activation_offload_fraction=self.config.activation_offload_fraction, - max_inflight_offloads=self.config.fine_grained_offloading_max_inflight_offloads, - ) - if self.disable_param_offloading: - for param in self.decoder.parameters(): - off_interface.mark_not_offload(param) - if self.mtp_process: - for param in self.mtp.parameters(): - off_interface.mark_not_offload(param) - if self.post_process: - for param in self.output_layer.parameters(): - off_interface.mark_not_offload(param) - self.disable_param_offloading = False - - def preprocess_for_paged_stash(self): - """Preprocess for paged stash.""" - return paged_stash_init_chunk_handler( - vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage - ) - - def _should_call_local_cudagraph(self, *args, **kwargs): - """ - Check if we should call the local cudagraph path. - """ - if ( - InferenceMode.is_active() - and hasattr(self, 'cudagraph_manager') - and ( - kwargs.get('inference_context') is not None - or kwargs.get('inference_params') is not None - ) - and self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block - ): - if kwargs['inference_context'].is_static_batching(): - using_cuda_graph = kwargs['inference_context'].is_decode_only() - else: - using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step() - - if using_cuda_graph: - return True - return False - - def __call__(self, *args, **kwargs): - if self._should_call_local_cudagraph(*args, **kwargs): - return super().__call__(*args, **kwargs)[0] - return super().__call__(*args, **kwargs) - - def create_mcore_cudagraph_manager(self, config): - """ - Create the cudagraph manager for the full iteration inference scope - """ - if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: - from megatron.core.transformer.cuda_graphs import CudaGraphManager - - self.cudagraph_manager = CudaGraphManager(config) - - def forward( + def _preprocess( self, input_ids: Tensor, position_ids: Tensor, - attention_mask: Tensor, decoder_input: Tensor = None, - labels: Tensor = None, inference_context: BaseInferenceContext = None, - runtime_gather_output: Optional[bool] = None, - *, - inference_params: Optional[BaseInferenceContext] = None, - loss_mask: Optional[Tensor] = None, - packed_seq_params: Optional[PackedSeqParams] = None, + packed_seq_params: PackedSeqParams = None, padding_mask: Optional[Tensor] = None, - ) -> Tensor: - """Forward function of the Hybrid model. This function passes the input tensors - through the embedding layer, and then the decoder and finally into the post - processing layer (optional). - - It either returns the Loss values if labels are given or the final hidden units + ): + """Preprocess inputs for HybridStack or combined-1F1B scheduling. + + Mirrors ``GPTModel._preprocess`` so the eager forward and the + EP-overlap ``PreProcessNode`` see the same embedding / rotary / + padding-mask code. Returns the canonical 6-tuple ``(decoder_input, + rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset, + padding_mask)`` — slots HybridModel does not compute (rotary cos/sin) + come back as ``None``. """ - # If decoder_input is provided (not None), then input_ids and position_ids are ignored. - # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input. - - if self.config.fine_grained_activation_offloading: - self.preprocess_for_fine_grained_offloading() - - if self.config.moe_paged_stash: - self.preprocess_for_paged_stash() - - inference_context = deprecate_inference_params(inference_context, inference_params) - in_inference_mode = InferenceMode.is_active() - if in_inference_mode: - assert runtime_gather_output, "Inference must always gather TP logits" - - # Decoder embedding. + # If decoder_input is provided, input_ids and position_ids are ignored; + # otherwise apply the embedding layer to get decoder_input. if decoder_input is not None: pass elif self.pre_process: + # Decoder embedding. decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) - # Clear the outputs for padding tokens when using dynamic batching with - # quantization scales to avoid corrupting amax calculations + # Clear the outputs for padding tokens when using dynamic batching + # with quantization scales to avoid corrupting amax calculations. if ( in_inference_mode and inference_context is not None @@ -472,8 +397,8 @@ def forward( decoder_input, group=self.pg_collection.tp ) else: - # intermediate stage of pipeline - # decoder will get hidden_states from encoder.input_tensor + # Intermediate stage of pipeline parallelism — the decoder will get + # hidden_states from encoder.input_tensor. decoder_input = None rotary_pos_emb = None @@ -489,54 +414,113 @@ def forward( rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( inference_context, self.decoder, decoder_input, self.config, packed_seq_params ) - # YarnRotaryEmbedding.forward returns (emb, mscale); discard mscale here + # YarnRotaryEmbedding.forward returns (emb, mscale); discard mscale here. rotary_pos_emb, _ = self.rotary_pos_emb( rotary_seq_len, packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd', ) - # Wrap decoder_input to allow the decoder (HybridStack) to delete the - # reference held by this caller function, enabling early garbage collection - # for inference. + # ``sequence_len_offset`` is only needed for flash-decode / local-cudagraph + # static-batching inference; otherwise leave it as ``None``. + if ( + in_inference_mode + and inference_context is not None + and (self.config.cuda_graph_impl == "local" or self.config.flash_decode) + and inference_context.is_static_batching() + ): + current_batch_size = input_ids.shape[0] + import torch + + sequence_len_offset = torch.tensor( + [inference_context.sequence_len_offset] * current_batch_size, + dtype=torch.int32, + device='cuda', + ) + else: + sequence_len_offset = None + + # Wrap decoder_input so the decoder (HybridStack) can drop its caller's + # reference for early garbage collection during inference. if in_inference_mode: decoder_input = WrappedTensor(decoder_input) - # The following assert will currently fail when running inference. - # Commented out for now. - # TODO (duncan/rwaleffe): (1) confirm that the externally-generated - # attention mask is not needed and is ignored by the model in - # inference mode, (2) reduce the size of the externally-generated - # attention mask to prevent CPU OOM (as we did for training), (3) - # force the attention mask passed to the model in inference mode to - # be None, so this assert will succeed. - # assert attention_mask is None, "The attention mask is ignored and should be set to None" + return decoder_input, rotary_pos_emb, None, None, sequence_len_offset, padding_mask - # Run decoder. - hidden_states = self.decoder( - hidden_states=decoder_input, - attention_mask=attention_mask, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb, - packed_seq_params=packed_seq_params, - padding_mask=padding_mask, + def preprocess_for_fine_grained_offloading(self): + """Preprocess for fine-grained activation offloading.""" + off_interface.init_chunk_handler( + pp_rank=self.pg_collection.pp.rank(), + vp_size=self.config.virtual_pipeline_model_parallel_size, + vp_stage=self.vp_stage, + min_offloaded_tensor_size=self.config.min_offloaded_tensor_size, + delta_offload_bytes_across_pp_ranks=self.config.delta_offload_bytes_across_pp_ranks, + activation_offload_fraction=self.config.activation_offload_fraction, + max_inflight_offloads=self.config.fine_grained_offloading_max_inflight_offloads, ) + if self.disable_param_offloading: + for param in self.decoder.parameters(): + off_interface.mark_not_offload(param) + if self.mtp_process: + for param in self.mtp.parameters(): + off_interface.mark_not_offload(param) + if self.post_process: + for param in self.output_layer.parameters(): + off_interface.mark_not_offload(param) + self.disable_param_offloading = False + + def _postprocess( + self, + hidden_states, + input_ids, + position_ids, + labels, + rotary_pos_emb, + rotary_pos_cos=None, + rotary_pos_sin=None, + mtp_in_postprocess=None, + loss_mask=None, + decoder_input=None, + attention_mask=None, + inference_params=None, + packed_seq_params=None, + sequence_len_offset=None, + runtime_gather_output=None, + extra_block_kwargs=None, + inference_context=None, + is_spec_decode=None, + output_processor=None, + output_processor_context=None, + ): + """Postprocess HybridStack hidden states into logits or language-model loss. + + Mirrors ``GPTModel._postprocess`` so the eager forward and the EP-overlap + ``PostProcessNode`` produce the same logits / loss / MTP outputs. + ``mtp_in_postprocess`` lets the EP-overlap path skip the inline MTP block + (it schedules MTP as separate layer nodes); the eager forward leaves it + ``True`` so the regular MTP forward runs here. + """ + in_inference_mode = inference_context is not None and not self.training + if in_inference_mode: + assert runtime_gather_output, "Inference must always gather TP logits" output_weight = None if self.share_embeddings_and_output_weights: output_weight = self.shared_embedding_or_output_weight() - # Check if speculative decoding is active. When it is, MTP must be - # computed *after* verification so that it is conditioned on verified - # tokens rather than stale speculative tokens from the previous step. - is_spec_decode = ( - in_inference_mode - and inference_context is not None - and inference_context.is_dynamic_batching() - and inference_context.num_speculative_tokens > 0 - ) + # Speculative decoding: when active, MTP must run *after* verification so + # it conditions on verified tokens rather than stale speculative ones. + if is_spec_decode is None: + is_spec_decode = ( + in_inference_mode + and inference_context.is_dynamic_batching() + and inference_context.num_speculative_tokens > 0 + ) - mtp_forward_ran = self.mtp_process and not (in_inference_mode or is_spec_decode) - if mtp_forward_ran: + # MTP forward inline (skipped when the EP-overlap plan schedules MTP + # separately, when running inference, or when speculative decoding is + # active). ``self.mtp_process`` guards against models built without an + # MTP block. + if mtp_in_postprocess and self.mtp_process and not (in_inference_mode or is_spec_decode): hidden_states = self.mtp( input_ids=input_ids, position_ids=position_ids, @@ -553,23 +537,12 @@ def forward( if self.config.mtp_num_layers is not None and self.mtp_process: assert self.config.mtp_num_layers > 0 - if is_spec_decode: - assert inference_context is not None - if self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: - # Block-scope CUDA graph mode: copy_() into the - # pre-allocated buffer so every graph replay writes to - # the same fixed GPU address regardless of batch size. - assert inference_context.mtp_decoder_hidden_states is not None - inference_context.mtp_decoder_hidden_states[: hidden_states.shape[0]].copy_( - hidden_states - ) - else: - # Non-block scope: direct assignment; the controller will set - # this back to None after reading to allow GC. - inference_context.mtp_decoder_hidden_states = hidden_states - elif not in_inference_mode: - # For RL (labels is None), process_mtp_loss derives labels from - # input_ids to match the SFT label format. + 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 + else: + # In training/eval, fold MTP loss into hidden_states. hidden_states = process_mtp_loss( hidden_states=hidden_states, labels=labels, @@ -581,24 +554,19 @@ def forward( compute_language_model_loss=self.compute_language_model_loss, config=self.config, cp_group=self.pg_collection.cp, - tp_group=self.tp_group, packed_seq_params=packed_seq_params, scale_logits_fn=self._scale_logits if self.config.use_mup else None, - input_ids=input_ids, ) + sequence_parallel_override = False - if ( - in_inference_mode - and inference_context is not None - and inference_context.config.materialize_only_last_token_logits - ): + if in_inference_mode and inference_context.config.materialize_only_last_token_logits: if inference_context.is_static_batching(): hidden_states = hidden_states[-1:, :, :] else: if self.output_layer.sequence_parallel: - # Perform the sequence parallel gather here instead of after the output layer - # because we need to slice the last token logits from the full view of the - # packed logits across all requests. + # Perform the sequence-parallel gather here instead of after + # the output layer so we can slice the last-token logits from + # the full view of the packed logits across all requests. hidden_states = gather_from_sequence_parallel_region( hidden_states, group=self.pg_collection.tp ) @@ -615,7 +583,6 @@ def forward( ) logits = self._scale_logits(logits) - # Restore sequence parallel execution to the output layer if necessary. if sequence_parallel_override: assert ( in_inference_mode @@ -625,9 +592,191 @@ def forward( self.output_layer.sequence_parallel = True if labels is None: - # [s b h] => [b s h] return logits.transpose(0, 1).contiguous() loss = self.compute_language_model_loss(labels, logits) - return loss + + def build_schedule_plan( + self, + input_ids: Tensor, + position_ids: Tensor, + attention_mask: Tensor, + decoder_input: Tensor = None, + labels: Tensor = None, + inference_context: BaseInferenceContext = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + runtime_gather_output: Optional[bool] = None, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, + ): + """Build the HybridModel combined-1F1B schedule plan.""" + if self.config.fine_grained_activation_offloading: + self.preprocess_for_fine_grained_offloading() + + from .model_chunk_schedule_plan import HybridStackModelChunkSchedulePlan + + return HybridStackModelChunkSchedulePlan( + self, + input_ids, + position_ids, + attention_mask, + decoder_input, + labels, + packed_seq_params, + extra_block_kwargs, + runtime_gather_output, + loss_mask, + padding_mask, + ) + + def sharded_state_dict( + self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None + ) -> ShardedStateDict: + """Return a Transformer-compatible sharded state dict for HybridModel.""" + sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata) + output_layer_extra_state_key = f'{prefix}output_layer._extra_state' + + # Match GPTModel checkpoint compatibility: old GPT checkpoints do not include + # output layer extra state, and the TE extra state should be empty. + output_extra_state = sharded_state_dict.pop(output_layer_extra_state_key, None) + assert not ( + output_extra_state and output_extra_state.data + ), f'Expected output layer extra state to be empty, got: {output_extra_state}' + + return sharded_state_dict + + def preprocess_for_paged_stash(self): + """Preprocess for paged stash.""" + return paged_stash_init_chunk_handler( + vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage + ) + + def _should_call_local_cudagraph(self, *args, **kwargs): + """ + Check if we should call the local cudagraph path. + """ + if ( + InferenceMode.is_active() + and hasattr(self, 'cudagraph_manager') + and ( + kwargs.get('inference_context') is not None + or kwargs.get('inference_params') is not None + ) + and self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + ): + if kwargs['inference_context'].is_static_batching(): + using_cuda_graph = kwargs['inference_context'].is_decode_only() + else: + using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step() + + if using_cuda_graph: + return True + return False + + def __call__(self, *args, **kwargs): + if self._should_call_local_cudagraph(*args, **kwargs): + return super().__call__(*args, **kwargs)[0] + return super().__call__(*args, **kwargs) + + def create_mcore_cudagraph_manager(self, config): + """ + Create the cudagraph manager for the full iteration inference scope + """ + if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: + from megatron.core.transformer.cuda_graphs import CudaGraphManager + + self.cudagraph_manager = CudaGraphManager(config) + + def forward( + self, + input_ids: Tensor, + position_ids: Tensor, + attention_mask: Tensor, + decoder_input: Tensor = None, + labels: Tensor = None, + inference_context: BaseInferenceContext = None, + runtime_gather_output: Optional[bool] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """Forward function of the Hybrid model. This function passes the input tensors + through the embedding layer, and then the decoder and finally into the post + processing layer (optional). + + It either returns the Loss values if labels are given or the final hidden units + """ + if self.config.fine_grained_activation_offloading: + self.preprocess_for_fine_grained_offloading() + + if self.config.moe_paged_stash: + self.preprocess_for_paged_stash() + + inference_context = deprecate_inference_params(inference_context, inference_params) + + in_inference_mode = InferenceMode.is_active() + if in_inference_mode: + assert runtime_gather_output, "Inference must always gather TP logits" + + # Mirror GPTModel.forward: delegate the embedding / rotary computation and + # the output-layer / MTP / loss computation to the same hooks the + # EP-overlap PreProcessNode / PostProcessNode call. Keeps the eager and + # combined-1F1B paths on the same code. + ( + decoder_input, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + sequence_len_offset, + padding_mask, + ) = self._preprocess( + input_ids=input_ids, + position_ids=position_ids, + decoder_input=decoder_input, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + ) + + # The following assert will currently fail when running inference. + # Commented out for now. + # TODO (duncan/rwaleffe): (1) confirm that the externally-generated + # attention mask is not needed and is ignored by the model in + # inference mode, (2) reduce the size of the externally-generated + # attention mask to prevent CPU OOM (as we did for training), (3) + # force the attention mask passed to the model in inference mode to + # be None, so this assert will succeed. + # assert attention_mask is None, "The attention mask is ignored and should be set to None" + + # Run decoder. + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + ) + + return self._postprocess( + hidden_states=hidden_states, + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + mtp_in_postprocess=True, + loss_mask=loss_mask, + attention_mask=attention_mask, + inference_params=inference_params, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + runtime_gather_output=runtime_gather_output, + inference_context=inference_context, + ) diff --git a/megatron/core/models/hybrid/model_chunk_schedule_plan.py b/megatron/core/models/hybrid/model_chunk_schedule_plan.py new file mode 100644 index 00000000000..20a5c027ea2 --- /dev/null +++ b/megatron/core/models/hybrid/model_chunk_schedule_plan.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Schedule-plan classes for HybridStack-based decoders. + +These extend the GPT-side ``TransformerLayerSchedulePlan`` / +``TransformerModelChunkSchedulePlan`` with the per-layer ``layer_type`` symbol +that HybridStack assigns to each entry of its ``layer_type_list`` (including +bracketed groups like ``[*-]``). The base classes remain GPT-only; this module +adds the hybrid-specific dispatch into ``build_hybrid_stack_callables`` and +uses ``HybridStackNode`` so the schedule node's free-input policy can diverge +from the GPT default. The pre/post-process nodes from +``core.models.common.utils`` are reused as-is — they already call +``model._preprocess`` / ``model._postprocess`` which work on a HybridModel. +""" + +from contextlib import nullcontext + +from megatron.core.models.common.model_chunk_schedule_plan import ( + TransformerLayerSchedulePlan, + TransformerModelChunkSchedulePlan, +) + + +class HybridStackSchedulePlan(TransformerLayerSchedulePlan): + """Per-layer schedule plan for HybridStack decoders. + + Adds the ``layer_type`` extra-arg propagation; routes through + ``build_hybrid_stack_callables`` when ``layer_type`` is set (i.e. the layer + is a HybridStack entry, possibly a bracketed group); falls back to the GPT + path for plain TransformerLayer / MTP layers when ``layer_type`` is None. + """ + + def __init__(self, layer, event, chunk_state, comp_stream, comm_stream, extra_args=None): + if extra_args is None: + extra_args = {} + self.layer_type = extra_args.get("layer_type", None) + super().__init__(layer, event, chunk_state, comp_stream, comm_stream, extra_args) + + def _build_callable_nodes(self, event, comp_stream, comm_stream, extra_args): + if self.layer_type is None: + return super()._build_callable_nodes(event, comp_stream, comm_stream, extra_args) + + # Hybrid grouped path. Imports are local because hybrid pulls in TE / SSM + # extensions that we don't want to load when only the GPT path is used. + from megatron.core.models.hybrid.fine_grained_callables import ( + HybridStackNode, + build_hybrid_stack_callables, + ) + from megatron.core.pipeline_parallel.utils import NoopScheduleNode + + fwd_callables, bwd_dw_callable_map, is_moe, num_local_experts = ( + build_hybrid_stack_callables(self.layer, layer_type=self.layer_type) + ) + + extra_args["config"] = self.layer.config + extra_args["is_moe"] = is_moe + extra_args["num_local_experts"] = num_local_experts + extra_args["delay_wgrad_compute"] = self.layer.config.delay_wgrad_compute + extra_args["is_mtp"] = False + + def create_node(stream, module, name): + bwd_dw_callables = bwd_dw_callable_map.get(name, None) + node_extra_args = dict(extra_args) + if bwd_dw_callables is None: + node_extra_args["delay_wgrad_compute"] = False + return HybridStackNode( + stream, + event, + self.layer_state, + self.chunk_state, + module, + name=name, + bwd_dw_callables=bwd_dw_callables, + extra_args=node_extra_args, + ) + + ( + pre_dispatch_module, + moe_dispatch_module, + mlp_module, + moe_combine_module, + mtp_post_process_module, + ) = fwd_callables + + self.pre_dispatch_computation = create_node( + comp_stream, pre_dispatch_module, "pre_dispatch_computation" + ) + self.mlp = create_node(comp_stream, mlp_module, "mlp") + if is_moe: + self.moe_dispatch = create_node(comm_stream, moe_dispatch_module, "moe_dispatch") + self.moe_combine = create_node(comm_stream, moe_combine_module, "moe_combine") + else: + self.moe_dispatch = NoopScheduleNode() + self.moe_combine = NoopScheduleNode() + + # HybridStack groups never carry an MTP terminal, so mtp_post_process is + # always a no-op here. + self.mtp_post_process = NoopScheduleNode() + + def get_fp8_context(self): + """Return an FP8 context only for plain transformer layers.""" + # Grouped hybrid layers (and inferred-layer-type entries that point at + # a HybridStack rather than a plain TransformerLayer) don't have a + # ``layer_number`` we can hand to ``get_fp8_context``; the inner layers + # manage their own per-layer fp8 context inside the hybrid callables. + if self.layer_type is not None or not hasattr(self.layer, "layer_number"): + return nullcontext() + return super().get_fp8_context() + + +class HybridStackModelChunkSchedulePlan(TransformerModelChunkSchedulePlan): + """Model-chunk schedule plan that builds ``HybridStackSchedulePlan`` layer plans. + + Threads HybridStack's ``layer_type_list[layer_idx]`` symbol into each + layer plan's ``extra_args`` so the per-layer plan can dispatch grouped + layers correctly. Ordinary GPT/MTP layers (no ``layer_type_list``) + default to ``layer_type=None`` and follow the GPT path. The pre/post + process nodes inherit from the GPT base class — they already dispatch + on ``model._preprocess`` / ``model._postprocess`` which a HybridModel + implements. + """ + + LAYER_SCHEDULE_PLAN_CLASS = HybridStackSchedulePlan + + def __init__(self, model, *args, **kwargs): + """Initialize the hybrid chunk plan after validating cuda graph support.""" + assert model.config.cuda_graph_impl == "none", ( + "EP A2A overlap with grouped HybridStack patterns (e.g. '[*E]') does not " + "support cuda graphs yet. Set cuda_graph_impl='none' or use an ungrouped pattern." + ) + super().__init__(model, *args, **kwargs) + + def _extra_args_for_layer(self, module, layer_idx, num_layers): + extra_args = super()._extra_args_for_layer(module, layer_idx, num_layers) + extra_args["layer_type"] = ( + module.layer_type_list[layer_idx] if hasattr(module, "layer_type_list") else None + ) + return extra_args diff --git a/megatron/core/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py index 68d41c56a31..53f452ec747 100644 --- a/megatron/core/ssm/mamba_layer.py +++ b/megatron/core/ssm/mamba_layer.py @@ -162,6 +162,17 @@ def forward( return hidden_states + def backward_dw(self): + """Compute weight gradients for the layer's linear projections. + + Delegates to the mixer; lets the hybrid EP-overlap schedule plan + register a Mamba pre-layer's wgrad alongside attention/GDN pre-layers + so the schedule node iterates a uniform set of callables. No-op when + the linears in the spec do not support delayed wgrad. + """ + if hasattr(self.mixer, "backward_dw"): + self.mixer.backward_dw() + def sharded_state_dict( self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None ) -> ShardedStateDict: diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index 7a44c8a493a..6631eaa9a81 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -1296,6 +1296,21 @@ def mamba_state_shapes_per_request(self) -> Tuple[Tuple[int], Tuple[int]]: ssm_states_shape = (self.nheads_local_tp, self.headdim, self.d_state) return (conv_states_shape, ssm_states_shape) + def backward_dw(self): + """Compute weight gradients for the linear layers wrapped by this mixer. + + Mirrors ``GatedDeltaNet.backward_dw``. The selective-scan kernel is a + single autograd function whose wgrad runs in the regular backward pass, + so only the input/output projections need delayed wgrad here. Each + ``backward_dw`` call is a no-op unless the underlying linear is built + from a TE primitive that supports delayed wgrad; if the spec uses + non-TE linears, ``backward_dw`` simply does nothing. + """ + if hasattr(self.in_proj, "backward_dw"): + self.in_proj.backward_dw() + if hasattr(self.out_proj, "backward_dw"): + self.out_proj.backward_dw() + def _get_states_from_cache(self, inference_context, batch_size, *, inference_params=None): """Initializes or retrieves the SSM state tensors from the cache. diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index e4591ce3acf..9b15446c771 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -580,16 +580,21 @@ def attach_and_log_load_balancing_loss( ): aux_loss = aux_loss / self.config.mtp_num_layers - # TODO (zijiey): fix the per_layer_logging for MTP, currently it will incorrectly - # add the aux loss logging value to other layer's since it is difficult to get the - # correct layer_number for MTP. It does not affect the correctness of the calculation - # results and the reduced load_balancing_loss logging value. + # The tracker has one slot per main decoder layer and one per MTP depth, so its + # size is (num_layers + mtp_num_layers). For MTP routers, the slot index must be + # in [num_layers + 1, num_layers + mtp_num_layers]. When the MTP block wraps a + # plain TransformerLayer, ``self.layer_number`` already equals the MTP depth + # (1..mtp_num_layers). When it wraps a HybridStack (e.g. ``*E`` for one depth), + # ``self.layer_number`` is the position within the inner HybridStack and can + # exceed mtp_num_layers, which would index past the tracker; collapse it to a + # valid MTP-depth slot via modulo so the aux loss lands at the right depth. num_layers = self.config.num_layers if self.config.mtp_num_layers is not None: num_layers += self.config.mtp_num_layers if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers + mtp_depth = ((self.layer_number - 1) % self.config.mtp_num_layers) + 1 + layer_number = mtp_depth + self.config.num_layers else: layer_number = self.layer_number diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 39bc7f30b57..890f962d79b 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -277,12 +277,13 @@ def loss_func( return loss, num_tokens, report -def forward_step(data_iterator, model: HybridModel): +def forward_step(data_iterator, model: HybridModel, return_schedule_plan: bool = False): """Forward training step. Args: data_iterator : Input data iterator model (HybridModel): The Hybrid Model + return_schedule_plan (bool): Whether to return the schedule plan instead of output tensor. """ timers = get_timers() @@ -332,14 +333,28 @@ def forward_step(data_iterator, model: HybridModel): timers('batch-generator').stop() with stimer: - output_tensor = model( - tokens, - position_ids, - attention_mask, - labels=labels, - packed_seq_params=packed_seq_params, - loss_mask=loss_mask, - ) + if return_schedule_plan: + args = get_args() + assert args.overlap_moe_expert_parallel_comm, ( + "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + ) + output_tensor = model.build_schedule_plan( + tokens, + position_ids, + attention_mask, + labels=labels, + packed_seq_params=packed_seq_params, + loss_mask=loss_mask, + ) + else: + output_tensor = model( + tokens, + position_ids, + attention_mask, + labels=labels, + packed_seq_params=packed_seq_params, + loss_mask=loss_mask, + ) # [ModelOpt]: model is needed to access ModelOpt distillation losses return output_tensor, partial(loss_func, loss_mask, model=model) diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index 95bcaa2d7d0..fc288d6293e 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -383,6 +383,29 @@ def test_save_load(self, tmp_path): self.model.load_state_dict(torch.load(path)) + def test_grouped_sharded_state_dict_uses_transformer_checkpoint_keys(self): + """Grouped HybridModel checkpoints should be load-compatible with GPTModel keys.""" + model_config = TransformerConfig( + num_layers=2, hidden_size=256, num_attention_heads=4, use_cpu_initialization=True + ) + model = HybridModel( + config=model_config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=100, + max_sequence_length=4, + hybrid_layer_pattern="[*-]", + ) + + sharded_state_dict = model.sharded_state_dict() + sharded_keys = {value.key for value in sharded_state_dict.values() if hasattr(value, "key")} + + assert "decoder.layers.0.self_attention.linear_qkv.weight" in sharded_keys + assert "decoder.layers.0.mlp.linear_fc1.weight" in sharded_keys + assert "decoder.layers.1.mlp.linear_fc1.weight" not in sharded_keys + assert "decoder.final_layernorm.weight" in sharded_keys + assert "decoder.final_norm.weight" not in sharded_keys + assert "output_layer._extra_state" not in sharded_state_dict + def test_layer_numbers(self): """ The layer numbers should start at one (for the embedding # layer) and go up diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 5d3c33264f4..793fded6c78 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -4,8 +4,13 @@ import torch from megatron.core.extensions.transformer_engine import TEDotProductAttention +from megatron.core.models.hybrid.fine_grained_callables import build_hybrid_stack_callables from megatron.core.models.hybrid.hybrid_block import HybridStack -from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols, validate_segment_layers +from megatron.core.models.hybrid.hybrid_layer_allocation import ( + Symbols, + get_layer_type_list_physical_count, + validate_segment_layers, +) from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net import GatedDeltaNet @@ -31,8 +36,10 @@ def setup_method(self, method): Utils.initialize_model_parallel(1, 1) model_parallel_cuda_manual_seed(123) - def get_pg_collection(self): - return ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'pp', 'cp']) + def get_pg_collection(self, required_pgs=None): + if required_pgs is None: + required_pgs = ['tp', 'pp', 'cp'] + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=required_pgs) def get_hybrid_block(self, layer_pattern, **config_kwargs): layer_type_list = validate_segment_layers(layer_pattern) @@ -116,6 +123,60 @@ def get_mla_hybrid_block(self, layer_pattern): pg_collection=self.get_pg_collection(), ) + def get_attention_mlp_block(self, layer_pattern): + layer_type_list = validate_segment_layers(layer_pattern) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=get_layer_type_list_physical_count(layer_type_list), + num_attention_heads=4, + hidden_dropout=0.0, + attention_dropout=0.0, + use_cpu_initialization=True, + ) + return HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + + def get_attention_moe_block(self, layer_pattern): + layer_type_list = validate_segment_layers(layer_pattern) + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=get_layer_type_list_physical_count(layer_type_list), + num_attention_heads=4, + ffn_hidden_size=256, + num_moe_experts=8, + expert_model_parallel_size=1, + moe_router_topk=2, + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + hidden_dropout=0.0, + attention_dropout=0.0, + use_cpu_initialization=True, + ) + return HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection( + required_pgs=[ + 'tp', + 'pp', + 'cp', + 'tp_cp', + 'tp_dp_cp', + 'ep', + 'expt_tp', + 'tp_ep', + 'expt_dp', + ] + ), + ) + def teardown_method(self, method): Utils.destroy_model_parallel() @@ -243,6 +304,102 @@ def test_layer_types(self): assert isinstance(layers[2], TransformerLayer) assert isinstance(layers[2].mlp, MLP) + def test_group_layer_type_builds_nested_hybrid_stack(self): + """Bracketed groups build an inner HybridStack with physical layer numbering.""" + layer_type_list = validate_segment_layers("M[M*]-") + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=get_layer_type_list_physical_count(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + assert isinstance(block.layers[0], MambaLayer) + assert isinstance(block.layers[1], HybridStack) + assert isinstance(block.layers[1].layers[0], MambaLayer) + assert isinstance(block.layers[1].layers[1], TransformerLayer) + assert isinstance(block.layers[2], TransformerLayer) + assert [layer.layer_number for layer in block.layers[1].layers] == [2, 3] + assert block.layers[2].layer_number == 4 + + def test_group_sharded_state_dict_uses_logical_layer_keys(self): + """Grouped attention+MLP layers share one Transformer-compatible checkpoint key.""" + layer_type_list = validate_segment_layers("[*-]") + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=get_layer_type_list_physical_count(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + logical_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + + sharded_state_dict = block.sharded_state_dict(prefix="decoder.") + sharded_keys = {value.key for value in sharded_state_dict.values() if hasattr(value, "key")} + + assert "decoder.layers.0.self_attention.linear_qkv.weight" in sharded_keys + assert "decoder.layers.0.mlp.linear_fc1.weight" in sharded_keys + assert "decoder.layers.1.mlp.linear_fc1.weight" not in sharded_keys + assert "decoder.final_layernorm.weight" in sharded_keys + assert "decoder.final_norm.weight" not in sharded_keys + + def test_group_forward_matches_equivalent_flat_layers(self): + """A bracket group is only a scheduling/checkpoint boundary, not new math.""" + flat_block = self.get_attention_mlp_block("*-") + group_block = self.get_attention_mlp_block("[*-]") + + group_block.layers[0].layers[0].load_state_dict(flat_block.layers[0].state_dict()) + group_block.layers[0].layers[1].load_state_dict(flat_block.layers[1].state_dict()) + group_block.final_norm.load_state_dict(flat_block.final_norm.state_dict()) + + flat_block.cuda().eval() + group_block.cuda().eval() + sequence_length = 16 + micro_batch_size = 2 + hidden_states = torch.randn( + sequence_length, micro_batch_size, flat_block.config.hidden_size, device="cuda" + ) + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool, device="cuda" + ) + + with torch.no_grad(): + flat_output = flat_block(hidden_states.clone(), attention_mask=attention_mask) + group_output = group_block(hidden_states.clone(), attention_mask=attention_mask) + + torch.testing.assert_close(group_output, flat_output, rtol=0, atol=0) + + def test_group_overlap_callables_keep_ep_moe_split_visible(self): + """EP-overlap scheduling still sees dispatch/experts/combine inside a group.""" + block = self.get_attention_moe_block("[*E]") + + forward_callables, bwd_dw_callable_map, is_moe, num_local_experts = ( + build_hybrid_stack_callables(block.layers[0], layer_type=block.layer_type_list[0]) + ) + + pre_dispatch, dispatch, experts, combine, mtp_post_process = forward_callables + assert callable(pre_dispatch) + assert callable(dispatch) + assert callable(experts) + assert callable(combine) + assert mtp_post_process is None + assert is_moe + assert num_local_experts == 8 + assert "pre_dispatch_computation" in bwd_dw_callable_map + assert "mlp" in bwd_dw_callable_map + def test_invalid_layer_types_cause_failure(self): invalid_symbol = 'X' assert invalid_symbol not in Symbols.VALID_LAYERS # sanity check. diff --git a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py index 8b4c181ee30..3220830f58c 100644 --- a/tests/unit_tests/ssm/test_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -15,6 +15,7 @@ parse_hybrid_pattern, pattern_from_ratios, select_pipeline_segment, + select_pipeline_segment_with_logical_offset, validate_segment_layers, ) @@ -71,6 +72,8 @@ def test_valid_patterns(self): """Test that valid segment patterns produce the correct layer type lists.""" test_cases = [ ("M*-M*-M*-", ['M', '*', '-', 'M', '*', '-', 'M', '*', '-']), + ("M[M*]-", ['M', ('M', '*'), '-']), + ("[M*E]", [('M', '*', 'E')]), ("MMMMMMMMM", ['M'] * 9), ("MM*-MM*-", ['M', 'M', '*', '-', 'M', 'M', '*', '-']), ("E", ['E']), @@ -99,6 +102,10 @@ def test_invalid_symbols_cause_failure(self): validate_segment_layers("M|M") # pipe not valid in a segment with pytest.raises(ValueError): validate_segment_layers("M/M") # MTP separator not valid in a segment + with pytest.raises(ValueError): + validate_segment_layers("M[[M]]") # nested groups are not valid + with pytest.raises(ValueError): + validate_segment_layers("M[EM]") # MoE must be last in a group with pytest.raises(ValueError): # Not allowed to have both standard Attention and MLA/DSA validate_segment_layers("MDM*-") @@ -116,6 +123,8 @@ def test_simple_patterns(self): assert get_hybrid_total_layer_count("M*M*") == 4 assert get_hybrid_total_layer_count("MMMM") == 4 assert get_hybrid_total_layer_count("M") == 1 + assert get_hybrid_total_layer_count("[M*E]") == 3 + assert get_hybrid_total_layer_count("M[M*]-") == 4 def test_with_pipe_separators(self): assert get_hybrid_total_layer_count("M-M-|M-M*-") == 9 @@ -161,6 +170,8 @@ def test_main_pattern_only(self): """Test patterns without MTP (no / separator).""" test_cases = [ ("M*M*", "M*M*"), + ("[M*E]", "[M*E]"), + ("M[M*]-", "M[M*]-"), ("MMMM", "MMMM"), ("*M*M", "*M*M"), ("MM-*", "MM-*"), @@ -239,11 +250,22 @@ def test_invalid_symbols_in_main_pattern(self): "M*X*", # X is not valid "MaMM", # a is not valid "M*M*1", # 1 is not valid + "M[M*]X", # X is not valid after a group ] for pattern in invalid_patterns: with pytest.raises(ValueError, match="not a valid layer symbol"): parse_hybrid_pattern(pattern) + def test_invalid_group_syntax(self): + with pytest.raises(ValueError, match="without a matching"): + parse_hybrid_pattern("M[M*") + with pytest.raises(ValueError, match="not supported"): + parse_hybrid_pattern("M[M[*]]") + with pytest.raises(ValueError, match="cannot be empty"): + parse_hybrid_pattern("M[]") + with pytest.raises(ValueError, match="must be the last"): + parse_hybrid_pattern("M[EM]") + def test_invalid_symbols_in_mtp_pattern(self): """Test that invalid symbols in MTP pattern raise ValueError.""" # Single MTP depth with invalid symbol - should raise "not a valid layer symbol" @@ -412,6 +434,17 @@ def test_moe_pattern(self): 'E': 2, } + def test_group_pattern(self): + assert get_hybrid_layer_counts("M[M*]E") == { + '*': 1, + 'D': 0, + 'G': 0, + 'M': 2, + '+': 0, + '-': 0, + 'E': 1, + } + def test_mtp_with_attention(self): # MTP pattern "*M" repeated 3 depths -> 3 attn + 3 mamba from MTP assert get_hybrid_layer_counts("MMMM/*M/*M/*M") == { @@ -520,6 +553,21 @@ def test_four_segments(self, mock_log): assert layer_types == expected_layers, f"Failed for vp_stage={vp_stage}" assert offset == expected_offset, f"Failed for vp_stage={vp_stage}" + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') + def test_group_segment_offsets(self, mock_log): + layer_types, offset = select_pipeline_segment("[M*E]|M-", pp_group=None, vp_stage=1) + assert layer_types == ['M', '-'] + assert offset == 3 + + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') + def test_group_segment_logical_offsets(self, mock_log): + layer_types, physical_offset, logical_offset = select_pipeline_segment_with_logical_offset( + "[*-][*-]|[*E][*E]", pp_group=None, vp_stage=1 + ) + assert layer_types == [('*', 'E'), ('*', 'E')] + assert physical_offset == 4 + assert logical_offset == 2 + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_empty_segment(self, mock_log): """Empty segments are allowed for pipeline balancing.""" @@ -840,3 +888,12 @@ def test_mixed_dsa_and_mla(self): assert mamba_map == {2: 0} assert mlp_map == {3: 0} assert moe_map == {} + + def test_grouped_layers_are_flattened(self): + maps = get_layer_maps_from_layer_type_list([("M", "*", "E"), "M"]) + attention_map, mamba_map, moe_map = operator.itemgetter( + Symbols.ATTENTION, Symbols.MAMBA, Symbols.MOE + )(maps) + assert attention_map == {1: 0} + assert mamba_map == {0: 0, 3: 1} + assert moe_map == {2: 0} From baa725688ce80eae959e27d98437e1f69efaa750 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Fri, 22 May 2026 09:50:49 -0700 Subject: [PATCH 02/11] [fix] Drop redundant pre_mlp_layernorm recompute hook in moe_combine Carries over upstream commit ce6e22987 from #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 #4798 (original changes by @Wohox). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Yan Xu --- .../models/hybrid/fine_grained_callables.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/megatron/core/models/hybrid/fine_grained_callables.py b/megatron/core/models/hybrid/fine_grained_callables.py index 7e45881518c..8e158b6f3c9 100644 --- a/megatron/core/models/hybrid/fine_grained_callables.py +++ b/megatron/core/models/hybrid/fine_grained_callables.py @@ -15,6 +15,9 @@ from megatron.core.models.hybrid.hybrid_layer_allocation import LayerPatternItem from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols from megatron.core.models.hybrid.hybrid_layer_allocation import is_layer_group +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) from megatron.core.pipeline_parallel.utils import ScheduleNode from megatron.core.transformer.transformer_layer import make_viewless_tensor @@ -211,7 +214,28 @@ def _run_moe_combine(layer, node: ScheduleNode, output: Tensor): shared_expert_output = getattr(node.layer_state, 'shared_expert_output', None) output = layer.mlp.combine(output) output = layer.mlp.postprocess(output, shared_expert_output) - output = layer._forward_post_mlp((output, None), residual) + # Inline bda instead of calling ``layer._forward_post_mlp`` so we can skip + # the redundant ``discard_output_and_register_recompute(mlp_output_with_bias[0])`` + # that ``_forward_post_mlp`` would otherwise issue. The pre_mlp_layernorm recompute + # is already registered on ``expert_output`` inside ``_run_moe_experts``; the second + # hook on the combine-slot ``mlp_output_with_bias[0]`` is not only unnecessary but + # harmful in the bracketed-hybrid case (``[*E]``): it fires during combine_bwd's + # autograd backward and triggers the LN recompute ahead of attention's backward + # in the same pre_dispatch slot, corrupting attention gradients (grad_norm explodes + # from iter 2). GPT's ``submodule_combine_forward`` likewise inlines bda and does + # not call ``_forward_post_mlp`` for the same reason. + mlp_output_with_bias = (output, None) + with layer.bias_dropout_add_exec_handler(): + output = layer.mlp_bda(layer.training, layer.config.bias_dropout_fusion)( + mlp_output_with_bias, residual, layer.hidden_dropout + ) + if layer.offload_mlp_norm: + output = off_interface.group_commit( + output, name="mlp_norm", forced_released_tensors=[residual] + ) + output = make_viewless_tensor( + inp=output, requires_grad=output.requires_grad, keep_graph=True + ) node.layer_state.residual.record_stream(torch.cuda.current_stream()) if shared_expert_output is not None: From dbcf0cb03dd68ba6f9bf859edc306019237b6ee0 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Sun, 24 May 2026 22:23:05 -0700 Subject: [PATCH 03/11] fix(hybrid): add missing enumerate index in HybridStack layer-build loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Signed-off-by: Yan Xu Co-authored-by: Pingtian Li --- megatron/core/models/hybrid/hybrid_block.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 073fd0385ce..51e0d48ddcb 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -132,7 +132,12 @@ def __init__( # Build layers from the pre-selected segment self.layers = nn.ModuleList() physical_layer_offset = pp_layer_offset - for layer_type in self.layer_type_list: + # ``i`` is the logical layer index within this stack, used only for the + # downstream ``name=...{i}`` argument introduced from main; the existing + # FP8/FP4 contexts and ``layer_number`` continue to use the physical- + # layer counter ``physical_layer_offset`` which advances by + # ``get_layer_type_physical_count`` (bracket groups count > 1). + for i, layer_type in enumerate(self.layer_type_list): layer_number = physical_layer_offset + 1 if is_layer_group(layer_type): quant_init_context = nullcontext() From b567967896f4b5f241e48628759b2ece9a0e48b9 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Wed, 3 Jun 2026 12:20:36 -0700 Subject: [PATCH 04/11] fix: apply black formatting to hybrid callables Signed-off-by: Yan Xu --- megatron/core/models/hybrid/fine_grained_callables.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/models/hybrid/fine_grained_callables.py b/megatron/core/models/hybrid/fine_grained_callables.py index 8e158b6f3c9..517b48a22cb 100644 --- a/megatron/core/models/hybrid/fine_grained_callables.py +++ b/megatron/core/models/hybrid/fine_grained_callables.py @@ -233,9 +233,7 @@ def _run_moe_combine(layer, node: ScheduleNode, output: Tensor): output = off_interface.group_commit( output, name="mlp_norm", forced_released_tensors=[residual] ) - output = make_viewless_tensor( - inp=output, requires_grad=output.requires_grad, keep_graph=True - ) + output = make_viewless_tensor(inp=output, requires_grad=output.requires_grad, keep_graph=True) node.layer_state.residual.record_stream(torch.cuda.current_stream()) if shared_expert_output is not None: From 0004229afd4e585ef03227e7516a01f36f412c73 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Mon, 20 Jul 2026 09:29:08 -0700 Subject: [PATCH 05/11] fix(hybrid): preserve MTP postprocess state 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 --- megatron/core/models/hybrid/hybrid_model.py | 25 +++- tests/unit_tests/models/test_hybrid_model.py | 139 ++++++++++++++++++- 2 files changed, 158 insertions(+), 6 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 7d0feb44ec8..3d8cff196e3 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -499,7 +499,7 @@ def _postprocess( (it schedules MTP as separate layer nodes); the eager forward leaves it ``True`` so the regular MTP forward runs here. """ - in_inference_mode = inference_context is not None and not self.training + in_inference_mode = InferenceMode.is_active() if in_inference_mode: assert runtime_gather_output, "Inference must always gather TP logits" @@ -512,6 +512,7 @@ def _postprocess( if is_spec_decode is None: is_spec_decode = ( in_inference_mode + and inference_context is not None and inference_context.is_dynamic_batching() and inference_context.num_speculative_tokens > 0 ) @@ -537,11 +538,19 @@ def _postprocess( if self.config.mtp_num_layers is not None and self.mtp_process: assert self.config.mtp_num_layers > 0 - if in_inference_mode or is_spec_decode: + if is_spec_decode: # Cache decoder hidden states for serial MTP computation after # speculative token verification. - self._decoder_hidden_states_cache = hidden_states - else: + assert inference_context is not None + if self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block: + # Block-scope CUDA graph mode replays into a fixed buffer. + 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: # In training/eval, fold MTP loss into hidden_states. hidden_states = process_mtp_loss( hidden_states=hidden_states, @@ -554,12 +563,18 @@ def _postprocess( compute_language_model_loss=self.compute_language_model_loss, config=self.config, cp_group=self.pg_collection.cp, + tp_group=self.tp_group, packed_seq_params=packed_seq_params, scale_logits_fn=self._scale_logits if self.config.use_mup else None, + input_ids=input_ids, ) sequence_parallel_override = False - if in_inference_mode and inference_context.config.materialize_only_last_token_logits: + if ( + in_inference_mode + and inference_context is not None + and inference_context.config.materialize_only_last_token_logits + ): if inference_context.is_static_batching(): hidden_states = hidden_states[-1:, :, :] else: diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index fc288d6293e..f0d269e414a 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -26,7 +26,7 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import MLATransformerConfig, TransformerConfig -from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.enums import AttnBackend, InferenceCudaGraphScope from megatron.core.transformer.module import Float16Module from megatron.core.utils import divide, is_fa_min_version, is_torch_min_version from tests.unit_tests.test_utilities import Utils @@ -95,6 +95,143 @@ def _assert_equal_with_partial_contents(left, right, path="root"): assert left == right, f"{path}: values differ" +def _make_postprocess_stub(config, training): + """Build the minimal model surface needed by ``HybridModel._postprocess``.""" + + def output_layer(hidden_states, weight=None, runtime_gather_output=None): + return hidden_states, None + + output_layer.sequence_parallel = False + pg_collection = SimpleNamespace(cp=object(), tp=object()) + return SimpleNamespace( + config=config, + training=training, + post_process=True, + mtp_process=True, + share_embeddings_and_output_weights=False, + output_layer=output_layer, + pg_collection=pg_collection, + tp_group=pg_collection.tp, + _scale_logits=lambda logits: logits, + compute_language_model_loss=lambda labels, logits: logits, + ) + + +@pytest.mark.parametrize( + "cuda_graph_scope", [InferenceCudaGraphScope.none, InferenceCudaGraphScope.block] +) +def test_hybrid_postprocess_caches_spec_decode_hidden_states(cuda_graph_scope): + """Speculative decoding publishes decoder states through the inference context.""" + hidden_states = torch.randn(3, 2, 8) + context_buffer = ( + torch.full((5, 2, 8), -1.0) if cuda_graph_scope == InferenceCudaGraphScope.block else None + ) + inference_context = SimpleNamespace( + config=SimpleNamespace(materialize_only_last_token_logits=False), + num_speculative_tokens=1, + mtp_decoder_hidden_states=context_buffer, + is_dynamic_batching=lambda: True, + is_static_batching=lambda: False, + ) + model = _make_postprocess_stub( + SimpleNamespace( + mtp_num_layers=1, inference_cuda_graph_scope=cuda_graph_scope, use_mup=False + ), + training=False, + ) + + with InferenceMode.active(): + output = HybridModel._postprocess( + model, + hidden_states=hidden_states, + input_ids=torch.zeros(2, 3, dtype=torch.long), + position_ids=torch.zeros(2, 3, dtype=torch.long), + labels=None, + rotary_pos_emb=None, + mtp_in_postprocess=False, + runtime_gather_output=True, + inference_context=inference_context, + ) + + torch.testing.assert_close(output, hidden_states.transpose(0, 1), rtol=0, atol=0) + if cuda_graph_scope == InferenceCudaGraphScope.block: + assert inference_context.mtp_decoder_hidden_states is context_buffer + torch.testing.assert_close(context_buffer[:3], hidden_states, rtol=0, atol=0) + assert torch.all(context_buffer[3:] == -1) + else: + assert inference_context.mtp_decoder_hidden_states is hidden_states + assert not hasattr(model, "_decoder_hidden_states_cache") + + +def test_hybrid_postprocess_does_not_cache_regular_inference_hidden_states(): + """Regular inference must not make the controller run serial MTP decoding.""" + hidden_states = torch.randn(3, 2, 8) + inference_context = SimpleNamespace( + config=SimpleNamespace(materialize_only_last_token_logits=False), + num_speculative_tokens=0, + mtp_decoder_hidden_states=None, + is_dynamic_batching=lambda: True, + is_static_batching=lambda: False, + ) + model = _make_postprocess_stub( + SimpleNamespace( + mtp_num_layers=1, inference_cuda_graph_scope=InferenceCudaGraphScope.none, use_mup=False + ), + training=False, + ) + + with InferenceMode.active(): + HybridModel._postprocess( + model, + hidden_states=hidden_states, + input_ids=torch.zeros(2, 3, dtype=torch.long), + position_ids=torch.zeros(2, 3, dtype=torch.long), + labels=None, + rotary_pos_emb=None, + mtp_in_postprocess=False, + runtime_gather_output=True, + inference_context=inference_context, + ) + + assert inference_context.mtp_decoder_hidden_states is None + assert not hasattr(model, "_decoder_hidden_states_cache") + + +def test_hybrid_postprocess_forwards_rl_mtp_inputs(monkeypatch): + """RL MTP loss receives token IDs and the TP group needed to derive labels.""" + captured_kwargs = {} + + def fake_process_mtp_loss(**kwargs): + captured_kwargs.update(kwargs) + return kwargs["hidden_states"][:2] + + monkeypatch.setattr( + "megatron.core.models.hybrid.hybrid_model.process_mtp_loss", fake_process_mtp_loss + ) + model = _make_postprocess_stub( + SimpleNamespace( + mtp_num_layers=1, inference_cuda_graph_scope=InferenceCudaGraphScope.none, use_mup=False + ), + training=True, + ) + input_ids = torch.arange(4, dtype=torch.long).reshape(1, 4) + + HybridModel._postprocess( + model, + hidden_states=torch.randn(4, 1, 8), + input_ids=input_ids, + position_ids=torch.arange(4, dtype=torch.long).reshape(1, 4), + labels=None, + rotary_pos_emb=None, + mtp_in_postprocess=False, + runtime_gather_output=False, + inference_context=None, + ) + + assert captured_kwargs["input_ids"] is input_ids + assert captured_kwargs["tp_group"] is model.tp_group + + def test_hybrid_logging_process_groups_are_paired(): tp_group = object() dp_cp_group = object() From f5796d0605b76e3ee4438faf8c44f5694a517fa2 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Mon, 20 Jul 2026 10:13:06 -0700 Subject: [PATCH 06/11] Fix grouped hybrid review issues Signed-off-by: Yan Xu --- megatron/core/models/hybrid/hybrid_block.py | 7 +++- .../models/hybrid/hybrid_layer_allocation.py | 2 +- megatron/core/models/hybrid/hybrid_model.py | 2 + megatron/core/recompute.py | 8 +++- megatron/core/transformer/moe/router.py | 3 +- pretrain_hybrid.py | 8 ++-- tests/unit_tests/ssm/test_hybrid_block.py | 3 +- .../transformer/moe/test_aux_loss.py | 37 +++++++++++++++++++ 8 files changed, 61 insertions(+), 9 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 51e0d48ddcb..99fb9c5dfaf 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -326,6 +326,7 @@ def forward( inference_params: Optional[BaseInferenceContext] = None, packed_seq_params: Optional[PackedSeqParams] = None, padding_mask=None, + _checkpointed_forward_in_parent: bool = False, ): """ Forward function of the HybridStack class. @@ -405,7 +406,11 @@ def get_inner_quant_context(config, layer_number): return nullcontext() with outer_fp8_context: - if self.config.recompute_granularity == 'full' and self.training: + if ( + self.config.recompute_granularity == 'full' + and self.training + and not _checkpointed_forward_in_parent + ): hidden_states = checkpointed_forward( self, hidden_states=hidden_states, diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py index 86ccd7dd1d6..cf1e0f06c5a 100644 --- a/megatron/core/models/hybrid/hybrid_layer_allocation.py +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -551,7 +551,7 @@ def select_pipeline_segment( last_stage_layers: Optional[int] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, -) -> Tuple[List[str], int]: +) -> Tuple[List[LayerPatternItem], int]: """Select and validate the pipeline segment for the given PP rank and VP stage. When the main pattern contains '|' pipe separators, splits by '|' into diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 3d8cff196e3..064f67c32c6 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -630,6 +630,8 @@ def build_schedule_plan( """Build the HybridModel combined-1F1B schedule plan.""" if self.config.fine_grained_activation_offloading: self.preprocess_for_fine_grained_offloading() + if self.config.moe_paged_stash: + self.preprocess_for_paged_stash() from .model_chunk_schedule_plan import HybridStackModelChunkSchedulePlan diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py index bd0d1bcb3b2..3f1bff25ab0 100644 --- a/megatron/core/recompute.py +++ b/megatron/core/recompute.py @@ -75,9 +75,10 @@ def custom_forward( # Use self.layers[index] (not self._get_layer) so this # function works for both TransformerBlock and HybridStack. layer = self.layers[index] + is_hybrid_group = getattr(layer, "is_layer_group_stack", False) # Get appropriate inner quantization context - if use_inner_quantization_context: + if use_inner_quantization_context and not is_hybrid_group: if self.config.fp8: inner_quantization_context = get_fp8_context( self.config, layer.layer_number - 1 @@ -109,6 +110,11 @@ def custom_forward( with inner_quantization_context: if isinstance(layer, TransformerLayer): hidden_states, context = layer(**layer_kwargs) + elif is_hybrid_group: + for k in ("context", "context_mask", "attention_bias"): + layer_kwargs.pop(k, None) + hidden_states = layer(**layer_kwargs, _checkpointed_forward_in_parent=True) + context = None else: # MambaLayer (HybridStack `M` slot) for k in ("context", "context_mask", "attention_bias", "padding_mask"): layer_kwargs.pop(k, None) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 9b15446c771..b6d323b9885 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -695,7 +695,8 @@ def apply_z_loss(self, logits, padding_mask: Optional[torch.Tensor] = None): num_layers += self.config.mtp_num_layers if self.is_mtp_layer: - layer_number = self.layer_number + self.config.num_layers + mtp_depth = ((self.layer_number - 1) % self.config.mtp_num_layers) + 1 + layer_number = mtp_depth + self.config.num_layers else: layer_number = self.layer_number diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index 890f962d79b..2a5d8579a1d 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -285,6 +285,7 @@ def forward_step(data_iterator, model: HybridModel, return_schedule_plan: bool = model (HybridModel): The Hybrid Model return_schedule_plan (bool): Whether to return the schedule plan instead of output tensor. """ + args = get_args() timers = get_timers() # Get the batch. @@ -334,10 +335,9 @@ def forward_step(data_iterator, model: HybridModel, return_schedule_plan: bool = with stimer: if return_schedule_plan: - args = get_args() - assert args.overlap_moe_expert_parallel_comm, ( - "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" - ) + assert ( + args.overlap_moe_expert_parallel_comm + ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" output_tensor = model.build_schedule_plan( tokens, position_ids, diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 793fded6c78..532c5284a81 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -47,7 +47,7 @@ def get_hybrid_block(self, layer_pattern, **config_kwargs): hidden_size=256, # The Mamba layer places several constraints on this # Need to specify num_attention_heads and num_layers or TransformerConfig # will generate errors. - num_layers=len(layer_type_list), + num_layers=get_layer_type_list_physical_count(layer_type_list), num_attention_heads=4, use_cpu_initialization=True, **config_kwargs, @@ -228,6 +228,7 @@ def _run_forward(self, block, sequence_length=32, micro_batch_size=2): Symbols.MLP * 5, Symbols.ATTENTION + Symbols.MLP + Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP, Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP, + "[*-]", ], ) def test_recompute(self, recompute_kwargs: dict, layer_pattern: str): diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py index c8c7bf0dd0f..b7d8a6fe4cc 100644 --- a/tests/unit_tests/transformer/moe/test_aux_loss.py +++ b/tests/unit_tests/transformer/moe/test_aux_loss.py @@ -133,6 +133,43 @@ def test_a2a_dispatcher(self, tp_size, ep_size, cp_size): container.aux_loss_test(self.input, self.baseline_grad, "load_balancing_loss") +@pytest.mark.internal +def test_z_loss_wraps_hybrid_mtp_layer_number_to_tracker_slot(): + """Hybrid MTP routers must record z-loss in one of the configured MTP slots.""" + + class DummyGroup: + @staticmethod + def size(): + return 1 + + class DummyConfig: + moe_z_loss_coeff = 1.0 + mtp_num_layers = 2 + mtp_use_repeated_layer = False + num_layers = 8 + + class DummyRouter: + config = DummyConfig() + tp_cp_group = DummyGroup() + tp_dp_cp_group = DummyGroup() + training = True + calculate_per_token_loss = False + is_mtp_layer = True + layer_number = 5 + + clear_aux_losses_tracker() + try: + logits = torch.randn(4, 3, requires_grad=True) + TopKRouter.apply_z_loss(DummyRouter(), logits) + + values = get_moe_layer_wise_logging_tracker()["z_loss"]["values"] + assert values.shape == (10,) + assert values[8] > 0 + assert torch.count_nonzero(values) == 1 + finally: + clear_aux_losses_tracker() + + class TestSeqAuxLoss: def setup_method(self, method): baseline_container = AuxlossTestContainer( From 2c75fea2c65888d2dd5c76252ce3295433b470c3 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Tue, 21 Jul 2026 14:43:29 -0700 Subject: [PATCH 07/11] Fix grouped hybrid unit test failures Signed-off-by: Yan Xu --- megatron/core/models/hybrid/hybrid_block.py | 10 +++++++--- tests/unit_tests/transformer/moe/test_aux_loss.py | 5 +++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index 99fb9c5dfaf..adc3b28d22e 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -283,8 +283,7 @@ def final_layernorm(self): Lets generic decoder consumers (e.g. ``GPTModel.PostProcessNode``) discover the final norm via the same attribute name they use for non-hybrid decoders, while - keeping ``final_norm`` as the registered submodule so existing hybrid checkpoint - keys are unchanged. + keeping ``final_norm`` as the registered submodule for local state-dict compatibility. """ return getattr(self, "final_norm", None) @@ -562,9 +561,14 @@ def _sharded_state_dict( # Add modules other than self.layers for name, module in self.named_children(): if not module is self.layers: + module_prefix = f'{prefix}{name}.' module_sharded_state_dict = sharded_state_dict_default( - module, f'{prefix}{name}.', sharded_offsets, metadata, tp_group=self.tp_group + module, module_prefix, sharded_offsets, metadata, tp_group=self.tp_group ) + if name == 'final_norm': + replace_prefix_for_sharding( + module_sharded_state_dict, module_prefix, f'{prefix}final_layernorm.' + ) sharded_state_dict.update(module_sharded_state_dict) return sharded_state_dict diff --git a/tests/unit_tests/transformer/moe/test_aux_loss.py b/tests/unit_tests/transformer/moe/test_aux_loss.py index b7d8a6fe4cc..8c440c00b33 100644 --- a/tests/unit_tests/transformer/moe/test_aux_loss.py +++ b/tests/unit_tests/transformer/moe/test_aux_loss.py @@ -11,6 +11,7 @@ get_cuda_rng_tracker, model_parallel_cuda_manual_seed, ) +from megatron.core.transformer.moe.moe_logging import destroy_moe_metrics_tracker from megatron.core.transformer.moe.moe_utils import ( clear_aux_losses_tracker, get_default_pg_collection, @@ -157,7 +158,7 @@ class DummyRouter: is_mtp_layer = True layer_number = 5 - clear_aux_losses_tracker() + destroy_moe_metrics_tracker() try: logits = torch.randn(4, 3, requires_grad=True) TopKRouter.apply_z_loss(DummyRouter(), logits) @@ -167,7 +168,7 @@ class DummyRouter: assert values[8] > 0 assert torch.count_nonzero(values) == 1 finally: - clear_aux_losses_tracker() + destroy_moe_metrics_tracker() class TestSeqAuxLoss: From de0edd4ad1a263519b3c5c168c260caa0a27c646 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Mon, 27 Jul 2026 15:30:36 -0700 Subject: [PATCH 08/11] Address review: scope hybrid checkpoint key rename, honor output_processor Two compatibility fixes from the strict review on #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 --- megatron/core/models/hybrid/hybrid_block.py | 16 +++++- megatron/core/models/hybrid/hybrid_model.py | 34 +++++++++++ tests/unit_tests/models/test_hybrid_model.py | 60 ++++++++++++++++++++ 3 files changed, 109 insertions(+), 1 deletion(-) diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py index adc3b28d22e..74a35c023c5 100644 --- a/megatron/core/models/hybrid/hybrid_block.py +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -91,6 +91,7 @@ def __init__( pp_layer_offset: int = 0, logical_layer_offset: int = 0, is_layer_group_stack: bool = False, + transformer_sharded_keys: bool = False, post_layer_norm: bool = True, post_process: bool = True, device=None, @@ -101,6 +102,12 @@ def __init__( ) -> None: """ Args: + transformer_sharded_keys (bool): emit ``TransformerBlock``-style sharded + checkpoint keys (``final_layernorm`` instead of ``final_norm``) so the + checkpoint is interchangeable with a ``GPTModel`` one. Only set for + bracketed-group patterns, whose logical layers map one-to-one onto + transformer layers; leaving it off keeps the historical hybrid keys so + existing non-grouped hybrid checkpoints stay loadable. name (str | None): module instance name passed top-down from its paranet module """ super().__init__(config=config) @@ -110,6 +117,7 @@ def __init__( self.is_mtp_layer = is_mtp_layer self.logical_layer_offset = logical_layer_offset self.is_layer_group_stack = is_layer_group_stack + self.transformer_sharded_keys = transformer_sharded_keys assert pg_collection is not None, "pg_collection must be provided for HybridStack" @@ -161,6 +169,7 @@ def __init__( pp_layer_offset=physical_layer_offset, logical_layer_offset=logical_layer_offset + len(self.layers), is_layer_group_stack=True, + transformer_sharded_keys=transformer_sharded_keys, post_layer_norm=False, post_process=False, device=device, @@ -565,7 +574,12 @@ def _sharded_state_dict( module_sharded_state_dict = sharded_state_dict_default( module, module_prefix, sharded_offsets, metadata, tp_group=self.tp_group ) - if name == 'final_norm': + # The registered submodule stays ``final_norm`` (local state-dict keys + # are unchanged), but grouped stacks publish the sharded key under + # TransformerBlock's ``final_layernorm`` name so their checkpoints + # cross-load with GPTModel. Non-grouped stacks keep ``final_norm`` so + # hybrid checkpoints written before this feature still load. + if name == 'final_norm' and self.transformer_sharded_keys: replace_prefix_for_sharding( module_sharded_state_dict, module_prefix, f'{prefix}final_layernorm.' ) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 064f67c32c6..394de959312 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -193,6 +193,7 @@ def __init__( # Parse unified pattern to extract main and MTP components, and # determine the pipeline segment for this model instance. from megatron.core.models.hybrid.hybrid_layer_allocation import ( + Symbols, parse_hybrid_pattern, select_pipeline_segment_with_logical_offset, ) @@ -201,6 +202,13 @@ def __init__( self.mtp_pattern = parsed.mtp_pattern self.mtp_num_depths = parsed.mtp_num_depths + # Bracketed-group patterns give every logical layer the structure of a + # transformer layer, so their checkpoints are made key-compatible with + # GPTModel. Derived from the full pattern rather than this rank's segment so + # every PP stage agrees on the naming. Non-grouped patterns keep the + # historical hybrid keys, which existing hybrid checkpoints were saved with. + transformer_sharded_keys = Symbols.GROUP_START in (parsed.main_pattern or '') + logging_pg_kwargs = _hybrid_logging_pg_kwargs(self.pg_collection) layer_type_list, layer_offset, logical_layer_offset = ( @@ -282,6 +290,7 @@ def __init__( layer_type_list=layer_type_list, pp_layer_offset=layer_offset, logical_layer_offset=logical_layer_offset, + transformer_sharded_keys=transformer_sharded_keys, post_process=self.post_process, dtype=config.params_dtype, pg_collection=self.pg_collection, @@ -498,6 +507,10 @@ def _postprocess( ``mtp_in_postprocess`` lets the EP-overlap path skip the inline MTP block (it schedules MTP as separate layer nodes); the eager forward leaves it ``True`` so the regular MTP forward runs here. + ``output_processor`` replaces the default logits / loss computation with a + caller-supplied hook (used by RL and other custom output paths); it is + forwarded by ``PostProcessNode`` and handled here exactly as in + ``GPTModel._postprocess``. """ in_inference_mode = InferenceMode.is_active() if in_inference_mode: @@ -570,6 +583,27 @@ def _postprocess( ) sequence_parallel_override = False + + if output_processor is not None: + return output_processor( + hidden_states=hidden_states, + output_layer=self.output_layer, + output_weight=output_weight, + labels=labels, + loss_mask=loss_mask, + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=decoder_input, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + runtime_gather_output=runtime_gather_output, + context=output_processor_context, + compute_language_model_loss=self.compute_language_model_loss, + scale_logits=self._scale_logits, + config=self.config, + ) + if ( in_inference_mode and inference_context is not None diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index f0d269e414a..725b404c28e 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -232,6 +232,47 @@ def fake_process_mtp_loss(**kwargs): assert captured_kwargs["tp_group"] is model.tp_group +def test_hybrid_postprocess_uses_output_processor_hook(): + """A caller-supplied output processor replaces the default logits / loss path.""" + captured_kwargs = {} + sentinel = torch.randn(2, 4, 8) + + def output_processor(**kwargs): + captured_kwargs.update(kwargs) + return sentinel + + model = _make_postprocess_stub( + SimpleNamespace( + mtp_num_layers=None, inference_cuda_graph_scope=InferenceCudaGraphScope.none + ), + training=True, + ) + hidden_states = torch.randn(4, 2, 8) + labels = torch.zeros(2, 4, dtype=torch.long) + context = object() + + output = HybridModel._postprocess( + model, + hidden_states=hidden_states, + input_ids=torch.zeros(2, 4, dtype=torch.long), + position_ids=torch.zeros(2, 4, dtype=torch.long), + labels=labels, + rotary_pos_emb=None, + mtp_in_postprocess=False, + runtime_gather_output=False, + inference_context=None, + output_processor=output_processor, + output_processor_context=context, + ) + + assert output is sentinel + assert captured_kwargs["hidden_states"] is hidden_states + assert captured_kwargs["labels"] is labels + assert captured_kwargs["context"] is context + assert captured_kwargs["output_layer"] is model.output_layer + assert captured_kwargs["compute_language_model_loss"] is model.compute_language_model_loss + + def test_hybrid_logging_process_groups_are_paired(): tp_group = object() dp_cp_group = object() @@ -543,6 +584,25 @@ def test_grouped_sharded_state_dict_uses_transformer_checkpoint_keys(self): assert "decoder.final_norm.weight" not in sharded_keys assert "output_layer._extra_state" not in sharded_state_dict + def test_ungrouped_sharded_state_dict_keeps_hybrid_final_norm_key(self): + """Non-grouped patterns keep ``final_norm`` so older hybrid checkpoints load.""" + model_config = TransformerConfig( + num_layers=2, hidden_size=256, num_attention_heads=4, use_cpu_initialization=True + ) + model = HybridModel( + config=model_config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=100, + max_sequence_length=4, + hybrid_layer_pattern="*-", + ) + + sharded_state_dict = model.sharded_state_dict() + sharded_keys = {value.key for value in sharded_state_dict.values() if hasattr(value, "key")} + + assert "decoder.final_norm.weight" in sharded_keys + assert "decoder.final_layernorm.weight" not in sharded_keys + def test_layer_numbers(self): """ The layer numbers should start at one (for the embedding # layer) and go up From 6b3772fdc33f84b21bf67540145926f1b36cc0d1 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Mon, 27 Jul 2026 16:19:53 -0700 Subject: [PATCH 09/11] Opt directly-constructed grouped HybridStack into transformer keys 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 --- tests/unit_tests/ssm/test_hybrid_block.py | 27 +++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 532c5284a81..54ee1bc94b2 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -344,6 +344,9 @@ def test_group_sharded_state_dict_uses_logical_layer_keys(self): layer_type_list=layer_type_list, pp_layer_offset=0, logical_layer_offset=0, + # HybridModel sets this from the full layer pattern; a directly + # constructed stack has to opt in itself. + transformer_sharded_keys=True, pg_collection=self.get_pg_collection(), ) @@ -356,6 +359,30 @@ def test_group_sharded_state_dict_uses_logical_layer_keys(self): assert "decoder.final_layernorm.weight" in sharded_keys assert "decoder.final_norm.weight" not in sharded_keys + def test_sharded_state_dict_keeps_final_norm_key_without_transformer_keys(self): + """Default (non-grouped) stacks keep the historical ``final_norm`` sharded key.""" + layer_type_list = validate_segment_layers("*-") + transformer_config = TransformerConfig( + hidden_size=256, + num_layers=get_layer_type_list_physical_count(layer_type_list), + num_attention_heads=4, + use_cpu_initialization=True, + ) + block = HybridStack( + transformer_config, + hybrid_stack_spec.submodules, + layer_type_list=layer_type_list, + pp_layer_offset=0, + logical_layer_offset=0, + pg_collection=self.get_pg_collection(), + ) + + sharded_state_dict = block.sharded_state_dict(prefix="decoder.") + sharded_keys = {value.key for value in sharded_state_dict.values() if hasattr(value, "key")} + + assert "decoder.final_norm.weight" in sharded_keys + assert "decoder.final_layernorm.weight" not in sharded_keys + def test_group_forward_matches_equivalent_flat_layers(self): """A bracket group is only a scheduling/checkpoint boundary, not new math.""" flat_block = self.get_attention_mlp_block("*-") From 343e029cbe25992e1b2256439d6d4c70a439fedf Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Thu, 21 May 2026 15:12:44 -0700 Subject: [PATCH 10/11] [feat] FSDP support for HybridStack EP-overlap 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 #4798 (original changes by @Wohox). Depends on the HybridStack changes in part 2/4 (#TBD). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Yan Xu --- .../distributed/fsdp/mcore_fsdp_adapter.py | 23 +- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 38 ++- .../a2a_overlap/test_fsdp_hybrid_overlap.py | 230 ++++++++++++++++++ 3 files changed, 281 insertions(+), 10 deletions(-) create mode 100644 tests/unit_tests/a2a_overlap/test_fsdp_hybrid_overlap.py diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 3ccd9f932c8..7be8a26cb16 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -181,7 +181,14 @@ def __init__( config.overlap_moe_expert_parallel_comm and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" ): - supported_fsdp_unit_modules = [TransformerLayer, MoETransformerLayer, MambaLayer] + from megatron.core.models.hybrid.hybrid_block import HybridStack + + supported_fsdp_unit_modules = [ + TransformerLayer, + MoETransformerLayer, + MambaLayer, + HybridStack, + ] assert self.fsdp_unit_modules and all( module in supported_fsdp_unit_modules for module in self.fsdp_unit_modules ), ( @@ -190,6 +197,19 @@ def __init__( f"{supported_fsdp_unit_modules}, " f"got {self.fsdp_unit_modules}." ) + + # HybridStack-specific filter: when bracketed hybrid patterns are used, + # the model has a nested layout -- an outer HybridStack root + # (``is_layer_group_stack=False``) whose ``layers`` are inner + # bracket-group HybridStacks (``is_layer_group_stack=True``). + # ``named_modules()`` walks root-first, so with ``[HybridStack]`` the + # outer matches first, the inner ones get skipped as its submodules, + # and the whole decoder becomes a single FSDP unit. We exclude the + # outer so each bracket group is its own unit. Modules without the + # attribute (TransformerLayer, etc.) keep the default ``True``. + def _fsdp_unit_filter(m): + return getattr(m, "is_layer_group_stack", True) + super().__init__( config=config, module=MegatronFSDP( @@ -197,6 +217,7 @@ def __init__( mixed_precision_policy=self.mp_policy, module=module, fsdp_unit_modules=self.fsdp_unit_modules, + fsdp_unit_filter=_fsdp_unit_filter, disable_bucketing=disable_bucketing, device=self.device, dist_index=self.megatron_fsdp_dist_index, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index f2df87ff256..3fea1de8123 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -18,7 +18,7 @@ from contextlib import contextmanager from enum import Enum, auto from functools import partial -from typing import Any, Dict, List, Literal, Optional, Tuple, Type +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Type import torch import torch.nn as nn @@ -211,6 +211,7 @@ def __init__( ddp_config: DistributedDataParallelConfig = None, mixed_precision_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(), fsdp_unit_modules: Optional[List[torch.nn.Module] | List[str]] = None, + fsdp_unit_filter: Optional[Callable[[torch.nn.Module], bool]] = None, disable_bucketing: bool = False, device: Optional[torch.device] = None, calculate_per_token_loss: bool = False, @@ -322,6 +323,11 @@ def __init__( if fsdp_unit_modules is not None else [] ) + # Optional caller-supplied filter run after the isinstance check; lets the + # adapter exclude specific class instances (e.g. an outer wrapper that + # shares its class with the actual FSDP-unit instances) without leaking + # model knowledge into this library. + self.fsdp_unit_filter = fsdp_unit_filter # Determine if we should delay the gradient reduction. self.is_delay_grad_reduce = self.data_parallel_sharding_strategy in ["no_shard", "optim"] @@ -423,7 +429,7 @@ def _init_fsdp_param_and_grad_buffer(self): total_param_elements = 0 total_fsdp_module = 0 for module in self.module.modules(): - if isinstance(module, tuple(self.fsdp_unit_modules)): + if self._is_fsdp_unit_module(module): total_fsdp_module += 1 total_param_elements += sum(p.numel() for p in module.parameters()) # The suggested size is twice the number of elements in the FSDP modules. @@ -455,6 +461,21 @@ def _import_class_from_path(self, class_path: str): cls = getattr(module, class_name) return cls + def _is_fsdp_unit_module(self, module: nn.Module) -> bool: + """Whether ``module`` should be treated as an FSDP unit. + + Default: ``isinstance(module, tuple(fsdp_unit_modules))``. When the + caller provides ``fsdp_unit_filter``, the filter runs after the + isinstance check and can exclude specific instances -- useful when a + wrapping container shares its class with the actual unit instances + (so a pure class-based match would register the wrong layer). + """ + if not isinstance(module, tuple(self.fsdp_unit_modules)): + return False + if self.fsdp_unit_filter is not None: + return self.fsdp_unit_filter(module) + return True + def all_gather_and_wait_parameters_ready( self, params, @@ -553,7 +574,6 @@ def _register_fsdp_hooks(self, root_module): `optim` and `optim_grads` do not require FSDP units because they do not shard model parameters. """ - fsdp_unit_modules = self.fsdp_unit_modules def _param_list_for_submodule_unshard( module: nn.Module, pass_direction: Literal["forward", "backward"] @@ -590,7 +610,7 @@ def _param_list_for_submodule_unshard( # recomputation on individual submodules. return list(module.parameters(recurse=False)) else: - if isinstance(module, tuple(fsdp_unit_modules)): + if self._is_fsdp_unit_module(module): # FSDP unit modules should be unsharded and communicated together. return list(module.parameters()) else: @@ -690,7 +710,7 @@ def _post_backward_release_module(module, *unused): - Releases the module's parameters for the backward phase to free memory. - Marks the module as IDLE in the training state machine. """ - assert isinstance(module, tuple(fsdp_unit_modules)) + assert self._is_fsdp_unit_module(module) assert self.data_parallel_sharding_strategy == "optim_grads_params" # Release parameters for this module after backward. @@ -967,8 +987,8 @@ def _post_forward(module: nn.Module, input: Any, output: Any): lazy_release = False module._training_state = TrainingState.IDLE - assert isinstance( - module, tuple(fsdp_unit_modules) + assert self._is_fsdp_unit_module( + module ), "_post_forward hook should only be registered on FSDP unit modules." # Release the module parameters after the forward pass to save memory. @@ -1063,7 +1083,7 @@ def _register_pre_backward_param_unshard_hook(module): if not self.enable_fine_grained_param_gather_hook: _register_pre_forward_param_unshard_hook(module) - if isinstance(module, tuple(fsdp_unit_modules)): + if self._is_fsdp_unit_module(module): fsdp_modules.append(module) # Register the forward post-hook to reshard FSDP unit module parameters # after the forward pass, except when recomputing forward activations, @@ -1087,7 +1107,7 @@ def _register_pre_backward_param_unshard_hook(module): # Register the post-backward hook to deallocate model parameters # and reduce-scatter gradients after the backward pass. - if isinstance(module, tuple(fsdp_unit_modules)): + if self._is_fsdp_unit_module(module): if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params": self.forward_pre_hooks[f"module {name} register post-backward hook"] = ( module.register_forward_pre_hook( diff --git a/tests/unit_tests/a2a_overlap/test_fsdp_hybrid_overlap.py b/tests/unit_tests/a2a_overlap/test_fsdp_hybrid_overlap.py new file mode 100644 index 00000000000..18d34c7dee3 --- /dev/null +++ b/tests/unit_tests/a2a_overlap/test_fsdp_hybrid_overlap.py @@ -0,0 +1,230 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""FSDP + EP overlap + Hybrid model integration test. + +End-to-end check that a HybridModel with a bracketed layer pattern (each +bracket builds an inner ``HybridStack`` that becomes the FSDP unit) trains +identically through two paths: + +- reference: standard FSDP forward/backward, no EP overlap +- test: ``combined_1f1b_schedule_for_no_pipelining`` with + ``overlap_moe_expert_parallel_comm=True`` + +Both paths use ``fsdp_unit_modules=[HybridStack]`` so meta-device +materialization traversal is identical and the seed lands on the same +parameter each draw -- a precondition for bit-exact comparison. The outer +HybridStack root (``is_layer_group_stack=False``) is filtered out of the +FSDP unit set by ``MegatronFSDP._is_fsdp_unit_module`` so each bracket +group's inner HybridStack is its own unit; this test also asserts that +property directly. +""" + +import gc + +import pytest +import torch + +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel +from megatron.core.distributed.fsdp.src.megatron_fsdp.fully_shard import fully_shard_optimizer +from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.pipeline_parallel.utils import set_streams +from megatron.core.ssm.mamba_mixer import HAVE_MAMBA_SSM +from megatron.core.transformer import TransformerConfig + +try: + import causal_conv1d # noqa: F401 + + HAVE_CAUSAL_CONV1D = True +except ImportError: + HAVE_CAUSAL_CONV1D = False +from megatron.core.utils import is_te_min_version, is_torch_min_version +from tests.unit_tests.a2a_overlap.utils import ( + assert_models_equal, + build_input_data, + deterministic_mode, + fsdp_train_step, + get_valid_flex_dispatcher_backend, + get_valid_token_dispatcher_types, + overlap_train_step, + reset_model, +) +from tests.unit_tests.test_utilities import Utils + +SEQ_LEN = 32 +VOCAB_SIZE = 128 +NUM_STEPS = 3 +LR = 0.01 + + +def _hybrid_config(hybrid_layer_pattern, num_moe_experts=8, extra_kwargs=None): + """Build a TransformerConfig usable by HybridModel + EP overlap.""" + extra_kwargs = dict(extra_kwargs or {}) + # HybridModel derives effective num_layers from the pattern; we still pass + # the flattened count so TransformerConfig.__post_init__ checks pass. + flat = hybrid_layer_pattern.replace("[", "").replace("]", "") + return TransformerConfig( + # ``deterministic_mode`` (utils.deterministic_mode) sets + # ``NVTE_FUSED_ATTN=0`` for reproducibility; the default attention + # backend ``auto`` asserts that env is unset, so pin it to ``unfused`` + # like the GPT-side a2a_overlap tests do. + attention_backend="unfused", + pipeline_model_parallel_size=1, + expert_model_parallel_size=4, + deterministic_mode=True, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + num_layers=len(flat), + hidden_size=512, + num_attention_heads=8, + num_query_groups=8, + ffn_hidden_size=512, + kv_channels=64, + hidden_dropout=0.0, + attention_dropout=0.0, + add_bias_linear=False, + num_moe_experts=num_moe_experts, + moe_grouped_gemm=True, + moe_router_dtype="fp32", + **extra_kwargs, + ) + + +def _hybrid_model(config, hybrid_layer_pattern): + return HybridModel( + config=config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=VOCAB_SIZE, + max_sequence_length=SEQ_LEN, + hybrid_layer_pattern=hybrid_layer_pattern, + pre_process=True, + post_process=True, + ).cuda() + + +def _make_ddp_config(): + return DistributedDataParallelConfig( + use_megatron_fsdp=True, + data_parallel_sharding_strategy="optim_grads_params", + overlap_grad_reduce=True, + overlap_param_gather=True, + megatron_fsdp_main_params_dtype=None, + ) + + +def _count_fsdp_units(fsdp_wrapper): + """Count HybridStack instances actually registered as FSDP units.""" + inner = fsdp_wrapper.module # MegatronFSDP wrapper + return sum( + 1 + for m in inner.module.modules() + if isinstance(m, HybridStack) and inner._is_fsdp_unit_module(m) + ) + + +class TestFSDPHybridOverlap: + """FSDP + EP overlap + hybrid model: per-step loss and final weights + must match the no-overlap reference (both using HybridStack as the + FSDP unit).""" + + def setup_method(self, method): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=4, + ) + set_streams() + + def teardown_method(self, method): + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not is_te_min_version("2.3.0"), reason="Requires TE >= 2.3.0") + @pytest.mark.skipif( + not is_torch_min_version("2.6.0"), reason="EP overlap hangs on torch < 2.6.0" + ) + @pytest.mark.parametrize("dispatcher_type", get_valid_token_dispatcher_types()) + @pytest.mark.parametrize("shared_expert_intermediate_size", [None, 512]) + @pytest.mark.parametrize("hybrid_layer_pattern", ["[*E][*E]", "[M*E][M*E]"]) + def test_fsdp_hybrid_overlap_training_step( + self, dispatcher_type, shared_expert_intermediate_size, hybrid_layer_pattern + ): + if "M" in hybrid_layer_pattern and not (HAVE_MAMBA_SSM and HAVE_CAUSAL_CONV1D): + pytest.skip( + "Mamba pattern requires both mamba-ssm and causal-conv1d " + "(`pip install mamba-ssm causal-conv1d`)." + ) + extra_kwargs = {"moe_token_dispatcher_type": dispatcher_type} + if dispatcher_type == "flex": + backend = get_valid_flex_dispatcher_backend() + if backend is None: + pytest.skip("No flex dispatcher backend available") + extra_kwargs["moe_flex_dispatcher_backend"] = backend + if shared_expert_intermediate_size is not None: + extra_kwargs["moe_shared_expert_intermediate_size"] = shared_expert_intermediate_size + + with deterministic_mode(): + data = build_input_data(seq_len=SEQ_LEN, vocab_size=VOCAB_SIZE) + + # Reference: no EP overlap, but same FSDP unit class so init + # consumes the seeded RNG in the same order as the test path. + ref_config = _hybrid_config(hybrid_layer_pattern, extra_kwargs=extra_kwargs) + ref_model = _hybrid_model(ref_config, hybrid_layer_pattern) + init_params = reset_model(ref_model) + + ref_fsdp = FullyShardedDataParallel( + config=ref_config, + ddp_config=_make_ddp_config(), + module=ref_model, + fsdp_unit_modules=[HybridStack], + ) + ref_opt = torch.optim.SGD(ref_fsdp.parameters(), lr=LR) + ref_opt = fully_shard_optimizer(optimizer=ref_opt) + + # Test: EP overlap on. + test_kwargs = {**extra_kwargs, "overlap_moe_expert_parallel_comm": True} + test_config = _hybrid_config(hybrid_layer_pattern, extra_kwargs=test_kwargs) + test_model = _hybrid_model(test_config, hybrid_layer_pattern) + reset_model(test_model, init_params) + + test_fsdp = FullyShardedDataParallel( + config=test_config, + ddp_config=_make_ddp_config(), + module=test_model, + fsdp_unit_modules=[HybridStack], + ) + test_opt = torch.optim.SGD(test_fsdp.parameters(), lr=LR) + test_opt = fully_shard_optimizer(optimizer=test_opt) + + # Lock in the FSDP-unit selection: each bracket group's inner + # HybridStack is its own unit, and the outer root is excluded. + # Pattern `[*E][*E]` has 2 bracket groups. + expected_units = hybrid_layer_pattern.count("[") + assert _count_fsdp_units(ref_fsdp) == expected_units, ( + f"reference: expected {expected_units} FSDP units, " + f"got {_count_fsdp_units(ref_fsdp)}" + ) + assert _count_fsdp_units(test_fsdp) == expected_units, ( + f"test: expected {expected_units} FSDP units, " + f"got {_count_fsdp_units(test_fsdp)}" + ) + + rank = torch.distributed.get_rank() + for step in range(NUM_STEPS): + if hasattr(ref_fsdp, "set_is_first_microbatch"): + ref_fsdp.set_is_first_microbatch() + ref_loss = fsdp_train_step(ref_fsdp, ref_opt, data) + test_loss = overlap_train_step(test_fsdp, test_opt, test_config, data) + + assert torch.equal(ref_loss, test_loss), ( + f"[rank {rank}] Loss mismatch at step {step}: " + f"ref={ref_loss.item()}, test={test_loss.item()}" + ) + + assert_models_equal(ref_fsdp, test_fsdp) + + del ref_fsdp, test_fsdp, ref_opt, test_opt + gc.collect() + torch.cuda.empty_cache() From bc5a0f707c9568987fcfba624809d9d2ec5e9928 Mon Sep 17 00:00:00 2001 From: Yan Xu Date: Mon, 25 May 2026 08:01:03 -0700 Subject: [PATCH 11/11] fix(mtp): detach mtp_hidden_states chunks at chunk_state boundary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the EP overlap schedule runs an MTP layer, `submodule_mtp_pre_dispatch_forward` stashes `torch.chunk(hidden_states, ...)` into ``node.chunk_state.mtp_hidden_states`` in the pre_dispatch slot, and ``submodule_mtp_postprocess_forward`` later does ``torch.cat(mtp_hidden_states, ...)`` in the mtp_post_process slot before feeding the LM head. Because ``chunk_state`` is a plain Python container shared across slots, the chunks carry their original grad_fn across the slot boundary — sidestepping the implicit ``detach()`` that ``ScheduleNode._forward`` applies to slot inputs. When the Bug 1 fix (commit f6ea23b1) added ``final_norm(hidden_states)`` ahead of the chunk on HybridModel-with-empty-decoder VPP chunks, the chunks' ``SplitBackward → final_norm`` chain became reachable from two independent ``run_backward`` calls: post_process → mtp_post_process traverses it via the cat, and pre_dispatch's own backward traverses it via the MTP forward chain (`chunks[offset]` is the same Tensor as ``mtp_hidden_states[offset]``). ``final_norm`` is a ``TENorm`` and therefore goes through TE's modular OpFuser, whose backward consumes ``ctx.tensor_objects`` and sets it to ``None``. The second traversal then trips the guard in ``transformer_engine/pytorch/quantized_tensor.py:restore_from_func_ctx`` and raises ``AttributeError: ctx must have .tensor_objects to restore saved tensors`` — observed on every rank of the PP stage that owns MTP on the DeepSeek-V3-Proxy-Hybrid-NoMLA 8-node EP-overlap run. GPT does not hit the same error because: - the Bug 1 final_norm branch is gated on ``isinstance(model, HybridModel)`` so GPT never inserts an OpFuser node into the MTP pre_dispatch slot's autograd chain, - GPT's mixed-VPP layout puts ``final_layernorm`` inside the *last decoder* layer's combine slot (see ``submodule_combine_forward``), and the ``ScheduleNode._forward`` input ``.detach()`` between that combine slot and MTP's pre_dispatch keeps the chunks' grad_fn rooted at a slot leaf rather than at ``final_layernorm`` itself. Fix: route the stored chunks through ``node.detach`` so the cross-slot view is a list of leaves. ``node.detach`` records the originals in ``before_detached`` and the detached copies in ``self.detached``, so ``TransformerLayerNode.backward_impl`` keeps pulling the LM-head-side grad (accumulated on the detached leaves by mtp_post_process / post_process backward) back into pre_dispatch's ``run_backward(outputs + before_detached, ...)`` call — gradient flow stays mathematically equivalent, just no longer shared across slots. ``hidden_states = chunks[offset]`` (the live tensor) remains the MTP input so the in-slot forward chain (eh_proj → attention → ...) still propagates grads correctly to the rest of the slot's graph. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Yan Xu Co-authored-by: Pingtian Li --- .../models/common/fine_grained_callables.py | 23 +++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/megatron/core/models/common/fine_grained_callables.py b/megatron/core/models/common/fine_grained_callables.py index 184e853d0c4..a3b6e0eb02b 100644 --- a/megatron/core/models/common/fine_grained_callables.py +++ b/megatron/core/models/common/fine_grained_callables.py @@ -64,8 +64,27 @@ def submodule_mtp_pre_dispatch_forward(node, hidden_states): ) offset = get_mtp_layer_offset(layer.config, node.chunk_state.model.vp_stage) - node.chunk_state.mtp_hidden_states = list(torch.chunk(hidden_states, 1 + offset, dim=0)) - hidden_states = node.chunk_state.mtp_hidden_states[offset] + chunks = list(torch.chunk(hidden_states, 1 + offset, dim=0)) + # Store DETACHED chunks in chunk_state. mtp_hidden_states is a + # ``chunk_state``-level Python list that crosses slot boundaries + # (set here in MTP's pre_dispatch slot, later torch.cat'd in MTP's + # mtp_post_process slot before feeding the LM head). Without an + # explicit detach the chunks keep their grad_fn from torch.chunk → + # whatever upstream node produced ``hidden_states`` (e.g. the + # final_norm we apply above for the HybridModel empty-decoder case), + # which means mtp_post_process.backward and pre_dispatch.backward + # both traverse that same grad_fn — for a TENorm-backed final_norm + # (an OpFuser op) the second traversal hits ``ctx.tensor_objects is + # None`` and raises ``ctx must have .tensor_objects to restore + # saved tensors``. Using ``node.detach`` records the originals in + # before_detached so pre_dispatch's backward_impl still pulls the + # LM-head-side grad (accumulated on the detached leaves by the + # post_process / mtp_post_process backward chain) back into the + # outputs+before_detached run_backward — i.e. the gradient flow + # remains mathematically equivalent, just no longer shared across + # slots. + node.chunk_state.mtp_hidden_states = [node.detach(c) for c in chunks] + hidden_states = chunks[offset] input_ids, position_ids, padding_mask, decoder_input, hidden_states = layer._get_embeddings( input_ids=node.chunk_state.input_ids,