diff --git a/megatron/core/models/hybrid/__init__.py b/megatron/core/models/hybrid/__init__.py index d8a0a817ee3..43672161637 100644 --- a/megatron/core/models/hybrid/__init__.py +++ b/megatron/core/models/hybrid/__init__.py @@ -1 +1,19 @@ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.models.hybrid.hybrid_architecture import ( + HybridLayerPattern, + HybridLayerSpec, + PipelineSplit, + ResolvedHybridArchitecture, + flatten_hybrid_layer_pattern, + resolve_hybrid_architecture, +) + +__all__ = [ + "HybridLayerPattern", + "HybridLayerSpec", + "PipelineSplit", + "ResolvedHybridArchitecture", + "flatten_hybrid_layer_pattern", + "resolve_hybrid_architecture", +] diff --git a/megatron/core/models/hybrid/hybrid_architecture.py b/megatron/core/models/hybrid/hybrid_architecture.py new file mode 100644 index 00000000000..e0c861d7340 --- /dev/null +++ b/megatron/core/models/hybrid/hybrid_architecture.py @@ -0,0 +1,535 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""First-class architecture descriptions for :class:`HybridModel`. + +This module deliberately keeps the public surface small. A layer occurrence is +an existing :class:`~megatron.core.transformer.spec_utils.ModuleSpec` paired with +the :class:`~megatron.core.transformer.transformer_config.TransformerConfig` +that should be passed to that layer. Nested Python lists provide composition, +and :class:`PipelineSplit` provides explicit PP/VPP chunk boundaries. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any, Sequence, TypeAlias + +from megatron.core.models.hybrid.hybrid_layer_allocation import ( + Symbols, + parse_hybrid_pattern, + validate_segment_layers, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + +HYBRID_LAYER_TYPE = "hybrid_layer_type" +VALID_HYBRID_LAYER_TYPES = frozenset({"mamba", "gdn", "attention", "dsa", "mla", "mlp", "moe"}) + + +@dataclass(frozen=True, slots=True) +class PipelineSplit: + """A boundary between two physical or virtual pipeline model chunks.""" + + +@dataclass(frozen=True, slots=True) +class HybridLayerSpec: + """An existing layer spec paired with its per-occurrence configuration. + + Args: + module_spec: Specification for an existing Megatron Core layer class. + config: Complete transformer configuration to pass to this occurrence. + """ + + module_spec: ModuleSpec + config: TransformerConfig + + def __post_init__(self) -> None: + if not isinstance(self.module_spec, ModuleSpec): + raise TypeError( + f"module_spec must be a ModuleSpec, got {type(self.module_spec).__name__}." + ) + if not isinstance(self.config, TransformerConfig): + raise TypeError( + f"config must be a TransformerConfig, got {type(self.config).__name__}." + ) + # Resolve eagerly so malformed public descriptors fail at recipe construction. + _get_layer_type(self.module_spec) + + @property + def layer_type(self) -> str: + """Return the stable semantic type carried by ``ModuleSpec.metainfo``.""" + + return _get_layer_type(self.module_spec) + + +HybridLayerPattern: TypeAlias = Sequence["HybridLayerSpec | PipelineSplit | HybridLayerPattern"] + + +@dataclass(frozen=True, slots=True) +class ResolvedHybridArchitecture: + """A validated global hybrid architecture shared by every PP/VPP chunk.""" + + segments: tuple[tuple[HybridLayerSpec, ...], ...] + mtp_layers: tuple[HybridLayerSpec, ...] = () + mtp_num_layers: int = 0 + source: str = "direct" + + @property + def main_layers(self) -> tuple[HybridLayerSpec, ...]: + """Return decoder occurrences in global model order.""" + + return tuple(layer for segment in self.segments for layer in segment) + + @property + def metric_num_layers(self) -> int: + """Return the number of explicit layer slots used by MoE metrics.""" + + return len(self.main_layers) + self.mtp_num_layers * len(self.mtp_layers) + + def select_segment( + self, *, pp_rank: int, pp_size: int, vp_stage: int | None + ) -> tuple[tuple[HybridLayerSpec, ...], int]: + """Select a local chunk and its global decoder-layer offset. + + Segments use the same VPP-major ordering as the legacy ``|`` syntax: + ``segment_index = vp_stage * pp_size + pp_rank``. + """ + + vp_rank = 0 if vp_stage is None else vp_stage + segment_index = vp_rank * pp_size + pp_rank + if segment_index >= len(self.segments): + raise ValueError( + f"Pipeline segment index {segment_index} is out of range for " + f"{len(self.segments)} resolved segments (pp_rank={pp_rank}, " + f"pp_size={pp_size}, vp_stage={vp_stage})." + ) + offset = sum(len(segment) for segment in self.segments[:segment_index]) + return self.segments[segment_index], offset + + @property + def has_heterogeneous_inference_shapes(self) -> bool: + """Whether dynamic inference would need non-uniform cache/buffer shapes.""" + + layers = self.main_layers + self.mtp_layers + signatures: dict[str, set[tuple[Any, ...]]] = { + "mamba": set(), + "attention": set(), + "moe": set(), + } + for layer in layers: + config = layer.config + if layer.layer_type == "mamba": + signatures["mamba"].add( + ( + config.mamba_state_dim, + config.mamba_head_dim, + config.mamba_num_heads, + config.mamba_num_groups, + ) + ) + elif layer.layer_type == "attention": + signatures["attention"].add( + (config.num_query_groups or config.num_attention_heads, config.kv_channels) + ) + elif layer.layer_type == "moe": + signatures["moe"].add((config.moe_router_topk,)) + return any(len(values) > 1 for values in signatures.values()) + + def has_incompatible_dynamic_inference_shapes(self, base_config: TransformerConfig) -> bool: + """Whether dynamic inference's model-global buffers fit every occurrence. + + Dynamic inference currently allocates attention KV and MoE routing buffers + from the model-wide config. Even a family-uniform direct override is + incompatible when it differs from that allocation source. Mamba state + shapes are discovered from the built layer, so only differing Mamba + occurrences are rejected. + """ + + if self.has_heterogeneous_inference_shapes: + return True + + base_attention_signature = ( + base_config.num_query_groups or base_config.num_attention_heads, + base_config.kv_channels, + ) + layers = self.main_layers + self.mtp_layers + for layer in layers: + if layer.layer_type == "attention": + occurrence_signature = ( + layer.config.num_query_groups or layer.config.num_attention_heads, + layer.config.kv_channels, + ) + if occurrence_signature != base_attention_signature: + return True + elif ( + layer.layer_type == "moe" + and layer.config.moe_router_topk != base_config.moe_router_topk + ): + return True + return False + + +def flatten_hybrid_layer_pattern( + pattern: HybridLayerPattern, *, allow_splits: bool = True +) -> tuple[tuple[HybridLayerSpec, ...], ...]: + """Recursively flatten a direct Python pattern while preserving split nodes. + + Args: + pattern: Nested lists or tuples containing layer descriptors and split nodes. + allow_splits: If false, encountering :class:`PipelineSplit` is an error. + + Returns: + A tuple of flat pipeline segments. Empty segments are intentionally retained. + """ + + segments: list[list[HybridLayerSpec]] = [[]] + + def visit(node: Any, path: tuple[int, ...]) -> None: + if isinstance(node, HybridLayerSpec): + segments[-1].append(node) + return + if isinstance(node, PipelineSplit): + if not allow_splits: + raise ValueError( + f"PipelineSplit at path {list(path)} is not allowed in an MTP pattern." + ) + segments.append([]) + return + if isinstance(node, (list, tuple)): + for index, child in enumerate(node): + visit(child, path + (index,)) + return + raise TypeError( + f"Hybrid layer pattern leaf at path {list(path)} has unsupported type " + f"{type(node).__name__}; expected HybridLayerSpec, PipelineSplit, list, or tuple." + ) + + visit(pattern, ()) + return tuple(tuple(segment) for segment in segments) + + +def resolve_hybrid_architecture( + *, + config: TransformerConfig, + hybrid_stack_spec: ModuleSpec, + layer_specs: HybridLayerPattern | None = None, + mtp_layer_specs: HybridLayerPattern | None = None, + hybrid_layer_pattern: str | None = None, +) -> ResolvedHybridArchitecture: + """Resolve direct descriptors or a legacy string into one global architecture. + + Direct split nodes are authoritative for PP/VPP chunking. Legacy strings retain + their existing pipe-free even/uneven PP compatibility behavior. + """ + + if layer_specs is not None and hybrid_layer_pattern is not None: + raise ValueError("layer_specs and hybrid_layer_pattern are mutually exclusive.") + if layer_specs is None and hybrid_layer_pattern is None: + raise ValueError("Exactly one of layer_specs or hybrid_layer_pattern must be provided.") + if mtp_layer_specs is not None and layer_specs is None: + raise ValueError("mtp_layer_specs requires direct layer_specs.") + + if layer_specs is not None: + if config.mtp_standalone: + raise ValueError("Direct hybrid architectures do not support standalone MTP placement.") + raw_segments = flatten_hybrid_layer_pattern(layer_specs) + _validate_direct_pipeline(config, raw_segments) + raw_mtp_segments = ( + flatten_hybrid_layer_pattern(mtp_layer_specs, allow_splits=False) + if mtp_layer_specs is not None + else ((),) + ) + raw_mtp_layers = raw_mtp_segments[0] + mtp_num_layers = config.mtp_num_layers or 0 + if raw_mtp_layers and mtp_num_layers <= 0: + raise ValueError("mtp_layer_specs requires config.mtp_num_layers > 0.") + if mtp_num_layers > 0 and not raw_mtp_layers: + raise ValueError("config.mtp_num_layers > 0 requires mtp_layer_specs.") + + segments = tuple( + tuple(_materialize_layer(layer, config) for layer in segment) + for segment in raw_segments + ) + mtp_layers = tuple(_materialize_layer(layer, config) for layer in raw_mtp_layers) + architecture = ResolvedHybridArchitecture( + segments=segments, mtp_layers=mtp_layers, mtp_num_layers=mtp_num_layers, source="direct" + ) + _validate_mtp_placement(architecture) + _validate_uniform_expert_count(architecture, config) + return architecture + + return _resolve_legacy_architecture(config, hybrid_stack_spec, hybrid_layer_pattern) + + +def _validate_direct_pipeline( + config: TransformerConfig, segments: tuple[tuple[HybridLayerSpec, ...], ...] +) -> None: + pp_size = config.pipeline_model_parallel_size + num_segments = len(segments) + num_layers = sum(len(segment) for segment in segments) + + if num_layers != config.num_layers: + raise ValueError( + f"Direct layer_specs contains {num_layers} decoder layers, but " + f"TransformerConfig.num_layers is {config.num_layers}." + ) + if pp_size <= 0: + raise ValueError(f"pipeline_model_parallel_size must be positive, got {pp_size}.") + if pp_size > 1 and num_segments == 1: + raise ValueError( + "Direct layer_specs with pipeline_model_parallel_size > 1 must contain " + "explicit PipelineSplit() boundaries." + ) + if num_segments % pp_size != 0: + raise ValueError( + f"Direct layer_specs defines {num_segments} pipeline segments, which is not " + f"divisible by pipeline_model_parallel_size={pp_size}." + ) + + inferred_vp_size = num_segments // pp_size + if pp_size == 1 and inferred_vp_size > 1: + raise ValueError("Virtual pipeline parallelism requires pipeline_model_parallel_size > 1.") + configured_vp_size = config.virtual_pipeline_model_parallel_size + if configured_vp_size is not None and configured_vp_size != inferred_vp_size: + raise ValueError( + f"PipelineSplit() nodes imply virtual_pipeline_model_parallel_size=" + f"{inferred_vp_size}, but the config specifies {configured_vp_size}." + ) + + if num_segments > 1: + conflicts = { + "pipeline_model_parallel_layout": config.pipeline_model_parallel_layout, + "num_layers_in_first_pipeline_stage": config.num_layers_in_first_pipeline_stage, + "num_layers_in_last_pipeline_stage": config.num_layers_in_last_pipeline_stage, + "account_for_embedding_in_pipeline_split": ( + config.account_for_embedding_in_pipeline_split + ), + "account_for_loss_in_pipeline_split": config.account_for_loss_in_pipeline_split, + } + active_conflicts = [name for name, value in conflicts.items() if value not in (None, False)] + if active_conflicts: + raise ValueError( + "PipelineSplit() already defines pipeline ownership and cannot be combined with: " + + ", ".join(active_conflicts) + ) + + config.virtual_pipeline_model_parallel_size = inferred_vp_size if inferred_vp_size > 1 else None + + +def _resolve_legacy_architecture( + config: TransformerConfig, hybrid_stack_spec: ModuleSpec, hybrid_layer_pattern: str | None +) -> ResolvedHybridArchitecture: + parsed = parse_hybrid_pattern(hybrid_layer_pattern) + submodules = hybrid_stack_spec.submodules + symbol_fields = { + Symbols.MAMBA: ("mamba", "mamba_layer"), + Symbols.GDN: ("gdn", "gdn_layer"), + Symbols.ATTENTION: ("attention", "attention_layer"), + Symbols.DS_ATTENTION: ("dsa", "dsa_layer"), + Symbols.MLA: ("mla", "mla_layer"), + Symbols.MLP: ("mlp", "mlp_layer"), + Symbols.MOE: ("moe", "moe_layer"), + } + + def descriptor(symbol: str) -> HybridLayerSpec: + layer_type, field_name = symbol_fields[symbol] + return HybridLayerSpec( + module_spec=_legacy_module_spec(getattr(submodules, field_name), layer_type), + config=config, + ) + + main_pattern = parsed.main_pattern or "" + # This is a read-only summary of the legacy pattern, not a replacement for + # select_pipeline_segment(). In particular, do not infer or mutate PP/VPP + # settings, or add validation that existing HybridModel callers did not see. + pattern_segments = main_pattern.split(Symbols.PIPE) + + segments = tuple( + tuple(descriptor(symbol) for symbol in validate_segment_layers(segment)) + for segment in pattern_segments + ) + + mtp_layers: tuple[HybridLayerSpec, ...] = () + if parsed.mtp_pattern: + mtp_layers = tuple( + descriptor(symbol) for symbol in validate_segment_layers(parsed.mtp_pattern) + ) + return ResolvedHybridArchitecture( + segments=segments, + mtp_layers=mtp_layers, + mtp_num_layers=parsed.mtp_num_depths, + source="legacy", + ) + + +_UNIFORM_CONFIG_FIELDS = ( + "hidden_size", + "params_dtype", + "pipeline_dtype", + "fp16", + "bf16", + "fp32_residual_connection", + "enable_autocast", + "autocast_dtype", + "apply_query_key_layer_scaling", + "attention_softmax_in_fp32", + "disable_bf16_reduced_precision_matmul", + "dsa_indexer_k_norm_fp32", + "fp8", + "fp8_recipe", + "fp8_param", + "fp8_quantizer_factory", + "fp8_margin", + "fp8_interval", + "fp8_amax_history_len", + "fp8_amax_compute_algo", + "fp8_wgrad", + "fp8_output_proj", + "fp8_dot_product_attention", + "fp8_multi_head_attention", + "tp_only_amax_red", + "activation_func_fp8_input_store", + "first_last_layers_bf16", + "num_layers_at_start_in_bf16", + "num_layers_at_end_in_bf16", + "moe_router_dtype", + "mamba_training_ssm_states_dtype", + "fp4", + "fp4_recipe", + "fp4_param", + "fp4_quantizer_factory", + "quant_recipe", + "normalization", + "layernorm_epsilon", + "layernorm_zero_centered_gamma", + "tensor_model_parallel_size", + "tensor_parallel_num_weight_shards", + "gtp_weight_remat_size", + "pipeline_model_parallel_size", + "pipeline_model_parallel_comm_backend", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "max_seqlen_per_dp_cp_rank", + "hybrid_context_parallel", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "expert_tensor_parallel_num_weight_shards", + "expert_gtp_weight_remat_size", + "sequence_parallel", +) + +_TOPOLOGY_CONFIG_FIELDS = ( + "num_layers", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "pipeline_model_parallel_layout", + "num_layers_in_first_pipeline_stage", + "num_layers_in_last_pipeline_stage", + "account_for_embedding_in_pipeline_split", + "account_for_loss_in_pipeline_split", + "mtp_num_layers", + "mtp_use_repeated_layer", + "mtp_standalone", +) + + +def _materialize_layer(layer: HybridLayerSpec, base_config: TransformerConfig) -> HybridLayerSpec: + for field_name in _UNIFORM_CONFIG_FIELDS: + layer_value = getattr(layer.config, field_name) + base_value = getattr(base_config, field_name) + if layer_value != base_value: + raise ValueError( + f"Per-layer config for {layer.layer_type!r} changes model-wide field " + f"{field_name!r}: {layer_value!r} != {base_value!r}." + ) + return HybridLayerSpec( + module_spec=layer.module_spec, config=_materialize_config(layer.config, base_config) + ) + + +def _materialize_config( + layer_config: TransformerConfig, base_config: TransformerConfig +) -> TransformerConfig: + materialized = copy.deepcopy(layer_config) + for field_name in _TOPOLOGY_CONFIG_FIELDS: + setattr(materialized, field_name, getattr(base_config, field_name)) + return materialized + + +def _validate_mtp_placement(architecture: ResolvedHybridArchitecture) -> None: + """Reject a prediction block on a logical chunk with no decoder layers.""" + + if ( + architecture.mtp_layers + and architecture.mtp_num_layers > 0 + and not architecture.segments[-1] + ): + raise ValueError( + "MTP must share the final logical PP/VPP chunk with at least one decoder layer; " + "standalone MTP placement is not supported." + ) + + +def _validate_uniform_expert_count( + architecture: ResolvedHybridArchitecture, base_config: TransformerConfig +) -> None: + expert_counts = { + layer.config.num_moe_experts + for layer in architecture.main_layers + architecture.mtp_layers + if layer.layer_type == "moe" + } + if len(expert_counts) > 1: + raise ValueError( + "All MoE occurrences must use one uniform num_moe_experts value; got " + f"{sorted(expert_counts)}." + ) + if expert_counts and next(iter(expert_counts)) != base_config.num_moe_experts: + occurrence_count = next(iter(expert_counts)) + raise ValueError( + "MoE occurrence num_moe_experts must match the model-wide config; " + f"got {occurrence_count!r} != {base_config.num_moe_experts!r}." + ) + + +def _get_layer_type(module_spec: ModuleSpec) -> str: + layer_type = module_spec.metainfo.get(HYBRID_LAYER_TYPE) + if layer_type not in VALID_HYBRID_LAYER_TYPES: + raise ValueError( + f"ModuleSpec.metainfo[{HYBRID_LAYER_TYPE!r}] must be one of " + f"{sorted(VALID_HYBRID_LAYER_TYPES)}, got {layer_type!r}." + ) + return layer_type + + +def _with_layer_type(module_spec: ModuleSpec, layer_type: str) -> ModuleSpec: + if not isinstance(module_spec, ModuleSpec): + raise TypeError( + f"Hybrid stack entry for {layer_type!r} must be a ModuleSpec, got " + f"{type(module_spec).__name__}." + ) + if module_spec.metainfo.get(HYBRID_LAYER_TYPE) == layer_type: + return module_spec + tagged = copy.copy(module_spec) + tagged.metainfo = dict(module_spec.metainfo) + tagged.metainfo[HYBRID_LAYER_TYPE] = layer_type + return tagged + + +def _legacy_module_spec(module_spec: ModuleSpec | type, layer_type: str) -> ModuleSpec: + """Tag a legacy stack entry without narrowing its historical accepted types.""" + + if isinstance(module_spec, ModuleSpec): + return _with_layer_type(module_spec, layer_type) + return ModuleSpec(module=module_spec, metainfo={HYBRID_LAYER_TYPE: layer_type}) + + +__all__ = [ + "HYBRID_LAYER_TYPE", + "HybridLayerPattern", + "HybridLayerSpec", + "PipelineSplit", + "ResolvedHybridArchitecture", + "flatten_hybrid_layer_pattern", + "resolve_hybrid_architecture", +] diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py index 03fef58159f..70bb2f965af 100755 --- a/megatron/core/models/hybrid/hybrid_layer_specs.py +++ b/megatron/core/models/hybrid/hybrid_layer_specs.py @@ -14,6 +14,7 @@ get_inference_optimized_moe_spec, get_moe_module_spec, ) +from megatron.core.models.hybrid.hybrid_architecture import HYBRID_LAYER_TYPE 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 @@ -92,6 +93,7 @@ submodules=HybridStackSubmodules( mamba_layer=ModuleSpec( module=MambaLayer, + metainfo={HYBRID_LAYER_TYPE: "mamba"}, submodules=MambaLayerSubmodules( mixer=ModuleSpec( module=MambaMixer, @@ -104,6 +106,7 @@ ), gdn_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "gdn"}, submodules=TransformerLayerSubmodules( self_attention=ModuleSpec( module=GatedDeltaNet, @@ -121,6 +124,7 @@ # working attention_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "attention"}, submodules=TransformerLayerSubmodules( self_attention=ModuleSpec( module=SelfAttention, @@ -136,6 +140,7 @@ ), dsa_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "dsa"}, submodules=TransformerLayerSubmodules( input_layernorm=TENorm, self_attention=ModuleSpec( @@ -171,6 +176,7 @@ ), mla_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "mla"}, submodules=TransformerLayerSubmodules( input_layernorm=TENorm, self_attention=ModuleSpec( @@ -196,6 +202,7 @@ # working mlp_layer=ModuleSpec( module=MLPLayer, + metainfo={HYBRID_LAYER_TYPE: "mlp"}, submodules=TransformerLayerSubmodules( mlp=partial( MLP.as_mlp_submodule, @@ -208,6 +215,7 @@ ), moe_layer=ModuleSpec( module=MoETransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "moe"}, submodules=TransformerLayerSubmodules( pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add ), @@ -222,6 +230,7 @@ submodules=HybridStackSubmodules( mamba_layer=ModuleSpec( module=MambaLayer, + metainfo={HYBRID_LAYER_TYPE: "mamba"}, submodules=MambaLayerSubmodules( mixer=ModuleSpec( module=MambaMixer, @@ -238,6 +247,7 @@ # working attention_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "attention"}, submodules=TransformerLayerSubmodules( self_attention=ModuleSpec( module=SelfAttention, @@ -253,6 +263,7 @@ ), dsa_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "dsa"}, submodules=TransformerLayerSubmodules( input_layernorm=TENorm, self_attention=ModuleSpec( @@ -288,6 +299,7 @@ ), mla_layer=ModuleSpec( module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "mla"}, submodules=TransformerLayerSubmodules( input_layernorm=TENorm, self_attention=ModuleSpec( @@ -313,6 +325,7 @@ # working mlp_layer=ModuleSpec( module=MLPLayer, + metainfo={HYBRID_LAYER_TYPE: "mlp"}, submodules=TransformerLayerSubmodules( mlp=partial( MLP.as_mlp_submodule, @@ -327,6 +340,7 @@ moe_layer=ModuleSpec( # Use inference-optimized MoE layer for end-to-end CUDA graph support module=TransformerLayer, + metainfo={HYBRID_LAYER_TYPE: "moe"}, submodules=TransformerLayerSubmodules( pre_mlp_layernorm=TENorm, mlp=moe_inference, mlp_bda=get_bias_dropout_add ), diff --git a/tests/unit_tests/models/hybrid/test_hybrid_architecture.py b/tests/unit_tests/models/hybrid/test_hybrid_architecture.py new file mode 100644 index 00000000000..444b1dccd9c --- /dev/null +++ b/tests/unit_tests/models/hybrid/test_hybrid_architecture.py @@ -0,0 +1,633 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for first-class HybridModel architecture descriptions.""" + +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.models.hybrid.hybrid_architecture import ( + HYBRID_LAYER_TYPE, + HybridLayerSpec, + PipelineSplit, + flatten_hybrid_layer_pattern, + resolve_hybrid_architecture, +) +from megatron.core.models.hybrid.hybrid_layer_specs import ( + hybrid_inference_stack_spec, + hybrid_stack_spec, +) +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.spec_utils import ModuleSpec + + +class _Layer: + """Stand-in module class; resolver tests only inspect ModuleSpec metadata.""" + + +class _Stack: + """Stand-in HybridStack class for the resolver's legacy symbol mapping.""" + + +def _tagged_spec(layer_type: str) -> ModuleSpec: + return ModuleSpec(module=_Layer, metainfo={HYBRID_LAYER_TYPE: layer_type}) + + +STACK_SPEC = ModuleSpec( + module=_Stack, + submodules=SimpleNamespace( + mamba_layer=_tagged_spec("mamba"), + gdn_layer=_tagged_spec("gdn"), + attention_layer=_tagged_spec("attention"), + dsa_layer=_tagged_spec("dsa"), + mla_layer=_tagged_spec("mla"), + mlp_layer=_tagged_spec("mlp"), + moe_layer=_tagged_spec("moe"), + ), +) + +SPEC_BY_TYPE = { + "mamba": STACK_SPEC.submodules.mamba_layer, + "gdn": STACK_SPEC.submodules.gdn_layer, + "attention": STACK_SPEC.submodules.attention_layer, + "dsa": STACK_SPEC.submodules.dsa_layer, + "mla": STACK_SPEC.submodules.mla_layer, + "mlp": STACK_SPEC.submodules.mlp_layer, + "moe": STACK_SPEC.submodules.moe_layer, +} + +STOCK_LAYER_TYPES = { + "mamba_layer": "mamba", + "gdn_layer": "gdn", + "attention_layer": "attention", + "dsa_layer": "dsa", + "mla_layer": "mla", + "mlp_layer": "mlp", + "moe_layer": "moe", +} + + +def _config(num_layers: int, *, pp_size: int = 1, mtp_num_layers: int | None = None): + return TransformerConfig( + num_layers=num_layers, + hidden_size=64, + num_attention_heads=8, + num_query_groups=4, + kv_channels=8, + ffn_hidden_size=128, + num_moe_experts=8, + moe_ffn_hidden_size=96, + moe_router_topk=2, + mamba_state_dim=16, + mamba_head_dim=8, + mamba_num_heads=8, + mamba_num_groups=2, + add_bias_linear=False, + pipeline_model_parallel_size=pp_size, + pipeline_dtype=torch.float32, + mtp_num_layers=mtp_num_layers, + ) + + +def _layer(layer_type: str, config: TransformerConfig) -> HybridLayerSpec: + return HybridLayerSpec(module_spec=SPEC_BY_TYPE[layer_type], config=config) + + +def _types(layers) -> list[str]: + return [layer.layer_type for layer in layers] + + +@pytest.mark.parametrize( + ("stack_spec", "expected_types"), + [ + pytest.param(hybrid_stack_spec, STOCK_LAYER_TYPES, id="training"), + pytest.param( + hybrid_inference_stack_spec, + {name: value for name, value in STOCK_LAYER_TYPES.items() if name != "gdn_layer"}, + id="inference", + ), + ], +) +def test_every_stock_hybrid_layer_module_spec_has_stable_semantic_tag(stack_spec, expected_types): + tagged_types = { + field_name: module_spec.metainfo.get(HYBRID_LAYER_TYPE) + for field_name in STOCK_LAYER_TYPES + if isinstance((module_spec := getattr(stack_spec.submodules, field_name)), ModuleSpec) + } + + assert tagged_types == expected_types + + +def test_flatten_preserves_nested_alias_multiplication_and_empty_segments(): + config = _config(6) + mamba = _layer("mamba", config) + attention = _layer("attention", config) + moe = _layer("moe", config) + + segments = flatten_hybrid_layer_pattern( + [[mamba, attention] * 2, PipelineSplit(), [], PipelineSplit(), (moe, [mamba])] + ) + + assert [_types(segment) for segment in segments] == [ + ["mamba", "attention", "mamba", "attention"], + [], + ["moe", "mamba"], + ] + assert segments[0][0] is mamba + assert segments[0][2] is mamba + + +def test_pp2_vpp2_selection_is_vpp_major_and_offsets_include_empty_chunks(): + config = _config(6, pp_size=2) + mamba = _layer("mamba", config) + attention = _layer("attention", config) + moe = _layer("moe", config) + mlp = _layer("mlp", config) + + architecture = resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[ + [mamba], + PipelineSplit(), + [attention, moe], + PipelineSplit(), + [], + PipelineSplit(), + [mlp, mamba, attention], + ], + ) + + assert config.virtual_pipeline_model_parallel_size == 2 + expected = { + (0, 0): (["mamba"], 0), + (1, 0): (["attention", "moe"], 1), + (0, 1): ([], 3), + (1, 1): (["mlp", "mamba", "attention"], 3), + } + for (pp_rank, vp_stage), (expected_types, expected_offset) in expected.items(): + layers, offset = architecture.select_segment(pp_rank=pp_rank, pp_size=2, vp_stage=vp_stage) + assert _types(layers) == expected_types + assert offset == expected_offset + + +def test_resolver_materializes_an_isolated_config_for_every_occurrence(): + config = _config(3) + alias_config = deepcopy(config) + alias_config.mamba_state_dim = 24 + alias = _layer("mamba", alias_config) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, layer_specs=[[alias] * 3] + ) + + resolved_configs = [layer.config for layer in architecture.main_layers] + assert len({id(layer_config) for layer_config in resolved_configs}) == 3 + assert all(layer_config is not alias_config for layer_config in resolved_configs) + assert [layer_config.mamba_state_dim for layer_config in resolved_configs] == [24, 24, 24] + + resolved_configs[0].mamba_state_dim = 32 + assert resolved_configs[1].mamba_state_dim == 24 + assert alias_config.mamba_state_dim == 24 + + +def test_resolver_materializes_pipeline_and_mtp_topology_from_base_config(): + config = _config(2, pp_size=2, mtp_num_layers=2) + config.mtp_use_repeated_layer = True + occurrence_config = deepcopy(config) + occurrence_config.pipeline_model_parallel_layout = object() + occurrence_config.mtp_num_layers = 7 + occurrence_config.mtp_use_repeated_layer = False + occurrence_config.mtp_standalone = True + + architecture = resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[ + _layer("mamba", occurrence_config), + PipelineSplit(), + _layer("mamba", occurrence_config), + ], + mtp_layer_specs=[_layer("attention", occurrence_config)], + ) + + for layer in architecture.main_layers + architecture.mtp_layers: + assert layer.config.pipeline_model_parallel_layout is None + assert layer.config.mtp_num_layers == 2 + assert layer.config.mtp_use_repeated_layer is True + assert layer.config.mtp_standalone is False + + +def test_resolver_preserves_permitted_per_layer_shape_differences(): + config = _config(8) + configs = [deepcopy(config) for _ in range(8)] + + configs[0].mamba_state_dim = 24 + configs[0].mamba_head_dim = 4 + configs[0].mamba_num_heads = 16 + configs[0].mamba_num_groups = 4 + configs[1].mamba_state_dim = 32 + + configs[2].num_attention_heads = 16 + configs[2].num_query_groups = 8 + configs[2].kv_channels = 4 + configs[3].num_attention_heads = 4 + configs[3].num_query_groups = 2 + configs[3].kv_channels = 16 + + configs[4].ffn_hidden_size = 160 + configs[5].ffn_hidden_size = 192 + + configs[6].moe_ffn_hidden_size = 112 + configs[6].moe_router_topk = 3 + configs[7].moe_ffn_hidden_size = 144 + configs[7].moe_router_topk = 4 + + architecture = resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[ + _layer("mamba", configs[0]), + _layer("mamba", configs[1]), + _layer("attention", configs[2]), + _layer("attention", configs[3]), + _layer("mlp", configs[4]), + _layer("mlp", configs[5]), + _layer("moe", configs[6]), + _layer("moe", configs[7]), + ], + ) + + layers = architecture.main_layers + assert ( + layers[0].config.mamba_state_dim, + layers[0].config.mamba_head_dim, + layers[0].config.mamba_num_heads, + layers[0].config.mamba_num_groups, + ) == (24, 4, 16, 4) + assert layers[1].config.mamba_state_dim == 32 + assert ( + layers[2].config.num_attention_heads, + layers[2].config.num_query_groups, + layers[2].config.kv_channels, + ) == (16, 8, 4) + assert ( + layers[3].config.num_attention_heads, + layers[3].config.num_query_groups, + layers[3].config.kv_channels, + ) == (4, 2, 16) + assert [layers[index].config.ffn_hidden_size for index in (4, 5)] == [160, 192] + assert [ + (layers[index].config.moe_ffn_hidden_size, layers[index].config.moe_router_topk) + for index in (6, 7) + ] == [(112, 3), (144, 4)] + + +@pytest.mark.parametrize( + ("layer_specs", "mtp_layer_specs", "legacy_pattern", "message"), + [ + pytest.param(["direct"], None, "M", "mutually exclusive", id="direct-and-legacy"), + pytest.param(None, ["mtp"], "M", "requires direct layer_specs", id="mtp-with-legacy"), + pytest.param(None, None, None, "Exactly one", id="missing-architecture"), + ], +) +def test_resolver_rejects_invalid_architecture_source_combinations( + layer_specs, mtp_layer_specs, legacy_pattern, message +): + config = _config(1) + direct_layer = _layer("mamba", config) + if layer_specs is not None: + layer_specs = [direct_layer] + if mtp_layer_specs is not None: + mtp_layer_specs = [direct_layer] + + with pytest.raises(ValueError, match=message): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=layer_specs, + mtp_layer_specs=mtp_layer_specs, + hybrid_layer_pattern=legacy_pattern, + ) + + +def test_direct_pp_requires_splits_and_segment_count_divisible_by_pp_size(): + config = _config(2, pp_size=2) + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match="must contain explicit PipelineSplit"): + resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, layer_specs=[mamba, mamba] + ) + + config = _config(2, pp_size=2) + mamba = _layer("mamba", config) + with pytest.raises(ValueError, match="not divisible"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[mamba, PipelineSplit(), mamba, PipelineSplit()], + ) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("pipeline_model_parallel_layout", object()), + ("num_layers_in_first_pipeline_stage", 1), + ("num_layers_in_last_pipeline_stage", 1), + ("account_for_embedding_in_pipeline_split", True), + ("account_for_loss_in_pipeline_split", True), + ], +) +def test_direct_splits_reject_other_pipeline_ownership_controls(field_name, value): + config = _config(2, pp_size=2) + setattr(config, field_name, value) + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match=field_name): + resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, layer_specs=[mamba, PipelineSplit(), mamba] + ) + + +def test_direct_splits_reject_configured_vpp_mismatch(): + config = _config(4, pp_size=2) + config.virtual_pipeline_model_parallel_size = 3 + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match="imply virtual_pipeline_model_parallel_size=2"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[ + mamba, + PipelineSplit(), + mamba, + PipelineSplit(), + mamba, + PipelineSplit(), + mamba, + ], + ) + + +def test_mtp_rejects_pipeline_splits(): + config = _config(1, mtp_num_layers=1) + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match="not allowed in an MTP pattern"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[mamba], + mtp_layer_specs=[mamba, PipelineSplit(), mamba], + ) + + +def test_direct_mtp_rejects_an_empty_final_decoder_chunk(): + config = _config(1, pp_size=2, mtp_num_layers=1) + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match="standalone MTP placement"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[mamba, PipelineSplit()], + mtp_layer_specs=[mamba], + ) + + +def test_direct_architecture_rejects_standalone_mtp_placement(): + config = _config(1, mtp_num_layers=1) + config.mtp_standalone = True + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match="do not support standalone MTP"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[mamba], + mtp_layer_specs=[mamba], + ) + + +@pytest.mark.parametrize( + ("num_layers", "layer_count"), + [pytest.param(2, 1, id="too-few"), pytest.param(1, 2, id="too-many")], +) +def test_direct_layer_count_must_match_transformer_config(num_layers, layer_count): + config = _config(num_layers) + mamba = _layer("mamba", config) + + with pytest.raises(ValueError, match=f"contains {layer_count} decoder layers"): + resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, layer_specs=[mamba] * layer_count + ) + + +def test_direct_mtp_presence_must_match_configured_depth(): + config = _config(1, mtp_num_layers=1) + mamba = _layer("mamba", config) + with pytest.raises(ValueError, match="requires mtp_layer_specs"): + resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, layer_specs=[mamba] + ) + + +def test_legacy_summary_does_not_validate_or_mutate_mtp_depth(): + config = _config(1, mtp_num_layers=1) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M/M/M" + ) + + assert architecture.mtp_num_layers == 2 + assert config.mtp_num_layers == 1 + + config = _config(1) + mamba = _layer("mamba", config) + with pytest.raises(ValueError, match="requires config.mtp_num_layers > 0"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[mamba], + mtp_layer_specs=[mamba], + ) + + +def test_legacy_summary_preserves_shared_config_identity(): + config = _config(4, mtp_num_layers=2) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M*E-/M*/M*" + ) + + assert all(layer.config is config for layer in architecture.main_layers) + assert all(layer.config is config for layer in architecture.mtp_layers) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("pipeline_model_parallel_layout", object()), + ("account_for_embedding_in_pipeline_split", True), + ("account_for_loss_in_pipeline_split", True), + ], +) +def test_legacy_summary_does_not_reject_new_direct_split_conflicts(field_name, value): + config = _config(2, pp_size=2) + setattr(config, field_name, value) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M|M" + ) + + assert [_types(segment) for segment in architecture.segments] == [["mamba"], ["mamba"]] + assert getattr(config, field_name) is value + + +def test_legacy_summary_does_not_infer_or_mutate_vpp(): + config = _config(4, pp_size=2) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M|M|M|M" + ) + + assert len(architecture.segments) == 4 + assert config.virtual_pipeline_model_parallel_size is None + + +def test_legacy_summary_does_not_add_a_layer_count_validation(): + config = _config(3) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M" + ) + + assert _types(architecture.main_layers) == ["mamba"] + assert config.num_layers == 3 + + +def test_legacy_summary_accepts_layer_classes_supported_by_hybrid_stack(): + stack_spec = deepcopy(STACK_SPEC) + stack_spec.submodules.mamba_layer = _Layer + config = _config(1) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=stack_spec, hybrid_layer_pattern="M" + ) + + assert architecture.main_layers[0].module_spec.module is _Layer + assert architecture.main_layers[0].config is config + + +def test_legacy_summary_only_reads_submodule_fields_used_by_the_pattern(): + stack_spec = ModuleSpec( + module=_Stack, submodules=SimpleNamespace(mamba_layer=_tagged_spec("mamba")) + ) + config = _config(1) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=stack_spec, hybrid_layer_pattern="M" + ) + + assert _types(architecture.main_layers) == ["mamba"] + + +def test_legacy_summary_preserves_empty_final_chunk_with_mtp(): + config = _config(1, pp_size=2, mtp_num_layers=1) + + architecture = resolve_hybrid_architecture( + config=config, hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M|/M" + ) + + assert [len(segment) for segment in architecture.segments] == [1, 0] + assert _types(architecture.mtp_layers) == ["mamba"] + + +def test_direct_architecture_rejects_nonuniform_total_expert_count(): + config = _config(1, mtp_num_layers=1) + first_config = deepcopy(config) + second_config = deepcopy(config) + second_config.num_moe_experts = 16 + + with pytest.raises(ValueError, match="uniform num_moe_experts"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[_layer("moe", first_config)], + mtp_layer_specs=[_layer("moe", second_config)], + ) + + +def test_direct_architecture_expert_count_must_match_base_config(): + config = _config(1) + occurrence_config = deepcopy(config) + occurrence_config.num_moe_experts = 16 + + with pytest.raises(ValueError, match="must match the model-wide config"): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[_layer("moe", occurrence_config)], + ) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("fp32_residual_connection", True), + ("attention_softmax_in_fp32", False), + ("enable_autocast", True), + ("fp8_margin", 1), + ("layernorm_zero_centered_gamma", True), + ("moe_router_dtype", "fp32"), + ("mamba_training_ssm_states_dtype", torch.float32), + ("dsa_indexer_k_norm_fp32", True), + ("tensor_parallel_num_weight_shards", 2), + ("expert_tensor_parallel_num_weight_shards", 2), + ("hierarchical_context_parallel_sizes", [1]), + ("hybrid_context_parallel", True), + ], +) +def test_direct_architecture_rejects_per_layer_model_wide_changes(field_name, value): + config = _config(1) + occurrence_config = deepcopy(config) + setattr(occurrence_config, field_name, value) + + with pytest.raises(ValueError, match=field_name): + resolve_hybrid_architecture( + config=config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[_layer("mamba", occurrence_config)], + ) + + +def test_equivalent_legacy_and_direct_architectures_have_same_types_and_order(): + direct_config = _config(4, pp_size=2) + direct_architecture = resolve_hybrid_architecture( + config=direct_config, + hybrid_stack_spec=STACK_SPEC, + layer_specs=[ + _layer("mamba", direct_config), + _layer("attention", direct_config), + PipelineSplit(), + _layer("moe", direct_config), + _layer("mlp", direct_config), + ], + ) + legacy_architecture = resolve_hybrid_architecture( + config=_config(4, pp_size=2), hybrid_stack_spec=STACK_SPEC, hybrid_layer_pattern="M*|E-" + ) + + assert [_types(segment) for segment in direct_architecture.segments] == [ + _types(segment) for segment in legacy_architecture.segments + ] + assert [layer.module_spec.module for layer in direct_architecture.main_layers] == [ + layer.module_spec.module for layer in legacy_architecture.main_layers + ] + assert direct_architecture.source == "direct" + assert legacy_architecture.source == "legacy"