From 423c19cc488e62d5bc362bb2045460f16fe9923d Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 17:23:02 +0000 Subject: [PATCH 01/23] Rename MambaModel/MambaStack/MambaStackSubmodules to Hybrid equivalents Rename the generic model classes that support multiple layer types (Mamba SSM, Attention, MoE, GDN, MLP) via hybrid_layer_pattern: - MambaModel -> HybridModel - MambaStack -> HybridStack - MambaStackSubmodules -> HybridStackSubmodules - mamba_stack_spec -> hybrid_stack_spec - mamba_inference_stack_spec -> hybrid_inference_stack_spec - get_mamba_stack_modelopt_spec -> get_hybrid_stack_modelopt_spec Move canonical files to megatron/core/models/hybrid/: - hybrid_model.py, hybrid_block.py, hybrid_layer_specs.py, hybrid_layer_allocation.py Backward-compatible re-export stubs at old import paths and class aliases (MambaModel is a thin subclass accepting mamba_stack_spec kwarg). Mamba-specific SSM classes (MambaLayer, MambaMixer, etc.) unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) --- megatron/core/inference/config.py | 4 +- .../inference/contexts/dynamic_context.py | 2 +- megatron/core/models/gpt/moe_module_specs.py | 6 +- megatron/core/models/hybrid/__init__.py | 1 + megatron/core/models/hybrid/hybrid_block.py | 425 +++++++++++++++ .../models/hybrid/hybrid_layer_allocation.py | 484 +++++++++++++++++ .../core/models/hybrid/hybrid_layer_specs.py | 226 ++++++++ megatron/core/models/hybrid/hybrid_model.py | 492 ++++++++++++++++++ megatron/core/models/mamba/__init__.py | 7 +- .../core/models/mamba/mamba_layer_specs.py | 224 +------- megatron/core/models/mamba/mamba_model.py | 479 +---------------- .../core/models/multimodal/llava_model.py | 8 +- .../post_training/modelopt/hybrid/__init__.py | 1 + .../modelopt/hybrid/model_specs.py | 145 ++++++ megatron/core/ssm/mamba_block.py | 431 +-------------- .../core/ssm/mamba_hybrid_layer_allocation.py | 487 +---------------- .../transformer/multi_token_prediction.py | 20 +- 17 files changed, 1823 insertions(+), 1619 deletions(-) create mode 100644 megatron/core/models/hybrid/__init__.py create mode 100644 megatron/core/models/hybrid/hybrid_block.py create mode 100644 megatron/core/models/hybrid/hybrid_layer_allocation.py create mode 100755 megatron/core/models/hybrid/hybrid_layer_specs.py create mode 100644 megatron/core/models/hybrid/hybrid_model.py mode change 100755 => 100644 megatron/core/models/mamba/mamba_layer_specs.py create mode 100644 megatron/core/post_training/modelopt/hybrid/__init__.py create mode 100755 megatron/core/post_training/modelopt/hybrid/model_specs.py diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index 590377979ae..e1a36ff1563 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -24,7 +24,7 @@ class MambaInferenceStateConfig: layer_type_list: List[str] """ A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer. - See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list of symbols. + See `megatron/core/models/hybrid/hybrid_layer_allocation.py` for the list of symbols. """ conv_states_shape: Tuple[int] @@ -50,7 +50,7 @@ def from_model( ssm_states_dtype: Optional[torch.dtype] = None, ) -> Optional["MambaInferenceStateConfig"]: """Returns Mamba inference state config from the model if it is a hybrid model.""" - from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols decoder = get_attr_wrapped_model(model, "decoder") layer_type_list = getattr(decoder, "layer_type_list", None) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 383d0ecbd95..5d3596c86dd 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -28,8 +28,8 @@ ) from megatron.core.inference.utils import device_memory_summary, tensor_swap from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb +from megatron.core.models.hybrid.hybrid_layer_allocation import get_layer_maps_from_layer_type_list from megatron.core.package_info import __version__ as mcore_version -from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list from megatron.core.transformer import MLATransformerConfig, TransformerConfig from megatron.core.utils import deprecate_args from megatron.core.utils import divide as core_divide diff --git a/megatron/core/models/gpt/moe_module_specs.py b/megatron/core/models/gpt/moe_module_specs.py index 53bca85f502..e82b6638f0b 100755 --- a/megatron/core/models/gpt/moe_module_specs.py +++ b/megatron/core/models/gpt/moe_module_specs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. from functools import partial from typing import Optional @@ -23,7 +23,7 @@ def get_moe_module_spec( ) -> ModuleSpec: """Helper function to get module spec for MoE. - Called by mamba_layer_specs.py for standard (non-inference) MoE specs. + Called by hybrid_layer_specs.py for standard (non-inference) MoE specs. The GPT layer specs call get_moe_module_spec_for_backend directly. Args: @@ -78,7 +78,7 @@ def get_inference_optimized_moe_spec() -> ModuleSpec: InferenceTopKRouter, InferenceGroupedMLP. MoELayer detects inference mode via config.transformer_impl and sets up the inference dispatcher internally. - Called by mamba_layer_specs.py and gpt_layer_specs.py. + Called by hybrid_layer_specs.py and gpt_layer_specs.py. """ backend = InferenceSpecProvider() activation_func = backend.activation_func() diff --git a/megatron/core/models/hybrid/__init__.py b/megatron/core/models/hybrid/__init__.py new file mode 100644 index 00000000000..d8a0a817ee3 --- /dev/null +++ b/megatron/core/models/hybrid/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py new file mode 100644 index 00000000000..02e3914c171 --- /dev/null +++ b/megatron/core/models/hybrid/hybrid_block.py @@ -0,0 +1,425 @@ +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024, Tri Dao, Albert Gu. + +# Some of this code was adopted from https://github.com/state-spaces/mamba/ +# This source code is licensed under the Apache license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor, nn + +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding +from megatron.core.enums import Fp8Recipe +from megatron.core.extensions.transformer_engine import TENorm +from megatron.core.fp4_utils import get_fp4_context +from megatron.core.fp8_utils import get_fp8_context +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.enums import CudaGraphScope +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_layer import TransformerLayer +from megatron.core.transformer.utils import sharded_state_dict_default +from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor + + +@dataclass +class HybridStackSubmodules: + """ + A class for the module specs for the HybridStack. + """ + + mamba_layer: Union[ModuleSpec, type] = IdentityOp + gdn_layer: Union[ModuleSpec, type] = IdentityOp + attention_layer: Union[ModuleSpec, type] = IdentityOp + mlp_layer: Union[ModuleSpec, type] = IdentityOp + moe_layer: Union[ModuleSpec, type] = IdentityOp + mtp_block_spec: Optional[ModuleSpec] = None + + +class HybridStack(GraphableMegatronModule, MegatronModule): + """ + Constructor for the HybridStack class. + + Args: + config (TransformerConfig): the model configuration + submodules (HybridStackSubmodules): the submodules for the stack + pre_process (bool, optional): whether to include an embedding layer. + Defaults to True. + layer_type_list (list, optional): pre-computed list of layer type symbols for + this pipeline segment. When provided (by HybridModel), pipeline stage + selection has already been done via '|' separators in the pattern. + pp_layer_offset (int, optional): the global layer offset for this pipeline + segment. Defaults to 0. + post_layer_norm (bool, optional): whether to include a final layer norm. + Defaults to True. + post_process (bool, optional): whether to include an output layer. + Defaults to True. + device (optional): the device to use. Defaults to None. + dtype (optional): the data type to use. Defaults to None. + pg_collection (ProcessGroupCollection): the required model communication + process groups to use. + is_mtp_layer (bool, optional): whether this is an MTP layer. Defaults to False. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: HybridStackSubmodules, + pre_process: bool = True, + layer_type_list: Optional[list[str]] = None, + pp_layer_offset: int = 0, + post_layer_norm: bool = True, + post_process: bool = True, + device=None, + dtype=None, + pg_collection: ProcessGroupCollection = None, + is_mtp_layer: bool = False, + ) -> None: + super().__init__(config=config) + self.pre_process = pre_process + self.post_layer_norm = post_layer_norm + self.post_process = post_process + self.is_mtp_layer = is_mtp_layer + + assert pg_collection is not None, "pg_collection must be provided for HybridStack" + + self.pp_group = pg_collection.pp + self.tp_group = pg_collection.tp + + # Required for pipeline parallel schedules + self.input_tensor = None + self.pg_collection = pg_collection + + assert layer_type_list is not None, ( + "layer_type_list must be provided. It should be pre-computed from " + "--hybrid-layer-pattern by HybridModel." + ) + self.layer_type_list = layer_type_list + + # 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) + elif self.config.fp4: + quant_init_context = get_fp4_context(self.config, i + pp_layer_offset, is_init=True) + else: + quant_init_context = nullcontext() + with quant_init_context: + if layer_type == LayerSymbols.MAMBA: + layer = build_module( + submodules.mamba_layer, + config=self.config, + layer_number=layer_number, + pp_layer_offset=pp_layer_offset, + pg_collection=pg_collection, + ) + elif layer_type == LayerSymbols.ATTENTION: + layer = build_module( + submodules.attention_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + add_layer_offset=False, + pp_layer_offset=pp_layer_offset, + ) + elif layer_type == LayerSymbols.MLP: + layer = build_module( + submodules.mlp_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + add_layer_offset=False, + ) + elif layer_type == LayerSymbols.MOE: + layer = build_module( + submodules.moe_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + add_layer_offset=False, + ) + elif layer_type == LayerSymbols.GDN: + layer = build_module( + submodules.gdn_layer, + config=self.config, + layer_number=layer_number, + pg_collection=pg_collection, + # Set to False as we do not want to change offset. + add_layer_offset=False, + ) + else: + assert False, "unexpected layer_type" + self.layers.append(layer) + + # Required for activation recomputation + self.num_layers_per_pipeline_rank = len(self.layers) + + if self.post_process and self.post_layer_norm: + # Final layer norm before output. + self.final_norm = TENorm( + config=self.config, + hidden_size=self.config.hidden_size, + eps=self.config.layernorm_epsilon, + ) + + def set_input_tensor(self, input_tensor: Tensor): + """Set input tensor to be used instead of forward()'s input. + + When doing pipeline parallelism the input from the previous + stage comes from communication, not from the input, so the + model's forward_step_func won't have it. This function is thus + used by internal code to bypass the input provided by the + forward_step_func""" + self.input_tensor = input_tensor + + def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: + """ + Returns the Mamba conv and ssm states shapes per input sequence + 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: + return layer.mamba_state_shapes_per_request() + return None + + def _should_call_local_cudagraph(self, *args, **kwargs): + """ + Check if we should call the local cudagraph path. + """ + if ( + not self.training + and hasattr(self, 'cudagraph_manager') + and kwargs['attention_mask'] is None + and ( + kwargs.get('inference_context') is not None + or kwargs.get('inference_params') is not None + ) + and CudaGraphScope.full_iteration_inference in self.config.cuda_graph_scope + ): + 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): + kwargs['hidden_states'] = ( + kwargs['hidden_states'].unwrap() + if isinstance(kwargs['hidden_states'], WrappedTensor) + else kwargs['hidden_states'] + ) + return super().__call__(*args, **kwargs)[0] + return super().__call__(*args, **kwargs) + + def forward( + self, + hidden_states: Union[Tensor, WrappedTensor], + attention_mask: Tensor, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Tensor] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask=None, + ): + """ + Forward function of the HybridStack class. + + It either returns the Loss values if labels are given or the + final hidden units + + Args: + hidden_states (Union[Tensor, WrappedTensor]): the input tensor. + Can be passed as a WrappedTensor during inference to avoid an obsolete + reference in the calling function. + attention_mask (Tensor): the attention mask. + inference_context (BaseInferenceContext): the inference parameters. + rotary_pos_emb (Tensor, optional): the rotary positional embeddings. + Defaults to None. + Returns: + Tensor: the output tensor. + """ + + inference_context = deprecate_inference_params(inference_context, inference_params) + + if not self.pre_process: + # See set_input_tensor() + hidden_states = self.input_tensor + + # Delete the obsolete reference to the initial input tensor if necessary + if isinstance(hidden_states, WrappedTensor): + hidden_states = hidden_states.unwrap() + + if inference_context and inference_context.is_static_batching(): + # NOTE(bnorick): match BaseInferenceContext attributes for + # mamba_ssm.utils.generation.BaseInferenceContext, + # this hack supports eval + inference_context.max_seqlen = inference_context.max_sequence_length + inference_context.seqlen_offset = inference_context.sequence_len_offset + + if ( + ( + ( + self.config.cuda_graph_impl == "local" + and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope + ) + or self.config.flash_decode + ) + and inference_context + and inference_context.is_static_batching() + and not self.training + ): + current_batch_size = hidden_states.shape[1] + sequence_len_offset = torch.tensor( + [inference_context.sequence_len_offset] * current_batch_size, + dtype=torch.int32, + device='cuda', + ) + else: + sequence_len_offset = None + + # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), + # otherwise do nothing extra at the outer level + # if we are using other fp8 recipes, then the context manager enter&exit are free + # we can wrap fp8_context within the for loop over layers, so that we can fine-grained + # control which layer will be fp8 or bf16 + use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed + use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed + use_fp4_context = self.config.fp4 is not None + outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext() + + if use_inner_fp8_context: + + def get_inner_quant_context(config, layer_number): + return get_fp8_context(config, layer_number) + + elif use_fp4_context: + + def get_inner_quant_context(config, layer_number): + return get_fp4_context(config, layer_number) + + else: + + def get_inner_quant_context(config, layer_number): + return nullcontext() + + with outer_fp8_context: + 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) + with inner_quant_context: + if isinstance(layer, TransformerLayer): + hidden_states, _ = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + ) + else: # MambaLayer, Expert, or MLP + hidden_states = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + ) + + # The attention layer (currently a simplified transformer layer) + # outputs a tuple of (hidden_states, context). Context is intended + # for cross-attention, and is not needed in our model. + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + # Final layer norm. + if self.post_process and self.post_layer_norm: + hidden_states = self.final_norm(hidden_states) + + # Ensure that the tensor passed between pipeline parallel stages is + # viewless. See related notes in TransformerBlock and TransformerLayer + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + + return hidden_states + + def sharded_state_dict( + self, + prefix: str = '', + sharded_offsets: Optional[tuple] = None, + metadata: Optional[dict] = None, + ) -> ShardedStateDict: + """ + Returns a sharded state dictionary for the current object. + + This function constructs a sharded state dictionary by iterating over the layers + in the current object, computing the sharded state dictionary for each layer, + and combining the results into a single dictionary. + + Parameters: + prefix (str): The prefix to use for the state dictionary keys. + sharded_offsets (tuple): The sharded offsets to use for the state dictionary. + metadata (dict): Additional metadata to use when computing the sharded state dictionary. + + Returns: + dict: The sharded state dictionary for the current object. + """ + + sharded_state_dict = {} + layer_prefix = f'{prefix}layers.' + + 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 + ) + + sharded_prefix = f'{layer_prefix}{global_layer_offset}.' + sharded_pp_offset = [] + + layer_sharded_state_dict = layer.sharded_state_dict( + state_dict_prefix, sharded_pp_offset, metadata + ) + + replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix) + + sharded_state_dict.update(layer_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, + ) + ) + + return sharded_state_dict + + +# Backward-compatible aliases +MambaStackSubmodules = HybridStackSubmodules +MambaStack = HybridStack diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py new file mode 100644 index 00000000000..023b59cd8ea --- /dev/null +++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py @@ -0,0 +1,484 @@ +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. + +import logging +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch + +from megatron.core.utils import log_on_each_pipeline_stage, log_single_rank + +logger = logging.getLogger(__name__) + + +class Symbols: + """Symbols for different layer types and pattern separators.""" + + MAMBA = "M" + GDN = 'G' + ATTENTION = "*" + MLP = "-" + MOE = 'E' + PIPE = '|' + MTP_SEPARATOR = "/" + VALID_LAYERS = {MAMBA, GDN, ATTENTION, MLP, MOE} + + +@dataclass +class ParsedHybridPattern: + """Result of parsing a unified hybrid pattern string. + + A unified pattern encodes both the main decoder pattern and the MTP pattern + in a single string using "/" as a separator. The main pattern may also + contain "|" pipe symbols to define pipeline stage boundaries for flexible + virtual pipeline parallelism (fVPP). + + Format: "///..." + + Examples: + - "M*M*" -> main="M*M*", mtp=None, depths=0 (no MTP) + - "M*M*/MM/MM" -> main="M*M*", mtp="MM", depths=2 + - "MMMM/*M/*M/*M" -> main="MMMM", mtp="*M", depths=3 + - "M-M-|M-M*-/MM/MM" -> main="M-M-|M-M*-" (2 PP stages), mtp="MM", depths=2 + + The "/" symbol introduces MTP patterns. Each repeated pattern after the main + decoder represents one MTP prediction depth. + + The "|" symbol in the main pattern defines pipeline stage boundaries. + + Attributes: + main_pattern: The main decoder layer pattern (e.g., "M*M*" or "M-M-|M-M*-") + mtp_pattern: The MTP layer pattern per depth (e.g., "MM"), or None if no MTP + mtp_num_depths: Number of MTP prediction depths (0 if no MTP) + """ + + main_pattern: Optional[str] + mtp_pattern: Optional[str] + mtp_num_depths: int + + +def pattern_from_ratios( + num_layers: int, attention_ratio: float = 0.0, mlp_ratio: float = 0.0 +) -> str: + """Convert deprecated ratio arguments to a layer pattern string. + + Generates an evenly-spaced hybrid layer pattern from target attention and MLP + ratios. This exists for backward compatibility with code that uses the deprecated + hybrid_attention_ratio and hybrid_mlp_ratio parameters. + + Args: + num_layers: Total number of layers. + attention_ratio: Target ratio of attention layers to total layers. + mlp_ratio: Target ratio of MLP layers to total layers. + + Returns: + A layer pattern string (e.g., "MMM*MMM*MM"). + """ + assert num_layers > 0 + assert 0.0 <= attention_ratio <= 1.0 + assert 0.0 <= mlp_ratio <= 1.0 + assert attention_ratio + mlp_ratio <= 1.0 + + # Allocate attention layers (evenly spaced, starting and ending with mamba) + attention_count = round(num_layers * attention_ratio) + mamba_count = num_layers - attention_count + sections = attention_count + 1 + section_len = mamba_count / sections + + layer_types = [Symbols.MAMBA] * num_layers + x = section_len + for i in range(num_layers): + if x < 0.5: + layer_types[i] = Symbols.ATTENTION + x += section_len + else: + x -= 1 + + # Allocate MLP layers (evenly distributed, not replacing attention) + mlp_count = round(num_layers * mlp_ratio) + if mlp_count > 0: + mamba_count -= mlp_count + ratio = mamba_count / mlp_count + x = ratio + for i in range(num_layers): + if layer_types[i] == Symbols.MAMBA: + if x < 0.5: + layer_types[i] = Symbols.MLP + x += ratio + else: + x -= 1 + + return ''.join(layer_types) + + +def get_hybrid_total_layer_count(pattern: str) -> int: + """Returns the total number of main decoder layers in a hybrid layer pattern. + + Extracts the main pattern (before the first MTP separator '/'), strips + pipeline stage separators '|', and returns the character count. + + Args: + pattern: Full hybrid layer pattern, possibly including MTP and pipe separators. + + Returns: + Total number of layers in the main decoder pattern. + """ + main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] + _validate_pattern(main_pattern, "main", allow_pipe=True) + return len(main_pattern.replace(Symbols.PIPE, '')) + + +def get_hybrid_total_pipeline_segment_count(pattern: str) -> int: + """Returns the number of pipeline segments in a hybrid layer pattern. + + Extracts the main pattern (before the first MTP separator '/') and counts + the number of segments delimited by '|'. + + Args: + pattern: Full hybrid layer pattern, possibly including MTP and pipe separators. + + Returns: + Number of pipeline segments (pipe count + 1). + """ + main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] + return main_pattern.count(Symbols.PIPE) + 1 + + +def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: + """Count layers by type across the full hybrid pattern (main + MTP). + + Parses the pattern to extract main and MTP components, then counts + each layer type. Main pattern '|' separators are skipped. MTP layers + are counted once per MTP depth. + + Args: + pattern: Full hybrid layer pattern string. + + Returns: + Dictionary mapping layer symbol to count. Keys are Symbols.MAMBA, + Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, and Symbols.MOE. + + Examples: + >>> get_hybrid_layer_counts("M*M*") + {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0} + + >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") + {'M': 8, 'G': 0, '*': 1, '-': 4, 'E': 0} + """ + parsed = parse_hybrid_pattern(pattern) + counts = { + Symbols.MAMBA: 0, + Symbols.GDN: 0, + Symbols.ATTENTION: 0, + Symbols.MLP: 0, + Symbols.MOE: 0, + } + + # Count main decoder layers (skip '|' pipe separators) + if parsed.main_pattern: + for char in parsed.main_pattern: + if char in counts: + 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 + + return counts + + +def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern: + """Parse a unified hybrid pattern string into main and MTP components. + + The pattern uses "/" as a separator between the main decoder pattern and + MTP patterns. Each MTP pattern after the separator represents one prediction + depth. The main pattern may contain "|" pipe symbols for pipeline stage + boundaries. + + Format: "///..." + + Args: + pattern: Unified pattern string, e.g., "M*M*/MM/MM" or just "M*M*" + + Returns: + ParsedHybridPattern with main_pattern, mtp_pattern, and mtp_num_depths + + Raises: + ValueError: If MTP patterns are inconsistent (all must be identical) + ValueError: If pattern contains invalid layer symbols + + Examples: + >>> parse_hybrid_pattern("M*M*") + ParsedHybridPattern(main_pattern="M*M*", mtp_pattern=None, mtp_num_depths=0) + + >>> parse_hybrid_pattern("M*M*/MM/MM") + ParsedHybridPattern(main_pattern="M*M*", mtp_pattern="MM", mtp_num_depths=2) + + >>> parse_hybrid_pattern("MMMM/*M/*M/*M") + ParsedHybridPattern(main_pattern="MMMM", mtp_pattern="*M", mtp_num_depths=3) + + >>> parse_hybrid_pattern("M-M-|M-M*-/MM/MM") + ParsedHybridPattern(main_pattern="M-M-|M-M*-", mtp_pattern="MM", mtp_num_depths=2) + """ + if pattern is None: + return ParsedHybridPattern(main_pattern=None, mtp_pattern=None, mtp_num_depths=0) + + parts = pattern.split(Symbols.MTP_SEPARATOR) + + if len(parts) == 1: + # No MTP separator found - pattern is main decoder only + main_pattern = parts[0] + _validate_pattern(main_pattern, "main", allow_pipe=True) + return ParsedHybridPattern(main_pattern=main_pattern, mtp_pattern=None, mtp_num_depths=0) + + # First part is main decoder pattern + main_pattern = parts[0] + if main_pattern: + _validate_pattern(main_pattern, "main", allow_pipe=True) + + # Remaining parts are MTP patterns (one per depth) + mtp_parts = parts[1:] + + if not mtp_parts or all(p == "" for p in mtp_parts): + # No MTP patterns after separator + return ParsedHybridPattern( + main_pattern=main_pattern if main_pattern else None, mtp_pattern=None, mtp_num_depths=0 + ) + + # Validate all MTP patterns are identical + mtp_pattern = mtp_parts[0] + for i, part in enumerate(mtp_parts[1:], start=2): + if part != mtp_pattern: + raise ValueError( + f"All MTP patterns must be identical. " + f"Pattern 1 is '{mtp_pattern}', but pattern {i} is '{part}'. " + f"Full pattern: '{pattern}'" + ) + + _validate_pattern(mtp_pattern, "MTP", allow_pipe=False) + + return ParsedHybridPattern( + main_pattern=main_pattern if main_pattern else None, + mtp_pattern=mtp_pattern, + mtp_num_depths=len(mtp_parts), + ) + + +def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) -> None: + """Validate that a pattern contains only valid layer symbols. + + Args: + pattern: Layer pattern string to validate + pattern_name: Name of pattern for error messages (e.g., "main" or "MTP") + allow_pipe: Whether to allow the pipe '|' separator (for main patterns) + + 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: + raise ValueError( + f"In {pattern_name} pattern, '{char}' is not a valid layer symbol. " + f"Valid symbols are: {valid_chars}" + ) + + +def validate_segment_layers(segment: str) -> List[str]: + """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. + Each segment should contain only valid layer symbols (no '|'). + + Args: + segment: A single pipeline segment pattern string (e.g., "M-M*-") + + Returns: + List of layer type characters. + + 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: + raise ValueError( + f"In hybrid layer pattern segment, '{layer_char}' is not " + f"one of {Symbols.VALID_LAYERS}" + ) + return layer_type_list + + +def select_pipeline_segment( + 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, +) -> Tuple[List[str], 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 + pipeline segments and selects the segment for the current PP rank / VP stage. + + When the pattern has no pipes but pp_size > 1, falls back to runtime layer + slicing (for backwards compatibility), supporting both even and uneven PP splits + via first_stage_layers / last_stage_layers. + + Args: + main_pattern: Main decoder pattern (may contain '|' separators). + Empty string is allowed (produces one empty segment). + pp_group: Pipeline parallel process group, or None if not using PP. + vp_stage: Virtual pipeline stage, or None if not using VPP. + first_stage_layers: Number of layers on the first pipeline stage for + uneven PP. Only valid when the pattern has no pipe separators. + last_stage_layers: Number of layers on the last pipeline stage for + uneven PP. Only valid when the pattern has no pipe separators. + + Returns: + Tuple of (layer_type_list, layer_offset) where layer_type_list is + the list of layer type characters for this segment, and layer_offset + is the sum of layer counts from all preceding segments. + + Raises: + ValueError: If the segment contains invalid layer symbols, if + first/last_stage_layers are used with pipe separators, if VPP is + requested without pipe separators, or if layer counts are not + evenly divisible across pipeline stages. + """ + segments = main_pattern.split(Symbols.PIPE) if main_pattern 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 + + if len(segments) > 1 and (first_stage_layers is not None or last_stage_layers is not None): + raise ValueError( + "Cannot specify num_layers_in_first_pipeline_stage or " + "num_layers_in_last_pipeline_stage when hybrid_layer_pattern " + "contains pipe ('|') separators. The pipeline layout is already " + "explicitly defined by the pipe separators." + ) + + if len(segments) == 1 and pp_size > 1: + if vp_stage is not None: + raise ValueError( + "Virtual pipeline parallelism (vp_stage != None) is not supported " + "when hybrid_layer_pattern has no pipe ('|') separators. " + "Add '|' separators to define explicit pipeline/virtual-pipeline " + "stage boundaries." + ) + log_single_rank( + logger, + logging.WARNING, + "DEPRECATION: Using hybrid_layer_pattern without pipe ('|') separators " + "with pipeline_model_parallel_size > 1 is deprecated. Please add '|' " + "separators to explicitly define pipeline stage boundaries. " + "Example: 'M*M*M*M*' with pp_size=2 should become 'M*M*|M*M*'.", + ) + full_pattern = segments[0] + layer_type_list = validate_segment_layers(full_pattern) + num_layers = len(layer_type_list) + + if first_stage_layers is not None or last_stage_layers is not None: + first = first_stage_layers or 0 + last = last_stage_layers or 0 + middle_num_layers = num_layers - first - last + middle_stages = pp_size - sum( + 1 for x in (first_stage_layers, last_stage_layers) if x is not None + ) + if middle_stages > 0: + if middle_num_layers % middle_stages != 0: + raise ValueError( + f"Middle layers ({middle_num_layers}) must be evenly divisible " + f"by middle pipeline stages ({middle_stages})." + ) + layers_per_middle = middle_num_layers // middle_stages + else: + layers_per_middle = 0 + + is_first = first_stage_layers is not None and pp_rank == 0 + is_last = last_stage_layers is not None and pp_rank == pp_size - 1 + + if is_first: + offset = 0 + count = first + elif is_last: + offset = num_layers - last + count = last + else: + middle_rank = pp_rank if first_stage_layers is None else pp_rank - 1 + offset = middle_rank * layers_per_middle + first + count = layers_per_middle + else: + if num_layers % pp_size != 0: + raise ValueError( + f"Number of layers ({num_layers}) must be evenly divisible " + f"by pipeline-model-parallel-size ({pp_size}) when no pipe " + f"separators are specified in the pattern." + ) + layers_per_rank = num_layers // pp_size + offset = pp_rank * layers_per_rank + count = layers_per_rank + + selected = layer_type_list[offset : 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"layer_offset={offset} (auto-split)", + ) + return selected, offset + + # Pipe-based segment selection + if len(segments) > 1 and len(segments) % pp_size != 0: + raise ValueError( + f"The number of pipe-delimited segments ({len(segments)}) in " + f"hybrid_layer_pattern must be evenly divisible by " + f"pipeline_model_parallel_size ({pp_size})." + ) + + vp_rel = vp_stage if vp_stage is not None else 0 + segment_index = vp_rel * pp_size + pp_rank + + if segment_index >= len(segments): + raise ValueError( + f"Pipeline segment index {segment_index} (pp_rank={pp_rank}, " + f"vp_stage={vp_rel}) is out of range for {len(segments)} segments. " + f"The pattern does not define enough pipe-delimited segments for " + f"the current PP/VPP configuration." + ) + + layer_offset = sum(len(segments[i]) for i in range(segment_index)) + my_segment = segments[segment_index] + + layer_type_list = validate_segment_layers(my_segment) + + log_on_each_pipeline_stage( + logger, + 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"layer_offset={layer_offset}", + ) + + return layer_type_list, layer_offset + + +def get_layer_maps_from_layer_type_list( + layer_type_list: List[str], +) -> Tuple[Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int]]: + """ + Returns maps from global layer index to the corresponding layer index + for each layer type in [Mamba, GDN, Attention, MLP, MoE] given a layer type list. + """ + layer_types = [Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE] + layer_maps = {layer_type: {} for layer_type in layer_types} + for global_layer_idx, layer_type in enumerate(layer_type_list): + layer_map = layer_maps[layer_type] + local_layer_idx = len(layer_map) + layer_map[global_layer_idx] = local_layer_idx + return [layer_maps[layer_type] for layer_type in layer_types] diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py new file mode 100755 index 00000000000..2cc6c8a8870 --- /dev/null +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -0,0 +1,226 @@ +# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TEDotProductAttention, + TELayerNormColumnParallelLinear, + TENorm, + TERowParallelLinear, +) +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.gpt.moe_module_specs import ( + get_inference_optimized_moe_spec, + get_moe_module_spec, +) +from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules +from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules +from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules +from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules +from megatron.core.ssm.mlp_layer import MLPLayer +from megatron.core.tensor_parallel import ( + InferenceColumnParallelLinear, + InferenceLayerNormColumnParallelLinear, + InferenceRowParallelLinear, +) +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionBlock, + MultiTokenPredictionBlockSubmodules, + MultiTokenPredictionLayer, + MultiTokenPredictionLayerSubmodules, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import ( + MoETransformerLayer, + TransformerLayer, + TransformerLayerSubmodules, +) + +# This should be private and should not be used outside of this file. +moe = get_moe_module_spec( + use_te=True, + num_experts=8, # Can be any positive integer (must not be None). + moe_grouped_gemm=True, +) + +# Inference-optimized MoE spec +moe_inference = get_inference_optimized_moe_spec() + + +# MTP block spec - provides norms and projection only. +# Inner layers are built by MultiTokenPredictionLayer using nested HybridStack +_hybrid_mtp_block_spec = ModuleSpec( + module=MultiTokenPredictionBlock, + submodules=MultiTokenPredictionBlockSubmodules( + layer_specs=[ + ModuleSpec( + module=MultiTokenPredictionLayer, + submodules=MultiTokenPredictionLayerSubmodules( + enorm=TENorm, + hnorm=TENorm, + eh_proj=TEColumnParallelLinear, + mtp_model_layer=None, # Built via pattern + mamba_submodules + layer_norm=TENorm, + ), + ) + ] + ), +) + + +hybrid_stack_spec = ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + mamba_layer=ModuleSpec( + module=MambaLayer, + submodules=MambaLayerSubmodules( + mixer=ModuleSpec( + module=MambaMixer, + submodules=MambaMixerSubmodules( + in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear + ), + ), + mamba_bda=get_bias_dropout_add, + ), + ), + gdn_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=GatedDeltaNet, + submodules=GatedDeltaNetSubmodules( + in_proj=TELayerNormColumnParallelLinear, + out_norm=TENorm, + out_proj=TERowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), + # Started with spec from gpt_layer_specs.py (with MLP removed) + # Using the TE spec because we had problems getting the non-TE spec + # working + attention_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=SelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=SelfAttentionSubmodules( + linear_qkv=TELayerNormColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=TERowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), + # Started with spec from gpt_layer_specs.py + # Using the TE spec because we had problems getting the non-TE spec + # working + mlp_layer=ModuleSpec( + module=MLPLayer, + submodules=TransformerLayerSubmodules( + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear + ), + ), + mlp_bda=get_bias_dropout_add, + ), + ), + moe_layer=ModuleSpec( + module=MoETransformerLayer, + submodules=TransformerLayerSubmodules( + pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add + ), + ), + mtp_block_spec=_hybrid_mtp_block_spec, + ), +) + + +hybrid_inference_stack_spec = ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + mamba_layer=ModuleSpec( + module=MambaLayer, + submodules=MambaLayerSubmodules( + mixer=ModuleSpec( + module=MambaMixer, + submodules=MambaMixerSubmodules( + in_proj=InferenceLayerNormColumnParallelLinear, + out_proj=InferenceRowParallelLinear, + ), + ), + mamba_bda=get_bias_dropout_add, + ), + ), + # Started with spec from gpt_layer_specs.py (with MLP removed) + # Using the TE spec because we had problems getting the non-TE spec + # working + attention_layer=ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=SelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=SelfAttentionSubmodules( + linear_qkv=InferenceLayerNormColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=InferenceRowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + ), + ), + # Started with spec from gpt_layer_specs.py + # Using the TE spec because we had problems getting the non-TE spec + # working + mlp_layer=ModuleSpec( + module=MLPLayer, + submodules=TransformerLayerSubmodules( + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=InferenceLayerNormColumnParallelLinear, + linear_fc2=InferenceRowParallelLinear, + ), + ), + mlp_bda=get_bias_dropout_add, + ), + ), + moe_layer=ModuleSpec( + # Use inference-optimized MoE layer for end-to-end CUDA graph support + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + pre_mlp_layernorm=TENorm, mlp=moe_inference, mlp_bda=get_bias_dropout_add + ), + ), + mtp_block_spec=ModuleSpec( + module=MultiTokenPredictionBlock, + submodules=MultiTokenPredictionBlockSubmodules( + layer_specs=[ + ModuleSpec( + module=MultiTokenPredictionLayer, + submodules=MultiTokenPredictionLayerSubmodules( + enorm=TENorm, + hnorm=TENorm, + eh_proj=InferenceColumnParallelLinear, + mtp_model_layer=None, # Built via pattern + mamba_submodules + layer_norm=TENorm, + ), + ) + ] + ), + ), + ), +) + + +# Backward-compatible aliases +mamba_stack_spec = hybrid_stack_spec +mamba_inference_stack_spec = hybrid_inference_stack_spec diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py new file mode 100644 index 00000000000..cbdd028d9ba --- /dev/null +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -0,0 +1,492 @@ +# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. + +import logging +from typing import 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.inference.contexts import BaseInferenceContext +from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding +from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding +from megatron.core.models.common.language_module.language_module import LanguageModule +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.quantization.utils import get_quant_config_or_none +from megatron.core.tensor_parallel import gather_from_sequence_parallel_region +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.enums import ModelType +from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionBlock, + mtp_on_this_rank, + process_mtp_loss, +) +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.utils import ( + WrappedTensor, + deprecate_inference_params, + is_using_quantization_scales, + log_single_rank, +) + +logger = logging.getLogger(__name__) + + +class HybridModel(LanguageModule): + """Hybrid language model. + + Args: + config (TransformerConfig): Model config + hybrid_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types + vocab_size (int): Vocabulary size + max_sequence_length (int): maximum size of sequence. + This is used for positional embedding + hybrid_layer_pattern (str): Unified hybrid layer pattern with optional MTP and + pipeline stage boundaries. + Format: "///..." + The main pattern may contain "|" to define pipeline stage boundaries. + Examples: + - "M*M*" -> main decoder only, no MTP + - "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths + - "M-M-|M-M*-|M-M-|M-M*-" -> 4 pipeline segments + hybrid_attention_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead. + If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be + generated from the ratio with a deprecation warning. + hybrid_mlp_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead. + If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be + generated from the ratio with a deprecation warning. + hybrid_override_pattern (str, optional): Deprecated. Use hybrid_layer_pattern instead. + If set and hybrid_layer_pattern is None, the value is copied to hybrid_layer_pattern + with a deprecation warning. + pre_process (bool, optional): Include embedding layer + (used with pipeline parallelism). Defaults to True. + post_process (bool, optional): Include an output layer (used with pipeline parallelism). + Defaults to True. + fp16_lm_cross_entropy (bool, optional): Defaults to False. + parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor + parallel ranks. Defaults to True. + share_embeddings_and_output_weights (bool, optional): When True, input embeddings and + output logit weights are shared. Defaults to False. + position_embedding_type (Literal[learned_absolute,rope,none], optional): Position + embedding type. Defaults to 'none'. + rotary_percent (float, optional): Percent of rotary dimension to use for rotary position + embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 1.0. + rotary_base (int, optional): Base period for rotary position embeddings. Ignored unless + position_embedding_type is 'rope'. Defaults to 10000. + seq_len_interpolation_factor (Optional[float], optional): scale of linearly + interpolating RoPE for longer sequences. The value must be a float larger than 1.0. + Defaults to None. + pg_collection (ProcessGroupCollection, optional): Model communication process groups. + vp_stage (Optional[int], optional): Virtual pipeline stage index. Defaults to None. + """ + + def __init__( + self, + config: TransformerConfig, + hybrid_stack_spec: ModuleSpec, + vocab_size: int, + max_sequence_length: int, + hybrid_layer_pattern: Optional[str] = None, + hybrid_attention_ratio: Optional[float] = None, + hybrid_mlp_ratio: Optional[float] = None, + hybrid_override_pattern: Optional[str] = None, + pre_process: bool = True, + post_process: bool = True, + fp16_lm_cross_entropy: bool = False, + parallel_output: bool = True, + share_embeddings_and_output_weights: bool = False, + # Mamba with no attention has no need for position embeddings, so none is default + position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'none', + rotary_percent: float = 1.0, + rotary_base: int = 10000, + scatter_embedding_sequence_parallel: bool = True, + seq_len_interpolation_factor: Optional[float] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + vp_stage: Optional[int] = None, + ) -> None: + super().__init__(config=config, pg_collection=pg_collection) + + if has_config_logger_enabled(config): + log_config_to_disk(config, locals(), prefix=type(self).__name__) + + if self.config.use_mup and not getattr(HybridModel, "mup_warning_printed", False): + log_single_rank( + logger, + logging.WARNING, + "MuP for HybridModel is experimental and not fully validated yet.", + ) + HybridModel.mup_warning_printed = True + + self.hybrid_stack_spec: ModuleSpec = hybrid_stack_spec + self.vocab_size = vocab_size + self.max_sequence_length = max_sequence_length + self.hybrid_layer_pattern = hybrid_layer_pattern + self.pre_process = pre_process + self.post_process = post_process + self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.parallel_output = parallel_output + self.share_embeddings_and_output_weights = share_embeddings_and_output_weights + self.position_embedding_type = position_embedding_type + self.vp_stage = vp_stage + + # Backward compatibility for deprecated hybrid parameters + if hybrid_override_pattern is not None: + if self.hybrid_layer_pattern is None: + log_single_rank( + logger, + logging.WARNING, + "hybrid_override_pattern has been deprecated. " + "Use hybrid_layer_pattern instead.", + ) + self.hybrid_layer_pattern = hybrid_override_pattern + else: + raise ValueError( + "hybrid_override_pattern and hybrid_layer_pattern cannot both be set. " + "hybrid_override_pattern has been deprecated; use hybrid_layer_pattern instead." + ) + if (hybrid_attention_ratio is not None and hybrid_attention_ratio > 0.0) or ( + hybrid_mlp_ratio is not None and hybrid_mlp_ratio > 0.0 + ): + if hybrid_layer_pattern is not None: + raise ValueError( + "hybrid_layer_pattern cannot be used together with " + "hybrid_attention_ratio or hybrid_mlp_ratio. " + "These ratios have been deprecated; use hybrid_layer_pattern alone." + ) + log_single_rank( + logger, + logging.WARNING, + "hybrid_attention_ratio and hybrid_mlp_ratio have been deprecated. " + "Use hybrid_layer_pattern instead.", + ) + if self.hybrid_layer_pattern is None: + from megatron.core.models.hybrid.hybrid_layer_allocation import pattern_from_ratios + + attn_ratio = hybrid_attention_ratio if hybrid_attention_ratio else 0.0 + mlp_ratio = hybrid_mlp_ratio if hybrid_mlp_ratio else 0.0 + self.hybrid_layer_pattern = pattern_from_ratios( + config.num_layers, attn_ratio, mlp_ratio + ) + + # 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 ( + parse_hybrid_pattern, + select_pipeline_segment, + ) + + parsed = parse_hybrid_pattern(self.hybrid_layer_pattern) + self.mtp_pattern = parsed.mtp_pattern + self.mtp_num_depths = parsed.mtp_num_depths + + 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, + ) + + # Determine if MTP is needed (based on pattern parsing) + self.mtp_process = ( + self.mtp_pattern is not None + and self.mtp_num_depths > 0 + # The following forces MTP to be on the final pipeline stage. It might be more optimal + # to split the hybrid layer pattern into pipeline stages before parsing the pattern for + # the current pipeline stage. This could also enable MTP standalone (MTP in a pipeline + # stage separate from loss) to be supported in the hybrid model. + and mtp_on_this_rank(self.config, ignore_virtual=False, vp_stage=self.vp_stage) + ) + + # megatron core pipelining currently depends on model type + # TODO: remove this dependency ? + self.model_type = ModelType.encoder_or_decoder + + if self.pre_process or self.mtp_process: + self.embedding = LanguageModelEmbedding( + config=self.config, + vocab_size=self.vocab_size, + max_sequence_length=self.max_sequence_length, + position_embedding_type=position_embedding_type, + scatter_to_sequence_parallel=scatter_embedding_sequence_parallel, + tp_group=self.pg_collection.tp, + ) + + if self.position_embedding_type == 'rope': + self.rotary_pos_emb = RotaryEmbedding( + kv_channels=self.config.kv_channels, + rotary_percent=rotary_percent, + seq_len_interpolation_factor=seq_len_interpolation_factor, + rotary_base=rotary_base, + use_cpu_initialization=self.config.use_cpu_initialization, + cp_group=self.pg_collection.cp, + ) + + self.decoder = build_module( + hybrid_stack_spec, + self.config, + pre_process=self.pre_process, + layer_type_list=layer_type_list, + pp_layer_offset=layer_offset, + post_process=self.post_process, + dtype=config.params_dtype, + pg_collection=self.pg_collection, + ) + + # MTP block - uses mtp_block_spec from hybrid_stack_spec.submodules + if self.mtp_process: + mamba_submodules = hybrid_stack_spec.submodules + mtp_block_spec = mamba_submodules.mtp_block_spec + assert mtp_block_spec is not None, ( + "MTP pattern specified but mtp_block_spec is None in hybrid_stack_spec.submodules. " + "Ensure hybrid_stack_spec includes mtp_block_spec for MTP support." + ) + + self.mtp = MultiTokenPredictionBlock( + config=self.config, + spec=mtp_block_spec, + pg_collection=self.pg_collection, + vp_stage=self.vp_stage, + mtp_layer_pattern=self.mtp_pattern, + mtp_num_depths=self.mtp_num_depths, + mamba_submodules=mamba_submodules, + ) + + # Output + if post_process or self.mtp_process: + self.output_layer = tensor_parallel.ColumnParallelLinear( + config.hidden_size, + self.vocab_size, + config=config, + init_method=( + config.embedding_init_method + if config.use_mup and not self.share_embeddings_and_output_weights + else config.init_method + ), + bias=False, + skip_bias_add=False, + gather_output=not self.parallel_output, + skip_weight_param_allocation=self.pre_process + and self.share_embeddings_and_output_weights, + tp_group=self.pg_collection.tp, + ) + + if self.pre_process or self.post_process or self.mtp_process: + self.setup_embeddings_and_output_layer() + + for name, module in self.named_modules(): + if hasattr(module, 'finish_init'): + quant_config = get_quant_config_or_none(name, self.config.quant_recipe) + module.finish_init(quant_config) + + def set_input_tensor(self, input_tensor: Tensor) -> None: + """Sets input tensor to the model. + + See megatron.model.transformer.set_input_tensor() + + Args: + input_tensor (Tensor): Sets the input tensor for the model. + """ + # This is usually handled in schedules.py but some inference code still + # gives us non-lists or None + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + + assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert' + self.decoder.set_input_tensor(input_tensor[0]) + + 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, + is_spec_decode: Optional[bool] = 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 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. + + inference_context = deprecate_inference_params(inference_context, inference_params) + + 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" + + # Decoder embedding. + if decoder_input is not None: + pass + elif self.pre_process: + 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 + if ( + in_inference_mode + and inference_context.is_dynamic_batching() + and is_using_quantization_scales(self.config) + ): + decoder_input[inference_context.padding_slice] = 0.0 + else: + # intermediate stage of pipeline + # decoder will get hidden_states from encoder.input_tensor + decoder_input = None + + rotary_pos_emb = None + if self.position_embedding_type == 'rope': + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, self.decoder, decoder_input, self.config, packed_seq_params + ) + 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. + 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" + + # 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, + ) + + 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. + 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: + hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + attention_mask=attention_mask, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + embedding=self.embedding, + ) + + if not self.post_process: + return hidden_states + + 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: + self._decoder_hidden_states_cache = hidden_states + else: + hidden_states = process_mtp_loss( + hidden_states=hidden_states, + labels=labels, + loss_mask=loss_mask, + output_layer=self.output_layer, + output_weight=output_weight, + runtime_gather_output=runtime_gather_output, + is_training=self.training, + compute_language_model_loss=self.compute_language_model_loss, + config=self.config, + cp_group=self.pg_collection.cp, + packed_seq_params=packed_seq_params, + scale_logits_fn=self._scale_logits if self.config.use_mup else None, + ) + sequence_parallel_override = False + 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. + hidden_states = gather_from_sequence_parallel_region( + hidden_states, group=self.pg_collection.tp + ) + self.output_layer.sequence_parallel = False + sequence_parallel_override = True + + # Reshape [S, B, H] (with B=1) to [1, S, H] for logit extraction, + # then back to [S', B, H] for the output layer. + reshaped = hidden_states.squeeze(1).unsqueeze(0) + hidden_states = inference_context.last_token_logits(reshaped).unsqueeze(1) + + logits, _ = self.output_layer( + hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output + ) + logits = self._scale_logits(logits) + + # Restore sequence parallel execution to the output layer if necessary. + if sequence_parallel_override: + assert ( + in_inference_mode + and inference_context.is_dynamic_batching() + and inference_context.config.materialize_only_last_token_logits + ) + 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 + + +class MambaModel(HybridModel): + """Backward-compatible wrapper that accepts the deprecated mamba_stack_spec kwarg.""" + + def __init__(self, *args, mamba_stack_spec: ModuleSpec = None, **kwargs): + log_single_rank( + logger, logging.WARNING, "MambaModel has been deprecated. Use HybridModel instead." + ) + if mamba_stack_spec is not None: + if 'hybrid_stack_spec' in kwargs or (args and len(args) >= 2): + raise ValueError( + "Cannot specify both hybrid_stack_spec and mamba_stack_spec. " + "mamba_stack_spec has been deprecated; use hybrid_stack_spec instead." + ) + kwargs['hybrid_stack_spec'] = mamba_stack_spec + super().__init__(*args, **kwargs) diff --git a/megatron/core/models/mamba/__init__.py b/megatron/core/models/mamba/__init__.py index 5aaf8524018..2a52cedd9b5 100644 --- a/megatron/core/models/mamba/__init__.py +++ b/megatron/core/models/mamba/__init__.py @@ -1,2 +1,5 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -from .mamba_model import MambaModel +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. + +# Backward-compatible re-exports. The canonical location is now +# megatron.core.models.hybrid. +from megatron.core.models.hybrid.hybrid_model import HybridModel, MambaModel diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py old mode 100755 new mode 100644 index 48f25bdbab9..5fb9e49a0dd --- a/megatron/core/models/mamba/mamba_layer_specs.py +++ b/megatron/core/models/mamba/mamba_layer_specs.py @@ -1,221 +1,5 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. -from megatron.core.extensions.transformer_engine import ( - TEColumnParallelLinear, - TEDotProductAttention, - TELayerNormColumnParallelLinear, - TENorm, - TERowParallelLinear, -) -from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add -from megatron.core.models.gpt.moe_module_specs import ( - get_inference_optimized_moe_spec, - get_moe_module_spec, -) -from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules -from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules -from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules -from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules -from megatron.core.ssm.mlp_layer import MLPLayer -from megatron.core.tensor_parallel import ( - InferenceColumnParallelLinear, - InferenceLayerNormColumnParallelLinear, - InferenceRowParallelLinear, -) -from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules -from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.mlp import MLP, MLPSubmodules -from megatron.core.transformer.multi_token_prediction import ( - MultiTokenPredictionBlock, - MultiTokenPredictionBlockSubmodules, - MultiTokenPredictionLayer, - MultiTokenPredictionLayerSubmodules, -) -from megatron.core.transformer.spec_utils import ModuleSpec -from megatron.core.transformer.transformer_layer import ( - MoETransformerLayer, - TransformerLayer, - TransformerLayerSubmodules, -) - -# This should be private and should not be used outside of this file. -moe = get_moe_module_spec( - use_te=True, - num_experts=8, # Can be any positive integer (must not be None). - moe_grouped_gemm=True, -) - -# Inference-optimized MoE spec -moe_inference = get_inference_optimized_moe_spec() - - -# MTP block spec for Mamba - provides norms and projection only. -# Inner layers are built by MultiTokenPredictionLayer using nested MambaStack -_mamba_mtp_block_spec = ModuleSpec( - module=MultiTokenPredictionBlock, - submodules=MultiTokenPredictionBlockSubmodules( - layer_specs=[ - ModuleSpec( - module=MultiTokenPredictionLayer, - submodules=MultiTokenPredictionLayerSubmodules( - enorm=TENorm, - hnorm=TENorm, - eh_proj=TEColumnParallelLinear, - mtp_model_layer=None, # Built via pattern + mamba_submodules - layer_norm=TENorm, - ), - ) - ] - ), -) - - -mamba_stack_spec = ModuleSpec( - module=MambaStack, - submodules=MambaStackSubmodules( - mamba_layer=ModuleSpec( - module=MambaLayer, - submodules=MambaLayerSubmodules( - mixer=ModuleSpec( - module=MambaMixer, - submodules=MambaMixerSubmodules( - in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear - ), - ), - mamba_bda=get_bias_dropout_add, - ), - ), - gdn_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - self_attention=ModuleSpec( - module=GatedDeltaNet, - submodules=GatedDeltaNetSubmodules( - in_proj=TELayerNormColumnParallelLinear, - out_norm=TENorm, - out_proj=TERowParallelLinear, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), - # Started with spec from gpt_layer_specs.py (with MLP removed) - # Using the TE spec because we had problems getting the non-TE spec - # working - attention_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - self_attention=ModuleSpec( - module=SelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, - submodules=SelfAttentionSubmodules( - linear_qkv=TELayerNormColumnParallelLinear, - core_attention=TEDotProductAttention, - linear_proj=TERowParallelLinear, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), - # Started with spec from gpt_layer_specs.py - # Using the TE spec because we had problems getting the non-TE spec - # working - mlp_layer=ModuleSpec( - module=MLPLayer, - submodules=TransformerLayerSubmodules( - mlp=ModuleSpec( - module=MLP, - submodules=MLPSubmodules( - linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear - ), - ), - mlp_bda=get_bias_dropout_add, - ), - ), - moe_layer=ModuleSpec( - module=MoETransformerLayer, - submodules=TransformerLayerSubmodules( - pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add - ), - ), - mtp_block_spec=_mamba_mtp_block_spec, - ), -) - - -mamba_inference_stack_spec = ModuleSpec( - module=MambaStack, - submodules=MambaStackSubmodules( - mamba_layer=ModuleSpec( - module=MambaLayer, - submodules=MambaLayerSubmodules( - mixer=ModuleSpec( - module=MambaMixer, - submodules=MambaMixerSubmodules( - in_proj=InferenceLayerNormColumnParallelLinear, - out_proj=InferenceRowParallelLinear, - ), - ), - mamba_bda=get_bias_dropout_add, - ), - ), - # Started with spec from gpt_layer_specs.py (with MLP removed) - # Using the TE spec because we had problems getting the non-TE spec - # working - attention_layer=ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - self_attention=ModuleSpec( - module=SelfAttention, - params={"attn_mask_type": AttnMaskType.causal}, - submodules=SelfAttentionSubmodules( - linear_qkv=InferenceLayerNormColumnParallelLinear, - core_attention=TEDotProductAttention, - linear_proj=InferenceRowParallelLinear, - ), - ), - self_attn_bda=get_bias_dropout_add, - ), - ), - # Started with spec from gpt_layer_specs.py - # Using the TE spec because we had problems getting the non-TE spec - # working - mlp_layer=ModuleSpec( - module=MLPLayer, - submodules=TransformerLayerSubmodules( - mlp=ModuleSpec( - module=MLP, - submodules=MLPSubmodules( - linear_fc1=InferenceLayerNormColumnParallelLinear, - linear_fc2=InferenceRowParallelLinear, - ), - ), - mlp_bda=get_bias_dropout_add, - ), - ), - moe_layer=ModuleSpec( - # Use inference-optimized MoE layer for end-to-end CUDA graph support - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - pre_mlp_layernorm=TENorm, mlp=moe_inference, mlp_bda=get_bias_dropout_add - ), - ), - mtp_block_spec=ModuleSpec( - module=MultiTokenPredictionBlock, - submodules=MultiTokenPredictionBlockSubmodules( - layer_specs=[ - ModuleSpec( - module=MultiTokenPredictionLayer, - submodules=MultiTokenPredictionLayerSubmodules( - enorm=TENorm, - hnorm=TENorm, - eh_proj=InferenceColumnParallelLinear, - mtp_model_layer=None, # Built via pattern + mamba_submodules - layer_norm=TENorm, - ), - ) - ] - ), - ), - ), -) +# Backward-compatible re-export. The canonical location is now +# megatron.core.models.hybrid.hybrid_layer_specs. +from megatron.core.models.hybrid.hybrid_layer_specs import * # noqa: F401,F403 diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py index 444d6b86398..a3abb80890c 100644 --- a/megatron/core/models/mamba/mamba_model.py +++ b/megatron/core/models/mamba/mamba_model.py @@ -1,475 +1,6 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. -import logging -from typing import 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.inference.contexts import BaseInferenceContext -from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding -from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding -from megatron.core.models.common.language_module.language_module import LanguageModule -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.quantization.utils import get_quant_config_or_none -from megatron.core.tensor_parallel import gather_from_sequence_parallel_region -from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.enums import ModelType -from megatron.core.transformer.multi_token_prediction import ( - MultiTokenPredictionBlock, - mtp_on_this_rank, - process_mtp_loss, -) -from megatron.core.transformer.spec_utils import ModuleSpec, build_module -from megatron.core.utils import ( - WrappedTensor, - deprecate_inference_params, - is_using_quantization_scales, - log_single_rank, -) - -logger = logging.getLogger(__name__) - - -class MambaModel(LanguageModule): - """Mamba language model. - - Args: - config (TransformerConfig): Model config - mamba_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types - vocab_size (int): Vocabulary size - max_sequence_length (int): maximum size of sequence. - This is used for positional embedding - hybrid_layer_pattern (str): Unified hybrid layer pattern with optional MTP and - pipeline stage boundaries. - Format: "///..." - The main pattern may contain "|" to define pipeline stage boundaries. - Examples: - - "M*M*" -> main decoder only, no MTP - - "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths - - "M-M-|M-M*-|M-M-|M-M*-" -> 4 pipeline segments - hybrid_attention_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead. - If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be - generated from the ratio with a deprecation warning. - hybrid_mlp_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead. - If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be - generated from the ratio with a deprecation warning. - hybrid_override_pattern (str, optional): Deprecated. Use hybrid_layer_pattern instead. - If set and hybrid_layer_pattern is None, the value is copied to hybrid_layer_pattern - with a deprecation warning. - pre_process (bool, optional): Include embedding layer - (used with pipeline parallelism). Defaults to True. - post_process (bool, optional): Include an output layer (used with pipeline parallelism). - Defaults to True. - fp16_lm_cross_entropy (bool, optional): Defaults to False. - parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor - parallel ranks. Defaults to True. - share_embeddings_and_output_weights (bool, optional): When True, input embeddings and - output logit weights are shared. Defaults to False. - position_embedding_type (Literal[learned_absolute,rope,none], optional): Position - embedding type. Defaults to 'none'. - rotary_percent (float, optional): Percent of rotary dimension to use for rotary position - embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 1.0. - rotary_base (int, optional): Base period for rotary position embeddings. Ignored unless - position_embedding_type is 'rope'. Defaults to 10000. - seq_len_interpolation_factor (Optional[float], optional): scale of linearly - interpolating RoPE for longer sequences. The value must be a float larger than 1.0. - Defaults to None. - pg_collection (ProcessGroupCollection, optional): Model communication process groups. - vp_stage (Optional[int], optional): Virtual pipeline stage index. Defaults to None. - """ - - def __init__( - self, - config: TransformerConfig, - mamba_stack_spec: ModuleSpec, - vocab_size: int, - max_sequence_length: int, - hybrid_layer_pattern: Optional[str] = None, - hybrid_attention_ratio: Optional[float] = None, - hybrid_mlp_ratio: Optional[float] = None, - hybrid_override_pattern: Optional[str] = None, - pre_process: bool = True, - post_process: bool = True, - fp16_lm_cross_entropy: bool = False, - parallel_output: bool = True, - share_embeddings_and_output_weights: bool = False, - # Mamba with no attention has no need for position embeddings, so none is default - position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'none', - rotary_percent: float = 1.0, - rotary_base: int = 10000, - scatter_embedding_sequence_parallel: bool = True, - seq_len_interpolation_factor: Optional[float] = None, - pg_collection: Optional[ProcessGroupCollection] = None, - vp_stage: Optional[int] = None, - ) -> None: - super().__init__(config=config, pg_collection=pg_collection) - - if has_config_logger_enabled(config): - log_config_to_disk(config, locals(), prefix=type(self).__name__) - - if self.config.use_mup and not getattr(MambaModel, "mup_warning_printed", False): - log_single_rank( - logger, - logging.WARNING, - "MuP for MambaModel is experimental and not fully validated yet.", - ) - MambaModel.mup_warning_printed = True - - self.mamba_stack_spec: ModuleSpec = mamba_stack_spec - self.vocab_size = vocab_size - self.max_sequence_length = max_sequence_length - self.hybrid_layer_pattern = hybrid_layer_pattern - self.pre_process = pre_process - self.post_process = post_process - self.fp16_lm_cross_entropy = fp16_lm_cross_entropy - self.parallel_output = parallel_output - self.share_embeddings_and_output_weights = share_embeddings_and_output_weights - self.position_embedding_type = position_embedding_type - self.vp_stage = vp_stage - - # Backward compatibility for deprecated hybrid parameters - if hybrid_override_pattern is not None: - if self.hybrid_layer_pattern is None: - log_single_rank( - logger, - logging.WARNING, - "hybrid_override_pattern has been deprecated. " - "Use hybrid_layer_pattern instead.", - ) - self.hybrid_layer_pattern = hybrid_override_pattern - else: - raise ValueError( - "hybrid_override_pattern and hybrid_layer_pattern cannot both be set. " - "hybrid_override_pattern has been deprecated; use hybrid_layer_pattern instead." - ) - if (hybrid_attention_ratio is not None and hybrid_attention_ratio > 0.0) or ( - hybrid_mlp_ratio is not None and hybrid_mlp_ratio > 0.0 - ): - if hybrid_layer_pattern is not None: - raise ValueError( - "hybrid_layer_pattern cannot be used together with " - "hybrid_attention_ratio or hybrid_mlp_ratio. " - "These ratios have been deprecated; use hybrid_layer_pattern alone." - ) - log_single_rank( - logger, - logging.WARNING, - "hybrid_attention_ratio and hybrid_mlp_ratio have been deprecated. " - "Use hybrid_layer_pattern instead.", - ) - if self.hybrid_layer_pattern is None: - from megatron.core.ssm.mamba_hybrid_layer_allocation import pattern_from_ratios - - attn_ratio = hybrid_attention_ratio if hybrid_attention_ratio else 0.0 - mlp_ratio = hybrid_mlp_ratio if hybrid_mlp_ratio else 0.0 - self.hybrid_layer_pattern = pattern_from_ratios( - config.num_layers, attn_ratio, mlp_ratio - ) - - # Parse unified pattern to extract main and MTP components, and - # determine the pipeline segment for this model instance. - from megatron.core.ssm.mamba_hybrid_layer_allocation import ( - parse_hybrid_pattern, - select_pipeline_segment, - ) - - parsed = parse_hybrid_pattern(self.hybrid_layer_pattern) - self.mtp_pattern = parsed.mtp_pattern - self.mtp_num_depths = parsed.mtp_num_depths - - 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, - ) - - # Determine if MTP is needed (based on pattern parsing) - self.mtp_process = ( - self.mtp_pattern is not None - and self.mtp_num_depths > 0 - # The following forces MTP to be on the final pipeline stage. It might be more optimal - # to split the hybrid layer pattern into pipeline stages before parsing the pattern for - # the current pipeline stage. This could also enable MTP standalone (MTP in a pipeline - # stage separate from loss) to be supported in the hybrid model. - and mtp_on_this_rank(self.config, ignore_virtual=False, vp_stage=self.vp_stage) - ) - - # megatron core pipelining currently depends on model type - # TODO: remove this dependency ? - self.model_type = ModelType.encoder_or_decoder - - if self.pre_process or self.mtp_process: - self.embedding = LanguageModelEmbedding( - config=self.config, - vocab_size=self.vocab_size, - max_sequence_length=self.max_sequence_length, - position_embedding_type=position_embedding_type, - scatter_to_sequence_parallel=scatter_embedding_sequence_parallel, - tp_group=self.pg_collection.tp, - ) - - if self.position_embedding_type == 'rope': - self.rotary_pos_emb = RotaryEmbedding( - kv_channels=self.config.kv_channels, - rotary_percent=rotary_percent, - seq_len_interpolation_factor=seq_len_interpolation_factor, - rotary_base=rotary_base, - use_cpu_initialization=self.config.use_cpu_initialization, - cp_group=self.pg_collection.cp, - ) - - self.decoder = build_module( - mamba_stack_spec, - self.config, - pre_process=self.pre_process, - layer_type_list=layer_type_list, - pp_layer_offset=layer_offset, - post_process=self.post_process, - dtype=config.params_dtype, - pg_collection=self.pg_collection, - ) - - # MTP block - uses mtp_block_spec from mamba_stack_spec.submodules - if self.mtp_process: - mamba_submodules = mamba_stack_spec.submodules - mtp_block_spec = mamba_submodules.mtp_block_spec - assert mtp_block_spec is not None, ( - "MTP pattern specified but mtp_block_spec is None in mamba_stack_spec.submodules. " - "Ensure mamba_stack_spec includes mtp_block_spec for MTP support." - ) - - self.mtp = MultiTokenPredictionBlock( - config=self.config, - spec=mtp_block_spec, - pg_collection=self.pg_collection, - vp_stage=self.vp_stage, - mtp_layer_pattern=self.mtp_pattern, - mtp_num_depths=self.mtp_num_depths, - mamba_submodules=mamba_submodules, - ) - - # Output - if post_process or self.mtp_process: - self.output_layer = tensor_parallel.ColumnParallelLinear( - config.hidden_size, - self.vocab_size, - config=config, - init_method=( - config.embedding_init_method - if config.use_mup and not self.share_embeddings_and_output_weights - else config.init_method - ), - bias=False, - skip_bias_add=False, - gather_output=not self.parallel_output, - skip_weight_param_allocation=self.pre_process - and self.share_embeddings_and_output_weights, - tp_group=self.pg_collection.tp, - ) - - if self.pre_process or self.post_process or self.mtp_process: - self.setup_embeddings_and_output_layer() - - for name, module in self.named_modules(): - if hasattr(module, 'finish_init'): - quant_config = get_quant_config_or_none(name, self.config.quant_recipe) - module.finish_init(quant_config) - - def set_input_tensor(self, input_tensor: Tensor) -> None: - """Sets input tensor to the model. - - See megatron.model.transformer.set_input_tensor() - - Args: - input_tensor (Tensor): Sets the input tensor for the model. - """ - # This is usually handled in schedules.py but some inference code still - # gives us non-lists or None - if not isinstance(input_tensor, list): - input_tensor = [input_tensor] - - assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert' - self.decoder.set_input_tensor(input_tensor[0]) - - 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, - is_spec_decode: Optional[bool] = None, - ) -> Tensor: - """Forward function of the Mamba 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 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. - - inference_context = deprecate_inference_params(inference_context, inference_params) - - 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" - - # Decoder embedding. - if decoder_input is not None: - pass - elif self.pre_process: - 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 - if ( - in_inference_mode - and inference_context.is_dynamic_batching() - and is_using_quantization_scales(self.config) - ): - decoder_input[inference_context.padding_slice] = 0.0 - else: - # intermediate stage of pipeline - # decoder will get hidden_states from encoder.input_tensor - decoder_input = None - - rotary_pos_emb = None - if self.position_embedding_type == 'rope': - rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( - inference_context, self.decoder, decoder_input, self.config, packed_seq_params - ) - 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 (MambaBlock) to delete the - # reference held by this caller function, enabling early garbage collection - # for 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" - - # 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, - ) - - 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. - 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: - hidden_states = self.mtp( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - attention_mask=attention_mask, - inference_params=inference_params, - rotary_pos_emb=rotary_pos_emb, - packed_seq_params=packed_seq_params, - embedding=self.embedding, - ) - - if not self.post_process: - return hidden_states - - 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: - self._decoder_hidden_states_cache = hidden_states - else: - hidden_states = process_mtp_loss( - hidden_states=hidden_states, - labels=labels, - loss_mask=loss_mask, - output_layer=self.output_layer, - output_weight=output_weight, - runtime_gather_output=runtime_gather_output, - is_training=self.training, - compute_language_model_loss=self.compute_language_model_loss, - config=self.config, - cp_group=self.pg_collection.cp, - packed_seq_params=packed_seq_params, - scale_logits_fn=self._scale_logits if self.config.use_mup else None, - ) - sequence_parallel_override = False - 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. - hidden_states = gather_from_sequence_parallel_region( - hidden_states, group=self.pg_collection.tp - ) - self.output_layer.sequence_parallel = False - sequence_parallel_override = True - - # Reshape [S, B, H] (with B=1) to [1, S, H] for logit extraction, - # then back to [S', B, H] for the output layer. - reshaped = hidden_states.squeeze(1).unsqueeze(0) - hidden_states = inference_context.last_token_logits(reshaped).unsqueeze(1) - - logits, _ = self.output_layer( - hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output - ) - logits = self._scale_logits(logits) - - # Restore sequence parallel execution to the output layer if necessary. - if sequence_parallel_override: - assert ( - in_inference_mode - and inference_context.is_dynamic_batching() - and inference_context.config.materialize_only_last_token_logits - ) - 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 +# Backward-compatible re-export. The canonical location is now +# megatron.core.models.hybrid.hybrid_model. +from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 +from megatron.core.models.hybrid.hybrid_model import HybridModel, MambaModel # noqa: F401 diff --git a/megatron/core/models/multimodal/llava_model.py b/megatron/core/models/multimodal/llava_model.py index 86ce04521a7..70f58216d20 100644 --- a/megatron/core/models/multimodal/llava_model.py +++ b/megatron/core/models/multimodal/llava_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import logging from collections import namedtuple from functools import partial @@ -11,7 +11,7 @@ from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.models.gpt import GPTModel -from megatron.core.models.mamba import MambaModel +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.models.vision.clip_vit_model import CLIPViTModel, get_num_image_embeddings from megatron.core.models.vision.multimodal_projector import MultimodalProjector from megatron.core.models.vision.radio import RADIOViTModel @@ -196,9 +196,9 @@ def __init__( ) self.language_model = build_hf_model(language_transformer_config) elif language_model_type.startswith('nemotron5-hybrid'): - self.language_model = MambaModel( + self.language_model = HybridModel( config=language_transformer_config, - mamba_stack_spec=language_transformer_layer_spec, + hybrid_stack_spec=language_transformer_layer_spec, vocab_size=language_vocab_size, max_sequence_length=language_max_sequence_length, parallel_output=parallel_output, diff --git a/megatron/core/post_training/modelopt/hybrid/__init__.py b/megatron/core/post_training/modelopt/hybrid/__init__.py new file mode 100644 index 00000000000..e76ed74857b --- /dev/null +++ b/megatron/core/post_training/modelopt/hybrid/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. diff --git a/megatron/core/post_training/modelopt/hybrid/model_specs.py b/megatron/core/post_training/modelopt/hybrid/model_specs.py new file mode 100755 index 00000000000..59a509be718 --- /dev/null +++ b/megatron/core/post_training/modelopt/hybrid/model_specs.py @@ -0,0 +1,145 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.extensions.transformer_engine import TEDotProductAttention +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec +from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules +from megatron.core.post_training.modelopt.layers import Norm +from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules +from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules +from megatron.core.transformer.dot_product_attention import DotProductAttention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules + + +# Use this spec for ModelOpt PTQ and TensorRT-LLM export +def get_hybrid_stack_modelopt_spec( + local_core_attention: bool = False, + remap_te_layernorm: bool = False, + use_default_te_spec: bool = False, +) -> ModuleSpec: + """Get the hybrid stack spec for ModelOpt PTQ and TensorRT-LLM export. + + When use_default_te_spec=False (default), this is the native local spec with TENorm + from Transformer-Engine for the layernorm implementation (since FusedLayerNorm from + apex has stopped supporting RMSNorm needed by llama). The remap_te_layernorm flag + can be used to add sharded state_dict key remapping for TE-compatible checkpoint + saving/loading. + + When use_default_te_spec=True, this returns the standard hybrid_stack_spec from + hybrid_layer_specs.py which uses full TE modules (TELayerNormColumnParallelLinear, + TERowParallelLinear, TEDotProductAttention, TENorm, moe_grouped_gemm=True). + + + Args: + local_core_attention: whether to use local DotProductAttention + (only for use_default_te_spec=False) + remap_te_layernorm: whether to perform sharded state_dict prefix mapping + on layernorm (only for use_default_te_spec=False) + use_default_te_spec: whether to use the default Transformer-Engine spec + """ + if use_default_te_spec: + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + + return hybrid_stack_spec + + return _get_hybrid_stack_local_spec( + local_core_attention=local_core_attention, remap_te_layernorm=remap_te_layernorm + ) + + +# Backward-compatible alias +get_mamba_stack_modelopt_spec = get_hybrid_stack_modelopt_spec + + +def _get_hybrid_stack_local_spec( + local_core_attention: bool = False, remap_te_layernorm: bool = False +) -> ModuleSpec: + """Get the hybrid stack spec with local (non-TE) modules. + + This is essentially the native local spec except for the layernorm implementation + is using TENorm from Transformer-Engine. + """ + mamba_state_dict_keys_map = {} + transformer_state_dict_keys_map = {} + if remap_te_layernorm: + mamba_state_dict_keys_map = {'norm.': 'mixer.in_proj.layer_norm_'} + transformer_state_dict_keys_map = { + 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_', + 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_', + } + + mamba_layer = ModuleSpec( + module=MambaLayer, + submodules=MambaLayerSubmodules( + norm=Norm, + mixer=ModuleSpec( + module=MambaMixer, + submodules=MambaMixerSubmodules( + in_proj=ColumnParallelLinear, out_proj=RowParallelLinear + ), + ), + mamba_bda=get_bias_dropout_add, + sharded_state_dict_keys_map=mamba_state_dict_keys_map, + ), + ) + + attn_mask_type = AttnMaskType.causal + core_attention = DotProductAttention if local_core_attention else TEDotProductAttention + attention_layer = ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=Norm, + self_attention=ModuleSpec( + module=SelfAttention, + params={"attn_mask_type": attn_mask_type}, + submodules=SelfAttentionSubmodules( + linear_qkv=ColumnParallelLinear, + core_attention=core_attention, + linear_proj=RowParallelLinear, + ), + ), + self_attn_bda=get_bias_dropout_add, + sharded_state_dict_keys_map=transformer_state_dict_keys_map, + ), + ) + + mlp_layer = ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + pre_mlp_layernorm=Norm, + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear + ), + ), + mlp_bda=get_bias_dropout_add, + sharded_state_dict_keys_map=transformer_state_dict_keys_map, + ), + ) + + moe_layer = ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + pre_mlp_layernorm=Norm, + mlp=get_moe_module_spec( + use_te=False, num_experts=8, moe_grouped_gemm=False # Can be anything non None + ), + mlp_bda=get_bias_dropout_add, + ), + ) + + return ModuleSpec( + module=HybridStack, + submodules=HybridStackSubmodules( + mamba_layer=mamba_layer, + attention_layer=attention_layer, + mlp_layer=mlp_layer, + moe_layer=moe_layer, + ), + ) diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py index f42f3542c3d..f37273f0c31 100644 --- a/megatron/core/ssm/mamba_block.py +++ b/megatron/core/ssm/mamba_block.py @@ -1,420 +1,11 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# Copyright (c) 2024, Tri Dao, Albert Gu. - -# Some of this code was adopted from https://github.com/state-spaces/mamba/ -# This source code is licensed under the Apache license found in the -# LICENSE file in the root directory of this source tree. - -from contextlib import nullcontext -from dataclasses import dataclass -from typing import Optional, Tuple, Union - -import torch -from torch import Tensor, nn - -from megatron.core.dist_checkpointing.mapping import ShardedStateDict -from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding -from megatron.core.enums import Fp8Recipe -from megatron.core.extensions.transformer_engine import TENorm -from megatron.core.fp4_utils import get_fp4_context -from megatron.core.fp8_utils import get_fp8_context -from megatron.core.inference.contexts import BaseInferenceContext -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols as LayerSymbols -from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.enums import CudaGraphScope -from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule -from megatron.core.transformer.spec_utils import ModuleSpec, build_module -from megatron.core.transformer.transformer_layer import TransformerLayer -from megatron.core.transformer.utils import sharded_state_dict_default -from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor - - -@dataclass -class MambaStackSubmodules: - """ - A class for the module specs for the MambaStack. - """ - - mamba_layer: Union[ModuleSpec, type] = IdentityOp - gdn_layer: Union[ModuleSpec, type] = IdentityOp - attention_layer: Union[ModuleSpec, type] = IdentityOp - mlp_layer: Union[ModuleSpec, type] = IdentityOp - moe_layer: Union[ModuleSpec, type] = IdentityOp - mtp_block_spec: Optional[ModuleSpec] = None - - -class MambaStack(GraphableMegatronModule, MegatronModule): - """ - Constructor for the MambaStack class. - - Args: - config (TransformerConfig): the model configuration - submodules (MambaStackSubmodules): the submodules for the stack - pre_process (bool, optional): whether to include an embedding layer. - Defaults to True. - layer_type_list (list, optional): pre-computed list of layer type symbols for - this pipeline segment. When provided (by MambaModel), pipeline stage - selection has already been done via '|' separators in the pattern. - pp_layer_offset (int, optional): the global layer offset for this pipeline - segment. Defaults to 0. - post_layer_norm (bool, optional): whether to include a final layer norm. - Defaults to True. - post_process (bool, optional): whether to include an output layer. - Defaults to True. - device (optional): the device to use. Defaults to None. - dtype (optional): the data type to use. Defaults to None. - pg_collection (ProcessGroupCollection): the required model communication - process groups to use. - is_mtp_layer (bool, optional): whether this is an MTP layer. Defaults to False. - """ - - def __init__( - self, - config: TransformerConfig, - submodules: MambaStackSubmodules, - pre_process: bool = True, - layer_type_list: Optional[list[str]] = None, - pp_layer_offset: int = 0, - post_layer_norm: bool = True, - post_process: bool = True, - device=None, - dtype=None, - pg_collection: ProcessGroupCollection = None, - is_mtp_layer: bool = False, - ) -> None: - super().__init__(config=config) - self.pre_process = pre_process - self.post_layer_norm = post_layer_norm - self.post_process = post_process - self.is_mtp_layer = is_mtp_layer - - assert pg_collection is not None, "pg_collection must be provided for MambaStack" - - self.pp_group = pg_collection.pp - self.tp_group = pg_collection.tp - - # Required for pipeline parallel schedules - self.input_tensor = None - self.pg_collection = pg_collection - - assert layer_type_list is not None, ( - "layer_type_list must be provided. It should be pre-computed from " - "--hybrid-layer-pattern by MambaModel." - ) - self.layer_type_list = layer_type_list - - # 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) - elif self.config.fp4: - quant_init_context = get_fp4_context(self.config, i + pp_layer_offset, is_init=True) - else: - quant_init_context = nullcontext() - with quant_init_context: - if layer_type == LayerSymbols.MAMBA: - layer = build_module( - submodules.mamba_layer, - config=self.config, - layer_number=layer_number, - pp_layer_offset=pp_layer_offset, - pg_collection=pg_collection, - ) - elif layer_type == LayerSymbols.ATTENTION: - layer = build_module( - submodules.attention_layer, - config=self.config, - layer_number=layer_number, - pg_collection=pg_collection, - is_mtp_layer=is_mtp_layer, - add_layer_offset=False, - pp_layer_offset=pp_layer_offset, - ) - elif layer_type == LayerSymbols.MLP: - layer = build_module( - submodules.mlp_layer, - config=self.config, - layer_number=layer_number, - pg_collection=pg_collection, - add_layer_offset=False, - ) - elif layer_type == LayerSymbols.MOE: - layer = build_module( - submodules.moe_layer, - config=self.config, - layer_number=layer_number, - pg_collection=pg_collection, - add_layer_offset=False, - ) - elif layer_type == LayerSymbols.GDN: - layer = build_module( - submodules.gdn_layer, - config=self.config, - layer_number=layer_number, - pg_collection=pg_collection, - # Set to False as we do not want to change offset. - add_layer_offset=False, - ) - else: - assert False, "unexpected layer_type" - self.layers.append(layer) - - # Required for activation recomputation - self.num_layers_per_pipeline_rank = len(self.layers) - - if self.post_process and self.post_layer_norm: - # Final layer norm before output. - self.final_norm = TENorm( - config=self.config, - hidden_size=self.config.hidden_size, - eps=self.config.layernorm_epsilon, - ) - - def set_input_tensor(self, input_tensor: Tensor): - """Set input tensor to be used instead of forward()'s input. - - When doing pipeline parallelism the input from the previous - stage comes from communication, not from the input, so the - model's forward_step_func won't have it. This function is thus - used by internal code to bypass the input provided by the - forward_step_func""" - self.input_tensor = input_tensor - - def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]: - """ - Returns the Mamba conv and ssm states shapes per input sequence - 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: - return layer.mamba_state_shapes_per_request() - return None - - def _should_call_local_cudagraph(self, *args, **kwargs): - """ - Check if we should call the local cudagraph path. - """ - if ( - not self.training - and hasattr(self, 'cudagraph_manager') - and kwargs['attention_mask'] is None - and ( - kwargs.get('inference_context') is not None - or kwargs.get('inference_params') is not None - ) - and CudaGraphScope.full_iteration_inference in self.config.cuda_graph_scope - ): - 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): - kwargs['hidden_states'] = ( - kwargs['hidden_states'].unwrap() - if isinstance(kwargs['hidden_states'], WrappedTensor) - else kwargs['hidden_states'] - ) - return super().__call__(*args, **kwargs)[0] - return super().__call__(*args, **kwargs) - - def forward( - self, - hidden_states: Union[Tensor, WrappedTensor], - attention_mask: Tensor, - inference_context: Optional[BaseInferenceContext] = None, - rotary_pos_emb: Optional[Tensor] = None, - *, - inference_params: Optional[BaseInferenceContext] = None, - packed_seq_params: Optional[PackedSeqParams] = None, - padding_mask=None, - ): - """ - Forward function of the MambaStack class. - - It either returns the Loss values if labels are given or the - final hidden units - - Args: - hidden_states (Union[Tensor, WrappedTensor]): the input tensor. - Can be passed as a WrappedTensor during inference to avoid an obsolete - reference in the calling function. - attention_mask (Tensor): the attention mask. - inference_context (BaseInferenceContext): the inference parameters. - rotary_pos_emb (Tensor, optional): the rotary positional embeddings. - Defaults to None. - Returns: - Tensor: the output tensor. - """ - - inference_context = deprecate_inference_params(inference_context, inference_params) - - if not self.pre_process: - # See set_input_tensor() - hidden_states = self.input_tensor - - # Delete the obsolete reference to the initial input tensor if necessary - if isinstance(hidden_states, WrappedTensor): - hidden_states = hidden_states.unwrap() - - if inference_context and inference_context.is_static_batching(): - # NOTE(bnorick): match BaseInferenceContext attributes for - # mamba_ssm.utils.generation.BaseInferenceContext, - # this hack supports eval - inference_context.max_seqlen = inference_context.max_sequence_length - inference_context.seqlen_offset = inference_context.sequence_len_offset - - if ( - ( - ( - self.config.cuda_graph_impl == "local" - and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope - ) - or self.config.flash_decode - ) - and inference_context - and inference_context.is_static_batching() - and not self.training - ): - current_batch_size = hidden_states.shape[1] - sequence_len_offset = torch.tensor( - [inference_context.sequence_len_offset] * current_batch_size, - dtype=torch.int32, - device='cuda', - ) - else: - sequence_len_offset = None - - # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), - # otherwise do nothing extra at the outer level - # if we are using other fp8 recipes, then the context manager enter&exit are free - # we can wrap fp8_context within the for loop over layers, so that we can fine-grained - # control which layer will be fp8 or bf16 - use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed - use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed - use_fp4_context = self.config.fp4 is not None - outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext() - - if use_inner_fp8_context: - - def get_inner_quant_context(config, layer_number): - return get_fp8_context(config, layer_number) - - elif use_fp4_context: - - def get_inner_quant_context(config, layer_number): - return get_fp4_context(config, layer_number) - - else: - - def get_inner_quant_context(config, layer_number): - return nullcontext() - - with outer_fp8_context: - 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) - with inner_quant_context: - if isinstance(layer, TransformerLayer): - hidden_states, _ = layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb, - sequence_len_offset=sequence_len_offset, - packed_seq_params=packed_seq_params, - padding_mask=padding_mask, - ) - else: # MambaLayer, Expert, or MLP - hidden_states = layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - inference_context=inference_context, - packed_seq_params=packed_seq_params, - ) - - # The attention layer (currently a simplified transformer layer) - # outputs a tuple of (hidden_states, context). Context is intended - # for cross-attention, and is not needed in our model. - if isinstance(hidden_states, tuple): - hidden_states = hidden_states[0] - - # Final layer norm. - if self.post_process and self.post_layer_norm: - hidden_states = self.final_norm(hidden_states) - - # Ensure that the tensor passed between pipeline parallel stages is - # viewless. See related notes in TransformerBlock and TransformerLayer - hidden_states = make_viewless_tensor( - inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True - ) - - return hidden_states - - def sharded_state_dict( - self, - prefix: str = '', - sharded_offsets: Optional[tuple] = None, - metadata: Optional[dict] = None, - ) -> ShardedStateDict: - """ - Returns a sharded state dictionary for the current object. - - This function constructs a sharded state dictionary by iterating over the layers - in the current object, computing the sharded state dictionary for each layer, - and combining the results into a single dictionary. - - Parameters: - prefix (str): The prefix to use for the state dictionary keys. - sharded_offsets (tuple): The sharded offsets to use for the state dictionary. - metadata (dict): Additional metadata to use when computing the sharded state dictionary. - - Returns: - dict: The sharded state dictionary for the current object. - """ - - sharded_state_dict = {} - layer_prefix = f'{prefix}layers.' - - 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 MambaBlock - ) - - sharded_prefix = f'{layer_prefix}{global_layer_offset}.' - sharded_pp_offset = [] - - layer_sharded_state_dict = layer.sharded_state_dict( - state_dict_prefix, sharded_pp_offset, metadata - ) - - replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix) - - sharded_state_dict.update(layer_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, - ) - ) - - return sharded_state_dict +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +# Backward-compatible re-export. The canonical location is now +# megatron.core.models.hybrid.hybrid_block. +from megatron.core.models.hybrid.hybrid_block import * # noqa: F401,F403 +from megatron.core.models.hybrid.hybrid_block import ( # noqa: F401 + HybridStack, + HybridStackSubmodules, + MambaStack, + MambaStackSubmodules, +) diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py index 1cb1e4a31d1..fd43ed2aab9 100644 --- a/megatron/core/ssm/mamba_hybrid_layer_allocation.py +++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py @@ -1,484 +1,5 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -import logging -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple - -import torch - -from megatron.core.utils import log_on_each_pipeline_stage, log_single_rank - -logger = logging.getLogger(__name__) - - -class Symbols: - """Symbols for different layer types and pattern separators.""" - - MAMBA = "M" - GDN = 'G' - ATTENTION = "*" - MLP = "-" - MOE = 'E' - PIPE = '|' - MTP_SEPARATOR = "/" - VALID_LAYERS = {MAMBA, GDN, ATTENTION, MLP, MOE} - - -@dataclass -class ParsedHybridPattern: - """Result of parsing a unified hybrid pattern string. - - A unified pattern encodes both the main decoder pattern and the MTP pattern - in a single string using "/" as a separator. The main pattern may also - contain "|" pipe symbols to define pipeline stage boundaries for flexible - virtual pipeline parallelism (fVPP). - - Format: "///..." - - Examples: - - "M*M*" -> main="M*M*", mtp=None, depths=0 (no MTP) - - "M*M*/MM/MM" -> main="M*M*", mtp="MM", depths=2 - - "MMMM/*M/*M/*M" -> main="MMMM", mtp="*M", depths=3 - - "M-M-|M-M*-/MM/MM" -> main="M-M-|M-M*-" (2 PP stages), mtp="MM", depths=2 - - The "/" symbol introduces MTP patterns. Each repeated pattern after the main - decoder represents one MTP prediction depth. - - The "|" symbol in the main pattern defines pipeline stage boundaries. - - Attributes: - main_pattern: The main decoder layer pattern (e.g., "M*M*" or "M-M-|M-M*-") - mtp_pattern: The MTP layer pattern per depth (e.g., "MM"), or None if no MTP - mtp_num_depths: Number of MTP prediction depths (0 if no MTP) - """ - - main_pattern: Optional[str] - mtp_pattern: Optional[str] - mtp_num_depths: int - - -def pattern_from_ratios( - num_layers: int, attention_ratio: float = 0.0, mlp_ratio: float = 0.0 -) -> str: - """Convert deprecated ratio arguments to a layer pattern string. - - Generates an evenly-spaced hybrid layer pattern from target attention and MLP - ratios. This exists for backward compatibility with code that uses the deprecated - hybrid_attention_ratio and hybrid_mlp_ratio parameters. - - Args: - num_layers: Total number of layers. - attention_ratio: Target ratio of attention layers to total layers. - mlp_ratio: Target ratio of MLP layers to total layers. - - Returns: - A layer pattern string (e.g., "MMM*MMM*MM"). - """ - assert num_layers > 0 - assert 0.0 <= attention_ratio <= 1.0 - assert 0.0 <= mlp_ratio <= 1.0 - assert attention_ratio + mlp_ratio <= 1.0 - - # Allocate attention layers (evenly spaced, starting and ending with mamba) - attention_count = round(num_layers * attention_ratio) - mamba_count = num_layers - attention_count - sections = attention_count + 1 - section_len = mamba_count / sections - - layer_types = [Symbols.MAMBA] * num_layers - x = section_len - for i in range(num_layers): - if x < 0.5: - layer_types[i] = Symbols.ATTENTION - x += section_len - else: - x -= 1 - - # Allocate MLP layers (evenly distributed, not replacing attention) - mlp_count = round(num_layers * mlp_ratio) - if mlp_count > 0: - mamba_count -= mlp_count - ratio = mamba_count / mlp_count - x = ratio - for i in range(num_layers): - if layer_types[i] == Symbols.MAMBA: - if x < 0.5: - layer_types[i] = Symbols.MLP - x += ratio - else: - x -= 1 - - return ''.join(layer_types) - - -def get_hybrid_total_layer_count(pattern: str) -> int: - """Returns the total number of main decoder layers in a hybrid layer pattern. - - Extracts the main pattern (before the first MTP separator '/'), strips - pipeline stage separators '|', and returns the character count. - - Args: - pattern: Full hybrid layer pattern, possibly including MTP and pipe separators. - - Returns: - Total number of layers in the main decoder pattern. - """ - main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] - _validate_pattern(main_pattern, "main", allow_pipe=True) - return len(main_pattern.replace(Symbols.PIPE, '')) - - -def get_hybrid_total_pipeline_segment_count(pattern: str) -> int: - """Returns the number of pipeline segments in a hybrid layer pattern. - - Extracts the main pattern (before the first MTP separator '/') and counts - the number of segments delimited by '|'. - - Args: - pattern: Full hybrid layer pattern, possibly including MTP and pipe separators. - - Returns: - Number of pipeline segments (pipe count + 1). - """ - main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] - return main_pattern.count(Symbols.PIPE) + 1 - - -def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]: - """Count layers by type across the full hybrid pattern (main + MTP). - - Parses the pattern to extract main and MTP components, then counts - each layer type. Main pattern '|' separators are skipped. MTP layers - are counted once per MTP depth. - - Args: - pattern: Full hybrid layer pattern string. - - Returns: - Dictionary mapping layer symbol to count. Keys are Symbols.MAMBA, - Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, and Symbols.MOE. - - Examples: - >>> get_hybrid_layer_counts("M*M*") - {'M': 2, 'G': 0, '*': 2, '-': 0, 'E': 0} - - >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM") - {'M': 8, 'G': 0, '*': 1, '-': 4, 'E': 0} - """ - parsed = parse_hybrid_pattern(pattern) - counts = { - Symbols.MAMBA: 0, - Symbols.GDN: 0, - Symbols.ATTENTION: 0, - Symbols.MLP: 0, - Symbols.MOE: 0, - } - - # Count main decoder layers (skip '|' pipe separators) - if parsed.main_pattern: - for char in parsed.main_pattern: - if char in counts: - 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 - - return counts - - -def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern: - """Parse a unified hybrid pattern string into main and MTP components. - - The pattern uses "/" as a separator between the main decoder pattern and - MTP patterns. Each MTP pattern after the separator represents one prediction - depth. The main pattern may contain "|" pipe symbols for pipeline stage - boundaries. - - Format: "///..." - - Args: - pattern: Unified pattern string, e.g., "M*M*/MM/MM" or just "M*M*" - - Returns: - ParsedHybridPattern with main_pattern, mtp_pattern, and mtp_num_depths - - Raises: - ValueError: If MTP patterns are inconsistent (all must be identical) - ValueError: If pattern contains invalid layer symbols - - Examples: - >>> parse_hybrid_pattern("M*M*") - ParsedHybridPattern(main_pattern="M*M*", mtp_pattern=None, mtp_num_depths=0) - - >>> parse_hybrid_pattern("M*M*/MM/MM") - ParsedHybridPattern(main_pattern="M*M*", mtp_pattern="MM", mtp_num_depths=2) - - >>> parse_hybrid_pattern("MMMM/*M/*M/*M") - ParsedHybridPattern(main_pattern="MMMM", mtp_pattern="*M", mtp_num_depths=3) - - >>> parse_hybrid_pattern("M-M-|M-M*-/MM/MM") - ParsedHybridPattern(main_pattern="M-M-|M-M*-", mtp_pattern="MM", mtp_num_depths=2) - """ - if pattern is None: - return ParsedHybridPattern(main_pattern=None, mtp_pattern=None, mtp_num_depths=0) - - parts = pattern.split(Symbols.MTP_SEPARATOR) - - if len(parts) == 1: - # No MTP separator found - pattern is main decoder only - main_pattern = parts[0] - _validate_pattern(main_pattern, "main", allow_pipe=True) - return ParsedHybridPattern(main_pattern=main_pattern, mtp_pattern=None, mtp_num_depths=0) - - # First part is main decoder pattern - main_pattern = parts[0] - if main_pattern: - _validate_pattern(main_pattern, "main", allow_pipe=True) - - # Remaining parts are MTP patterns (one per depth) - mtp_parts = parts[1:] - - if not mtp_parts or all(p == "" for p in mtp_parts): - # No MTP patterns after separator - return ParsedHybridPattern( - main_pattern=main_pattern if main_pattern else None, mtp_pattern=None, mtp_num_depths=0 - ) - - # Validate all MTP patterns are identical - mtp_pattern = mtp_parts[0] - for i, part in enumerate(mtp_parts[1:], start=2): - if part != mtp_pattern: - raise ValueError( - f"All MTP patterns must be identical. " - f"Pattern 1 is '{mtp_pattern}', but pattern {i} is '{part}'. " - f"Full pattern: '{pattern}'" - ) - - _validate_pattern(mtp_pattern, "MTP", allow_pipe=False) - - return ParsedHybridPattern( - main_pattern=main_pattern if main_pattern else None, - mtp_pattern=mtp_pattern, - mtp_num_depths=len(mtp_parts), - ) - - -def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) -> None: - """Validate that a pattern contains only valid layer symbols. - - Args: - pattern: Layer pattern string to validate - pattern_name: Name of pattern for error messages (e.g., "main" or "MTP") - allow_pipe: Whether to allow the pipe '|' separator (for main patterns) - - 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: - raise ValueError( - f"In {pattern_name} pattern, '{char}' is not a valid layer symbol. " - f"Valid symbols are: {valid_chars}" - ) - - -def validate_segment_layers(segment: str) -> List[str]: - """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. - Each segment should contain only valid layer symbols (no '|'). - - Args: - segment: A single pipeline segment pattern string (e.g., "M-M*-") - - Returns: - List of layer type characters. - - 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: - raise ValueError( - f"In hybrid layer pattern segment, '{layer_char}' is not " - f"one of {Symbols.VALID_LAYERS}" - ) - return layer_type_list - - -def select_pipeline_segment( - 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, -) -> Tuple[List[str], 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 - pipeline segments and selects the segment for the current PP rank / VP stage. - - When the pattern has no pipes but pp_size > 1, falls back to runtime layer - slicing (for backwards compatibility), supporting both even and uneven PP splits - via first_stage_layers / last_stage_layers. - - Args: - main_pattern: Main decoder pattern (may contain '|' separators). - Empty string is allowed (produces one empty segment). - pp_group: Pipeline parallel process group, or None if not using PP. - vp_stage: Virtual pipeline stage, or None if not using VPP. - first_stage_layers: Number of layers on the first pipeline stage for - uneven PP. Only valid when the pattern has no pipe separators. - last_stage_layers: Number of layers on the last pipeline stage for - uneven PP. Only valid when the pattern has no pipe separators. - - Returns: - Tuple of (layer_type_list, layer_offset) where layer_type_list is - the list of layer type characters for this segment, and layer_offset - is the sum of layer counts from all preceding segments. - - Raises: - ValueError: If the segment contains invalid layer symbols, if - first/last_stage_layers are used with pipe separators, if VPP is - requested without pipe separators, or if layer counts are not - evenly divisible across pipeline stages. - """ - segments = main_pattern.split(Symbols.PIPE) if main_pattern 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 - - if len(segments) > 1 and (first_stage_layers is not None or last_stage_layers is not None): - raise ValueError( - "Cannot specify num_layers_in_first_pipeline_stage or " - "num_layers_in_last_pipeline_stage when hybrid_layer_pattern " - "contains pipe ('|') separators. The pipeline layout is already " - "explicitly defined by the pipe separators." - ) - - if len(segments) == 1 and pp_size > 1: - if vp_stage is not None: - raise ValueError( - "Virtual pipeline parallelism (vp_stage != None) is not supported " - "when hybrid_layer_pattern has no pipe ('|') separators. " - "Add '|' separators to define explicit pipeline/virtual-pipeline " - "stage boundaries." - ) - log_single_rank( - logger, - logging.WARNING, - "DEPRECATION: Using hybrid_layer_pattern without pipe ('|') separators " - "with pipeline_model_parallel_size > 1 is deprecated. Please add '|' " - "separators to explicitly define pipeline stage boundaries. " - "Example: 'M*M*M*M*' with pp_size=2 should become 'M*M*|M*M*'.", - ) - full_pattern = segments[0] - layer_type_list = validate_segment_layers(full_pattern) - num_layers = len(layer_type_list) - - if first_stage_layers is not None or last_stage_layers is not None: - first = first_stage_layers or 0 - last = last_stage_layers or 0 - middle_num_layers = num_layers - first - last - middle_stages = pp_size - sum( - 1 for x in (first_stage_layers, last_stage_layers) if x is not None - ) - if middle_stages > 0: - if middle_num_layers % middle_stages != 0: - raise ValueError( - f"Middle layers ({middle_num_layers}) must be evenly divisible " - f"by middle pipeline stages ({middle_stages})." - ) - layers_per_middle = middle_num_layers // middle_stages - else: - layers_per_middle = 0 - - is_first = first_stage_layers is not None and pp_rank == 0 - is_last = last_stage_layers is not None and pp_rank == pp_size - 1 - - if is_first: - offset = 0 - count = first - elif is_last: - offset = num_layers - last - count = last - else: - middle_rank = pp_rank if first_stage_layers is None else pp_rank - 1 - offset = middle_rank * layers_per_middle + first - count = layers_per_middle - else: - if num_layers % pp_size != 0: - raise ValueError( - f"Number of layers ({num_layers}) must be evenly divisible " - f"by pipeline-model-parallel-size ({pp_size}) when no pipe " - f"separators are specified in the pattern." - ) - layers_per_rank = num_layers // pp_size - offset = pp_rank * layers_per_rank - count = layers_per_rank - - selected = layer_type_list[offset : offset + count] - log_on_each_pipeline_stage( - logger, - logging.INFO, - f"MambaModel: pp_rank={pp_rank}/{pp_size}, vp_stage={vp_stage}, " - f"layers='{''.join(selected)}' ({len(selected)} layers), " - f"layer_offset={offset} (auto-split)", - ) - return selected, offset - - # Pipe-based segment selection - if len(segments) > 1 and len(segments) % pp_size != 0: - raise ValueError( - f"The number of pipe-delimited segments ({len(segments)}) in " - f"hybrid_layer_pattern must be evenly divisible by " - f"pipeline_model_parallel_size ({pp_size})." - ) - - vp_rel = vp_stage if vp_stage is not None else 0 - segment_index = vp_rel * pp_size + pp_rank - - if segment_index >= len(segments): - raise ValueError( - f"Pipeline segment index {segment_index} (pp_rank={pp_rank}, " - f"vp_stage={vp_rel}) is out of range for {len(segments)} segments. " - f"The pattern does not define enough pipe-delimited segments for " - f"the current PP/VPP configuration." - ) - - layer_offset = sum(len(segments[i]) for i in range(segment_index)) - my_segment = segments[segment_index] - - layer_type_list = validate_segment_layers(my_segment) - - log_on_each_pipeline_stage( - logger, - logging.INFO, - f"MambaModel: 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"layer_offset={layer_offset}", - ) - - return layer_type_list, layer_offset - - -def get_layer_maps_from_layer_type_list( - layer_type_list: List[str], -) -> Tuple[Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int], Dict[int, int]]: - """ - Returns maps from global layer index to the corresponding layer index - for each layer type in [Mamba, GDN, Attention, MLP, MoE] given a layer type list. - """ - layer_types = [Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE] - layer_maps = {layer_type: {} for layer_type in layer_types} - for global_layer_idx, layer_type in enumerate(layer_type_list): - layer_map = layer_maps[layer_type] - local_layer_idx = len(layer_map) - layer_map[global_layer_idx] = local_layer_idx - return [layer_maps[layer_type] for layer_type in layer_types] +# Backward-compatible re-export. The canonical location is now +# megatron.core.models.hybrid.hybrid_layer_allocation. +from megatron.core.models.hybrid.hybrid_layer_allocation import * # noqa: F401,F403 diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 8fe7a2636b0..fb940289cb2 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. from __future__ import annotations import warnings @@ -37,7 +37,7 @@ ) if TYPE_CHECKING: - from megatron.core.ssm.mamba_block import MambaStackSubmodules + from megatron.core.models.hybrid.hybrid_block import HybridStackSubmodules if is_torch_min_version("1.13.0"): dist_all_gather_func = torch.distributed.all_gather_into_tensor @@ -738,7 +738,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, # For Mamba path - pattern and submodules to build inner layers directly mtp_layer_pattern: Optional[str] = None, - mamba_submodules: Optional[MambaStackSubmodules] = None, + mamba_submodules: Optional[HybridStackSubmodules] = None, ): super().__init__(config=config) self.sequence_parallel = config.sequence_parallel @@ -753,11 +753,11 @@ def __init__( if self.submodules.mtp_model_layer is not None and hasattr( self.submodules.mtp_model_layer, 'submodules' ): - from megatron.core.ssm.mamba_block import MambaStackSubmodules + from megatron.core.models.hybrid.hybrid_block import HybridStackSubmodules from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules layer_submodules = None - if isinstance(self.submodules.mtp_model_layer.submodules, MambaStackSubmodules): + if isinstance(self.submodules.mtp_model_layer.submodules, HybridStackSubmodules): attention_layer_spec = self.submodules.mtp_model_layer.submodules.attention_layer if hasattr(attention_layer_spec, 'submodules'): assert isinstance(attention_layer_spec.submodules, TransformerLayerSubmodules) @@ -809,13 +809,13 @@ def __init__( ) # Build inner layers: two possible paths - # 1. Mamba path: use MambaStack for hybrid pattern support + # 1. Hybrid path: use HybridStack for hybrid pattern support # 2. GPT path: single TransformerLayer if mtp_layer_pattern is not None and mamba_submodules is not None: - from megatron.core.ssm.mamba_block import MambaStack - from megatron.core.ssm.mamba_hybrid_layer_allocation import validate_segment_layers + from megatron.core.models.hybrid.hybrid_block import HybridStack + from megatron.core.models.hybrid.hybrid_layer_allocation import validate_segment_layers - self.mtp_model_layer = MambaStack( + self.mtp_model_layer = HybridStack( config=self.config, submodules=mamba_submodules, layer_type_list=validate_segment_layers(mtp_layer_pattern), @@ -1272,7 +1272,7 @@ def __init__( # New: For Mamba path with unified pattern syntax mtp_layer_pattern: Optional[str] = None, mtp_num_depths: int = 0, - mamba_submodules: Optional["MambaStackSubmodules"] = None, + mamba_submodules: Optional["HybridStackSubmodules"] = None, ): super().__init__(config=config) self.submodules = _get_mtp_block_submodules(config, spec) From f8aaa60a5af4e5a0016bafc29701bd33a604be5b Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 17:29:09 +0000 Subject: [PATCH 02/23] Fix pylint unused-import warnings in backward-compat stubs Remove redundant explicit named imports that duplicate the wildcard import and add pylint disable comments. Co-Authored-By: Claude Opus 4.6 (1M context) --- megatron/core/models/mamba/mamba_model.py | 3 +-- megatron/core/ssm/mamba_block.py | 8 +------- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py index a3abb80890c..f0494e802cb 100644 --- a/megatron/core/models/mamba/mamba_model.py +++ b/megatron/core/models/mamba/mamba_model.py @@ -2,5 +2,4 @@ # Backward-compatible re-export. The canonical location is now # megatron.core.models.hybrid.hybrid_model. -from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 -from megatron.core.models.hybrid.hybrid_model import HybridModel, MambaModel # noqa: F401 +from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 # pylint: disable=unused-import diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py index f37273f0c31..bbfc3c62561 100644 --- a/megatron/core/ssm/mamba_block.py +++ b/megatron/core/ssm/mamba_block.py @@ -2,10 +2,4 @@ # Backward-compatible re-export. The canonical location is now # megatron.core.models.hybrid.hybrid_block. -from megatron.core.models.hybrid.hybrid_block import * # noqa: F401,F403 -from megatron.core.models.hybrid.hybrid_block import ( # noqa: F401 - HybridStack, - HybridStackSubmodules, - MambaStack, - MambaStackSubmodules, -) +from megatron.core.models.hybrid.hybrid_block import * # noqa: F401,F403 # pylint: disable=unused-import From 19e6688e9ba9b9386f2ccb0856b035034f550b1e Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 18:32:27 +0000 Subject: [PATCH 03/23] Add corresponding unit tests for renamed core classes Include the test files that exercise the renamed megatron/core classes (HybridModel, HybridStack, hybrid_layer_allocation, etc.) and update their imports to use the new canonical paths. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../contexts/test_dynamic_context.py | 2 +- .../inference/engines/test_dynamic_engine.py | 44 +++++++++---------- ...e.py => test_hybrid_prefix_caching_e2e.py} | 8 ++-- .../test_prefix_caching_cuda_graphs.py | 12 ++--- ...st_mamba_model.py => test_hybrid_model.py} | 28 ++++++------ ...hybrid_model_expert_parallel_inference.py} | 14 +++--- ..._moe_model.py => test_hybrid_moe_model.py} | 14 +++--- .../test_modelopt_module_spec.py | 32 +++++++------- .../unit_tests/resharding/test_model_swap.py | 10 ++--- ...st_mamba_block.py => test_hybrid_block.py} | 18 ++++---- ...ion.py => test_hybrid_layer_allocation.py} | 38 ++++++++-------- tests/unit_tests/ssm/test_mamba_layer.py | 12 ++--- tests/unit_tests/ssm/test_mamba_mixer.py | 22 +++++----- .../transformer/test_cuda_graphs.py | 12 ++--- .../test_multi_token_prediction.py | 18 ++++---- 15 files changed, 143 insertions(+), 141 deletions(-) rename tests/unit_tests/inference/engines/{test_mamba_prefix_caching_e2e.py => test_hybrid_prefix_caching_e2e.py} (99%) rename tests/unit_tests/models/{test_mamba_model.py => test_hybrid_model.py} (95%) rename tests/unit_tests/models/{test_mamba_model_expert_parallel_inference.py => test_hybrid_model_expert_parallel_inference.py} (96%) rename tests/unit_tests/models/{test_mamba_moe_model.py => test_hybrid_moe_model.py} (97%) rename tests/unit_tests/ssm/{test_mamba_block.py => test_hybrid_block.py} (91%) rename tests/unit_tests/ssm/{test_mamba_hybrid_layer_allocation.py => test_hybrid_layer_allocation.py} (93%) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 385ea09e345..cfff78a1780 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -16,7 +16,7 @@ ) from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams -from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index a3b2ce71e60..0b476ad3ae3 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio import gc @@ -46,8 +46,8 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -65,7 +65,7 @@ def skip_if_mamba_sequence_packing_not_available(model_provider: str): - if model_provider == "mamba": + if model_provider in ("hybrid", "mamba"): sequence_packing_available, reason_for_no_sequence_packing = ( _check_mamba_sequence_packing_support() ) @@ -366,7 +366,7 @@ def _build_test_env(cls, test_config): post_process=parallel_state.is_pipeline_last_stage(), mtp_block_spec=mtp_block_spec, ).cuda() - elif test_config.model_provider == "mamba": + elif test_config.model_provider in ("hybrid", "mamba"): pp_size = test_config.pipeline_model_parallel_size # Transformer config. transformer_config = TransformerConfig( @@ -406,9 +406,9 @@ def _build_test_env(cls, test_config): ) # Mamba model. - model = MambaModel( + model = HybridModel( config=transformer_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=test_config.vocab_size, max_sequence_length=test_config.max_sequence_length, parallel_output=True, @@ -567,7 +567,7 @@ def teardown_class(cls): @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @pytest.mark.parametrize("num_cuda_graphs", [None, 1, 4, -1]) @pytest.mark.parametrize("cuda_graph_scope", [[], [CudaGraphScope.full_iteration_inference]]) def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None: @@ -625,7 +625,7 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None if model_provider == "gpt": expected_generated_tokens_list = gpt_expected_generated_tokens - elif model_provider == "mamba": + elif model_provider in ("hybrid", "mamba"): expected_generated_tokens_list = mamba_expected_generated_tokens else: raise ValueError(f"Invalid model_provider {model_provider}") @@ -686,7 +686,7 @@ def test_token_overflow_nontransient(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_block_overflow(self, model_provider: str) -> None: """Test block overflow.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -732,7 +732,7 @@ def test_block_overflow_insufficient_kv_cache(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_multi_add(self, model_provider: str) -> None: """Test adding multiple requests simultaneously.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -742,7 +742,7 @@ def test_multi_add(self, model_provider: str) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_fixed_output_lengths(self, model_provider: str) -> None: """Test generating a fixed number of output tokens.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -785,7 +785,7 @@ def test_cuda_graph_token_counts(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @torch.inference_mode() def test_generate_function(self, model_provider: str) -> None: """Test the generate function that processes multiple prompts at once.""" @@ -879,7 +879,7 @@ async def test_run_engine(self): not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @pytest.mark.skipif(not is_te_min_version("2.2.0"), reason="TE 2.2.0 is required") - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_fp8_inference(self, model_provider: str): skip_if_mamba_sequence_packing_not_available(model_provider) @@ -1085,7 +1085,7 @@ def test_log_probs_token_correspondence(self): @pytest.mark.parametrize("ep_size", [1, 2]) @pytest.mark.parametrize("pp_size", [1, 2]) @pytest.mark.parametrize("tp_size", [1, 2]) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @pytest.mark.parametrize("transformer_impl", ["local", "inference_optimized"]) @torch.inference_mode() def test_parallel_inference( @@ -1124,7 +1124,7 @@ def test_parallel_inference( "when tp_size > 1." ) ) - if model_provider == "mamba": + if model_provider in ("hybrid", "mamba"): pytest.skip( reason="Mamba model is not supported with the inference optimized transformer." ) @@ -1292,11 +1292,11 @@ def test_mamba_chunked_prefill(self): """ Test chunked prefill with a Mamba model. """ - skip_if_mamba_sequence_packing_not_available("mamba") + skip_if_mamba_sequence_packing_not_available("hybrid") # Context max tokens = 50. test_config = DynamicEngineTestConfig( - model_provider="mamba", + model_provider="hybrid", num_requests=0, num_tokens_to_generate=None, num_tokens_total=200, @@ -3058,7 +3058,7 @@ def _create_model(self, model_provider, num_cuda_graphs): pre_process=parallel_state.is_pipeline_first_stage(), post_process=parallel_state.is_pipeline_last_stage(), ).cuda() - elif model_provider == "mamba": + elif model_provider in ("hybrid", "mamba"): config = TransformerConfig( params_dtype=torch.bfloat16, num_layers=3, @@ -3074,9 +3074,9 @@ def _create_model(self, model_provider, num_cuda_graphs): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=CHUNKED_CG_VOCAB_SIZE, max_sequence_length=CHUNKED_CG_MAX_SEQ_LEN, parallel_output=True, @@ -3162,7 +3162,7 @@ def _run_to_completion(self, engine, prompts, num_tokens_to_generate): return finished, step_count - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @pytest.mark.parametrize("chunked_prefill", [False, True]) @pytest.mark.parametrize("num_cuda_graphs", [None, 2]) @torch.inference_mode() diff --git a/tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py similarity index 99% rename from tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py rename to tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py index ce21c775b73..303cf76d122 100644 --- a/tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py +++ b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py @@ -54,8 +54,8 @@ from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -131,9 +131,9 @@ def _create_model(self, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=transformer_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, diff --git a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py index 52a05f7f80f..26a81c5baef 100644 --- a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py +++ b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py @@ -37,8 +37,8 @@ ) from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -121,9 +121,9 @@ def _create_model(self, model_type, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, @@ -343,9 +343,9 @@ def _create_hybrid_model(self, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, diff --git a/tests/unit_tests/models/test_mamba_model.py b/tests/unit_tests/models/test_hybrid_model.py similarity index 95% rename from tests/unit_tests/models/test_mamba_model.py rename to tests/unit_tests/models/test_hybrid_model.py index 6f1c27ca70c..ade1fa524be 100644 --- a/tests/unit_tests/models/test_mamba_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import os from datetime import timedelta @@ -15,8 +15,8 @@ from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel 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 TransformerConfig @@ -37,9 +37,9 @@ def setup_method(self, method): num_attention_heads=4, use_cpu_initialization=True, ) - self.model = MambaModel( + self.model = HybridModel( config=model_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=100, max_sequence_length=4, hybrid_layer_pattern="M*-", # 1 Mamba, 1 attention, 1 MLP @@ -49,7 +49,7 @@ def teardown_method(self, method): Utils.destroy_model_parallel() def test_constructor(self): - assert isinstance(self.model, MambaModel) + assert isinstance(self.model, HybridModel) assert self.model.max_sequence_length == 4 @@ -105,9 +105,9 @@ def test_forward_packed_sequence(self): attention_backend=AttnBackend.flash, # Needed for packed sequence ) vocab_size = 100 - model = MambaModel( + model = HybridModel( config=model_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=vocab_size, max_sequence_length=12, hybrid_layer_pattern="M*-", # 1 Mamba, 1 attention, 1 MLP @@ -212,7 +212,7 @@ def test_layer_numbers(self): ) @pytest.mark.parametrize("tp_size,cp_size,pp_size", [(2, 1, 4), (1, 1, 8), (8, 1, 1)]) def test_with_custom_process_groups(self, tmp_path, tp_size, cp_size, pp_size): - """Test MambaModel with custom process groups.""" + """Test HybridModel with custom process groups.""" Utils.initialize_model_parallel( tensor_model_parallel_size=tp_size, context_parallel_size=cp_size, @@ -261,9 +261,9 @@ def test_with_custom_process_groups(self, tmp_path, tp_size, cp_size, pp_size): pipeline_dtype=torch.bfloat16, ) - model = MambaModel( + model = HybridModel( config=model_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=128, max_sequence_length=4, hybrid_layer_pattern=hybrid_layer_pattern, @@ -293,7 +293,7 @@ def test_with_custom_process_groups(self, tmp_path, tp_size, cp_size, pp_size): class TestMambaWithDynamicInference: - """Tests MambaModel with dynamic inference.""" + """Tests HybridModel with dynamic inference.""" @torch.inference_mode() def setup_method(self, method): @@ -315,9 +315,9 @@ def setup_method(self, method): fp8_recipe="tensorwise", ) - self.model = MambaModel( + self.model = HybridModel( config=model_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=128, max_sequence_length=DynamicInferenceContext.TOKEN_ROUNDER, hybrid_layer_pattern="M*", # 1 Mamba, 1 attention diff --git a/tests/unit_tests/models/test_mamba_model_expert_parallel_inference.py b/tests/unit_tests/models/test_hybrid_model_expert_parallel_inference.py similarity index 96% rename from tests/unit_tests/models/test_mamba_model_expert_parallel_inference.py rename to tests/unit_tests/models/test_hybrid_model_expert_parallel_inference.py index 090a0fe7877..b90f7d65132 100644 --- a/tests/unit_tests/models/test_mamba_model_expert_parallel_inference.py +++ b/tests/unit_tests/models/test_hybrid_model_expert_parallel_inference.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Tests for full MambaModel inference with expert-parallel batch dimension sync. +"""Tests for full HybridModel inference with expert-parallel batch dimension sync. When expert parallelism > 1 with strict matching (hybrid models), batch dimensions are MAX-reduced across EP ranks. Different EP ranks can be in @@ -11,7 +11,7 @@ - PREFILL: 0 decode requests, >0 prefill requests - MIXED: >0 decode requests, >0 prefill requests -These tests verify that the full MambaModel produces correct output shapes +These tests verify that the full HybridModel produces correct output shapes for every combination of these states across EP ranks, using the real EP synchronization path (strict matching + MAX-reduce on batch dimensions). """ @@ -24,8 +24,8 @@ from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig @@ -56,7 +56,7 @@ @pytest.mark.internal class TestDynamicInference: - """Verify full MambaModel output shapes under EP strict matching scenarios.""" + """Verify full HybridModel output shapes under EP strict matching scenarios.""" HIDDEN_SIZE = 256 NUM_ATTN_HEADS = 4 @@ -95,9 +95,9 @@ def _build_model(self): num_moe_experts=2, moe_token_dispatcher_type="alltoall", ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=self.VOCAB_SIZE, max_sequence_length=self.MAX_SEQ_LEN, hybrid_layer_pattern="M*", diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py similarity index 97% rename from tests/unit_tests/models/test_mamba_moe_model.py rename to tests/unit_tests/models/test_hybrid_moe_model.py index ce594591503..5d14c2bebce 100644 --- a/tests/unit_tests/models/test_mamba_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import hashlib import inspect @@ -10,8 +10,8 @@ import pytest # type: ignore[import] import torch -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig @@ -397,7 +397,7 @@ def create_test_args(self): destroy_global_vars() destroy_num_microbatches_calculator() - sys.argv = ['test_mamba_moe_model.py'] + sys.argv = ['test_hybrid_moe_model.py'] args = parse_args() # The following args would be set from the nano v3 checkpoint. @@ -425,7 +425,7 @@ def create_test_args(self): args.hidden_dropout = 0.0 args.hybrid_layer_pattern = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" args.hybrid_override_pattern = None - args.spec = ["megatron.core.models.mamba.mamba_layer_specs", "mamba_stack_spec"] + args.spec = ["megatron.core.models.hybrid.hybrid_layer_specs", "hybrid_stack_spec"] args.num_experts = 128 args.moe_layer_freq = 1 args.moe_ffn_hidden_size = 1856 @@ -497,9 +497,9 @@ def setup_method(self, method): model_config = core_transformer_config_from_args(args, TransformerConfig) - self.model = MambaModel( + self.model = HybridModel( config=model_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=args.vocab_size, max_sequence_length=args.seq_length, hybrid_layer_pattern=args.hybrid_layer_pattern, diff --git a/tests/unit_tests/post_training/test_modelopt_module_spec.py b/tests/unit_tests/post_training/test_modelopt_module_spec.py index 585be52f944..a5f7d2ca943 100644 --- a/tests/unit_tests/post_training/test_modelopt_module_spec.py +++ b/tests/unit_tests/post_training/test_modelopt_module_spec.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import inspect import tempfile @@ -13,13 +13,13 @@ get_gpt_layer_with_transformer_engine_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.post_training.modelopt.gpt.model_specs import get_gpt_modelopt_spec from megatron.core.post_training.modelopt.gpt.state_dict_hooks import ( mcore_gpt_load_te_state_dict_pre_hook, ) -from megatron.core.post_training.modelopt.mamba.model_specs import get_mamba_stack_modelopt_spec +from megatron.core.post_training.modelopt.hybrid.model_specs import get_hybrid_stack_modelopt_spec from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -206,19 +206,19 @@ def setup_method(self, method): num_layers=3, hidden_size=256, num_attention_heads=4, use_cpu_initialization=True ) - # A Hybrid MambaModel using fused-TE spec (default) - self.default_model = MambaModel( + # A Hybrid HybridModel using fused-TE spec (default) + self.default_model = HybridModel( config=transformer_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=100, max_sequence_length=4, hybrid_layer_pattern="M*-", ) - # A Hybrid MambaModel using ModelOpt spec (local + TENorm). - self.modelopt_model = MambaModel( + # A Hybrid HybridModel using ModelOpt spec (local + TENorm). + self.modelopt_model = HybridModel( config=transformer_config, - mamba_stack_spec=get_mamba_stack_modelopt_spec(remap_te_layernorm=True), + hybrid_stack_spec=get_hybrid_stack_modelopt_spec(remap_te_layernorm=True), vocab_size=100, max_sequence_length=4, hybrid_layer_pattern="M*-", @@ -264,9 +264,9 @@ def test_get_gpt_modelopt_spec_interface(): ), f"Default value of {sig_defaults[k]} does not match the expected value of {v} for parameter {k}." -def test_get_mamba_stack_modelopt_spec_interface(): +def test_get_hybrid_stack_modelopt_spec_interface(): # Get the function signature - sig = inspect.signature(get_mamba_stack_modelopt_spec) + sig = inspect.signature(get_hybrid_stack_modelopt_spec) # Define the expected signature expected_params = { @@ -298,7 +298,7 @@ def test_get_mamba_stack_modelopt_spec_interface(): ), f"Default value of {sig_defaults[k]} does not match the expected value of {v} for parameter {k}." -def test_get_mamba_stack_modelopt_spec_use_default_te_spec(): - """Test that use_default_te_spec=True returns the standard mamba_stack_spec.""" - spec = get_mamba_stack_modelopt_spec(use_default_te_spec=True) - assert spec is mamba_stack_spec +def test_get_hybrid_stack_modelopt_spec_use_default_te_spec(): + """Test that use_default_te_spec=True returns the standard hybrid_stack_spec.""" + spec = get_hybrid_stack_modelopt_spec(use_default_te_spec=True) + assert spec is hybrid_stack_spec diff --git a/tests/unit_tests/resharding/test_model_swap.py b/tests/unit_tests/resharding/test_model_swap.py index 70d81d97829..e2d6a2bd096 100644 --- a/tests/unit_tests/resharding/test_model_swap.py +++ b/tests/unit_tests/resharding/test_model_swap.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import copy import gc import os @@ -37,8 +37,8 @@ try: import mamba_ssm # noqa: F401 - from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec - from megatron.core.models.mamba.mamba_model import MambaModel + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + from megatron.core.models.hybrid.hybrid_model import HybridModel has_mamba_deps = True except Exception: @@ -203,9 +203,9 @@ def _build_mamba( parallel_output: bool = True, ): pre_process, post_process = _pp_flags(pg_collection) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=vocab_size, max_sequence_length=seq_len, hybrid_layer_pattern=hybrid_layer_pattern, diff --git a/tests/unit_tests/ssm/test_mamba_block.py b/tests/unit_tests/ssm/test_hybrid_block.py similarity index 91% rename from tests/unit_tests/ssm/test_mamba_block.py rename to tests/unit_tests/ssm/test_hybrid_block.py index 7b743afbfad..bcf2791a118 100644 --- a/tests/unit_tests/ssm/test_mamba_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -1,13 +1,13 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import pytest import torch -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +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_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.ssm.gated_delta_net import GatedDeltaNet -from megatron.core.ssm.mamba_block import MambaStack -from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, validate_segment_layers from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig @@ -37,8 +37,8 @@ def get_mamba_block(self, layer_pattern): num_attention_heads=4, use_cpu_initialization=True, ) - modules = mamba_stack_spec.submodules - return MambaStack( + modules = hybrid_stack_spec.submodules + return HybridStack( transformer_config, modules, layer_type_list=layer_type_list, @@ -87,7 +87,7 @@ def test_invalid_layer_types_cause_failure(self): invalid_symbol = '+' assert invalid_symbol not in Symbols.VALID_LAYERS # sanity check. layer_pattern = Symbols.MAMBA + Symbols.ATTENTION + Symbols.MLP + invalid_symbol - # validate_segment_layers() in mamba_hybrid_layer_allocation.py throws a ValueError. + # validate_segment_layers() in hybrid_layer_allocation.py throws a ValueError. with pytest.raises(ValueError): block = self.get_mamba_block(layer_pattern) @@ -116,8 +116,8 @@ def test_gdn_gpu_forward(self): use_cpu_initialization=True, activation_func=torch.nn.functional.silu, ) - modules = mamba_stack_spec.submodules - block = MambaStack( + modules = hybrid_stack_spec.submodules + block = HybridStack( transformer_config, modules, layer_type_list=layer_type_list, diff --git a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py similarity index 93% rename from tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py rename to tests/unit_tests/ssm/test_hybrid_layer_allocation.py index 440c843bc27..6b36b12a542 100644 --- a/tests/unit_tests/ssm/test_mamba_hybrid_layer_allocation.py +++ b/tests/unit_tests/ssm/test_hybrid_layer_allocation.py @@ -4,7 +4,7 @@ import pytest -from megatron.core.ssm.mamba_hybrid_layer_allocation import ( +from megatron.core.models.hybrid.hybrid_layer_allocation import ( ParsedHybridPattern, Symbols, get_hybrid_layer_counts, @@ -346,28 +346,28 @@ class TestSelectPipelineSegment: is simply the vp_stage value. """ - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_single_segment_no_vp(self, mock_log): """Single segment, no VPP.""" layer_types, offset = select_pipeline_segment("M*M*", pp_group=None, vp_stage=None) assert layer_types == ['M', '*', 'M', '*'] assert offset == 0 - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_two_segments_vp0(self, mock_log): """Two segments, select first (vp_stage=0).""" layer_types, offset = select_pipeline_segment("M-M-|M-M*-", pp_group=None, vp_stage=0) assert layer_types == ['M', '-', 'M', '-'] assert offset == 0 - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_two_segments_vp1(self, mock_log): """Two segments, select second (vp_stage=1).""" layer_types, offset = select_pipeline_segment("M-M-|M-M*-", pp_group=None, vp_stage=1) assert layer_types == ['M', '-', 'M', '*', '-'] assert offset == 4 - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_four_segments(self, mock_log): """Four segments, verify each vp_stage selects correctly.""" pattern = "MM|M*|M-|ME" @@ -377,7 +377,7 @@ 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.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @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.""" layer_types, offset = select_pipeline_segment("||M*", pp_group=None, vp_stage=0) @@ -388,7 +388,7 @@ def test_empty_segment(self, mock_log): assert layer_types == ['M', '*'] assert offset == 0 - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_uneven_segments(self, mock_log): """Segments of different lengths.""" pattern = "MMM|M|MMMMM" @@ -404,44 +404,44 @@ def test_uneven_segments(self, mock_log): assert len(layer_types) == 5 assert offset == 4 - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_empty_main_pattern(self, mock_log): """Empty main pattern produces one empty segment.""" layer_types, offset = select_pipeline_segment("", pp_group=None, vp_stage=None) assert layer_types == [] assert offset == 0 - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_invalid_segment_raises(self, mock_log): """Invalid layer symbols in a segment should raise ValueError.""" with pytest.raises(ValueError): select_pipeline_segment("MX|M*", pp_group=None, vp_stage=0) - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_out_of_range_segment_raises(self, mock_log): """Segment index out of range should raise ValueError.""" with pytest.raises(ValueError, match="out of range"): select_pipeline_segment("M*|M*", pp_group=None, vp_stage=5) - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_logging_is_called(self, mock_log): """Verify that log_on_each_pipeline_stage is called.""" select_pipeline_segment("M*M*", pp_group=None, vp_stage=None) mock_log.assert_called_once() - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_mutual_exclusivity_pipes_with_first_stage(self, mock_log): """Pipe separators + first_stage_layers should raise ValueError.""" with pytest.raises(ValueError, match="Cannot specify"): select_pipeline_segment("M*|M*", pp_group=None, vp_stage=0, first_stage_layers=1) - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_mutual_exclusivity_pipes_with_last_stage(self, mock_log): """Pipe separators + last_stage_layers should raise ValueError.""" with pytest.raises(ValueError, match="Cannot specify"): select_pipeline_segment("M*|M*", pp_group=None, vp_stage=0, last_stage_layers=1) - @patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage') + @patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage') def test_segment_count_not_divisible_by_pp_size(self, mock_log): """Segment count not divisible by pp_size should raise ValueError.""" mock_group = object() @@ -475,8 +475,8 @@ def _call_for_rank( with ( patch('torch.distributed.get_rank', return_value=pp_rank), patch('torch.distributed.get_world_size', return_value=pp_size), - patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage'), - patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_single_rank'), + patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage'), + patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_single_rank'), ): return select_pipeline_segment( pattern, @@ -578,8 +578,10 @@ def test_deprecation_warning_logged(self): with ( patch('torch.distributed.get_rank', return_value=0), patch('torch.distributed.get_world_size', return_value=2), - patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_on_each_pipeline_stage'), - patch('megatron.core.ssm.mamba_hybrid_layer_allocation.log_single_rank') as mock_warn, + patch('megatron.core.models.hybrid.hybrid_layer_allocation.log_on_each_pipeline_stage'), + patch( + 'megatron.core.models.hybrid.hybrid_layer_allocation.log_single_rank' + ) as mock_warn, ): select_pipeline_segment("M*M*", pp_group=mock_group, vp_stage=None) mock_warn.assert_called_once() diff --git a/tests/unit_tests/ssm/test_mamba_layer.py b/tests/unit_tests/ssm/test_mamba_layer.py index 26c7ba4d92f..8d6e0ab8c91 100644 --- a/tests/unit_tests/ssm/test_mamba_layer.py +++ b/tests/unit_tests/ssm/test_mamba_layer.py @@ -1,11 +1,11 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import pytest import torch -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.models.hybrid.hybrid_block import HybridStackSubmodules +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.mamba_block import MambaStackSubmodules from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig @@ -26,12 +26,12 @@ def setup_method(self, method): num_attention_heads=1, use_cpu_initialization=True, ) - assert isinstance(mamba_stack_spec.submodules, MambaStackSubmodules) - assert isinstance(mamba_stack_spec.submodules.mamba_layer.submodules, MambaLayerSubmodules) + assert isinstance(hybrid_stack_spec.submodules, HybridStackSubmodules) + assert isinstance(hybrid_stack_spec.submodules.mamba_layer.submodules, MambaLayerSubmodules) pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) self.layer = MambaLayer( transformer_config, - mamba_stack_spec.submodules.mamba_layer.submodules, + hybrid_stack_spec.submodules.mamba_layer.submodules, pg_collection=pg_collection, ) diff --git a/tests/unit_tests/ssm/test_mamba_mixer.py b/tests/unit_tests/ssm/test_mamba_mixer.py index f32b62146f8..184a00fcadd 100644 --- a/tests/unit_tests/ssm/test_mamba_mixer.py +++ b/tests/unit_tests/ssm/test_mamba_mixer.py @@ -1,12 +1,12 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import pytest import torch from megatron.core.inference.contexts.static_context import StaticInferenceContext -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.models.hybrid.hybrid_block import HybridStackSubmodules +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.mamba_block import MambaStackSubmodules from megatron.core.ssm.mamba_layer import MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed @@ -39,16 +39,16 @@ def get_mixer(self, tp_size=1, cp_size=1, use_mem_eff_path=True): use_cpu_initialization=True, use_mamba_mem_eff_path=use_mem_eff_path, ) - assert isinstance(mamba_stack_spec.submodules, MambaStackSubmodules) - assert isinstance(mamba_stack_spec.submodules.mamba_layer.submodules, MambaLayerSubmodules) + assert isinstance(hybrid_stack_spec.submodules, HybridStackSubmodules) + assert isinstance(hybrid_stack_spec.submodules.mamba_layer.submodules, MambaLayerSubmodules) assert isinstance( - mamba_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + hybrid_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, MambaMixerSubmodules, ) pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) mixer = MambaMixer( transformer_config, - mamba_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + hybrid_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, transformer_config.hidden_size, layer_number=1, pg_collection=pg_collection, @@ -126,17 +126,17 @@ def test_error_check(self, hidden_size, ngroups, tp_size, expected_error_message use_cpu_initialization=True, mamba_num_groups=ngroups, ) - assert isinstance(mamba_stack_spec.submodules, MambaStackSubmodules) - assert isinstance(mamba_stack_spec.submodules.mamba_layer.submodules, MambaLayerSubmodules) + assert isinstance(hybrid_stack_spec.submodules, HybridStackSubmodules) + assert isinstance(hybrid_stack_spec.submodules.mamba_layer.submodules, MambaLayerSubmodules) assert isinstance( - mamba_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + hybrid_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, MambaMixerSubmodules, ) pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) with pytest.raises(AssertionError, match=expected_error_message): MambaMixer( transformer_config, - mamba_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, + hybrid_stack_spec.submodules.mamba_layer.submodules.mixer.submodules, transformer_config.hidden_size, pg_collection=pg_collection, ) diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index bfde9ff9cf1..f6a88ed1c78 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc import os @@ -15,15 +15,15 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.models.hybrid.hybrid_block import HybridStack +from megatron.core.models.hybrid.hybrid_layer_allocation import validate_segment_layers +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.num_microbatches_calculator import ( destroy_num_microbatches_calculator, init_num_microbatches_calculator, ) from megatron.core.pipeline_parallel.schedules import set_current_microbatch from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.mamba_block import MambaStack -from megatron.core.ssm.mamba_hybrid_layer_allocation import validate_segment_layers from megatron.core.tensor_parallel.random import ( HAVE_TE, initialize_rng_tracker, @@ -486,8 +486,8 @@ def get_mamba_block(hybrid_layer_pattern): use_cpu_initialization=True, cuda_graph_impl="local", ) - modules = mamba_stack_spec.submodules - return MambaStack( + modules = hybrid_stack_spec.submodules + return HybridStack( transformer_config, modules, layer_type_list=layer_type_list, diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 57423da335b..fc75affc22f 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os import sys @@ -14,8 +14,8 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import get_context_parallel_group @@ -712,10 +712,10 @@ def model_provider(self, pre_process=True, post_process=True, **config_kwargs): config = core_transformer_config_from_args(args) # MTP is configured via unified pattern in hybrid_layer_pattern - # MambaModel creates the MTP block internally based on the parsed pattern - model = MambaModel( + # HybridModel creates the MTP block internally based on the parsed pattern + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=args.vocab_size, max_sequence_length=args.max_position_embeddings, pre_process=pre_process, @@ -735,7 +735,7 @@ def create_test_args( destroy_global_vars() destroy_num_microbatches_calculator() - sys.argv = ['test_multi_token_prediction_mamba.py'] + sys.argv = ['test_multi_token_prediction_hybrid.py'] args = parse_args() args.mtp_num_layers = 2 args.mtp_loss_scaling_factor = 0.1 @@ -763,7 +763,7 @@ def create_test_args( args.bf16 = True # Unified pattern: "main/mtp/mtp" - main decoder "M*M*", MTP pattern "M*" with 2 depths args.hybrid_layer_pattern = "M*M*/M*/M*" - args.spec = "megatron.core.models.mamba.mamba_layer_specs.mamba_stack_spec" + args.spec = "megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec" if fp8 is not None: args.fp8 = 'e4m3' @@ -920,7 +920,7 @@ def test_attention_mask_validation_mamba(self): try: mamba_model = get_model(self.model_provider, ModelType.encoder_or_decoder) mamba_model = unwrap_model(mamba_model) - assert isinstance(mamba_model[0], MambaModel) + assert isinstance(mamba_model[0], HybridModel) assert mamba_model[0].mtp is not None except AssertionError as e: if "Multi-Token Prediction (MTP) is not yet supported" in str(e): From d6a3854d31510802cb879ef2908ca1623631c98f Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 17:24:38 +0000 Subject: [PATCH 04/23] Rename class-named Mamba references to Hybrid outside megatron/core Updates all files outside megatron/core/ to use the new Hybrid names introduced in the core rename PR. Depends on the core changes which provide backward-compat stubs at old import paths. File renames: - mamba_builders.py -> hybrid_builders.py - pretrain_mamba.py -> pretrain_hybrid.py - tools/run_mamba_text_generation_server*.py -> run_hybrid_* - test_mamba_model.py -> test_hybrid_model.py (and similar test renames) Function renames: - mamba_builder() -> hybrid_builder() - modelopt_gpt_mamba_builder() -> modelopt_gpt_hybrid_builder() Also updates --spec paths, --export-model-type, --model-provider, deprecation warnings, and documentation. Model-named files stay as-is: examples/mamba/, recipes/h100/mamba*.yaml Co-Authored-By: Claude Opus 4.6 (1M context) --- examples/mamba/run_text_gen_server_8b.sh | 4 +- examples/mamba/train.sh | 4 +- examples/multimodal/layer_specs.py | 10 +- examples/multimodal/model.py | 6 +- .../NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh | 2 +- .../NVIDIA-Nemotron-3-Super-120B-A12B-BF16.sh | 2 +- .../conf/nvidia/NVIDIA-Nemotron-Nano-9B-v2.sh | 2 +- .../nvidia/Nemotron-H-47B-Reasoning-128K.sh | 2 +- .../conf/nvidia/Nemotron-H-4B-Instruct.sh | 2 +- .../conf/nvidia/Nemotron-H-56B-Base-8K.sh | 2 +- .../conf/nvidia/Nemotron-H-8B-Base-8K.sh | 2 +- .../post_training/modelopt/convert_model.py | 4 +- .../post_training/modelopt/distillation.md | 2 +- examples/post_training/modelopt/export.py | 4 +- examples/post_training/modelopt/finetune.py | 4 +- examples/post_training/modelopt/generate.py | 4 +- examples/post_training/modelopt/mmlu.py | 4 +- .../modelopt/offline_feature_extract.py | 4 +- examples/post_training/modelopt/prune.py | 4 +- examples/post_training/modelopt/quantize.py | 4 +- examples/post_training/modelopt/train.sh | 4 +- examples/post_training/modelopt/validate.py | 4 +- examples/rl/model_configs/nemotron5_56b.sh | 2 +- examples/rl/model_configs/nemotron5_8b.sh | 2 +- .../rl/model_configs/nemotron5p5_12b_H.sh | 2 +- hybrid_builders.py | 54 +++ megatron/inference/utils.py | 20 +- megatron/post_training/arguments.py | 7 +- megatron/post_training/model_builder.py | 35 +- megatron/training/arguments.py | 9 +- megatron/training/training.py | 11 +- model_provider.py | 14 +- pretrain_gpt.py | 2 +- pretrain_hybrid.py | 366 ++++++++++++++++++ .../model_config.yaml | 4 +- .../model_config.yaml | 4 +- .../model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../model_config.yaml | 2 +- .../model_config.yaml | 4 +- .../model_config.yaml | 2 +- .../model_config.yaml | 4 +- .../model_config.yaml | 4 +- tests/test_utils/recipes/h100/mamba.yaml | 2 +- .../test_modelopt_model_builder.py | 2 +- ...y => run_hybrid_text_generation_server.py} | 2 +- ...rid_text_generation_server_completions.py} | 2 +- tools/run_inference_performance_test.py | 2 +- tools/run_text_generation_server.py | 14 +- train_rl.py | 4 +- 50 files changed, 555 insertions(+), 105 deletions(-) create mode 100644 hybrid_builders.py create mode 100644 pretrain_hybrid.py rename tools/{run_mamba_text_generation_server.py => run_hybrid_text_generation_server.py} (89%) rename tools/{run_mamba_text_generation_server_completions.py => run_hybrid_text_generation_server_completions.py} (89%) diff --git a/examples/mamba/run_text_gen_server_8b.sh b/examples/mamba/run_text_gen_server_8b.sh index d228e0c0edb..f183dea4ad1 100755 --- a/examples/mamba/run_text_gen_server_8b.sh +++ b/examples/mamba/run_text_gen_server_8b.sh @@ -22,7 +22,7 @@ export NCCL_IB_QPS_PER_CONNECTION=4 export TRITON_CACHE_DIR="./triton-cache/" export TRITON_CACHE_MANAGER="megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager" -torchrun $DISTRIBUTED_ARGS ../../tools/run_mamba_text_generation_server.py \ +torchrun $DISTRIBUTED_ARGS ../../tools/run_hybrid_text_generation_server.py \ --tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 \ --untie-embeddings-and-output-weights \ @@ -46,5 +46,5 @@ torchrun $DISTRIBUTED_ARGS ../../tools/run_mamba_text_generation_server.py \ --bf16 \ --micro-batch-size 1 \ --use-mcore-models \ - --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \ + --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --seed 42 diff --git a/examples/mamba/train.sh b/examples/mamba/train.sh index ba83f0d4e33..f971242ff0b 100755 --- a/examples/mamba/train.sh +++ b/examples/mamba/train.sh @@ -96,8 +96,8 @@ options=" \ --eval-iters 32 \ --bf16 \ --use-mcore-models \ - --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \ + --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --no-create-attention-mask-in-dataloader \ --tensorboard-dir ${TENSORBOARD_DIR}" -torchrun --nproc_per_node 8 ../../pretrain_mamba.py ${options} +torchrun --nproc_per_node 8 ../../pretrain_hybrid.py ${options} diff --git a/examples/multimodal/layer_specs.py b/examples/multimodal/layer_specs.py index ad24850b631..acced15eeb6 100644 --- a/examples/multimodal/layer_specs.py +++ b/examples/multimodal/layer_specs.py @@ -1,8 +1,8 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import torch from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add -from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules +from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules from megatron.core.ssm.mlp_layer import MLPLayer @@ -125,15 +125,15 @@ def get_layer_spec_te(is_vit=False, padding=False) -> ModuleSpec: ) -def get_mamba_layer_spec_te(padding=False) -> ModuleSpec: +def get_hybrid_layer_spec_te(padding=False) -> ModuleSpec: attn_mask_type = AttnMaskType.causal # Padding mask is needed for e.g. Context Parallel. if padding: attn_mask_type = AttnMaskType.padding_causal return ModuleSpec( - module=MambaStack, - submodules=MambaStackSubmodules( + module=HybridStack, + submodules=HybridStackSubmodules( mamba_layer=ModuleSpec( module=MambaLayer, submodules=MambaLayerSubmodules( diff --git a/examples/multimodal/model.py b/examples/multimodal/model.py index 494a854099e..a2d83428338 100644 --- a/examples/multimodal/model.py +++ b/examples/multimodal/model.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import warnings import logging from copy import deepcopy @@ -6,7 +6,7 @@ import torch from config import get_language_model_config, get_vision_model_config, get_vision_projection_config from layer_specs import (get_layer_spec, get_layer_spec_te, get_mlp_module_spec, get_norm_mlp_module_spec_te, - get_mamba_layer_spec_te) + get_hybrid_layer_spec_te) from megatron.core.models.multimodal.llava_model import IMAGE_TOKEN, LLaVAModel from megatron.core.models.vision.clip_vit_model import get_num_image_embeddings @@ -99,7 +99,7 @@ def model_provider( # Padding mask needed for SP/CP. padding = args.context_parallel_size > 1 and args.sequence_parallel if args.language_model_type.startswith('nemotron5-hybrid'): - language_transformer_layer_spec = get_mamba_layer_spec_te(padding=padding) + language_transformer_layer_spec = get_hybrid_layer_spec_te(padding=padding) else: language_transformer_layer_spec = get_layer_spec_te( is_vit=False, padding=padding diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh index 1fa00889e99..805302498fc 100644 --- a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.sh @@ -51,5 +51,5 @@ MODEL_ARGS=" \ --bf16 \ --seq-length 8192 \ --max-position-embeddings 8192 \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ " diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16.sh index 977be033df0..b9da9429eb5 100644 --- a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16.sh +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16.sh @@ -58,5 +58,5 @@ MODEL_ARGS=" \ --bf16 \ --seq-length 8192 \ --max-position-embeddings 8192 \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ " diff --git a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-9B-v2.sh b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-9B-v2.sh index 83867430a97..51aff10a22a 100644 --- a/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-9B-v2.sh +++ b/examples/post_training/modelopt/conf/nvidia/NVIDIA-Nemotron-Nano-9B-v2.sh @@ -35,6 +35,6 @@ MODEL_ARGS=" \ --tokenizer-type HuggingFaceTokenizer \ --make-vocab-size-divisible-by 1 \ --use-mcore-models \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ --padded-vocab-size 131072 \ " diff --git a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-47B-Reasoning-128K.sh b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-47B-Reasoning-128K.sh index 901e607f298..e2da6a3c33d 100644 --- a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-47B-Reasoning-128K.sh +++ b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-47B-Reasoning-128K.sh @@ -33,5 +33,5 @@ MODEL_ARGS=" \ --max-position-embeddings 8192 \ --tokenizer-type HuggingFaceTokenizer \ --use-mcore-models \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ " diff --git a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-4B-Instruct.sh b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-4B-Instruct.sh index 084db49e0eb..523f7d521b0 100644 --- a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-4B-Instruct.sh +++ b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-4B-Instruct.sh @@ -38,5 +38,5 @@ MODEL_ARGS=" \ --make-vocab-size-divisible-by 1 \ --use-mcore-models \ --rotary-base 10000 \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ " diff --git a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-56B-Base-8K.sh b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-56B-Base-8K.sh index 645a159d075..be80d8a9a19 100644 --- a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-56B-Base-8K.sh +++ b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-56B-Base-8K.sh @@ -35,5 +35,5 @@ MODEL_ARGS=" \ --max-position-embeddings 8192 \ --tokenizer-type HuggingFaceTokenizer \ --bf16 \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ " diff --git a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-8B-Base-8K.sh b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-8B-Base-8K.sh index 66f3ad368b4..36b242e36dd 100644 --- a/examples/post_training/modelopt/conf/nvidia/Nemotron-H-8B-Base-8K.sh +++ b/examples/post_training/modelopt/conf/nvidia/Nemotron-H-8B-Base-8K.sh @@ -37,6 +37,6 @@ MODEL_ARGS=" \ --use-mcore-models \ --rotary-percent 0.5 \ --rotary-base 500000 \ - --export-model-type MambaModel \ + --export-model-type HybridModel \ " # --rotary-base 10000 \ diff --git a/examples/post_training/modelopt/convert_model.py b/examples/post_training/modelopt/convert_model.py index eaec9789e1e..136b0273724 100644 --- a/examples/post_training/modelopt/convert_model.py +++ b/examples/post_training/modelopt/convert_model.py @@ -19,7 +19,7 @@ from megatron.core.parallel_state import destroy_model_parallel from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.utils import ( report_current_memory_info, to_empty_if_meta, @@ -129,7 +129,7 @@ def check_arguments(): ) model = get_model( - functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False + functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False ) report_current_memory_info() diff --git a/examples/post_training/modelopt/distillation.md b/examples/post_training/modelopt/distillation.md index 49f73c4edde..9946723364e 100644 --- a/examples/post_training/modelopt/distillation.md +++ b/examples/post_training/modelopt/distillation.md @@ -53,7 +53,7 @@ Without this configuration file, the default logits-only distillation with scale ### Training -Distillation is triggered by calling `pretrain_gpt.py` or `pretrain_mamba.py` with the following arguments: +Distillation is triggered by calling `pretrain_gpt.py` or `pretrain_hybrid.py` with the following arguments: ```bash --export-kd-teacher-load diff --git a/examples/post_training/modelopt/export.py b/examples/post_training/modelopt/export.py index 5e3b2a1716e..31c72d95eab 100755 --- a/examples/post_training/modelopt/export.py +++ b/examples/post_training/modelopt/export.py @@ -15,7 +15,7 @@ from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.training import get_args, get_model from megatron.training.initialize import initialize_megatron from megatron.training.utils import unwrap_model @@ -74,7 +74,7 @@ def add_modelopt_export_args(parser): ) model = get_model( - functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False + functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False ) # Materialize the model from meta device to cpu before loading the checkpoint. diff --git a/examples/post_training/modelopt/finetune.py b/examples/post_training/modelopt/finetune.py index f7f7c24f970..2efd3cde6a4 100755 --- a/examples/post_training/modelopt/finetune.py +++ b/examples/post_training/modelopt/finetune.py @@ -19,7 +19,7 @@ from megatron.core.models.gpt import GPTModel from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.loss_func import loss_func -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.non_loss_data_func import report_draft_acceptance_length from megatron.training import get_args, get_timers, pretrain from megatron.training.utils import ( @@ -486,7 +486,7 @@ def forward_step(data_iterator, model: GPTModel): if __name__ == "__main__": pretrain( train_valid_test_sft_datasets_provider, - partial(model_provider, modelopt_gpt_mamba_builder), + partial(model_provider, modelopt_gpt_hybrid_builder), ModelType.encoder_or_decoder, forward_step, extra_args_provider=add_finetune_args, diff --git a/examples/post_training/modelopt/generate.py b/examples/post_training/modelopt/generate.py index 3d3f6571b34..cc4c4e37a80 100644 --- a/examples/post_training/modelopt/generate.py +++ b/examples/post_training/modelopt/generate.py @@ -14,7 +14,7 @@ from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint from megatron.post_training.generate import simple_generate -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.utils import report_current_memory_info, to_empty_if_meta from megatron.training import get_args, get_model, initialize_megatron from utils import get_hf_tokenizer @@ -100,7 +100,7 @@ def get_conversations(example): UserWarning, ) - model = get_model(functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False) + model = get_model(functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False) report_current_memory_info() unwrapped_model = unwrap_model(model)[0] diff --git a/examples/post_training/modelopt/mmlu.py b/examples/post_training/modelopt/mmlu.py index 5aa5d1c24c7..466d5052b50 100644 --- a/examples/post_training/modelopt/mmlu.py +++ b/examples/post_training/modelopt/mmlu.py @@ -17,7 +17,7 @@ from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint from megatron.post_training.generate import simple_generate -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.utils import report_current_memory_info from megatron.training import get_args, get_model, initialize_megatron from utils import get_hf_tokenizer @@ -158,7 +158,7 @@ def generate_prompt(test_example, dev_examples, few_shots=0, no_subject_prompt=F UserWarning, ) - model = get_model(functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False) + model = get_model(functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False) report_current_memory_info() # Materialize the model from meta device to gpu before loading the checkpoint. diff --git a/examples/post_training/modelopt/offline_feature_extract.py b/examples/post_training/modelopt/offline_feature_extract.py index 80207faf2b2..92500b2950e 100644 --- a/examples/post_training/modelopt/offline_feature_extract.py +++ b/examples/post_training/modelopt/offline_feature_extract.py @@ -14,7 +14,7 @@ from megatron.core import mpu from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.training import get_args, get_model, get_tokenizer, initialize_megatron from megatron.training.utils import print_rank_0, unwrap_model from model_provider import model_provider @@ -53,7 +53,7 @@ def extract_feature(dataset, model, output_dir, idx_start, idx_end): args = get_args() tokenizer = get_tokenizer() - model = get_model(functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False) + model = get_model(functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False) load_modelopt_checkpoint(model, strict=not args.untie_embeddings_and_output_weights) print_rank_0("Done loading checkpoint") diff --git a/examples/post_training/modelopt/prune.py b/examples/post_training/modelopt/prune.py index 56bbffa0cd0..99e351a6198 100644 --- a/examples/post_training/modelopt/prune.py +++ b/examples/post_training/modelopt/prune.py @@ -28,7 +28,7 @@ from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint from megatron.post_training.generate import simple_generate -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.utils import ( report_current_memory_info, ) @@ -163,7 +163,7 @@ def get_params(model): tokenizer = get_hf_tokenizer() model = get_model( - functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False + functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False ) unwrapped_model = unwrap_model(model)[0] diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py index dc4947038e5..0c10696df84 100644 --- a/examples/post_training/modelopt/quantize.py +++ b/examples/post_training/modelopt/quantize.py @@ -39,7 +39,7 @@ from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint from megatron.post_training.generate import simple_generate -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.utils import ( print_distributed_quant_summary, report_current_memory_info, @@ -362,7 +362,7 @@ def get_calib_dataloader( tokenizer = get_hf_tokenizer() model = get_model( - functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False + functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False ) report_current_memory_info() diff --git a/examples/post_training/modelopt/train.sh b/examples/post_training/modelopt/train.sh index 1ebb8bf3d76..3afcd4f5be7 100755 --- a/examples/post_training/modelopt/train.sh +++ b/examples/post_training/modelopt/train.sh @@ -69,8 +69,8 @@ fi export HF_TOKEN=${HF_TOKEN} -if [[ ${MODEL_ARGS} == *"MambaModel"* ]]; then - PRETRAIN_EXE=${SCRIPT_DIR}/../../../pretrain_mamba.py +if [[ ${MODEL_ARGS} == *"HybridModel"* ]] || [[ ${MODEL_ARGS} == *"MambaModel"* ]]; then + PRETRAIN_EXE=${SCRIPT_DIR}/../../../pretrain_hybrid.py else PRETRAIN_EXE=${SCRIPT_DIR}/../../../pretrain_gpt.py fi diff --git a/examples/post_training/modelopt/validate.py b/examples/post_training/modelopt/validate.py index 8b8f1ffc9dd..4d9757da00c 100644 --- a/examples/post_training/modelopt/validate.py +++ b/examples/post_training/modelopt/validate.py @@ -14,7 +14,7 @@ from megatron.post_training.arguments import add_modelopt_args from megatron.post_training.checkpointing import load_modelopt_checkpoint -from megatron.post_training.model_builder import modelopt_gpt_mamba_builder +from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder from megatron.post_training.utils import get_mtbench_chat_data from megatron.training import get_args, get_model, initialize_megatron from utils import get_hf_tokenizer @@ -116,7 +116,7 @@ def report_current_memory_info(): ground_truth = [None for _ in range(len(prompts))] tokenizer = get_hf_tokenizer() - model = get_model(functools.partial(model_provider, modelopt_gpt_mamba_builder), wrap_with_ddp=False) + model = get_model(functools.partial(model_provider, modelopt_gpt_hybrid_builder), wrap_with_ddp=False) report_current_memory_info() diff --git a/examples/rl/model_configs/nemotron5_56b.sh b/examples/rl/model_configs/nemotron5_56b.sh index 23b9f99a72a..b4fcee17a8e 100644 --- a/examples/rl/model_configs/nemotron5_56b.sh +++ b/examples/rl/model_configs/nemotron5_56b.sh @@ -69,7 +69,7 @@ MODEL_OPTIONS="\ \ --fp8-recipe tensorwise \ --hybrid-layer-pattern M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- \ - --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \ + --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --mamba-state-dim 256 \ --per-split-data-args-path ${BLEND_PATH} \ --tiktoken-pattern v2 \ diff --git a/examples/rl/model_configs/nemotron5_8b.sh b/examples/rl/model_configs/nemotron5_8b.sh index c18149f03d6..198efd2a163 100644 --- a/examples/rl/model_configs/nemotron5_8b.sh +++ b/examples/rl/model_configs/nemotron5_8b.sh @@ -61,7 +61,7 @@ MODEL_OPTIONS="\ --inference-max-requests $MAX_INFERENCE_BS \ --pretrained-checkpoint $CHECKPOINT \ --hybrid-layer-pattern M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- \ - --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \ + --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --tiktoken-pattern v2 \ --distributed-timeout-minutes 60 \ --use-mcore-models \ diff --git a/examples/rl/model_configs/nemotron5p5_12b_H.sh b/examples/rl/model_configs/nemotron5p5_12b_H.sh index 1826d57e913..bfb4c7e4727 100644 --- a/examples/rl/model_configs/nemotron5p5_12b_H.sh +++ b/examples/rl/model_configs/nemotron5p5_12b_H.sh @@ -76,7 +76,7 @@ MODEL_OPTIONS="\ --disable-gloo-process-groups \ --mamba-head-dim 80 \ --hybrid-layer-pattern M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M- \ - --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \ + --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ --tiktoken-pattern v2 \ --distributed-timeout-minutes 10 \ --use-mcore-models \ diff --git a/hybrid_builders.py b/hybrid_builders.py new file mode 100644 index 00000000000..36a87a3940b --- /dev/null +++ b/hybrid_builders.py @@ -0,0 +1,54 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. + +from model_provider import count_parameters_in_layer +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.spec_utils import import_module +from megatron.training import print_rank_0 +from megatron.training.arguments import core_transformer_config_from_args +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec + + +def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None): + print_rank_0('building Hybrid model ...') + if config is None: + config = core_transformer_config_from_args(args, TransformerConfig) + assert args.use_legacy_models is False, "Hybrid model only supported in Mcore!" + + if config.transformer_impl == "inference_optimized": + hybrid_stack_spec = hybrid_inference_stack_spec + assert ( + not config.inference_fuse_tp_communication + ), "inference_fuse_tp_communication is not supported for HybridModel" + elif args.spec is not None: + hybrid_stack_spec = import_module(args.spec) + else: + raise ValueError("You must provide a valid hybrid layer spec via --spec") + + model = HybridModel( + config=config, + hybrid_stack_spec=hybrid_stack_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + hybrid_layer_pattern=args.hybrid_layer_pattern, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + rotary_base=args.rotary_base, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + + for l in range(model.decoder.num_layers_per_pipeline_rank): + layer_params = count_parameters_in_layer(model, f'decoder.layers.{l}.') + print_rank_0(f" == params layer {l}: {layer_params}") + + return model + + +# Backward-compatible alias +mamba_builder = hybrid_builder diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index f0fde214d4a..00687627b19 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -7,7 +7,7 @@ import torch from gpt_builders import gpt_builder -from mamba_builders import mamba_builder +from hybrid_builders import hybrid_builder from megatron.core.inference.config import ( InferenceConfig, KVCacheManagementMode, @@ -43,8 +43,16 @@ def get_model_for_inference() -> MegatronModule: if args.model_provider == "gpt": model_builder = gpt_builder - elif args.model_provider == "mamba": - model_builder = mamba_builder + elif args.model_provider in ("hybrid", "mamba"): + if args.model_provider == "mamba": + import warnings + + warnings.warn( + '--model-provider "mamba" is deprecated. Use --model-provider "hybrid" instead.', + DeprecationWarning, + stacklevel=2, + ) + model_builder = hybrid_builder else: raise ValueError(f"Invalid model provider {args.model_provider}") @@ -158,7 +166,11 @@ def add_inference_args(parser: ArgumentParser) -> ArgumentParser: "total number of requests. Set to -1 to add all requests together.", ) group.add_argument( - "--model-provider", choices=["mamba", "gpt"], default="gpt", help="Model provider" + "--model-provider", + choices=["hybrid", "mamba", "gpt"], + default="gpt", + help='Model provider. Use "hybrid" for HybridModel (formerly MambaModel). ' + '"mamba" is accepted for backward compatibility but deprecated.', ) group.add_argument( "--skip-prompt-log-probs", action='store_true', default=False, help='Skip prompt log probs.' diff --git a/megatron/post_training/arguments.py b/megatron/post_training/arguments.py index dc98c6d28e4..47c667b4d0a 100644 --- a/megatron/post_training/arguments.py +++ b/megatron/post_training/arguments.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. def add_modelopt_args(parser): @@ -10,8 +10,9 @@ def add_modelopt_args(parser): "--export-model-type", type=str, default="GPTModel", - choices=["GPTModel", "MambaModel"], - help="Model type to use in model_provider.", + choices=["GPTModel", "HybridModel", "MambaModel"], + help='Model type to use in model_provider. Use "HybridModel" for hybrid models ' + '(formerly MambaModel). "MambaModel" is accepted for backward compatibility but deprecated.', ) group.add_argument( "--export-legacy-megatron", diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py index 085d188e811..383ae6ec8aa 100644 --- a/megatron/post_training/model_builder.py +++ b/megatron/post_training/model_builder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. """ModelOpt GPT model provider.""" @@ -16,7 +16,7 @@ from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import ( get_gpt_heterogeneous_layer_spec, ) -from megatron.core.models.mamba import MambaModel as MCoreMambaModel +from megatron.core.models.hybrid.hybrid_model import HybridModel as MCoreHybridModel from megatron.core.post_training.modelopt.gpt.model_specs import get_gpt_modelopt_spec from megatron.core.post_training.modelopt.gpt.state_dict_hooks import ( mcore_gpt_load_te_state_dict_pre_hook, @@ -124,7 +124,7 @@ def _load_teacher_model(config, config_raw: Namespace, model_kwargs: Dict[str, A # _load_teacher_model_config, so config_raw.hybrid_layer_pattern is always set here. model_kwargs["hybrid_layer_pattern"] = config_raw.hybrid_layer_pattern - teacher = MCoreMambaModel(config=config, **model_kwargs) + teacher = MCoreHybridModel(config=config, **model_kwargs) else: # GPT layer spec needs re-creation since it depends on number of model layers. if config.heterogeneous_block_specs: @@ -158,14 +158,14 @@ def _load_teacher_model(config, config_raw: Namespace, model_kwargs: Dict[str, A return teacher -def modelopt_gpt_mamba_builder( +def modelopt_gpt_hybrid_builder( args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None, -) -> MCoreGPTModel | MCoreMambaModel: +) -> MCoreGPTModel | MCoreHybridModel: """Builds the model. Args: @@ -179,7 +179,7 @@ def modelopt_gpt_mamba_builder( attached to the returned model for downstream routing/resharding utilities. Returns: - MCoreGPTModel | MCoreMambaModel: The returned model + MCoreGPTModel | MCoreHybridModel: The returned model """ print_rank_0("building GPT model ...") @@ -259,8 +259,17 @@ def modelopt_gpt_mamba_builder( "pg_collection": pg_collection, } model = MCoreGPTModel(config=config, **model_kwargs) - elif args.export_model_type == "MambaModel" or getattr(args, 'hybrid_layer_pattern', None) is not None: - from megatron.core.post_training.modelopt.mamba.model_specs import get_mamba_stack_modelopt_spec + elif args.export_model_type in ("HybridModel", "MambaModel") or getattr(args, 'hybrid_layer_pattern', None) is not None: + if args.export_model_type == "MambaModel": + import warnings + + warnings.warn( + '--export-model-type "MambaModel" is deprecated. ' + 'Use --export-model-type "HybridModel" instead.', + DeprecationWarning, + stacklevel=2, + ) + from megatron.core.post_training.modelopt.hybrid.model_specs import get_hybrid_stack_modelopt_spec if args.export_default_te_spec and args.export_te_mcore_model: logging.getLogger(__name__).warning( @@ -269,12 +278,12 @@ def modelopt_gpt_mamba_builder( ) args.export_te_mcore_model = False - mamba_stack_spec = get_mamba_stack_modelopt_spec( + hybrid_stack_spec = get_hybrid_stack_modelopt_spec( remap_te_layernorm=args.export_te_mcore_model, use_default_te_spec=args.export_default_te_spec, ) model_kwargs = { - "mamba_stack_spec": mamba_stack_spec, + "hybrid_stack_spec": hybrid_stack_spec, "vocab_size": args.padded_vocab_size, "max_sequence_length": args.max_position_embeddings, "hybrid_layer_pattern": args.hybrid_layer_pattern, @@ -289,7 +298,7 @@ def modelopt_gpt_mamba_builder( "pg_collection": pg_collection, } - model = MCoreMambaModel(config=config, **model_kwargs) + model = MCoreHybridModel(config=config, **model_kwargs) for l in range(model.decoder.num_layers_per_pipeline_rank): layer_params = count_parameters_in_layer(model, f'decoder.layers.{l}.') @@ -352,3 +361,7 @@ def modelopt_gpt_mamba_builder( print_distributed_quant_summary(model) return model + + +# Backward-compatible alias +modelopt_gpt_mamba_builder = modelopt_gpt_hybrid_builder diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0bfe6142d01..6c326fd7f16 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -612,7 +612,7 @@ def validate_args(args, defaults={}): args.rank, ) - from megatron.core.ssm.mamba_hybrid_layer_allocation import ( + from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, parse_hybrid_pattern, get_hybrid_total_layer_count, get_hybrid_total_pipeline_segment_count, ) @@ -2482,14 +2482,12 @@ def _add_training_args(parser): help='use FlashAttention implementation of attention. ' 'https://arxiv.org/abs/2205.14135') group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], + choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap'], help='Optimizer function. ' 'Note: dist_muon is deprecated; use --optimizer muon ' 'with --use-distributed-optimizer instead.') group.add_argument('--optimizer-cpu-offload', action='store_true', help='Offload optimizer state to CPU') - group.add_argument('--optimizer-cuda-graph', action='store_true', - help='Enable CUDA graph for optimizer step') group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, help='Ratio of optimizer state to offload to CPU') group.add_argument('--use-torch-optimizer-for-cpu-offload', action='store_true', @@ -2696,6 +2694,9 @@ def _add_distributed_args(parser): help='If not set, all PP stages will launch param all-gathers simultaneously. ' 'Otherwise, each PP stage will independently launch as needed.', dest='align_param_gather') + group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false', + help='If not set, use scatter/gather to optimize communication of tensors in pipeline.', + dest='scatter_gather_tensors_in_pipeline') group.add_argument('--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.') group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', diff --git a/megatron/training/training.py b/megatron/training/training.py index b8707041695..e218efab1c5 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -156,7 +156,6 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.training.checkpointing import checkpoint_exists from megatron.training.checkpointing import get_loaded_iteration from megatron.core.full_cuda_graph import FullCudaGraphWrapper -from megatron.core.optimizer.optimizer_cuda_graph import OptimizerCudaGraphWrapper from megatron.core.transformer.cuda_graphs import TECudaGraphHelper from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.module import Float16Module @@ -667,7 +666,7 @@ def transformer_flops(): # Calculate the number of each type of layer. from operator import itemgetter - from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, get_hybrid_layer_counts + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols, get_hybrid_layer_counts num_mamba_layers, num_gdn_layers, num_attn_layers, num_mlp_layers, num_moe_layers = ( itemgetter(Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)( get_hybrid_layer_counts(args.hybrid_layer_pattern) @@ -2122,7 +2121,7 @@ def training_log( if is_hybrid_model(args): from operator import itemgetter - from megatron.core.ssm.mamba_hybrid_layer_allocation import ( + from megatron.core.models.hybrid.hybrid_layer_allocation import ( Symbols, get_hybrid_layer_counts, ) layers = itemgetter(Symbols.MOE)(get_hybrid_layer_counts(args.hybrid_layer_pattern)) @@ -2785,8 +2784,6 @@ def train( forward_backward_func = get_forward_backward_func() if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) - if args.optimizer_cuda_graph: - optimizer.step = OptimizerCudaGraphWrapper(optimizer.step, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) def get_e2e_base_metrics(): """Get base metrics values for one-logger to calculate E2E tracking metrics.""" @@ -3209,10 +3206,6 @@ def trace_handler(p): if args.cuda_graph_impl == "transformer_engine" and cuda_graph_helper.graphs_created(): cuda_graph_helper.delete_cuda_graphs() - # Call OptimizerCudaGraph destructor to destroy optimizer CUDA graph - if args.optimizer_cuda_graph: - del optimizer.step - one_logger_utils.track_e2e_metrics() # Flush TensorBoard, WandB writers and one-logger. diff --git a/model_provider.py b/model_provider.py index 0c80c54dfdb..3e61343bb0d 100644 --- a/model_provider.py +++ b/model_provider.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. """Common functions used in train_*.py and pretrain_*.py scripts.""" @@ -7,11 +7,11 @@ import torch from megatron.core.models.gpt import GPTModel -from megatron.core.models.mamba import MambaModel +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.training import get_args, print_rank_0 try: - from megatron.post_training.model_builder import modelopt_gpt_mamba_builder + from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder has_nvidia_modelopt = True except ImportError: has_nvidia_modelopt = False @@ -23,18 +23,18 @@ def model_provider( model_builder: Callable, pre_process=True, post_process=True, vp_stage: Optional[int] = None, config=None, pg_collection=None, -) -> Union[GPTModel, megatron.legacy.model.GPTModel, MambaModel]: +) -> Union[GPTModel, megatron.legacy.model.GPTModel, HybridModel]: """Builds the model. If you set the use_legacy_models to True, it will return the legacy GPT model and if not the mcore GPT model. Args: - model_builder: A callable that builds the actual model, its signature is the same as model_provider's with an exception of the first argument which is a builder itself. In addition might take a config passed from outside to skip its own config loading. See gpt_builder or mamba_builder for an example, see _gpt_model_builder in train_rl.py to see how to augment a default gpt builder and pass the config from outside + model_builder: A callable that builds the actual model, its signature is the same as model_provider's with an exception of the first argument which is a builder itself. In addition might take a config passed from outside to skip its own config loading. See gpt_builder or hybrid_builder for an example, see _gpt_model_builder in train_rl.py to see how to augment a default gpt builder and pass the config from outside pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True. post_process (bool, optional): Set to true if you need to compute output logits/loss. Defaults to True. Returns: - Union[GPTModel, megatron.legacy.model.GPTModel, MambaModel]: The returned model + Union[GPTModel, megatron.legacy.model.GPTModel, HybridModel]: The returned model """ args = get_args() @@ -58,7 +58,7 @@ def oom_observer(device, alloc, device_alloc, device_free): if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt]: Use custom builder + spec when modelopt is enabled - model_builder = modelopt_gpt_mamba_builder + model_builder = modelopt_gpt_hybrid_builder return model_builder(args, pre_process, post_process, vp_stage, config=config, pg_collection=pg_collection) diff --git a/pretrain_gpt.py b/pretrain_gpt.py index 31eee0f4dc6..c6291306a15 100644 --- a/pretrain_gpt.py +++ b/pretrain_gpt.py @@ -90,7 +90,7 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): - MTP ranks (``mtp_on_this_rank``) also receive the full batch, regardless of pipeline stage. - Difference from ``pretrain_mamba.py``: + Difference from ``pretrain_hybrid.py``: - Return format: GPT returns a 6-tuple ``(tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params)`` where ``packed_seq_params`` is a diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py new file mode 100644 index 00000000000..f60552c3523 --- /dev/null +++ b/pretrain_hybrid.py @@ -0,0 +1,366 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +"""Pretrain and SFT Hybrid.""" + +# Capture the true program start time BEFORE any heavy imports. +import time +_PROGRAM_START_TIME = time.time() + +import json + +# Suppress warnings on all ranks but rank 0. +import os +import warnings +rank = int(os.environ.get('RANK', 0)) +if rank != 0: + warnings.filterwarnings("ignore", category=UserWarning) + warnings.filterwarnings("ignore", category=FutureWarning) + +from functools import partial +from typing import List, Optional, Tuple + +import torch + +from hybrid_builders import hybrid_builder +from megatron.core import mpu +from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder +from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset +from megatron.core.enums import ModelType +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parallel_state import ( + get_context_parallel_rank, + get_context_parallel_world_size, +) +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.rerun_state_machine import get_rerun_state_machine +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer +from megatron.core.utils import get_attr_wrapped_model, is_te_min_version, StragglerDetector +from megatron.training import ( + get_args, + get_timers, + inprocess_restart, + pretrain, + print_rank_0, + set_startup_timestamps, +) +from megatron.training.datasets.sft_dataset import SFTDataset +from megatron.training.utils import ( + get_batch_on_this_cp_rank, + get_batch_on_this_tp_rank, + get_blend_and_blend_per_split, + is_first_or_last_pipeline_stage, +) +from model_provider import model_provider + +try: + from megatron.post_training.arguments import add_modelopt_args + from megatron.post_training.loss_func import loss_func as loss_func_modelopt + has_nvidia_modelopt = True +except ImportError: + has_nvidia_modelopt = False + +try: + # Register the TE CUDA kernels + import transformer_engine # pylint: disable=unused-import + + # Alias the PyTorch wrapper so we can call tex.* APIs + import transformer_engine_torch as tex +except ImportError: + # TE isn’t installed or the torch wrapper is missing + tex = None + +stimer = StragglerDetector() + + +def get_batch(data_iterator, vp_stage=None): + """Generate a batch.""" + + empty_batch = { + 'tokens': None, + 'labels': None, + 'loss_mask': None, + 'attention_mask': None, + 'position_ids': None, + 'cu_seqlens': None, + 'max_seqlen': None, + } + + # TODO(duncan): Is there a more efficient way to access is_packed_sequence here? + is_packed_sequence = get_args().sft # SFT always uses packed sequence + if not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence: + return empty_batch.values() + + batch = get_batch_on_this_tp_rank(data_iterator) + + cu_seqlens = batch['cu_seqlens'] + # Unused at the moment + cu_seqlens_padded = batch.pop('cu_seqlens_padded', None) + # Support for Hybrid Context Parallel (Unused in this script) + local_cp_size = batch.pop('local_cp_size', None) + + if cu_seqlens is not None: + assert ( + cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 + ), "micro-batch-size must be 1 for packing" + cu_seqlens = cu_seqlens[0] + batch['cu_seqlens'] = cu_seqlens + + max_seqlen = batch['max_seqlen'] + assert max_seqlen.dim() == 1 + # TODO(duncan): can this be kept as a 0-D tensor? + batch['max_seqlen'] = int(max_seqlen[0].item()) + + if mpu.is_pipeline_first_stage(ignore_virtual=(vp_stage is None), vp_stage=vp_stage): + total_tokens = batch['tokens'].size(1) + elif mpu.is_pipeline_last_stage(ignore_virtual=(vp_stage is None), vp_stage=vp_stage): + total_tokens = batch['labels'].size(1) + else: # packed sequence + empty_batch['cu_seqlens'] = cu_seqlens + empty_batch['max_seqlen'] = max_seqlen + return empty_batch.values() + + if cu_seqlens is None: + # slice batch along sequence dimension for context parallelism + batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore + else: # Packed THD format + cp_size = get_context_parallel_world_size() + if cp_size > 1: # slice batch along sequence dimension for context parallelism + assert tex is not None and is_te_min_version("1.10.0"), ( + "Please update Transformer Engine to >= 1.10 to use " + "Context Parallel with THD format data" + ) + cp_rank = get_context_parallel_rank() + index = tex.thd_get_partitioned_indices( + cu_seqlens, + total_tokens, + cp_size, + cp_rank, + ) + for key, data in batch.items(): + if key in {'attention_mask', 'cu_seqlens', 'max_seqlen'}: + continue + if data is not None: + # On first PP rank, labels and loss_mask can be None. + # On last PP rank, tokens and position_ids can be None. + batch[key] = data.index_select(1, index) + + return batch.values() + + +# define spiky loss as a loss that's 10x the max loss observed +SPIKY_LOSS_FACTOR = 10 + +def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[HybridModel] = None): + """Loss function. + + Args: + loss_mask (torch.Tensor): Used to mask out some portions of the loss + output_tensor (torch.Tensor): The tensor with the losses + + Returns: + the loss scalar for this micro-batch + the number of non-padded tokens in this microbatch + a dict containing reporting metrics on the loss and number of tokens across + the data parallel ranks + """ + args = get_args() + if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + loss, num_tokens, report = loss_func_modelopt(loss_mask, output_tensor, model=model) + else: + losses = output_tensor.view(-1).float() + loss_mask = loss_mask.view(-1).float() + loss = torch.sum(losses * loss_mask) + + num_tokens = loss_mask.sum().clone().detach().to(torch.int) + report = {'lm loss': torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])} + + # Check individual rank losses are not NaN prior to DP all-reduce. + rerun_state_machine = get_rerun_state_machine() + if args.check_for_nan_in_loss_and_grad: + rerun_state_machine.validate_result( + result=loss, + rejection_func=torch.isnan, + message="found NaN in local forward loss calculation", + tolerance=0.0, # forward pass calculations are deterministic + fatal=True, + ) + rerun_state_machine.validate_result( + result=loss, + rejection_func=torch.isinf, + message="found Inf in local forward loss calculation", + tolerance=0.0, # forward pass calculations are deterministic + fatal=True, + ) + # Check for spiky loss + if args.check_for_spiky_loss: + rerun_state_machine.validate_result( + result=loss, + rejection_func=partial( + rerun_state_machine.is_unexpectedly_large, + threshold=SPIKY_LOSS_FACTOR, + context="loss", + ), + message="Spiky loss", + tolerance=0.0, # forward pass calculations are deterministic + fatal=False, + ) + + return loss, num_tokens, report + + +def forward_step(data_iterator, model: HybridModel): + """Forward training step. + + Args: + data_iterator : Input data iterator + model (HybridModel): The Model + """ + timers = get_timers() + + # Get the batch. + timers('batch-generator', log_level=2).start() + + global stimer + + with stimer(bdata=True): + vp_stage = get_attr_wrapped_model(model, "vp_stage") + ( + tokens, + labels, + loss_mask, + attention_mask, + position_ids, + cu_seqlens, + max_seqlen, + ) = get_batch(data_iterator, vp_stage) + + if cu_seqlens is None: + packed_seq_params = None + else: + total_tokens = tokens.size(1) if tokens is not None else labels.size(1) + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=None, + cu_seqlens_kv_padded=None, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + total_tokens=total_tokens, + ) + + 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 + ) + + # [ModelOpt]: model is needed to access ModelOpt distillation losses + return output_tensor, partial(loss_func, loss_mask, model=model) + + +def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): + if mpu.get_tensor_model_parallel_rank() != 0: + return False + elif is_packed_sequence: + return True + else: + return is_first_or_last_pipeline_stage(vp_stage) + + +def core_gpt_dataset_config_from_args(args): + tokenizer = build_tokenizer(args) + + # Sometimes --data-path is too long, instead we parse it from a file. + blend: Optional[Tuple[List[str], Optional[List[float]]]] + blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] + blend, blend_per_split = get_blend_and_blend_per_split(args) + + sequences_per_dataset = None + if args.per_dataset_sequences_path is not None: + with open(args.per_dataset_sequences_path, "r") as f: + sequences_per_dataset = json.load(f) + + return GPTDatasetConfig( + random_seed=args.seed, + sequence_length=args.seq_length, + blend=blend, + blend_per_split=blend_per_split, + split=args.split, + num_dataset_builder_threads=args.num_dataset_builder_threads, + path_to_cache=args.data_cache_path, + mmap_bin_files=args.mmap_bin_files, + tokenizer=tokenizer, + reset_position_ids=args.reset_position_ids, + reset_attention_mask=args.reset_attention_mask, + eod_mask_loss=args.eod_mask_loss, + create_attention_mask=args.create_attention_mask_in_dataloader, + object_storage_cache_path=args.object_storage_cache_path, + mid_level_dataset_surplus=args.mid_level_dataset_surplus, + allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens, + fast_cache_load=args.dataloader_fast_cache_load, + sequences_per_dataset=sequences_per_dataset, + defer_npy_index_mmap=args.dataloader_defer_npy_index_mmap, + context_parallel_size=args.context_parallel_size, + ) + + +def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None): + """Build the train test and validation datasets. + + Args: + train_val_test_num_samples : A list containing the number of samples in train test and validation. + """ + args = get_args() + config = core_gpt_dataset_config_from_args(args) + + is_packed_sequence = False + if args.sft: + dataset_type = SFTDataset + is_packed_sequence = True # SFT always uses packed sequence + else: + if args.mock_data: + dataset_type = MockGPTDataset + else: + dataset_type = GPTDataset + + print_rank_0("> building train, validation, and test datasets for GPT ...") + + train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( + dataset_type, + train_val_test_num_samples, + partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence), + config + ).build() + + print_rank_0("> finished creating GPT datasets ...") + + return train_ds, valid_ds, test_ds + + +if __name__ == "__main__": + # Timestamp right after entering __main__ block (after all imports/library setup) + _MAIN_ENTRY_TIME = time.time() + + # Register startup timestamps for timing report in pretrain() + set_startup_timestamps(program_start=_PROGRAM_START_TIME, main_entry=_MAIN_ENTRY_TIME) + + # Temporary for transition to core datasets + train_valid_test_datasets_provider.is_distributed = True + + # Optionally enable inprocess restart on pretrain + pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) + + pretrain(train_valid_test_datasets_provider, + partial(model_provider, hybrid_builder), + ModelType.encoder_or_decoder, + forward_step, + args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + store=store, + extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, + ) diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml index f5de6eaac72..4b258afe0d6 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m/model_config.yaml @@ -24,7 +24,7 @@ MODEL_ARGS: --pipeline-model-parallel-size: 1 --expert-model-parallel-size: 1 --use-mcore-models: true - --model-provider: mamba + --model-provider: hybrid --init-method-std: 0.0198 --untie-embeddings-and-output-weights: true --disable-bias-linear: true @@ -35,7 +35,7 @@ MODEL_ARGS: --num-attention-heads: 16 --kv-channels: 128 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- - --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec + --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec --normalization: RMSNorm --swiglu: true --attention-dropout: 0.0 diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml index b10698d521f..bd86d2faa44 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_chunked_prefill/model_config.yaml @@ -24,7 +24,7 @@ MODEL_ARGS: --pipeline-model-parallel-size: 1 --expert-model-parallel-size: 1 --use-mcore-models: true - --model-provider: mamba + --model-provider: hybrid --init-method-std: 0.0198 --untie-embeddings-and-output-weights: true --disable-bias-linear: true @@ -35,7 +35,7 @@ MODEL_ARGS: --num-attention-heads: 16 --kv-channels: 128 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- - --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec + --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec --normalization: RMSNorm --swiglu: true --attention-dropout: 0.0 diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_cp1_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_cp1_dgx_a100_1N8G/model_config.yaml index 6d40098499d..9add53f8a49 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_cp1_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp1_cp1_dgx_a100_1N8G/model_config.yaml @@ -9,7 +9,7 @@ MODEL_ARGS: --group-query-attention: true --num-query-groups: 8 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M- - --spec: "[megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec]" + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec]" --log-params-norm: true --log-num-zeros-in-grad: true --log-validation-ppl-to-tensorboard: true diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_vpp2_cp1_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_vpp2_cp1_dgx_a100_1N8G/model_config.yaml index 51492f98c6e..25df6aa0359 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_vpp2_cp1_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp2_vpp2_cp1_dgx_a100_1N8G/model_config.yaml @@ -9,7 +9,7 @@ MODEL_ARGS: --group-query-attention: true --num-query-groups: 8 --hybrid-layer-pattern: M-M-M-M*-M-|M-M-M*-M-M-|M-M*-M-M-M-|M*-M-M-M-M- - --spec: "[megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec]" + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec]" --log-params-norm: true --log-num-zeros-in-grad: true --log-validation-ppl-to-tensorboard: true diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp4_cp1_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp4_cp1_dgx_a100_1N8G/model_config.yaml index 6eff846884a..fe4f9e63714 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp4_cp1_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp1_pp4_cp1_dgx_a100_1N8G/model_config.yaml @@ -9,7 +9,7 @@ MODEL_ARGS: --group-query-attention: true --num-query-groups: 8 --hybrid-layer-pattern: M-M-M-M*-M-|M-M-M*-M-M-|M-M*-M-M-M-|M*-M-M-M-M- - --spec: "[megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec]" + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec]" --log-params-norm: true --log-num-zeros-in-grad: true --log-validation-ppl-to-tensorboard: true diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml index 8c655bc135c..b89d305dc63 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml @@ -9,7 +9,7 @@ MODEL_ARGS: --group-query-attention: true --num-query-groups: 8 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M- - --spec: "[megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec]" + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec]" --log-params-norm: true --log-num-zeros-in-grad: true --log-validation-ppl-to-tensorboard: true @@ -55,4 +55,6 @@ MODEL_ARGS: --bf16: true --attention-backend: unfused --log-memory-to-tensorboard: true + --async-save: true + --use-persistent-ckpt-worker: true TEST_TYPE: regular diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp4_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp4_dgx_a100_1N8G/model_config.yaml index 44b588ee140..3efc155949f 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp4_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp4_dgx_a100_1N8G/model_config.yaml @@ -9,7 +9,7 @@ MODEL_ARGS: --group-query-attention: true --num-query-groups: 8 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M*-M-M-M-M- - --spec: "[megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec]" + --spec: "[megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec]" --log-params-norm: true --log-num-zeros-in-grad: true --log-validation-ppl-to-tensorboard: true diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml index 26708b32a60..02c5cc3055c 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_cudagraphs/model_config.yaml @@ -22,7 +22,7 @@ MODEL_ARGS: --pipeline-model-parallel-size: 1 --expert-model-parallel-size: 1 --use-mcore-models: true - --model-provider: mamba + --model-provider: hybrid --init-method-std: 0.0198 --untie-embeddings-and-output-weights: true --disable-bias-linear: true @@ -33,7 +33,7 @@ MODEL_ARGS: --num-attention-heads: 16 --kv-channels: 128 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- - --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec + --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec --normalization: RMSNorm --swiglu: true --attention-dropout: 0.0 diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml index 3964bcb8ecb..2543f59e668 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_static_inference_tp1_pp1_2B_logitsmatch/model_config.yaml @@ -22,7 +22,7 @@ MODEL_ARGS: --pipeline-model-parallel-size: 1 --expert-model-parallel-size: 1 --use-mcore-models: true - --model-provider: mamba + --model-provider: hybrid --init-method-std: 0.0198 --untie-embeddings-and-output-weights: true --disable-bias-linear: true @@ -33,7 +33,7 @@ MODEL_ARGS: --num-attention-heads: 16 --kv-channels: 128 --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- - --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec + --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec --normalization: RMSNorm --swiglu: true --attention-dropout: 0.0 diff --git a/tests/test_utils/recipes/h100/mamba.yaml b/tests/test_utils/recipes/h100/mamba.yaml index 703fb53160f..72b44495617 100644 --- a/tests/test_utils/recipes/h100/mamba.yaml +++ b/tests/test_utils/recipes/h100/mamba.yaml @@ -44,7 +44,7 @@ spec: "TENSORBOARD_PATH={assets_dir}/tensorboard" "CHECKPOINT_SAVE_PATH={artifacts_dir}/checkpoints" "CHECKPOINT_LOAD_PATH=/mnt/artifacts/model/{name}" - "TRAINING_SCRIPT_PATH=pretrain_mamba.py" + "TRAINING_SCRIPT_PATH=pretrain_hybrid.py" "TRAINING_PARAMS_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/model_config.yaml" "GOLDEN_VALUES_PATH=./tests/functional_tests/test_cases/{model}/{test_case}/golden_values_{environment}_{platforms}.json" "N_REPEAT={n_repeat}" diff --git a/tests/unit_tests/post_training/test_modelopt_model_builder.py b/tests/unit_tests/post_training/test_modelopt_model_builder.py index b489d659ec4..2ab8ebfe947 100644 --- a/tests/unit_tests/post_training/test_modelopt_model_builder.py +++ b/tests/unit_tests/post_training/test_modelopt_model_builder.py @@ -39,7 +39,7 @@ def test_model_provider_switches_to_modelopt_builder(monkeypatch): monkeypatch.setattr(mp, "has_nvidia_modelopt", True) monkeypatch.setattr(mp, "get_args", lambda: args) monkeypatch.setattr( - mp, "modelopt_gpt_mamba_builder", _sentinel_builder(modelopt_result, modelopt_calls) + mp, "modelopt_gpt_hybrid_builder", _sentinel_builder(modelopt_result, modelopt_calls) ) # original_builder should be ignored when ModelOpt is enabled. diff --git a/tools/run_mamba_text_generation_server.py b/tools/run_hybrid_text_generation_server.py similarity index 89% rename from tools/run_mamba_text_generation_server.py rename to tools/run_hybrid_text_generation_server.py index 33465f1bb4a..e70e5389e88 100644 --- a/tools/run_mamba_text_generation_server.py +++ b/tools/run_hybrid_text_generation_server.py @@ -8,4 +8,4 @@ from run_text_generation_server import main if __name__ == "__main__": - main(model_type="mamba") + main(model_type="hybrid") diff --git a/tools/run_mamba_text_generation_server_completions.py b/tools/run_hybrid_text_generation_server_completions.py similarity index 89% rename from tools/run_mamba_text_generation_server_completions.py rename to tools/run_hybrid_text_generation_server_completions.py index 33465f1bb4a..e70e5389e88 100644 --- a/tools/run_mamba_text_generation_server_completions.py +++ b/tools/run_hybrid_text_generation_server_completions.py @@ -8,4 +8,4 @@ from run_text_generation_server import main if __name__ == "__main__": - main(model_type="mamba") + main(model_type="hybrid") diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py index a1feef380ef..759fd67cbbc 100644 --- a/tools/run_inference_performance_test.py +++ b/tools/run_inference_performance_test.py @@ -9,7 +9,7 @@ import torch from gpt_builders import gpt_builder -from mamba_builders import mamba_builder +from hybrid_builders import hybrid_builder from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.engines import DynamicInferenceEngine, StaticInferenceEngine from megatron.core.inference.engines.abstract_engine import AbstractEngine diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py index 83a0bbc2369..a1f9a710505 100644 --- a/tools/run_text_generation_server.py +++ b/tools/run_text_generation_server.py @@ -15,7 +15,7 @@ import torch from gpt_builders import gpt_builder -from mamba_builders import mamba_builder +from hybrid_builders import hybrid_builder from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.engines import AbstractEngine, StaticInferenceEngine from megatron.core.inference.engines.abstract_engine import AbstractEngine @@ -138,8 +138,16 @@ def main(model_type: str = "gpt"): # Set up model and load checkpoint if model_type == "gpt": model_builder = gpt_builder - elif model_type == "mamba": - model_builder = mamba_builder + elif model_type in ("hybrid", "mamba"): + if model_type == "mamba": + import warnings + + warnings.warn( + 'model_type="mamba" is deprecated. Use model_type="hybrid" instead.', + DeprecationWarning, + stacklevel=2, + ) + model_builder = hybrid_builder else: raise ValueError(f"Invalid model provider {model_type}") model = get_model(partial(model_provider, model_builder), wrap_with_ddp=False) diff --git a/train_rl.py b/train_rl.py index f75bba12997..1afa9c79b25 100644 --- a/train_rl.py +++ b/train_rl.py @@ -8,7 +8,7 @@ import torch from gpt_builders import gpt_builder -from mamba_builders import mamba_builder +from hybrid_builders import hybrid_builder from megatron.core import mpu from megatron.core.enums import ModelType from megatron.core.models.gpt import GPTModel @@ -392,7 +392,7 @@ def _model_builder( args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None ): if is_hybrid_model(args): - return mamba_builder( + return hybrid_builder( args, pre_process, post_process, From 6f5cf1d1610fe81876a0701971a7287e2fb81cca Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 18:37:02 +0000 Subject: [PATCH 05/23] Move non-core test files back out of the core PR Tests under tests/unit_tests/inference/ and tests/unit_tests/resharding/ test the inference engine and resharding infrastructure, not megatron/core classes. They belong in the companion non-core PR (#4159). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../contexts/test_dynamic_context.py | 2 +- .../inference/engines/test_dynamic_engine.py | 44 +++++++++---------- ...2e.py => test_mamba_prefix_caching_e2e.py} | 8 ++-- .../test_prefix_caching_cuda_graphs.py | 12 ++--- .../unit_tests/resharding/test_model_swap.py | 10 ++--- 5 files changed, 38 insertions(+), 38 deletions(-) rename tests/unit_tests/inference/engines/{test_hybrid_prefix_caching_e2e.py => test_mamba_prefix_caching_e2e.py} (99%) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index cfff78a1780..385ea09e345 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -16,7 +16,7 @@ ) from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams -from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 0b476ad3ae3..a3b2ce71e60 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio import gc @@ -46,8 +46,8 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec -from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.models.mamba.mamba_model import MambaModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -65,7 +65,7 @@ def skip_if_mamba_sequence_packing_not_available(model_provider: str): - if model_provider in ("hybrid", "mamba"): + if model_provider == "mamba": sequence_packing_available, reason_for_no_sequence_packing = ( _check_mamba_sequence_packing_support() ) @@ -366,7 +366,7 @@ def _build_test_env(cls, test_config): post_process=parallel_state.is_pipeline_last_stage(), mtp_block_spec=mtp_block_spec, ).cuda() - elif test_config.model_provider in ("hybrid", "mamba"): + elif test_config.model_provider == "mamba": pp_size = test_config.pipeline_model_parallel_size # Transformer config. transformer_config = TransformerConfig( @@ -406,9 +406,9 @@ def _build_test_env(cls, test_config): ) # Mamba model. - model = HybridModel( + model = MambaModel( config=transformer_config, - hybrid_stack_spec=hybrid_stack_spec, + mamba_stack_spec=mamba_stack_spec, vocab_size=test_config.vocab_size, max_sequence_length=test_config.max_sequence_length, parallel_output=True, @@ -567,7 +567,7 @@ def teardown_class(cls): @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) @pytest.mark.parametrize("num_cuda_graphs", [None, 1, 4, -1]) @pytest.mark.parametrize("cuda_graph_scope", [[], [CudaGraphScope.full_iteration_inference]]) def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None: @@ -625,7 +625,7 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None if model_provider == "gpt": expected_generated_tokens_list = gpt_expected_generated_tokens - elif model_provider in ("hybrid", "mamba"): + elif model_provider == "mamba": expected_generated_tokens_list = mamba_expected_generated_tokens else: raise ValueError(f"Invalid model_provider {model_provider}") @@ -686,7 +686,7 @@ def test_token_overflow_nontransient(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) def test_block_overflow(self, model_provider: str) -> None: """Test block overflow.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -732,7 +732,7 @@ def test_block_overflow_insufficient_kv_cache(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) def test_multi_add(self, model_provider: str) -> None: """Test adding multiple requests simultaneously.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -742,7 +742,7 @@ def test_multi_add(self, model_provider: str) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) def test_fixed_output_lengths(self, model_provider: str) -> None: """Test generating a fixed number of output tokens.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -785,7 +785,7 @@ def test_cuda_graph_token_counts(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) @torch.inference_mode() def test_generate_function(self, model_provider: str) -> None: """Test the generate function that processes multiple prompts at once.""" @@ -879,7 +879,7 @@ async def test_run_engine(self): not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @pytest.mark.skipif(not is_te_min_version("2.2.0"), reason="TE 2.2.0 is required") - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) def test_fp8_inference(self, model_provider: str): skip_if_mamba_sequence_packing_not_available(model_provider) @@ -1085,7 +1085,7 @@ def test_log_probs_token_correspondence(self): @pytest.mark.parametrize("ep_size", [1, 2]) @pytest.mark.parametrize("pp_size", [1, 2]) @pytest.mark.parametrize("tp_size", [1, 2]) - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) @pytest.mark.parametrize("transformer_impl", ["local", "inference_optimized"]) @torch.inference_mode() def test_parallel_inference( @@ -1124,7 +1124,7 @@ def test_parallel_inference( "when tp_size > 1." ) ) - if model_provider in ("hybrid", "mamba"): + if model_provider == "mamba": pytest.skip( reason="Mamba model is not supported with the inference optimized transformer." ) @@ -1292,11 +1292,11 @@ def test_mamba_chunked_prefill(self): """ Test chunked prefill with a Mamba model. """ - skip_if_mamba_sequence_packing_not_available("hybrid") + skip_if_mamba_sequence_packing_not_available("mamba") # Context max tokens = 50. test_config = DynamicEngineTestConfig( - model_provider="hybrid", + model_provider="mamba", num_requests=0, num_tokens_to_generate=None, num_tokens_total=200, @@ -3058,7 +3058,7 @@ def _create_model(self, model_provider, num_cuda_graphs): pre_process=parallel_state.is_pipeline_first_stage(), post_process=parallel_state.is_pipeline_last_stage(), ).cuda() - elif model_provider in ("hybrid", "mamba"): + elif model_provider == "mamba": config = TransformerConfig( params_dtype=torch.bfloat16, num_layers=3, @@ -3074,9 +3074,9 @@ def _create_model(self, model_provider, num_cuda_graphs): add_bias_linear=True, is_hybrid_model=True, ) - model = HybridModel( + model = MambaModel( config=config, - hybrid_stack_spec=hybrid_stack_spec, + mamba_stack_spec=mamba_stack_spec, vocab_size=CHUNKED_CG_VOCAB_SIZE, max_sequence_length=CHUNKED_CG_MAX_SEQ_LEN, parallel_output=True, @@ -3162,7 +3162,7 @@ def _run_to_completion(self, engine, prompts, num_tokens_to_generate): return finished, step_count - @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) + @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) @pytest.mark.parametrize("chunked_prefill", [False, True]) @pytest.mark.parametrize("num_cuda_graphs", [None, 2]) @torch.inference_mode() diff --git a/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py b/tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py similarity index 99% rename from tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py rename to tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py index 303cf76d122..ce21c775b73 100644 --- a/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py +++ b/tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py @@ -54,8 +54,8 @@ from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec -from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.models.mamba.mamba_model import MambaModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -131,9 +131,9 @@ def _create_model(self, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = HybridModel( + model = MambaModel( config=transformer_config, - hybrid_stack_spec=hybrid_stack_spec, + mamba_stack_spec=mamba_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, diff --git a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py index 26a81c5baef..52a05f7f80f 100644 --- a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py +++ b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py @@ -37,8 +37,8 @@ ) from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec -from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.models.mamba.mamba_model import MambaModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -121,9 +121,9 @@ def _create_model(self, model_type, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = HybridModel( + model = MambaModel( config=config, - hybrid_stack_spec=hybrid_stack_spec, + mamba_stack_spec=mamba_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, @@ -343,9 +343,9 @@ def _create_hybrid_model(self, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = HybridModel( + model = MambaModel( config=config, - hybrid_stack_spec=hybrid_stack_spec, + mamba_stack_spec=mamba_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, diff --git a/tests/unit_tests/resharding/test_model_swap.py b/tests/unit_tests/resharding/test_model_swap.py index e2d6a2bd096..70d81d97829 100644 --- a/tests/unit_tests/resharding/test_model_swap.py +++ b/tests/unit_tests/resharding/test_model_swap.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. import copy import gc import os @@ -37,8 +37,8 @@ try: import mamba_ssm # noqa: F401 - from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec - from megatron.core.models.hybrid.hybrid_model import HybridModel + from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec + from megatron.core.models.mamba.mamba_model import MambaModel has_mamba_deps = True except Exception: @@ -203,9 +203,9 @@ def _build_mamba( parallel_output: bool = True, ): pre_process, post_process = _pp_flags(pg_collection) - model = HybridModel( + model = MambaModel( config=config, - hybrid_stack_spec=hybrid_stack_spec, + mamba_stack_spec=mamba_stack_spec, vocab_size=vocab_size, max_sequence_length=seq_len, hybrid_layer_pattern=hybrid_layer_pattern, From 834c84dcb3894d9f3b7785cae571e54f240e3b51 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 18:39:01 +0000 Subject: [PATCH 06/23] Remove unrelated upstream change from core PR .claude/skills/respond-to-issue/SKILL.md is not part of this rename. Co-Authored-By: Claude Opus 4.6 (1M context) --- .claude/skills/respond-to-issue/SKILL.md | 64 ++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 .claude/skills/respond-to-issue/SKILL.md diff --git a/.claude/skills/respond-to-issue/SKILL.md b/.claude/skills/respond-to-issue/SKILL.md new file mode 100644 index 00000000000..1d513da2e63 --- /dev/null +++ b/.claude/skills/respond-to-issue/SKILL.md @@ -0,0 +1,64 @@ +--- +name: respond-to-issue +description: Research and draft a response to a GitHub issue or question from an external contributor. Use when the user shares a GitHub issue URL or asks to respond to a community question. +user_invocable: true +argument: "" +--- + +# Respond to GitHub Issue + +Help a maintainer draft a high-quality response to a GitHub issue from an external contributor. + +## Workflow + +### 1. Understand the issue + +- Fetch the issue using `gh issue view --repo NVIDIA/Megatron-LM --json title,body,comments,labels,state`. +- Read the title, body, and all existing comments to understand the full context. +- Identify the type: bug report, feature request, question, or discussion. + +### 2. Research the codebase + +- Based on what the issue is asking, search the Megatron-LM codebase for the relevant code. +- Read the relevant source files to understand the current behavior. +- If the issue references specific files or functions, read those directly. +- Check `git log --oneline -20 -- ` to see if there have been recent changes that address or relate to the issue. +- Use `git log -S "" --oneline` to trace when code was added or removed — this is especially useful for questions about unused/deprecated code or missing features. +- If the issue is about a bug, try to confirm whether the reported behavior matches the code. +- Check whether an existing PR already addresses the issue: `gh pr list --repo NVIDIA/Megatron-LM --search "" --limit 5`. + +### 3. Verify before citing + +Before including specific details in the response, verify them: +- If citing a commit hash, confirm the commit message and diff match what you're claiming (`git show --stat`). +- If citing a file path and line number, re-read the file to confirm the line content is correct. +- If claiming code is unused or missing, do a thorough grep to make sure you haven't missed a reference. + +### 4. Draft the response + +Write a response that: +- Is technically accurate and grounded in the actual code (cite file paths and line numbers where helpful). +- Is respectful and welcoming to external contributors. +- Directly addresses the question or concern raised. +- If the contributor identified a real gap or bug, acknowledge it clearly. +- If there's a workaround, mention it. +- If work is planned or a fix would be welcome, say so and suggest next steps (e.g., "a PR to address this would be welcome"). +- Keeps the tone professional but friendly. +- Is concise -- contributors appreciate direct answers, not walls of text. + +### 5. Suggest follow-up actions + +If the issue identifies something cleanly actionable (dead code to remove, a small bug fix, a missing feature), tell the maintainer and offer to create a branch and PR to address it — don't just draft a comment. + +### 6. Present to the maintainer + +Show the drafted response to the user (the maintainer) for review. Do NOT post it to GitHub automatically. The maintainer will decide whether to post it, edit it, or ask for changes. + +Format the draft as a quoted markdown block so it's easy to copy. + +## Important guidelines + +- Never post comments to GitHub without explicit approval from the user. +- If you're unsure about the answer, say so clearly in your draft and flag the uncertainty for the maintainer. +- If the issue is outside the scope of what you can determine from the code, tell the maintainer what you found and what remains unclear. +- Check whether similar issues exist that might be relevant: `gh issue list --repo NVIDIA/Megatron-LM --search "" --limit 5`. From 2b827837051f69695d050153b5401f7891e2f812 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 18:40:20 +0000 Subject: [PATCH 07/23] Add missing inference and resharding test updates These test files were accidentally dropped during the core/non-core split. They test inference engines and resharding (not megatron/core classes), so they belong in this non-core PR. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../contexts/test_dynamic_context.py | 2 +- .../inference/engines/test_dynamic_engine.py | 44 +++++++++---------- ...e.py => test_hybrid_prefix_caching_e2e.py} | 8 ++-- .../test_prefix_caching_cuda_graphs.py | 12 ++--- .../unit_tests/resharding/test_model_swap.py | 10 ++--- 5 files changed, 38 insertions(+), 38 deletions(-) rename tests/unit_tests/inference/engines/{test_mamba_prefix_caching_e2e.py => test_hybrid_prefix_caching_e2e.py} (99%) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 385ea09e345..cfff78a1780 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -16,7 +16,7 @@ ) from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams -from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.transformer_config import TransformerConfig from tests.unit_tests.test_utilities import Utils diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index a3b2ce71e60..0b476ad3ae3 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import asyncio import gc @@ -46,8 +46,8 @@ get_gpt_mtp_block_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -65,7 +65,7 @@ def skip_if_mamba_sequence_packing_not_available(model_provider: str): - if model_provider == "mamba": + if model_provider in ("hybrid", "mamba"): sequence_packing_available, reason_for_no_sequence_packing = ( _check_mamba_sequence_packing_support() ) @@ -366,7 +366,7 @@ def _build_test_env(cls, test_config): post_process=parallel_state.is_pipeline_last_stage(), mtp_block_spec=mtp_block_spec, ).cuda() - elif test_config.model_provider == "mamba": + elif test_config.model_provider in ("hybrid", "mamba"): pp_size = test_config.pipeline_model_parallel_size # Transformer config. transformer_config = TransformerConfig( @@ -406,9 +406,9 @@ def _build_test_env(cls, test_config): ) # Mamba model. - model = MambaModel( + model = HybridModel( config=transformer_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=test_config.vocab_size, max_sequence_length=test_config.max_sequence_length, parallel_output=True, @@ -567,7 +567,7 @@ def teardown_class(cls): @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @pytest.mark.parametrize("num_cuda_graphs", [None, 1, 4, -1]) @pytest.mark.parametrize("cuda_graph_scope", [[], [CudaGraphScope.full_iteration_inference]]) def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None: @@ -625,7 +625,7 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None if model_provider == "gpt": expected_generated_tokens_list = gpt_expected_generated_tokens - elif model_provider == "mamba": + elif model_provider in ("hybrid", "mamba"): expected_generated_tokens_list = mamba_expected_generated_tokens else: raise ValueError(f"Invalid model_provider {model_provider}") @@ -686,7 +686,7 @@ def test_token_overflow_nontransient(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_block_overflow(self, model_provider: str) -> None: """Test block overflow.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -732,7 +732,7 @@ def test_block_overflow_insufficient_kv_cache(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_multi_add(self, model_provider: str) -> None: """Test adding multiple requests simultaneously.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -742,7 +742,7 @@ def test_multi_add(self, model_provider: str) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_fixed_output_lengths(self, model_provider: str) -> None: """Test generating a fixed number of output tokens.""" skip_if_mamba_sequence_packing_not_available(model_provider) @@ -785,7 +785,7 @@ def test_cuda_graph_token_counts(self) -> None: @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @torch.inference_mode() def test_generate_function(self, model_provider: str) -> None: """Test the generate function that processes multiple prompts at once.""" @@ -879,7 +879,7 @@ async def test_run_engine(self): not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @pytest.mark.skipif(not is_te_min_version("2.2.0"), reason="TE 2.2.0 is required") - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) def test_fp8_inference(self, model_provider: str): skip_if_mamba_sequence_packing_not_available(model_provider) @@ -1085,7 +1085,7 @@ def test_log_probs_token_correspondence(self): @pytest.mark.parametrize("ep_size", [1, 2]) @pytest.mark.parametrize("pp_size", [1, 2]) @pytest.mark.parametrize("tp_size", [1, 2]) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @pytest.mark.parametrize("transformer_impl", ["local", "inference_optimized"]) @torch.inference_mode() def test_parallel_inference( @@ -1124,7 +1124,7 @@ def test_parallel_inference( "when tp_size > 1." ) ) - if model_provider == "mamba": + if model_provider in ("hybrid", "mamba"): pytest.skip( reason="Mamba model is not supported with the inference optimized transformer." ) @@ -1292,11 +1292,11 @@ def test_mamba_chunked_prefill(self): """ Test chunked prefill with a Mamba model. """ - skip_if_mamba_sequence_packing_not_available("mamba") + skip_if_mamba_sequence_packing_not_available("hybrid") # Context max tokens = 50. test_config = DynamicEngineTestConfig( - model_provider="mamba", + model_provider="hybrid", num_requests=0, num_tokens_to_generate=None, num_tokens_total=200, @@ -3058,7 +3058,7 @@ def _create_model(self, model_provider, num_cuda_graphs): pre_process=parallel_state.is_pipeline_first_stage(), post_process=parallel_state.is_pipeline_last_stage(), ).cuda() - elif model_provider == "mamba": + elif model_provider in ("hybrid", "mamba"): config = TransformerConfig( params_dtype=torch.bfloat16, num_layers=3, @@ -3074,9 +3074,9 @@ def _create_model(self, model_provider, num_cuda_graphs): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=CHUNKED_CG_VOCAB_SIZE, max_sequence_length=CHUNKED_CG_MAX_SEQ_LEN, parallel_output=True, @@ -3162,7 +3162,7 @@ def _run_to_completion(self, engine, prompts, num_tokens_to_generate): return finished, step_count - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) + @pytest.mark.parametrize("model_provider", ["gpt", "hybrid"]) @pytest.mark.parametrize("chunked_prefill", [False, True]) @pytest.mark.parametrize("num_cuda_graphs", [None, 2]) @torch.inference_mode() diff --git a/tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py similarity index 99% rename from tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py rename to tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py index ce21c775b73..303cf76d122 100644 --- a/tests/unit_tests/inference/engines/test_mamba_prefix_caching_e2e.py +++ b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py @@ -54,8 +54,8 @@ from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -131,9 +131,9 @@ def _create_model(self, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=transformer_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, diff --git a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py index 52a05f7f80f..26a81c5baef 100644 --- a/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py +++ b/tests/unit_tests/inference/engines/test_prefix_caching_cuda_graphs.py @@ -37,8 +37,8 @@ ) from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.ssm.mamba_mixer import _check_mamba_sequence_packing_support from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord @@ -121,9 +121,9 @@ def _create_model(self, model_type, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, @@ -343,9 +343,9 @@ def _create_hybrid_model(self, num_cuda_graphs=None): add_bias_linear=True, is_hybrid_model=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=VOCAB_SIZE, max_sequence_length=MAX_SEQ_LEN, parallel_output=True, diff --git a/tests/unit_tests/resharding/test_model_swap.py b/tests/unit_tests/resharding/test_model_swap.py index 70d81d97829..e2d6a2bd096 100644 --- a/tests/unit_tests/resharding/test_model_swap.py +++ b/tests/unit_tests/resharding/test_model_swap.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. import copy import gc import os @@ -37,8 +37,8 @@ try: import mamba_ssm # noqa: F401 - from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec - from megatron.core.models.mamba.mamba_model import MambaModel + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + from megatron.core.models.hybrid.hybrid_model import HybridModel has_mamba_deps = True except Exception: @@ -203,9 +203,9 @@ def _build_mamba( parallel_output: bool = True, ): pre_process, post_process = _pp_flags(pg_collection) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=vocab_size, max_sequence_length=seq_len, hybrid_layer_pattern=hybrid_layer_pattern, From f3533a1e95827c884b8c2d7b73a74b10205d37c9 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 18:44:06 +0000 Subject: [PATCH 08/23] Revert "Remove unrelated upstream change from core PR" This reverts commit 834c84dcb3894d9f3b7785cae571e54f240e3b51. --- .claude/skills/respond-to-issue/SKILL.md | 64 ------------------------ 1 file changed, 64 deletions(-) delete mode 100644 .claude/skills/respond-to-issue/SKILL.md diff --git a/.claude/skills/respond-to-issue/SKILL.md b/.claude/skills/respond-to-issue/SKILL.md deleted file mode 100644 index 1d513da2e63..00000000000 --- a/.claude/skills/respond-to-issue/SKILL.md +++ /dev/null @@ -1,64 +0,0 @@ ---- -name: respond-to-issue -description: Research and draft a response to a GitHub issue or question from an external contributor. Use when the user shares a GitHub issue URL or asks to respond to a community question. -user_invocable: true -argument: "" ---- - -# Respond to GitHub Issue - -Help a maintainer draft a high-quality response to a GitHub issue from an external contributor. - -## Workflow - -### 1. Understand the issue - -- Fetch the issue using `gh issue view --repo NVIDIA/Megatron-LM --json title,body,comments,labels,state`. -- Read the title, body, and all existing comments to understand the full context. -- Identify the type: bug report, feature request, question, or discussion. - -### 2. Research the codebase - -- Based on what the issue is asking, search the Megatron-LM codebase for the relevant code. -- Read the relevant source files to understand the current behavior. -- If the issue references specific files or functions, read those directly. -- Check `git log --oneline -20 -- ` to see if there have been recent changes that address or relate to the issue. -- Use `git log -S "" --oneline` to trace when code was added or removed — this is especially useful for questions about unused/deprecated code or missing features. -- If the issue is about a bug, try to confirm whether the reported behavior matches the code. -- Check whether an existing PR already addresses the issue: `gh pr list --repo NVIDIA/Megatron-LM --search "" --limit 5`. - -### 3. Verify before citing - -Before including specific details in the response, verify them: -- If citing a commit hash, confirm the commit message and diff match what you're claiming (`git show --stat`). -- If citing a file path and line number, re-read the file to confirm the line content is correct. -- If claiming code is unused or missing, do a thorough grep to make sure you haven't missed a reference. - -### 4. Draft the response - -Write a response that: -- Is technically accurate and grounded in the actual code (cite file paths and line numbers where helpful). -- Is respectful and welcoming to external contributors. -- Directly addresses the question or concern raised. -- If the contributor identified a real gap or bug, acknowledge it clearly. -- If there's a workaround, mention it. -- If work is planned or a fix would be welcome, say so and suggest next steps (e.g., "a PR to address this would be welcome"). -- Keeps the tone professional but friendly. -- Is concise -- contributors appreciate direct answers, not walls of text. - -### 5. Suggest follow-up actions - -If the issue identifies something cleanly actionable (dead code to remove, a small bug fix, a missing feature), tell the maintainer and offer to create a branch and PR to address it — don't just draft a comment. - -### 6. Present to the maintainer - -Show the drafted response to the user (the maintainer) for review. Do NOT post it to GitHub automatically. The maintainer will decide whether to post it, edit it, or ask for changes. - -Format the draft as a quoted markdown block so it's easy to copy. - -## Important guidelines - -- Never post comments to GitHub without explicit approval from the user. -- If you're unsure about the answer, say so clearly in your draft and flag the uncertainty for the maintainer. -- If the issue is outside the scope of what you can determine from the code, tell the maintainer what you found and what remains unclear. -- Check whether similar issues exist that might be relevant: `gh issue list --repo NVIDIA/Megatron-LM --search "" --limit 5`. From 420f5c9590688268363c71fe0da3dd38e8cfe993 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 18:54:49 +0000 Subject: [PATCH 09/23] Move MambaModel subclass from hybrid_model.py to mamba_model.py The backward-compat MambaModel wrapper belongs in the legacy mamba module, not in the canonical hybrid_model.py. This keeps hybrid_model.py clean with only HybridModel. Co-Authored-By: Claude Opus 4.6 (1M context) --- megatron/core/models/hybrid/hybrid_model.py | 17 -------------- megatron/core/models/mamba/__init__.py | 3 ++- megatron/core/models/mamba/mamba_model.py | 26 +++++++++++++++++++-- 3 files changed, 26 insertions(+), 20 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index cbdd028d9ba..6f7dd29b5b8 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -473,20 +473,3 @@ def forward( loss = self.compute_language_model_loss(labels, logits) return loss - - -class MambaModel(HybridModel): - """Backward-compatible wrapper that accepts the deprecated mamba_stack_spec kwarg.""" - - def __init__(self, *args, mamba_stack_spec: ModuleSpec = None, **kwargs): - log_single_rank( - logger, logging.WARNING, "MambaModel has been deprecated. Use HybridModel instead." - ) - if mamba_stack_spec is not None: - if 'hybrid_stack_spec' in kwargs or (args and len(args) >= 2): - raise ValueError( - "Cannot specify both hybrid_stack_spec and mamba_stack_spec. " - "mamba_stack_spec has been deprecated; use hybrid_stack_spec instead." - ) - kwargs['hybrid_stack_spec'] = mamba_stack_spec - super().__init__(*args, **kwargs) diff --git a/megatron/core/models/mamba/__init__.py b/megatron/core/models/mamba/__init__.py index 2a52cedd9b5..4de391a62c9 100644 --- a/megatron/core/models/mamba/__init__.py +++ b/megatron/core/models/mamba/__init__.py @@ -2,4 +2,5 @@ # Backward-compatible re-exports. The canonical location is now # megatron.core.models.hybrid. -from megatron.core.models.hybrid.hybrid_model import HybridModel, MambaModel +from megatron.core.models.hybrid.hybrid_model import HybridModel +from megatron.core.models.mamba.mamba_model import MambaModel diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py index f0494e802cb..236865e67eb 100644 --- a/megatron/core/models/mamba/mamba_model.py +++ b/megatron/core/models/mamba/mamba_model.py @@ -1,5 +1,27 @@ # Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. -# Backward-compatible re-export. The canonical location is now -# megatron.core.models.hybrid.hybrid_model. +import logging + +from megatron.core.models.hybrid.hybrid_model import HybridModel # noqa: F401 from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 # pylint: disable=unused-import +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.utils import log_single_rank + +logger = logging.getLogger(__name__) + + +class MambaModel(HybridModel): + """Backward-compatible wrapper that accepts the deprecated mamba_stack_spec kwarg.""" + + def __init__(self, *args, mamba_stack_spec: ModuleSpec = None, **kwargs): + log_single_rank( + logger, logging.WARNING, "MambaModel has been deprecated. Use HybridModel instead." + ) + if mamba_stack_spec is not None: + if 'hybrid_stack_spec' in kwargs or (args and len(args) >= 2): + raise ValueError( + "Cannot specify both hybrid_stack_spec and mamba_stack_spec. " + "mamba_stack_spec has been deprecated; use hybrid_stack_spec instead." + ) + kwargs['hybrid_stack_spec'] = mamba_stack_spec + super().__init__(*args, **kwargs) From 642664637dc96bd066ffc94e00e295ce33f4b52e Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 6 Apr 2026 20:55:17 +0000 Subject: [PATCH 10/23] Fix imports --- megatron/core/models/mamba/mamba_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py index 236865e67eb..fa92f8685ba 100644 --- a/megatron/core/models/mamba/mamba_model.py +++ b/megatron/core/models/mamba/mamba_model.py @@ -2,8 +2,8 @@ import logging -from megatron.core.models.hybrid.hybrid_model import HybridModel # noqa: F401 from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 # pylint: disable=unused-import +from megatron.core.models.hybrid.hybrid_model import HybridModel # noqa: F401 from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.utils import log_single_rank From a1fb376c4afc1e927b3d40fbdb7ef910669b9f2c Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 9 Apr 2026 00:14:02 +0000 Subject: [PATCH 11/23] Add CODEOWNERS entry for megatron/core/models/hybrid/ Co-Authored-By: Claude Opus 4.6 (1M context) --- .github/CODEOWNERS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d9b619f9559..d2560a8d3ed 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -7,6 +7,8 @@ megatron/core/models/multimodal/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/mul megatron/core/models/mamba/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/hybrid-mamba megatron/core/ssm/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/hybrid-mamba +megatron/core/models/hybrid/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/hybrid-model + megatron/core/datasets/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/datasets megatron/core/tokenizers/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/tokenizers From 91897dcf1711ec21833c97f1c2891019685f9ee7 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 13 Apr 2026 18:00:59 +0000 Subject: [PATCH 12/23] Rename mamba_submodules parameter to hybrid_submodules This parameter passes HybridStackSubmodules (formerly MambaStackSubmodules), so its name should match the renamed class. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../core/models/hybrid/hybrid_layer_specs.py | 4 ++-- megatron/core/models/hybrid/hybrid_model.py | 6 ++--- .../transformer/multi_token_prediction.py | 22 +++++++++---------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 2cc6c8a8870..c690a7ba53a 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -61,7 +61,7 @@ enorm=TENorm, hnorm=TENorm, eh_proj=TEColumnParallelLinear, - mtp_model_layer=None, # Built via pattern + mamba_submodules + mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), ) @@ -210,7 +210,7 @@ enorm=TENorm, hnorm=TENorm, eh_proj=InferenceColumnParallelLinear, - mtp_model_layer=None, # Built via pattern + mamba_submodules + mtp_model_layer=None, # Built via pattern + hybrid_submodules layer_norm=TENorm, ), ) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index 6f9d6793501..eb8a2bf7ee4 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -240,8 +240,8 @@ def __init__( # MTP block - uses mtp_block_spec from hybrid_stack_spec.submodules if self.mtp_process: - mamba_submodules = hybrid_stack_spec.submodules - mtp_block_spec = mamba_submodules.mtp_block_spec + hybrid_submodules = hybrid_stack_spec.submodules + mtp_block_spec = hybrid_submodules.mtp_block_spec assert mtp_block_spec is not None, ( "MTP pattern specified but mtp_block_spec is None in hybrid_stack_spec.submodules. " "Ensure hybrid_stack_spec includes mtp_block_spec for MTP support." @@ -254,7 +254,7 @@ def __init__( vp_stage=self.vp_stage, mtp_layer_pattern=self.mtp_pattern, mtp_num_depths=self.mtp_num_depths, - mamba_submodules=mamba_submodules, + hybrid_submodules=hybrid_submodules, ) # Output diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index fb940289cb2..37be871e0c3 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -738,7 +738,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, # For Mamba path - pattern and submodules to build inner layers directly mtp_layer_pattern: Optional[str] = None, - mamba_submodules: Optional[HybridStackSubmodules] = None, + hybrid_submodules: Optional[HybridStackSubmodules] = None, ): super().__init__(config=config) self.sequence_parallel = config.sequence_parallel @@ -811,13 +811,13 @@ def __init__( # Build inner layers: two possible paths # 1. Hybrid path: use HybridStack for hybrid pattern support # 2. GPT path: single TransformerLayer - if mtp_layer_pattern is not None and mamba_submodules is not None: + if mtp_layer_pattern is not None and hybrid_submodules is not None: from megatron.core.models.hybrid.hybrid_block import HybridStack from megatron.core.models.hybrid.hybrid_layer_allocation import validate_segment_layers self.mtp_model_layer = HybridStack( config=self.config, - submodules=mamba_submodules, + submodules=hybrid_submodules, layer_type_list=validate_segment_layers(mtp_layer_pattern), pp_layer_offset=0, pre_process=True, # Always receives input from eh_proj @@ -1272,7 +1272,7 @@ def __init__( # New: For Mamba path with unified pattern syntax mtp_layer_pattern: Optional[str] = None, mtp_num_depths: int = 0, - mamba_submodules: Optional["HybridStackSubmodules"] = None, + hybrid_submodules: Optional["HybridStackSubmodules"] = None, ): super().__init__(config=config) self.submodules = _get_mtp_block_submodules(config, spec) @@ -1280,7 +1280,7 @@ def __init__( self.vp_stage = vp_stage self.mtp_layer_pattern = mtp_layer_pattern self.mtp_num_depths = mtp_num_depths - self.mamba_submodules = mamba_submodules + self.hybrid_submodules = hybrid_submodules self.mtp_use_repeated_layer = self.config.mtp_use_repeated_layer vp_size = config.virtual_pipeline_model_parallel_size @@ -1327,7 +1327,7 @@ def build_layer_legacy(layer_spec, layer_number): ) return module - def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_submodules): + def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, hybrid_submodules): """Build layer using pattern-based approach (new Mamba path).""" fp8_init_context = get_fp8_context(self.config, is_init=True) with fp8_init_context: @@ -1338,12 +1338,12 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_ vp_stage=self.vp_stage, pg_collection=pg_collection, mtp_layer_pattern=mtp_layer_pattern, - mamba_submodules=mamba_submodules, + hybrid_submodules=hybrid_submodules, ) return module - # New Mamba path: use mtp_layer_pattern and mamba_submodules - if self.mtp_layer_pattern is not None and self.mamba_submodules is not None: + # New Mamba path: use mtp_layer_pattern and hybrid_submodules + if self.mtp_layer_pattern is not None and self.hybrid_submodules is not None: if self.mtp_use_repeated_layer: # Shared/repeated layer: build one layer, use it for all depths layer_spec = self.submodules.layer_specs[0] @@ -1351,7 +1351,7 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_ layer_spec, layer_number=1, mtp_layer_pattern=self.mtp_layer_pattern, - mamba_submodules=self.mamba_submodules, + hybrid_submodules=self.hybrid_submodules, ) self.layers = torch.nn.ModuleList([shared_layer]) else: @@ -1364,7 +1364,7 @@ def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, mamba_ ], layer_number=i + 1, mtp_layer_pattern=self.mtp_layer_pattern, - mamba_submodules=self.mamba_submodules, + hybrid_submodules=self.hybrid_submodules, ) for i in range(num_depths) ] From f85130584bf4ee2b3c573cf382b63537c0f568c5 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 13 Apr 2026 19:24:34 +0000 Subject: [PATCH 13/23] Fix black formatting in multi_token_prediction.py Co-Authored-By: Claude Opus 4.6 (1M context) --- megatron/core/transformer/multi_token_prediction.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 37be871e0c3..1044efe94c5 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -1327,7 +1327,9 @@ def build_layer_legacy(layer_spec, layer_number): ) return module - def build_layer_with_pattern(layer_spec, layer_number, mtp_layer_pattern, hybrid_submodules): + def build_layer_with_pattern( + layer_spec, layer_number, mtp_layer_pattern, hybrid_submodules + ): """Build layer using pattern-based approach (new Mamba path).""" fp8_init_context = get_fp8_context(self.config, is_init=True) with fp8_init_context: From b88d93b09f2ec80e71450ea4f44f6166658f4301 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 18 Apr 2026 00:50:18 +0000 Subject: [PATCH 14/23] Rename test classes that reference the renamed model classes Test classes named after MambaModel/MambaStack/MambaStackSubmodules are renamed to match the new Hybrid class names: - TestMambaModel -> TestHybridModel - TestMambaQKLayernorm -> TestHybridQKLayernorm - TestMambaWithDynamicInference -> TestHybridWithDynamicInference - TestMambaMoEModel -> TestHybridMoEModel - TestMambaBlock -> TestHybridBlock - TestModelOptMambaModel -> TestModelOptHybridModel - TestMultiTokenPredictionMamba -> TestMultiTokenPredictionHybrid - TestParallelMambaBlockCudagraphs -> TestParallelHybridBlockCudagraphs Also updates the TestHybridQKLayernorm class (added by upstream merge) to use HybridModel/hybrid_stack_spec consistently. Test classes for Mamba-specific SSM components (MambaLayer, MambaMixer, MambaContextParallel, MambaMetadata, MambaSlotAllocator, etc.) are unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/unit_tests/models/test_hybrid_model.py | 18 +++++++++--------- .../unit_tests/models/test_hybrid_moe_model.py | 4 ++-- .../post_training/test_modelopt_module_spec.py | 2 +- tests/unit_tests/ssm/test_hybrid_block.py | 2 +- .../unit_tests/transformer/test_cuda_graphs.py | 2 +- .../transformer/test_multi_token_prediction.py | 2 +- 6 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/unit_tests/models/test_hybrid_model.py b/tests/unit_tests/models/test_hybrid_model.py index e6356e89ba9..c4bdc147621 100644 --- a/tests/unit_tests/models/test_hybrid_model.py +++ b/tests/unit_tests/models/test_hybrid_model.py @@ -26,7 +26,7 @@ from tests.unit_tests.test_utilities import Utils -class TestMambaModel: +class TestHybridModel: def setup_method(self, method): Utils.initialize_model_parallel(1, 1) @@ -292,7 +292,7 @@ def test_with_custom_process_groups(self, tmp_path, tp_size, cp_size, pp_size): assert logits.shape[2] == divide(model.vocab_size, tp_size) -class TestMambaQKLayernorm: +class TestHybridQKLayernorm: def setup_method(self, method): Utils.initialize_model_parallel(1, 1) @@ -309,9 +309,9 @@ def _build_model(self, **config_overrides): use_cpu_initialization=True, **config_overrides, ) - return MambaModel( + return HybridModel( config=config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=100, max_sequence_length=4, hybrid_layer_pattern="M*-", @@ -371,7 +371,7 @@ def test_spec_provided_norm_not_overwritten(self): ) # Build a spec that explicitly sets q/k layernorm to IdentityOp - spec = copy.deepcopy(mamba_stack_spec) + spec = copy.deepcopy(hybrid_stack_spec) spec.submodules.attention_layer.submodules.self_attention.submodules.q_layernorm = ( IdentityOp ) @@ -386,9 +386,9 @@ def test_spec_provided_norm_not_overwritten(self): use_cpu_initialization=True, qk_layernorm=True, ) - model = MambaModel( + model = HybridModel( config=config, - mamba_stack_spec=spec, + hybrid_stack_spec=spec, vocab_size=100, max_sequence_length=4, hybrid_layer_pattern="M*-", @@ -399,7 +399,7 @@ def test_spec_provided_norm_not_overwritten(self): assert isinstance(attn.k_layernorm, IdentityOp) def test_forward_with_qk_layernorm(self): - """MambaModel forward pass works with qk_layernorm enabled.""" + """HybridModel forward pass works with qk_layernorm enabled.""" model = self._build_model(qk_layernorm=True) model.cuda() @@ -421,7 +421,7 @@ def test_forward_with_qk_layernorm(self): assert logits.shape[2] == 100 -class TestMambaWithDynamicInference: +class TestHybridWithDynamicInference: """Tests HybridModel with dynamic inference.""" @torch.inference_mode() diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 5d14c2bebce..01a46efe083 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -390,8 +390,8 @@ def _diff_configs(expected: Mapping[str, Any], actual: Mapping[str, Any]) -> Tup return added, removed, changed -class TestMambaMoEModel: - """Test the initialization and use of an MoE Mamba model.""" +class TestHybridMoEModel: + """Test the initialization and use of an MoE Hybrid model.""" def create_test_args(self): destroy_global_vars() diff --git a/tests/unit_tests/post_training/test_modelopt_module_spec.py b/tests/unit_tests/post_training/test_modelopt_module_spec.py index a5f7d2ca943..732dc08eca1 100644 --- a/tests/unit_tests/post_training/test_modelopt_module_spec.py +++ b/tests/unit_tests/post_training/test_modelopt_module_spec.py @@ -197,7 +197,7 @@ def setup_method(self, method): ) -class TestModelOptMambaModel(TestModelOptGPTModel): +class TestModelOptHybridModel(TestModelOptGPTModel): def setup_method(self, method): Utils.initialize_model_parallel(1, 1) diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 0869cc1a1f2..f21ae68c06c 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -21,7 +21,7 @@ @pytest.mark.internal -class TestMambaBlock: +class TestHybridBlock: def setup_method(self, method): Utils.initialize_model_parallel(1, 1) diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py index f6a88ed1c78..69f12e4d983 100644 --- a/tests/unit_tests/transformer/test_cuda_graphs.py +++ b/tests/unit_tests/transformer/test_cuda_graphs.py @@ -462,7 +462,7 @@ def test_llava_cudagraph_is_last_layer_logic(self): del layer.cudagraph_manager.cudagraph_runners[0].bwd_graph -class TestParallelMambaBlockCudagraphs: +class TestParallelHybridBlockCudagraphs: def setup_method(self, method): # initialize parallel state initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index fc75affc22f..d4d7edfe44b 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -686,7 +686,7 @@ def log(self, metrics, iteration): assert MTPLossLoggingHelper.tracker["avg_group"] is None -class TestMultiTokenPredictionMamba: +class TestMultiTokenPredictionHybrid: """Test Multi-Token Prediction with Mamba hybrid models.""" def setup_method(self, method): From 3375b1060b133d8bb8e56a645397224aae3c07b0 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 18 Apr 2026 01:03:46 +0000 Subject: [PATCH 15/23] Convert modelopt/mamba/model_specs.py to backward-compat stub This stale file was missed in the original directory rename. The canonical location is now modelopt/hybrid/model_specs.py. This stub re-exports from the new canonical location to preserve backward compatibility with any external imports from the old path. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../modelopt/mamba/model_specs.py | 144 +----------------- 1 file changed, 4 insertions(+), 140 deletions(-) diff --git a/megatron/core/post_training/modelopt/mamba/model_specs.py b/megatron/core/post_training/modelopt/mamba/model_specs.py index e9f83b49f71..68907ed8bd9 100755 --- a/megatron/core/post_training/modelopt/mamba/model_specs.py +++ b/megatron/core/post_training/modelopt/mamba/model_specs.py @@ -1,141 +1,5 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. -from megatron.core.extensions.transformer_engine import TEDotProductAttention -from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add -from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec -from megatron.core.post_training.modelopt.layers import Norm -from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules -from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules -from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules -from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear -from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules -from megatron.core.transformer.dot_product_attention import DotProductAttention -from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.mlp import MLP, MLPSubmodules -from megatron.core.transformer.spec_utils import ModuleSpec -from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules - - -# Use this spec for ModelOpt PTQ and TensorRT-LLM export -def get_mamba_stack_modelopt_spec( - local_core_attention: bool = False, - remap_te_layernorm: bool = False, - use_default_te_spec: bool = False, -) -> ModuleSpec: - """Get the Mamba stack spec for ModelOpt PTQ and TensorRT-LLM export. - - When use_default_te_spec=False (default), this is the native local spec with TENorm - from Transformer-Engine for the layernorm implementation (since FusedLayerNorm from - apex has stopped supporting RMSNorm needed by llama). The remap_te_layernorm flag - can be used to add sharded state_dict key remapping for TE-compatible checkpoint - saving/loading. - - When use_default_te_spec=True, this returns the standard mamba_stack_spec from - mamba_layer_specs.py which uses full TE modules (TELayerNormColumnParallelLinear, - TERowParallelLinear, TEDotProductAttention, TENorm, moe_grouped_gemm=True). - - - Args: - local_core_attention: whether to use local DotProductAttention - (only for use_default_te_spec=False) - remap_te_layernorm: whether to perform sharded state_dict prefix mapping - on layernorm (only for use_default_te_spec=False) - use_default_te_spec: whether to use the default Transformer-Engine spec - """ - if use_default_te_spec: - from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec - - return mamba_stack_spec - - return _get_mamba_stack_local_spec( - local_core_attention=local_core_attention, remap_te_layernorm=remap_te_layernorm - ) - - -def _get_mamba_stack_local_spec( - local_core_attention: bool = False, remap_te_layernorm: bool = False -) -> ModuleSpec: - """Get the Mamba stack spec with local (non-TE) modules. - - This is essentially the native local spec except for the layernorm implementation - is using TENorm from Transformer-Engine. - """ - mamba_state_dict_keys_map = {} - transformer_state_dict_keys_map = {} - if remap_te_layernorm: - mamba_state_dict_keys_map = {'norm.': 'mixer.in_proj.layer_norm_'} - transformer_state_dict_keys_map = { - 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_', - 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_', - } - - mamba_layer = ModuleSpec( - module=MambaLayer, - submodules=MambaLayerSubmodules( - norm=Norm, - mixer=ModuleSpec( - module=MambaMixer, - submodules=MambaMixerSubmodules( - in_proj=ColumnParallelLinear, out_proj=RowParallelLinear - ), - ), - mamba_bda=get_bias_dropout_add, - sharded_state_dict_keys_map=mamba_state_dict_keys_map, - ), - ) - - attn_mask_type = AttnMaskType.causal - core_attention = DotProductAttention if local_core_attention else TEDotProductAttention - attention_layer = ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - input_layernorm=Norm, - self_attention=ModuleSpec( - module=SelfAttention, - params={"attn_mask_type": attn_mask_type}, - submodules=SelfAttentionSubmodules( - linear_qkv=ColumnParallelLinear, - core_attention=core_attention, - linear_proj=RowParallelLinear, - ), - ), - self_attn_bda=get_bias_dropout_add, - sharded_state_dict_keys_map=transformer_state_dict_keys_map, - ), - ) - - mlp_layer = ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - pre_mlp_layernorm=Norm, - mlp=ModuleSpec( - module=MLP, - submodules=MLPSubmodules( - linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear - ), - ), - mlp_bda=get_bias_dropout_add, - sharded_state_dict_keys_map=transformer_state_dict_keys_map, - ), - ) - - moe_layer = ModuleSpec( - module=TransformerLayer, - submodules=TransformerLayerSubmodules( - pre_mlp_layernorm=Norm, - mlp=get_moe_module_spec( - use_te=False, num_experts=8, moe_grouped_gemm=False # Can be anything non None - ), - mlp_bda=get_bias_dropout_add, - ), - ) - - return ModuleSpec( - module=MambaStack, - submodules=MambaStackSubmodules( - mamba_layer=mamba_layer, - attention_layer=attention_layer, - mlp_layer=mlp_layer, - moe_layer=moe_layer, - ), - ) +# Backward-compatible re-export. The canonical location is now +# megatron.core.post_training.modelopt.hybrid.model_specs. +from megatron.core.post_training.modelopt.hybrid.model_specs import * # noqa: F401,F403 From 556eaa5557931155c6aa1888d894582732e10c42 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 18 Apr 2026 01:16:23 +0000 Subject: [PATCH 16/23] Add mamba_submodules backward-compat alias in MTP MultiTokenPredictionLayer and MultiTokenPredictionBlock now accept mamba_submodules as a deprecated alias for hybrid_submodules, emitting DeprecationWarning and forwarding the value (raises if both are set). Also drops the redundant explicit HybridModel import in mamba_model.py. Co-Authored-By: Claude Opus 4.7 (1M context) --- megatron/core/models/mamba/mamba_model.py | 1 - .../transformer/multi_token_prediction.py | 30 +++++++++++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py index fa92f8685ba..13964286daf 100644 --- a/megatron/core/models/mamba/mamba_model.py +++ b/megatron/core/models/mamba/mamba_model.py @@ -3,7 +3,6 @@ import logging from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 # pylint: disable=unused-import -from megatron.core.models.hybrid.hybrid_model import HybridModel # noqa: F401 from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.utils import log_single_rank diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 1044efe94c5..b70fdc3f28b 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -736,11 +736,24 @@ def __init__( layer_number: int = 1, vp_stage: Optional[int] = None, pg_collection: Optional[ProcessGroupCollection] = None, - # For Mamba path - pattern and submodules to build inner layers directly + # For hybrid path - pattern and submodules to build inner layers directly mtp_layer_pattern: Optional[str] = None, hybrid_submodules: Optional[HybridStackSubmodules] = None, + mamba_submodules: Optional[HybridStackSubmodules] = None, ): super().__init__(config=config) + if mamba_submodules is not None: + if hybrid_submodules is not None: + raise ValueError( + "Cannot specify both hybrid_submodules and mamba_submodules. " + "mamba_submodules has been deprecated; use hybrid_submodules instead." + ) + warnings.warn( + "mamba_submodules has been deprecated. Use hybrid_submodules instead.", + DeprecationWarning, + stacklevel=2, + ) + hybrid_submodules = mamba_submodules self.sequence_parallel = config.sequence_parallel self.submodules = submodules self.layer_number = layer_number + get_mtp_layer_offset(self.config, vp_stage) @@ -1269,12 +1282,25 @@ def __init__( spec: Union[TransformerBlockSubmodules, ModuleSpec], vp_stage: Optional[int] = None, pg_collection: Optional[ProcessGroupCollection] = None, - # New: For Mamba path with unified pattern syntax + # New: For hybrid path with unified pattern syntax mtp_layer_pattern: Optional[str] = None, mtp_num_depths: int = 0, hybrid_submodules: Optional["HybridStackSubmodules"] = None, + mamba_submodules: Optional["HybridStackSubmodules"] = None, ): super().__init__(config=config) + if mamba_submodules is not None: + if hybrid_submodules is not None: + raise ValueError( + "Cannot specify both hybrid_submodules and mamba_submodules. " + "mamba_submodules has been deprecated; use hybrid_submodules instead." + ) + warnings.warn( + "mamba_submodules has been deprecated. Use hybrid_submodules instead.", + DeprecationWarning, + stacklevel=2, + ) + hybrid_submodules = mamba_submodules self.submodules = _get_mtp_block_submodules(config, spec) self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor self.vp_stage = vp_stage From 5202abca9a924063ef40eea1464b79cf3ebd4ee4 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 18 Apr 2026 05:04:01 +0000 Subject: [PATCH 17/23] Fix mamba_stack_spec/MambaStack refs in test_dsa_layer_types The test_dsa_layer_types test was added to test_hybrid_block.py by an upstream merge and used bare references to mamba_stack_spec and MambaStack (not via the backward-compat import). The file imports hybrid_stack_spec and HybridStack, so those bare references NameError'd. Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/unit_tests/ssm/test_hybrid_block.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index f21ae68c06c..08bf7f2bc28 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -72,8 +72,8 @@ def get_dsa_mamba_block(self, layer_pattern): dsa_indexer_head_dim=64, dsa_indexer_topk=32, ) - modules = mamba_stack_spec.submodules - return MambaStack( + modules = hybrid_stack_spec.submodules + return HybridStack( transformer_config, modules, layer_type_list=layer_type_list, From 118c8b5eb77824302231202ed1827359ec63917b Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sun, 19 Apr 2026 20:53:23 +0000 Subject: [PATCH 18/23] Remove internal references to deprecated Mamba class names Backward-compat stubs and aliases remain for external libraries, but all internal Megatron-LM code now uses the canonical Hybrid names. Changes: - mamba_builders.py: backward-compat re-export stub with deprecation warning (delegates to hybrid_builders) - pretrain_mamba.py: backward-compat wrapper that runpy's pretrain_hybrid with deprecation warning - pretrain_hybrid.py: port upstream parse_and_validate_args refactor (from upstream #4225 which modified pretrain_mamba.py) - megatron/training/arguments.py: import Symbols from canonical hybrid_layer_allocation path instead of backward-compat stub - nemotron3_super_release_g200/model_config.yaml: update --spec and --mtp-spec to use hybrid_layer_specs path - test_dsa_gpt_mamba_equivalence.py: update all class references and imports to HybridModel / hybrid_stack_spec - tools/checkpoint/remap_gpt_dsa_to_mamba.py: update docstrings and help text from MambaModel/MambaStack to HybridModel/HybridStack Co-Authored-By: Claude Opus 4.7 (1M context) --- mamba_builders.py | 59 +-- megatron/training/arguments.py | 2 +- pretrain_hybrid.py | 7 +- pretrain_mamba.py | 375 +----------------- .../model_config.yaml | 4 +- .../models/test_dsa_gpt_mamba_equivalence.py | 50 +-- tools/checkpoint/remap_gpt_dsa_to_mamba.py | 22 +- 7 files changed, 68 insertions(+), 451 deletions(-) diff --git a/mamba_builders.py b/mamba_builders.py index 650ea4a719f..f824fce9be3 100644 --- a/mamba_builders.py +++ b/mamba_builders.py @@ -1,50 +1,15 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +"""Backward-compatible re-export of hybrid_builders. -from model_provider import count_parameters_in_layer -from megatron.core.models.mamba import MambaModel -from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.spec_utils import import_module -from megatron.training import print_rank_0 -from megatron.training.arguments import core_transformer_config_from_args -from megatron.core.models.mamba.mamba_layer_specs import mamba_inference_stack_spec +Deprecated. Use hybrid_builders instead. +""" +import warnings +warnings.warn( + "mamba_builders has been deprecated. Use hybrid_builders instead.", + DeprecationWarning, + stacklevel=2, +) -def mamba_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None): - print_rank_0('building MAMBA model ...') - if config is None: - config = core_transformer_config_from_args(args, TransformerConfig) - assert args.use_legacy_models is False, "Mamba only supported in Mcore!" - - if config.transformer_impl == "inference_optimized": - mamba_stack_spec = mamba_inference_stack_spec - assert ( - not config.inference_fuse_tp_communication - ), "inference_fuse_tp_communication is not supported for Mamba" - elif args.spec is not None: - mamba_stack_spec = import_module(args.spec) - else: - raise ValueError("You must provide a valid Mamba layer spec via --spec") - - model = MambaModel( - config=config, - mamba_stack_spec=mamba_stack_spec, - vocab_size=args.padded_vocab_size, - max_sequence_length=args.max_position_embeddings, - hybrid_layer_pattern=args.hybrid_layer_pattern, - pre_process=pre_process, - post_process=post_process, - fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, - parallel_output=True, - share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, - position_embedding_type=args.position_embedding_type, - rotary_percent=args.rotary_percent, - rotary_base=args.rotary_base, - pg_collection=pg_collection, - vp_stage=vp_stage, - ) - - for l in range(model.decoder.num_layers_per_pipeline_rank): - layer_params = count_parameters_in_layer(model, f'decoder.layers.{l}.') - print_rank_0(f" == params layer {l}: {layer_params}") - - return model +from hybrid_builders import * # noqa: F401,F403 +from hybrid_builders import hybrid_builder as mamba_builder # noqa: F401 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index a02fb86037c..8c6d484dc4d 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1770,7 +1770,7 @@ def core_transformer_config_from_args(args, config_class=None): kw_args['cp_comm_type'] = args.cp_comm_type[0] if args.hybrid_layer_pattern is not None: kw_args['is_hybrid_model'] = True - from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols + from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols if Symbols.DS_ATTENTION in args.hybrid_layer_pattern: kw_args['experimental_attention_variant'] = 'dsa' diff --git a/pretrain_hybrid.py b/pretrain_hybrid.py index f60552c3523..f073e8e9ab3 100644 --- a/pretrain_hybrid.py +++ b/pretrain_hybrid.py @@ -42,6 +42,7 @@ print_rank_0, set_startup_timestamps, ) +from megatron.training.arguments import parse_and_validate_args from megatron.training.datasets.sft_dataset import SFTDataset from megatron.training.utils import ( get_batch_on_this_cp_rank, @@ -356,11 +357,13 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None # Optionally enable inprocess restart on pretrain pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) + args = parse_and_validate_args( + extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, + args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + ) pretrain(train_valid_test_datasets_provider, partial(model_provider, hybrid_builder), ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, store=store, - extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, ) diff --git a/pretrain_mamba.py b/pretrain_mamba.py index 590eb92ab28..a316b861b63 100644 --- a/pretrain_mamba.py +++ b/pretrain_mamba.py @@ -1,369 +1,18 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -"""Pretrain and SFT Mamba.""" +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +"""Backward-compatible wrapper for pretrain_hybrid.py. -# Capture the true program start time BEFORE any heavy imports. -import time -_PROGRAM_START_TIME = time.time() - -import json - -# Suppress warnings on all ranks but rank 0. +Deprecated. Use pretrain_hybrid.py instead. +""" import os +import runpy import warnings -rank = int(os.environ.get('RANK', 0)) -if rank != 0: - warnings.filterwarnings("ignore", category=UserWarning) - warnings.filterwarnings("ignore", category=FutureWarning) - -from functools import partial -from typing import List, Optional, Tuple - -import torch -from mamba_builders import mamba_builder -from megatron.core import mpu -from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder -from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset -from megatron.core.enums import ModelType -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.parallel_state import ( - get_context_parallel_rank, - get_context_parallel_world_size, +warnings.warn( + "pretrain_mamba.py has been deprecated. Use pretrain_hybrid.py instead.", + DeprecationWarning, + stacklevel=2, ) -from megatron.core.models.mamba import MambaModel -from megatron.core.rerun_state_machine import get_rerun_state_machine -from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer -from megatron.core.utils import get_attr_wrapped_model, is_te_min_version, StragglerDetector -from megatron.training import ( - get_args, - get_timers, - inprocess_restart, - pretrain, - print_rank_0, - set_startup_timestamps, -) -from megatron.training.arguments import parse_and_validate_args -from megatron.training.datasets.sft_dataset import SFTDataset -from megatron.training.utils import ( - get_batch_on_this_cp_rank, - get_batch_on_this_tp_rank, - get_blend_and_blend_per_split, - is_first_or_last_pipeline_stage, -) -from model_provider import model_provider - -try: - from megatron.post_training.arguments import add_modelopt_args - from megatron.post_training.loss_func import loss_func as loss_func_modelopt - has_nvidia_modelopt = True -except ImportError: - has_nvidia_modelopt = False - -try: - # Register the TE CUDA kernels - import transformer_engine # pylint: disable=unused-import - - # Alias the PyTorch wrapper so we can call tex.* APIs - import transformer_engine_torch as tex -except ImportError: - # TE isn’t installed or the torch wrapper is missing - tex = None - -stimer = StragglerDetector() - - -def get_batch(data_iterator, vp_stage=None): - """Generate a batch.""" - - empty_batch = { - 'tokens': None, - 'labels': None, - 'loss_mask': None, - 'attention_mask': None, - 'position_ids': None, - 'cu_seqlens': None, - 'max_seqlen': None, - } - - # TODO(duncan): Is there a more efficient way to access is_packed_sequence here? - is_packed_sequence = get_args().sft # SFT always uses packed sequence - if not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence: - return empty_batch.values() - - batch = get_batch_on_this_tp_rank(data_iterator) - - cu_seqlens = batch['cu_seqlens'] - # Unused at the moment - cu_seqlens_padded = batch.pop('cu_seqlens_padded', None) - # Support for Hybrid Context Parallel (Unused in this script) - local_cp_size = batch.pop('local_cp_size', None) - - if cu_seqlens is not None: - assert ( - cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 - ), "micro-batch-size must be 1 for packing" - cu_seqlens = cu_seqlens[0] - batch['cu_seqlens'] = cu_seqlens - - max_seqlen = batch['max_seqlen'] - assert max_seqlen.dim() == 1 - # TODO(duncan): can this be kept as a 0-D tensor? - batch['max_seqlen'] = int(max_seqlen[0].item()) - - if mpu.is_pipeline_first_stage(ignore_virtual=(vp_stage is None), vp_stage=vp_stage): - total_tokens = batch['tokens'].size(1) - elif mpu.is_pipeline_last_stage(ignore_virtual=(vp_stage is None), vp_stage=vp_stage): - total_tokens = batch['labels'].size(1) - else: # packed sequence - empty_batch['cu_seqlens'] = cu_seqlens - empty_batch['max_seqlen'] = max_seqlen - return empty_batch.values() - - if cu_seqlens is None: - # slice batch along sequence dimension for context parallelism - batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore - else: # Packed THD format - cp_size = get_context_parallel_world_size() - if cp_size > 1: # slice batch along sequence dimension for context parallelism - assert tex is not None and is_te_min_version("1.10.0"), ( - "Please update Transformer Engine to >= 1.10 to use " - "Context Parallel with THD format data" - ) - cp_rank = get_context_parallel_rank() - index = tex.thd_get_partitioned_indices( - cu_seqlens, - total_tokens, - cp_size, - cp_rank, - ) - for key, data in batch.items(): - if key in {'attention_mask', 'cu_seqlens', 'max_seqlen'}: - continue - if data is not None: - # On first PP rank, labels and loss_mask can be None. - # On last PP rank, tokens and position_ids can be None. - batch[key] = data.index_select(1, index) - - return batch.values() - - -# define spiky loss as a loss that's 10x the max loss observed -SPIKY_LOSS_FACTOR = 10 - -def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model: Optional[MambaModel] = None): - """Loss function. - - Args: - loss_mask (torch.Tensor): Used to mask out some portions of the loss - output_tensor (torch.Tensor): The tensor with the losses - - Returns: - the loss scalar for this micro-batch - the number of non-padded tokens in this microbatch - a dict containing reporting metrics on the loss and number of tokens across - the data parallel ranks - """ - args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] - loss, num_tokens, report = loss_func_modelopt(loss_mask, output_tensor, model=model) - else: - losses = output_tensor.view(-1).float() - loss_mask = loss_mask.view(-1).float() - loss = torch.sum(losses * loss_mask) - - num_tokens = loss_mask.sum().clone().detach().to(torch.int) - report = {'lm loss': torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])} - - # Check individual rank losses are not NaN prior to DP all-reduce. - rerun_state_machine = get_rerun_state_machine() - if args.check_for_nan_in_loss_and_grad: - rerun_state_machine.validate_result( - result=loss, - rejection_func=torch.isnan, - message="found NaN in local forward loss calculation", - tolerance=0.0, # forward pass calculations are deterministic - fatal=True, - ) - rerun_state_machine.validate_result( - result=loss, - rejection_func=torch.isinf, - message="found Inf in local forward loss calculation", - tolerance=0.0, # forward pass calculations are deterministic - fatal=True, - ) - # Check for spiky loss - if args.check_for_spiky_loss: - rerun_state_machine.validate_result( - result=loss, - rejection_func=partial( - rerun_state_machine.is_unexpectedly_large, - threshold=SPIKY_LOSS_FACTOR, - context="loss", - ), - message="Spiky loss", - tolerance=0.0, # forward pass calculations are deterministic - fatal=False, - ) - - return loss, num_tokens, report - - -def forward_step(data_iterator, model: MambaModel): - """Forward training step. - - Args: - data_iterator : Input data iterator - model (MambaModel): The GPT Model - """ - timers = get_timers() - - # Get the batch. - timers('batch-generator', log_level=2).start() - - global stimer - - with stimer(bdata=True): - vp_stage = get_attr_wrapped_model(model, "vp_stage") - ( - tokens, - labels, - loss_mask, - attention_mask, - position_ids, - cu_seqlens, - max_seqlen, - ) = get_batch(data_iterator, vp_stage) - - if cu_seqlens is None: - packed_seq_params = None - else: - total_tokens = tokens.size(1) if tokens is not None else labels.size(1) - packed_seq_params = PackedSeqParams( - qkv_format="thd", - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - cu_seqlens_q_padded=None, - cu_seqlens_kv_padded=None, - max_seqlen_q=max_seqlen, - max_seqlen_kv=max_seqlen, - total_tokens=total_tokens, - ) - - 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 - ) - - # [ModelOpt]: model is needed to access ModelOpt distillation losses - return output_tensor, partial(loss_func, loss_mask, model=model) - - -def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): - if mpu.get_tensor_model_parallel_rank() != 0: - return False - elif is_packed_sequence: - return True - else: - return is_first_or_last_pipeline_stage(vp_stage) - - -def core_gpt_dataset_config_from_args(args): - tokenizer = build_tokenizer(args) - - # Sometimes --data-path is too long, instead we parse it from a file. - blend: Optional[Tuple[List[str], Optional[List[float]]]] - blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] - blend, blend_per_split = get_blend_and_blend_per_split(args) - - sequences_per_dataset = None - if args.per_dataset_sequences_path is not None: - with open(args.per_dataset_sequences_path, "r") as f: - sequences_per_dataset = json.load(f) - - return GPTDatasetConfig( - random_seed=args.seed, - sequence_length=args.seq_length, - blend=blend, - blend_per_split=blend_per_split, - split=args.split, - num_dataset_builder_threads=args.num_dataset_builder_threads, - path_to_cache=args.data_cache_path, - mmap_bin_files=args.mmap_bin_files, - tokenizer=tokenizer, - reset_position_ids=args.reset_position_ids, - reset_attention_mask=args.reset_attention_mask, - eod_mask_loss=args.eod_mask_loss, - create_attention_mask=args.create_attention_mask_in_dataloader, - object_storage_cache_path=args.object_storage_cache_path, - mid_level_dataset_surplus=args.mid_level_dataset_surplus, - allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens, - fast_cache_load=args.dataloader_fast_cache_load, - sequences_per_dataset=sequences_per_dataset, - defer_npy_index_mmap=args.dataloader_defer_npy_index_mmap, - context_parallel_size=args.context_parallel_size, - ) - - -def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None): - """Build the train test and validation datasets. - - Args: - train_val_test_num_samples : A list containing the number of samples in train test and validation. - """ - args = get_args() - config = core_gpt_dataset_config_from_args(args) - - is_packed_sequence = False - if args.sft: - dataset_type = SFTDataset - is_packed_sequence = True # SFT always uses packed sequence - else: - if args.mock_data: - dataset_type = MockGPTDataset - else: - dataset_type = GPTDataset - - print_rank_0("> building train, validation, and test datasets for GPT ...") - - train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( - dataset_type, - train_val_test_num_samples, - partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence), - config - ).build() - - print_rank_0("> finished creating GPT datasets ...") - - return train_ds, valid_ds, test_ds - - -if __name__ == "__main__": - # Timestamp right after entering __main__ block (after all imports/library setup) - _MAIN_ENTRY_TIME = time.time() - - # Register startup timestamps for timing report in pretrain() - set_startup_timestamps(program_start=_PROGRAM_START_TIME, main_entry=_MAIN_ENTRY_TIME) - - # Temporary for transition to core datasets - train_valid_test_datasets_provider.is_distributed = True - - # Optionally enable inprocess restart on pretrain - pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) - args = parse_and_validate_args( - extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, - ) - pretrain(train_valid_test_datasets_provider, - partial(model_provider, mamba_builder), - ModelType.encoder_or_decoder, - forward_step, - store=store, - ) +# Execute pretrain_hybrid.py as if it were invoked directly. +_this_dir = os.path.dirname(os.path.abspath(__file__)) +runpy.run_path(os.path.join(_this_dir, "pretrain_hybrid.py"), run_name="__main__") diff --git a/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_g200/model_config.yaml b/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_g200/model_config.yaml index 1147dda6118..9c5f1807c2d 100644 --- a/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_g200/model_config.yaml +++ b/tests/functional_tests/test_cases/nemotron/nemotron3_super_release_g200/model_config.yaml @@ -42,7 +42,7 @@ MODEL_ARGS: # Network size args --use-mcore-models: true - --spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec + --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec --is-hybrid-model: true --mamba-num-heads: 128 --num-layers: 88 @@ -90,7 +90,7 @@ MODEL_ARGS: --moe-shared-expert-compute-before-router: true # MTP args - --mtp-spec: megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec + --mtp-spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec --mtp-num-layers: 2 --mtp-hybrid-override-pattern: \"*E\" --calculate-per-token-loss: true diff --git a/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py b/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py index 96b782fad85..229af268a79 100644 --- a/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py +++ b/tests/unit_tests/models/test_dsa_gpt_mamba_equivalence.py @@ -1,6 +1,6 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. """ -Equivalence tests: GPTModel with DSA vs MambaModel with DSA pattern. +Equivalence tests: GPTModel with DSA vs HybridModel with DSA pattern. A small DeepSeek-V3.2 proxy model (4 GPT layers / 8 Mamba layers) is built, weights are remapped GPT→Mamba, and logprobs are compared to verify they are @@ -9,8 +9,8 @@ Architecture equivalence ------------------------ GPTModel layer N (combined attention + MLP in one TransformerLayer) - ≡ MambaModel layer 2N (D, DSA TransformerLayer: input_layernorm + MLASelfAttention) - + MambaModel layer 2N+1 (-, MLPLayer: fused-norm MLP) + ≡ HybridModel layer 2N (D, DSA TransformerLayer: input_layernorm + MLASelfAttention) + + HybridModel layer 2N+1 (-, MLPLayer: fused-norm MLP) Run with:: @@ -35,10 +35,10 @@ get_transformer_block_with_experimental_attention_variant_spec, ) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec -from megatron.core.models.mamba.mamba_model import MambaModel +from megatron.core.models.hybrid.hybrid_layer_allocation import validate_segment_layers +from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec +from megatron.core.models.hybrid.hybrid_model import HybridModel from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.mamba_hybrid_layer_allocation import validate_segment_layers from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer.transformer_config import MLATransformerConfig from megatron.rl.rl_utils import selective_log_softmax @@ -193,15 +193,15 @@ def _build_mamba_model( layer_pattern: str, pre_process: bool = True, post_process: bool = True, -) -> MambaModel: - """Build a MambaModel with the given hybrid layer pattern.""" +) -> HybridModel: + """Build a HybridModel with the given hybrid layer pattern.""" layer_type_list = validate_segment_layers(layer_pattern) mamba_config = copy.deepcopy(config) mamba_config.num_layers = len(layer_type_list) assert mamba_config.num_layers == _NUM_GPT_LAYERS * 2 - model = MambaModel( + model = HybridModel( config=mamba_config, - mamba_stack_spec=mamba_stack_spec, + hybrid_stack_spec=hybrid_stack_spec, vocab_size=_VOCAB_SIZE, max_sequence_length=_MAX_SEQ_LEN, pre_process=pre_process, @@ -221,14 +221,14 @@ def _build_mamba_model( def _remap_gpt_to_mamba_state_dict( gpt_sd: Dict[str, torch.Tensor], num_local_gpt_layers: int ) -> Dict[str, torch.Tensor]: - """Remap a GPTModel state_dict to a MambaModel state_dict. + """Remap a GPTModel state_dict to a HybridModel state_dict. GPTModel layer N (combined attention + MLP) maps to: - * MambaModel layer 2N – DSA attention (input_layernorm + self_attention) - * MambaModel layer 2N+1 – MLP (mlp.*) + * HybridModel layer 2N – DSA attention (input_layernorm + self_attention) + * HybridModel layer 2N+1 – MLP (mlp.*) Additionally, ``decoder.final_layernorm.*`` (TransformerBlock naming) is - remapped to ``decoder.final_norm.*`` (MambaStack naming). + remapped to ``decoder.final_norm.*`` (HybridStack naming). All other keys (embedding, output_layer, rotary_pos_emb, …) are unchanged. @@ -238,7 +238,7 @@ def _remap_gpt_to_mamba_state_dict( pipeline stage (i.e. ``len(gpt_model.decoder.layers)``). Returns: - Remapped state dict ready for MambaModel.load_state_dict(strict=True). + Remapped state dict ready for HybridModel.load_state_dict(strict=True). """ mamba_sd: Dict[str, torch.Tensor] = {} layer_prefix = "decoder.layers." @@ -380,12 +380,12 @@ def _compare_against_golden_values( @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("tp,pp", [(1, 1), (2, 1), (1, 2)]) class TestDSAGPTMambaEquivalence: - """Verify logprob equivalence between GPTModel+DSA and MambaModel+DSA. + """Verify logprob equivalence between GPTModel+DSA and HybridModel+DSA. For each distributed configuration (TP, PP), the test: 1. Builds a GPTModel with 4 DSA layers. - 2. Builds a MambaModel with pattern "D-D-D-D-" (8 layers). - 3. Remaps and loads GPT weights into MambaModel (strict=True). + 2. Builds a HybridModel with pattern "D-D-D-D-" (8 layers). + 3. Remaps and loads GPT weights into HybridModel (strict=True). 4. Runs the same random tokens through both models. 5. Asserts logprob tensors are numerically close. """ @@ -416,7 +416,7 @@ def test_dsa_logprobs_match(self, tp: int, pp: int) -> None: num_local_gpt_layers = len(gpt_model.decoder.layers) gpt_sd = gpt_model.state_dict() - # ---- Build MambaModel ---- + # ---- Build HybridModel ---- mamba_model = _build_mamba_model( gpt_config, _MAMBA_PATTERN, pre_process=pre_process, post_process=post_process ) @@ -481,7 +481,7 @@ def test_weight_loading_strict(self, tp: int, pp: int) -> None: assert not unexpected, f"Unexpected keys: {unexpected}" def test_record_and_compare_golden_values(self, tp: int, pp: int) -> None: - """Record GPTModel logprobs as golden values, then compare MambaModel against them. + """Record GPTModel logprobs as golden values, then compare HybridModel against them. Golden values are written to the functional test directory so they can be committed and used by the CI inference golden-value tests. @@ -508,7 +508,7 @@ def test_record_and_compare_golden_values(self, tp: int, pp: int) -> None: gpt_logprobs = _forward_logprobs_pp1(gpt_model, tokens) mamba_logprobs = _forward_logprobs_pp1(mamba_model, tokens) - # Verify MambaModel matches golden values + # Verify HybridModel matches golden values _compare_against_golden_values(mamba_logprobs, gpt_logprobs, abs_tol=1e-3) @@ -520,7 +520,7 @@ def test_record_and_compare_golden_values(self, tp: int, pp: int) -> None: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("tp,pp", [(1, 1), (2, 1), (1, 2)]) class TestDSAMoEGPTMambaEquivalence: - """Verify logprob equivalence between GPTModel+DSA+MoE and MambaModel+DSA+MoE. + """Verify logprob equivalence between GPTModel+DSA+MoE and HybridModel+DSA+MoE. Architecture: 4 GPT layers with moe_layer_freq=[0,0,1,1] (first 2 dense, last 2 MoE) maps to 8 Mamba layers with pattern "D-D-DEDE": @@ -556,7 +556,7 @@ def test_dsa_moe_logprobs_match(self, tp: int, pp: int) -> None: num_local_gpt_layers = len(gpt_model.decoder.layers) gpt_sd = gpt_model.state_dict() - # ---- Build MambaModel with MoE pattern ---- + # ---- Build HybridModel with MoE pattern ---- mamba_model = _build_mamba_model( gpt_config, _MOE_MAMBA_PATTERN, pre_process=pre_process, post_process=post_process ) @@ -618,7 +618,7 @@ def test_moe_weight_loading_strict(self, tp: int, pp: int) -> None: assert not unexpected, f"Unexpected keys: {unexpected}" def test_moe_record_and_compare_golden_values(self, tp: int, pp: int) -> None: - """Record GPTModel+MoE logprobs as golden values, then compare MambaModel+MoE.""" + """Record GPTModel+MoE logprobs as golden values, then compare HybridModel+MoE.""" self._skip_if_insufficient_gpus(tp, pp) if tp != 1 or pp != 1: pytest.skip("Golden-value recording only runs for tp=1, pp=1") @@ -640,5 +640,5 @@ def test_moe_record_and_compare_golden_values(self, tp: int, pp: int) -> None: gpt_logprobs = _forward_logprobs_pp1(gpt_model, tokens) mamba_logprobs = _forward_logprobs_pp1(mamba_model, tokens) - # Verify MambaModel matches golden values + # Verify HybridModel matches golden values _compare_against_golden_values(mamba_logprobs, gpt_logprobs, abs_tol=1e-3) diff --git a/tools/checkpoint/remap_gpt_dsa_to_mamba.py b/tools/checkpoint/remap_gpt_dsa_to_mamba.py index 8a6888d1dc7..3d11c981c25 100644 --- a/tools/checkpoint/remap_gpt_dsa_to_mamba.py +++ b/tools/checkpoint/remap_gpt_dsa_to_mamba.py @@ -1,16 +1,16 @@ #!/usr/bin/env python3 # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -"""Convert a GPTModel DSA checkpoint to a MambaModel-compatible checkpoint. +"""Convert a GPTModel DSA checkpoint to a HybridModel-compatible checkpoint. A GPTModel with ``--experimental-attention-variant dsa`` uses one combined -TransformerLayer per model layer (attention + MLP). The equivalent MambaModel +TransformerLayer per model layer (attention + MLP). The equivalent HybridModel with pattern ``D-D-...`` stores them as two separate layers: * Layer 2N – DSA attention (TransformerLayer: input_layernorm + MLASelfAttention) * Layer 2N+1 – MLP (MLPLayer: fused-norm MLP) This script loads a GPTModel Distributed Checkpoint (DCP), remaps the state-dict -keys, and saves a new DCP that can be loaded by MambaModel. +keys, and saves a new DCP that can be loaded by HybridModel. Usage ----- @@ -43,14 +43,14 @@ def _remap_key(key: str, num_gpt_layers: int) -> str: - """Return the MambaModel state-dict key corresponding to *key* from GPTModel. + """Return the HybridModel state-dict key corresponding to *key* from GPTModel. Args: key: A key from the GPTModel state dict. num_gpt_layers: Total number of GPT decoder layers (across all PP stages). Returns: - The remapped key for MambaModel. + The remapped key for HybridModel. Raises: ValueError: If an unexpected sub-key is encountered in a decoder layer. @@ -58,7 +58,7 @@ def _remap_key(key: str, num_gpt_layers: int) -> str: layer_prefix = "decoder.layers." final_ln_prefix = "decoder.final_layernorm." - # Final layernorm name differs between TransformerBlock and MambaStack + # Final layernorm name differs between TransformerBlock and HybridStack if key.startswith(final_ln_prefix): return "decoder.final_norm." + key[len(final_ln_prefix):] @@ -96,11 +96,11 @@ def _remap_state_dict( def convert(input_path: Path, output_path: Path, num_gpt_layers: int) -> None: - """Load a GPTModel DCP checkpoint, remap keys, and save as MambaModel DCP. + """Load a GPTModel DCP checkpoint, remap keys, and save as HybridModel DCP. Args: input_path: Path to the GPTModel DCP checkpoint directory. - output_path: Destination directory for the MambaModel DCP checkpoint. + output_path: Destination directory for the HybridModel DCP checkpoint. num_gpt_layers: Number of GPT decoder layers in the original model. """ try: @@ -139,7 +139,7 @@ def convert(input_path: Path, output_path: Path, num_gpt_layers: int) -> None: output_path.mkdir(parents=True, exist_ok=True) torch_save_to_dcp(str(tmp_mamba), str(output_path)) - print(f"MambaModel DCP checkpoint saved to: {output_path}") + print(f"HybridModel DCP checkpoint saved to: {output_path}") finally: for tmp in (tmp_flat, output_path.parent / "_tmp_mamba_flat.pt"): @@ -149,7 +149,7 @@ def convert(input_path: Path, output_path: Path, num_gpt_layers: int) -> None: def main() -> None: parser = argparse.ArgumentParser( - description="Convert GPTModel DSA checkpoint to MambaModel-compatible format." + description="Convert GPTModel DSA checkpoint to HybridModel-compatible format." ) parser.add_argument( "--input", required=True, type=Path, @@ -157,7 +157,7 @@ def main() -> None: ) parser.add_argument( "--output", required=True, type=Path, - help="Destination path for the MambaModel DCP checkpoint.", + help="Destination path for the HybridModel DCP checkpoint.", ) parser.add_argument( "--num-gpt-layers", required=True, type=int, From b415521c80f8009e8102f46f9b29e400718158b6 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 20 Apr 2026 23:02:58 +0000 Subject: [PATCH 19/23] Restore upstream logic accidentally reverted by rename commit The rename commit (d6a3854d3) and subsequent merges accidentally reverted unrelated upstream changes. Restore upstream versions of: - megatron/training/arguments.py: restore 'adaptive_muon' optimizer choice and --optimizer-cuda-graph argument; remove stray --no-scatter-gather-tensors-in-pipeline - megatron/training/training.py: restore OptimizerCudaGraphWrapper import, save_checkpoint_and_time() usage, and timer logic from upstream - megatron/rl/sequence_packing_utils.py: restore removal of unused packed_attention_mask parameter (upstream #3859) - .github/actions/action.yml: restore retry loop around uv install (upstream #4387) Re-applied only the intended rename-related import path changes to arguments.py (2 lines) and training.py (2 lines). Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/actions/action.yml | 6 ++- megatron/rl/sequence_packing_utils.py | 63 +++------------------------ megatron/training/arguments.py | 7 ++- megatron/training/training.py | 21 ++++++--- 4 files changed, 29 insertions(+), 68 deletions(-) diff --git a/.github/actions/action.yml b/.github/actions/action.yml index d500a5896d0..3913414a1c7 100644 --- a/.github/actions/action.yml +++ b/.github/actions/action.yml @@ -83,7 +83,11 @@ runs: echo "apt attempt $i failed, retrying..." sleep 10 done - curl -LsSf https://astral.sh/uv/install.sh | UV_INSTALL_DIR=/usr/local/bin sh + for i in 1 2 3; do + curl -LsSf https://astral.sh/uv/install.sh | UV_INSTALL_DIR=/usr/local/bin sh && break + echo "uv install attempt $i failed, retrying..." + sleep 10 + done - name: Create run-script (unit test) shell: bash -x -e -u -o pipefail {0} diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index 6f3e4711ed7..50ab2ca27de 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -51,7 +51,6 @@ class PackingContext: original_trajs: All trajectories before packing packed_trajs: Packed trajectories tensor [num_bins, bin_size] packed_position_ids: Position IDs for packed sequences [num_bins, bin_size] - packed_attention_mask: Attention mask for packed sequences [num_bins, 1, bin_size, bin_size] packed_loss_mask: Loss mask for packed sequences [num_bins, bin_size] original_inference_logprobs: Inference logprobs for all sequences before packing (optional) bin_advantages: List of advantage tensors for each bin @@ -64,7 +63,6 @@ class PackingContext: original_trajs: torch.Tensor packed_trajs: torch.Tensor packed_position_ids: torch.Tensor - packed_attention_mask: torch.Tensor packed_loss_mask: torch.Tensor original_inference_logprobs: Optional[torch.Tensor] = None bin_advantages: List[torch.Tensor] = field(default_factory=list) @@ -314,9 +312,8 @@ def create_empty_bins( packed_trajs : torch.Tensor, packed_position_ids : torch.Tensor, packed_loss_mask : torch.Tensor, - packed_attention_mask : torch.Tensor, tokenizer, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, List[Dict[str, Any]]]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[Dict[str, Any]]]: """Create empty bins for padding to ensure all ranks have the same number of bins. Args: @@ -325,11 +322,10 @@ def create_empty_bins( packed_trajs: Packed trajectories tensor (for dtype/device reference) packed_position_ids: Packed position IDs tensor (for dtype/device reference) packed_loss_mask: Packed loss mask tensor (for dtype/device reference) - packed_attention_mask: Packed attention mask tensor (can be None) tokenizer: Tokenizer for pad token Returns: - Tuple of (empty_trajs, empty_position_ids, empty_loss_mask, empty_attention_mask, empty_packing_info_entries) + Tuple of (empty_trajs, empty_position_ids, empty_loss_mask, empty_packing_info_entries) """ device = packed_trajs.device @@ -337,7 +333,6 @@ def create_empty_bins( empty_bins = [] empty_position_ids_list = [] empty_loss_mask_list = [] - empty_attention_mask_list = [] empty_packing_info_entries = [] for i in range(num_empty_bins): @@ -355,14 +350,6 @@ def create_empty_bins( empty_loss = torch.zeros(1, bin_size, dtype=packed_loss_mask.dtype, device=device) empty_loss_mask_list.append(empty_loss) - # Zero attention mask if needed - if packed_attention_mask is not None: - # Attention mask is always 4D: [num_bins, 1, bin_size, bin_size] - empty_attn = torch.zeros( - 1, 1, bin_size, bin_size, dtype=packed_attention_mask.dtype, device=device - ) - empty_attention_mask_list.append(empty_attn) - # Empty packing info entries empty_packing_info_entries.append( { @@ -376,22 +363,15 @@ def create_empty_bins( empty_trajs = torch.cat(empty_bins, dim=0) empty_position_ids = torch.cat(empty_position_ids_list, dim=0) empty_loss_mask = torch.cat(empty_loss_mask_list, dim=0) - empty_attention_mask = ( - torch.cat(empty_attention_mask_list, dim=0) - if packed_attention_mask is not None - else None - ) else: empty_trajs = None empty_position_ids = None empty_loss_mask = None - empty_attention_mask = None return ( empty_trajs, empty_position_ids, empty_loss_mask, - empty_attention_mask, empty_packing_info_entries, ) @@ -708,9 +688,6 @@ def pack_sequences( position_ids = torch.zeros( (num_bins, self.bin_size), dtype=torch.long, device=device, requires_grad=False ) - attention_mask = torch.zeros( - (num_bins, 1, self.bin_size, self.bin_size), dtype=torch.bool, device=device - ) loss_mask = torch.zeros((num_bins, self.bin_size), dtype=torch.float, device=device) # Track packing information for unpacking later @@ -741,12 +718,6 @@ def pack_sequences( len(seq), device=device, requires_grad=False ) - # Causal attention mask within each sequence - seq_len = end - start - attention_mask[bin_idx, 0, start:end, start:end] = torch.tril( - torch.ones(seq_len, seq_len, dtype=torch.bool, device=device) - ) - # Loss mask (excluding padding) loss_mask[bin_idx, start:end] = 1.0 @@ -761,12 +732,6 @@ def pack_sequences( seq_starts.append(current_pos) seq_starts_dict[bin_idx] = seq_starts - # Note: We'll store the actual padded length later when we know it - # (it depends on the original trajectories passed to pack_sequences) - - # Invert attention mask, before inversion: (True = attend, False = mask) - attention_mask.bitwise_not_() - # Create the PackingInfo dataclass packing_info = PackingInfo( bin_seq_indices=bin_seq_indices, @@ -795,15 +760,14 @@ def pack_sequences( ) log_single_rank(logger, logging.DEBUG, f" - First 20 bins: {seq_per_bin[:20]}") - return packed_sequences, position_ids, attention_mask, loss_mask, packing_info + return packed_sequences, position_ids, loss_mask, packing_info def distribute_packed_bins( packed_trajs: torch.Tensor, packed_position_ids: torch.Tensor, - packed_attention_mask: torch.Tensor, packed_loss_mask: torch.Tensor, packing_info: PackingInfo, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, PackingInfo]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, PackingInfo]: """Distribute packed bins across the data parallel ranks.""" rank = mpu.get_data_parallel_rank() world_size = mpu.get_data_parallel_world_size() @@ -840,7 +804,6 @@ def distribute_packed_bins( # Extract this rank's bins my_packed_trajs = [] my_packed_position_ids = [] - my_packed_attention_mask = [] my_packed_loss_mask = [] my_bin_seq_indices = [] my_seq_starts = {} @@ -850,8 +813,6 @@ def distribute_packed_bins( for new_idx, old_idx in enumerate(my_bin_indices): my_packed_trajs.append(packed_trajs[old_idx]) my_packed_position_ids.append(packed_position_ids[old_idx]) - if packed_attention_mask is not None: - my_packed_attention_mask.append(packed_attention_mask[old_idx]) my_packed_loss_mask.append(packed_loss_mask[old_idx]) my_bin_seq_indices.append(packing_info.bin_seq_indices[old_idx]) my_seq_starts[new_idx] = packing_info.seq_starts[old_idx] @@ -877,9 +838,6 @@ def distribute_packed_bins( device=packed_position_ids.device, ) ) - packed_attention_mask = ( - torch.stack(my_packed_attention_mask) if my_packed_attention_mask else None - ) packed_loss_mask = ( torch.stack(my_packed_loss_mask) if my_packed_loss_mask @@ -937,7 +895,6 @@ def distribute_packed_bins( empty_trajs, empty_position_ids, empty_loss_mask, - empty_attention_mask, empty_packing_entries, ) = create_empty_bins( num_empty_bins, @@ -945,7 +902,6 @@ def distribute_packed_bins( packed_trajs, packed_position_ids, packed_loss_mask, - packed_attention_mask, tokenizer, ) @@ -956,18 +912,13 @@ def distribute_packed_bins( ) packed_loss_mask = torch.cat([packed_loss_mask, empty_loss_mask], dim=0) - if packed_attention_mask is not None and empty_attention_mask is not None: - packed_attention_mask = torch.cat( - [packed_attention_mask, empty_attention_mask], dim=0 - ) - # Add empty entries to packing_info for i, entry in enumerate(empty_packing_entries): bin_idx = current_bins + i new_packing_info.bin_seq_indices.append(entry['bin_seq_indices']) new_packing_info.seq_starts[bin_idx] = entry['seq_starts'] - return packed_trajs, packed_position_ids, packed_attention_mask, packed_loss_mask, new_packing_info + return packed_trajs, packed_position_ids, packed_loss_mask, new_packing_info def pack_all_trajectories(trajs, generation_masks, inference_logprobs, global_advantages, bin_size, max_sequences_per_bin, packing_algo): @@ -1000,7 +951,6 @@ def _gather(data): ( packed_trajs, packed_position_ids, - packed_attention_mask, packed_loss_mask, packing_info, ) = packer.pack_sequences(trajs, generation_masks) @@ -1010,13 +960,11 @@ def _gather(data): ( packed_trajs, packed_position_ids, - packed_attention_mask, packed_loss_mask, packing_info, ) = distribute_packed_bins( packed_trajs, packed_position_ids, - packed_attention_mask, packed_loss_mask, packing_info, ) @@ -1053,7 +1001,6 @@ def _gather(data): original_trajs=trajs, packed_trajs=packed_trajs, packed_position_ids=packed_position_ids, - packed_attention_mask=packed_attention_mask, packed_loss_mask=packed_loss_mask, original_inference_logprobs=inference_logprobs, bin_advantages=bin_advantages, diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 8c6d484dc4d..eaa5b15ee53 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2556,12 +2556,14 @@ def _add_training_args(parser): help='use FlashAttention implementation of attention. ' 'https://arxiv.org/abs/2205.14135') group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap'], + choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap', 'adaptive_muon'], help='Optimizer function. ' 'Note: dist_muon is deprecated; use --optimizer muon ' 'with --use-distributed-optimizer instead.') group.add_argument('--optimizer-cpu-offload', action='store_true', help='Offload optimizer state to CPU') + group.add_argument('--optimizer-cuda-graph', action='store_true', + help='Enable CUDA graph for optimizer step') group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, help='Ratio of optimizer state to offload to CPU') group.add_argument('--use-torch-optimizer-for-cpu-offload', action='store_true', @@ -2768,9 +2770,6 @@ def _add_distributed_args(parser): help='If not set, all PP stages will launch param all-gathers simultaneously. ' 'Otherwise, each PP stage will independently launch as needed.', dest='align_param_gather') - group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false', - help='If not set, use scatter/gather to optimize communication of tensors in pipeline.', - dest='scatter_gather_tensors_in_pipeline') group.add_argument('--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.') group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', diff --git a/megatron/training/training.py b/megatron/training/training.py index 331f058237b..8172c6a0cd5 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -157,6 +157,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.training.checkpointing import checkpoint_exists from megatron.training.checkpointing import get_loaded_iteration from megatron.core.full_cuda_graph import FullCudaGraphWrapper +from megatron.core.optimizer.optimizer_cuda_graph import OptimizerCudaGraphWrapper from megatron.core.transformer.cuda_graphs import TECudaGraphHelper from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.module import Float16Module @@ -1202,15 +1203,14 @@ def pretrain( print_datetime('after training is done') if not args.skip_train and args.save and iteration != 0 and iteration % args.save_interval != 0: - save_checkpoint( + save_checkpoint_and_time( iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far, checkpointing_context, - train_data_iterator=train_data_iterator, - preprocess_common_state_dict_fn=preprocess_common_state_dict, + train_data_iterator=train_data_iterator ) one_logger and one_logger.log_metrics( @@ -2425,6 +2425,11 @@ def save_checkpoint_and_time( train_data_iterator=train_data_iterator, preprocess_common_state_dict_fn=preprocess_common_state_dict, ) + + # Stop timer and compute time elapsed to save checkpoint. Stop timer before timers.log() call as it resets the timer. + timers(timer_key).stop(barrier=True) + save_checkpoint_duration = timers(timer_key).elapsed(reset=False) + if should_report_memory: # Track memory after checkpoint save. report_memory(f"(after save_checkpoint for iteration {iteration})") @@ -2435,12 +2440,12 @@ def save_checkpoint_and_time( # dequantized bf16 tensors that were temporarily created during fp8 # model checkpoint saving. gc.collect() - timers(timer_key).stop(barrier=True) + timers.log([timer_key]) # Log E2E metrics after save-checkpoint one_logger_utils.track_e2e_metrics() - save_checkpoint_duration = timers(timer_key).elapsed() + one_logger_utils.on_save_checkpoint_end(save_checkpoint_duration, iteration, args.async_save) if args.log_progress and not non_persistent_ckpt: @@ -2835,6 +2840,8 @@ def train( forward_backward_func = get_forward_backward_func() if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) + if args.optimizer_cuda_graph: + optimizer.step = OptimizerCudaGraphWrapper(optimizer.step, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) def get_e2e_base_metrics(): """Get base metrics values for one-logger to calculate E2E tracking metrics.""" @@ -3261,6 +3268,10 @@ def trace_handler(p): if args.cuda_graph_impl == "transformer_engine" and cuda_graph_helper.graphs_created(): cuda_graph_helper.delete_cuda_graphs() + # Call OptimizerCudaGraph destructor to destroy optimizer CUDA graph + if args.optimizer_cuda_graph: + del optimizer.step + one_logger_utils.track_e2e_metrics() # Flush TensorBoard, WandB writers and one-logger. From 9828cc5eb5f1468549d67284cd0bde4059a12416 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 20 Apr 2026 23:03:48 +0000 Subject: [PATCH 20/23] Restore two more files accidentally reverted by rename commit - tests/test_utils/python_scripts/notify.py: restore WEBHOOK_URL check and copyright header from upstream - tests/unit_tests/rl/test_sequence_packing_utils.py: restore removal of packed_attention_mask parameter (companion to the sequence_packing_utils.py revert) Co-Authored-By: Claude Opus 4.7 (1M context) --- tests/test_utils/python_scripts/notify.py | 6 ++++++ .../rl/test_sequence_packing_utils.py | 17 +++++++---------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/tests/test_utils/python_scripts/notify.py b/tests/test_utils/python_scripts/notify.py index 7da00dc401a..cdad04644a7 100644 --- a/tests/test_utils/python_scripts/notify.py +++ b/tests/test_utils/python_scripts/notify.py @@ -1,3 +1,5 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import logging import os @@ -94,6 +96,10 @@ def main(pipeline_id: int, check_for: str, pipeline_context: str, pipeline_creat messages.append("===============================================") + if not WEBHOOK_URL: + logger.info("No webhook URL configured, skipping Slack notification") + return + for message in messages: response = slack_sdk.webhook.WebhookClient(WEBHOOK_URL).send(text=message) logger.info(response.status_code) diff --git a/tests/unit_tests/rl/test_sequence_packing_utils.py b/tests/unit_tests/rl/test_sequence_packing_utils.py index 06e63adf217..0a4f4cd03f9 100644 --- a/tests/unit_tests/rl/test_sequence_packing_utils.py +++ b/tests/unit_tests/rl/test_sequence_packing_utils.py @@ -98,13 +98,12 @@ def test_sequence_packing_basic(): rewards = torch.tensor([1.0, 2.0, 3.0, 4.0]) sequences_tensor = torch.stack(sequences) - packed_trajs, packed_position_ids, packed_attention_mask, packed_loss_mask, packing_info = ( - packer.pack_sequences(sequences_tensor, generation_masks) + packed_trajs, packed_position_ids, packed_loss_mask, packing_info = packer.pack_sequences( + sequences_tensor, generation_masks ) assert packed_trajs is not None assert packed_position_ids is not None - assert packed_attention_mask is not None assert packed_loss_mask is not None assert packing_info is not None @@ -140,8 +139,8 @@ def test_sequence_packing_with_generation_masks(): ) padded_sequences_tensor = torch.stack(padded_sequences) - packed_trajs, packed_position_ids, packed_attention_mask, packed_loss_mask, packing_info = ( - packer.pack_sequences(padded_sequences_tensor, generation_masks) + packed_trajs, packed_position_ids, packed_loss_mask, packing_info = packer.pack_sequences( + padded_sequences_tensor, generation_masks ) assert packed_trajs.shape[0] == 1 @@ -162,16 +161,14 @@ def test_sequence_packing_empty_bins(): ) packed_position_ids = torch.tensor([[0, 1, 2, 3, 0, 0, 0, 0]]) packed_loss_mask = torch.tensor([[1, 1, 1, 1, 0, 0, 0, 0]], dtype=torch.float) - packed_attention_mask = torch.ones(1, bin_size, bin_size) - empty_trajs, empty_position_ids, empty_loss_mask, empty_attention_mask, empty_packing_info = ( + empty_trajs, empty_position_ids, empty_loss_mask, empty_packing_info = ( sequence_packing_utils.create_empty_bins( num_empty_bins=num_empty_bins, bin_size=bin_size, packed_trajs=packed_trajs, packed_position_ids=packed_position_ids, packed_loss_mask=packed_loss_mask, - packed_attention_mask=packed_attention_mask, tokenizer=tokenizer, ) ) @@ -220,8 +217,8 @@ def test_sequence_packing_integration(): ] sequences_tensor = torch.stack(sequences) - packed_trajs, packed_position_ids, packed_attention_mask, packed_loss_mask, packing_info = ( - packer.pack_sequences(sequences_tensor, generation_masks) + packed_trajs, packed_position_ids, packed_loss_mask, packing_info = packer.pack_sequences( + sequences_tensor, generation_masks ) assert packed_trajs is not None From 4e38d27d8b20923f8ab1a8e57de436cc44ffdd13 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Tue, 21 Apr 2026 00:16:30 +0000 Subject: [PATCH 21/23] Remove stray --async-save / --use-persistent-ckpt-worker from hybrid test config Upstream removed these flags from hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml in commit 2697b82ab ("base strategy simplification #4001"), but they were accidentally re-introduced by a merge conflict resolution. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../model_config.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml index b89d305dc63..2339f7a7ce9 100644 --- a/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml +++ b/tests/functional_tests/test_cases/hybrid/hybrid_mr_mcore_te_tp2_pp1_cp1_dgx_a100_1N8G/model_config.yaml @@ -55,6 +55,4 @@ MODEL_ARGS: --bf16: true --attention-backend: unfused --log-memory-to-tensorboard: true - --async-save: true - --use-persistent-ckpt-worker: true TEST_TYPE: regular From 30606dadf3bc7696d628d7807e3f6fe8e1fc3b0e Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Tue, 21 Apr 2026 00:28:11 +0000 Subject: [PATCH 22/23] Remove mamba CLI backward-compat from test_dynamic_engine.py Test files don't need to accept the old "mamba" CLI value since the backward-compat shim is only for external library consumers. Simplify the defensive ("hybrid", "mamba") checks to just "hybrid". Co-Authored-By: Claude Opus 4.7 (1M context) --- .../inference/engines/test_dynamic_engine.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 02943304bc1..b23e9562242 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -65,7 +65,7 @@ def skip_if_mamba_sequence_packing_not_available(model_provider: str): - if model_provider in ("hybrid", "mamba"): + if model_provider == "hybrid": sequence_packing_available, reason_for_no_sequence_packing = ( _check_mamba_sequence_packing_support() ) @@ -368,7 +368,7 @@ def _build_test_env(cls, test_config): mtp_block_spec=mtp_block_spec, position_embedding_type=test_config.position_embedding_type, ).cuda() - elif test_config.model_provider in ("hybrid", "mamba"): + elif test_config.model_provider == "hybrid": pp_size = test_config.pipeline_model_parallel_size # Transformer config. transformer_config = TransformerConfig( @@ -632,7 +632,7 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None if model_provider == "gpt": expected_generated_tokens_list = gpt_expected_generated_tokens - elif model_provider in ("hybrid", "mamba"): + elif model_provider == "hybrid": expected_generated_tokens_list = mamba_expected_generated_tokens else: raise ValueError(f"Invalid model_provider {model_provider}") @@ -1131,7 +1131,7 @@ def test_parallel_inference( "when tp_size > 1." ) ) - if model_provider in ("hybrid", "mamba"): + if model_provider == "hybrid": pytest.skip( reason="Mamba model is not supported with the inference optimized transformer." ) @@ -4319,7 +4319,7 @@ def test_speculative_decoding_mamba_hybrid(self, rejection_mode): Two requests run simultaneously to exercise batched rewind indexing where mamba_metadata.request_to_mamba_state_idx differs per request. """ - skip_if_mamba_sequence_packing_not_available("mamba") + skip_if_mamba_sequence_packing_not_available("hybrid") num_tokens_to_generate = 8 test_config = DynamicEngineTestConfig( @@ -4329,7 +4329,7 @@ def test_speculative_decoding_mamba_hybrid(self, rejection_mode): num_tokens_to_generate=num_tokens_to_generate, num_speculative_tokens=2, materialize_only_last_token_logits=False, - model_provider="mamba", + model_provider="hybrid", ) env = self._build_test_env(test_config) @@ -4460,7 +4460,7 @@ def _create_model(self, model_provider, num_cuda_graphs): pre_process=parallel_state.is_pipeline_first_stage(), post_process=parallel_state.is_pipeline_last_stage(), ).cuda() - elif model_provider in ("hybrid", "mamba"): + elif model_provider == "hybrid": config = TransformerConfig( params_dtype=torch.bfloat16, num_layers=3, From 018fa5902c32050bd328402daa3354e06ba5ac5b Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 23 Apr 2026 02:38:33 +0000 Subject: [PATCH 23/23] Guard pretrain_mamba.py launch with __main__ check Wrap the runpy.run_path call in `if __name__ == "__main__":` so that importing pretrain_mamba (e.g. to reuse model_provider / get_batch) does not launch distributed training as an import side-effect. Matches the pattern used in pretrain_gpt.py and pretrain_hybrid.py. Co-Authored-By: Claude Opus 4.7 (1M context) --- pretrain_mamba.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pretrain_mamba.py b/pretrain_mamba.py index a316b861b63..7eb7f461cab 100644 --- a/pretrain_mamba.py +++ b/pretrain_mamba.py @@ -13,6 +13,7 @@ stacklevel=2, ) -# Execute pretrain_hybrid.py as if it were invoked directly. -_this_dir = os.path.dirname(os.path.abspath(__file__)) -runpy.run_path(os.path.join(_this_dir, "pretrain_hybrid.py"), run_name="__main__") +if __name__ == "__main__": + # Execute pretrain_hybrid.py as if it were invoked directly. + _this_dir = os.path.dirname(os.path.abspath(__file__)) + runpy.run_path(os.path.join(_this_dir, "pretrain_hybrid.py"), run_name="__main__")