Refactor hybrid layers to use per-layer configs

megatron/core/dist_checkpointing/gpt_checkpoint_interop.py

@@ -56,20 +56,16 @@
     ShardedTensor,
     ShardedTensorFactory,
 )
-from megatron.core.models.hybrid.hybrid_layer_allocation import (
-    Symbols,
-    get_layer_maps_from_layer_type_list,
-    parse_hybrid_pattern,
-)
+from megatron.core.models.hybrid.hybrid_layer_allocation import PIPE_SEPARATOR, parse_hybrid_pattern
 
 # Hybrid layer symbols that have a GPT-side source of weights ('*', '-', 'E')
 # or that are explicitly initialized from scratch ('M'). Attention layers map
 # onto GPT ``self_attention`` sub-modules; dense and MoE MLP layers map onto
 # GPT ``mlp`` sub-modules (both models keep MoE tensors under ``mlp.*``).
 # GDN ('G') and DS-attention ('D') use different weight layouts than GPT
 # attention and are rejected rather than silently mistranslated.
-_GPT_SOURCED_SYMBOLS = (Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)
-_FRESH_INIT_SYMBOLS = (Symbols.MAMBA,)
+_GPT_SOURCED_PATTERN_CHARACTERS = ('*', '-', 'E')
+_FRESH_INIT_PATTERN_CHARACTERS = ('M',)
 
 _DECODER_LAYER_KEY_RE = re.compile(r'decoder\.layers\.(\d+)\.')
 
@@ -119,13 +115,12 @@ def gpt_compatible_layer_maps(hybrid_layer_pattern: str) -> GPTCompatLayerMaps:
             f"('/{parsed.mtp_pattern}'), which have no source weights in a GPT "
             f"checkpoint. Remove the MTP part of the pattern to load a GPT checkpoint."
         )
-    main_pattern = (parsed.main_pattern or '').replace(Symbols.PIPE, '')
+    main_pattern = (parsed.main_pattern or '').replace(PIPE_SEPARATOR, '')
     if not main_pattern:
         raise ValueError("Hybrid layer pattern is empty; set --hybrid-layer-pattern.")
 
-    layer_type_list = list(main_pattern)
-    translatable = set(_GPT_SOURCED_SYMBOLS) | set(_FRESH_INIT_SYMBOLS)
-    unknown = sorted(set(layer_type_list) - translatable)
+    translatable = set(_GPT_SOURCED_PATTERN_CHARACTERS) | set(_FRESH_INIT_PATTERN_CHARACTERS)
+    unknown = sorted(set(main_pattern) - translatable)
     if unknown:
         raise ValueError(
             f"Hybrid layer pattern {hybrid_layer_pattern!r} contains layer types "
@@ -134,17 +129,27 @@ def gpt_compatible_layer_maps(hybrid_layer_pattern: str) -> GPTCompatLayerMaps:
             f"initialization)."
         )
 
-    layer_maps = get_layer_maps_from_layer_type_list(layer_type_list)
-    dense_map = layer_maps[Symbols.MLP]
-    moe_map = layer_maps[Symbols.MOE]
+    attention_map = {}
+    dense_map = {}
+    moe_map = {}
+    fresh_init = set()
+    for global_layer_idx, pattern_char in enumerate(main_pattern):
+        if pattern_char == '*':
+            attention_map[global_layer_idx] = len(attention_map)
+        elif pattern_char == '-':
+            dense_map[global_layer_idx] = len(dense_map)
+        elif pattern_char == 'E':
+            moe_map[global_layer_idx] = len(moe_map)
+        elif pattern_char == 'M':
+            fresh_init.add(global_layer_idx)
+
     if dense_map and moe_map:
         raise ValueError(
             f"Hybrid layer pattern {hybrid_layer_pattern!r} mixes dense ('-') and "
             f"MoE ('E') MLP positions. GPT checkpoints have one MLP kind on every "
             f"layer, so the pattern must use only one of '-' or 'E'."
         )
     mlp_map = moe_map if moe_map else dense_map
-    attention_map = layer_maps[Symbols.ATTENTION]
 
     if len(attention_map) != len(mlp_map) or not attention_map:
         raise ValueError(
@@ -157,7 +162,7 @@ def gpt_compatible_layer_maps(hybrid_layer_pattern: str) -> GPTCompatLayerMaps:
     return GPTCompatLayerMaps(
         attention_to_gpt=dict(attention_map),
         mlp_to_gpt=dict(mlp_map),
-        fresh_init=frozenset(layer_maps[Symbols.MAMBA]),
+        fresh_init=frozenset(fresh_init),
         num_gpt_layers=len(attention_map),
     )
 

megatron/core/inference/config.py

@@ -7,7 +7,9 @@
 
 import torch
 
+from megatron.core.models.hybrid.hybrid_layer_allocation import HybridLayerConfig
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.ssm.mamba_layer_config import MambaLayerConfig
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_attr_wrapped_model
 
@@ -22,10 +24,9 @@ class MambaInferenceStateConfig:
     these. Once the kernels have been updated we can simplify this code.
     """
 
-    layer_type_list: List[str]
+    layer_config_list: List[HybridLayerConfig]
     """
-    A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer.
-    See `megatron/core/models/hybrid/hybrid_layer_allocation.py` for the list of symbols.
+    Per-layer configs used to derive dynamic inference cache indexing.
     """
 
     conv_states_shape: Tuple[int]
@@ -51,12 +52,12 @@ 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.models.hybrid.hybrid_layer_allocation import Symbols
-
         decoder = get_attr_wrapped_model(model, "decoder")
-        layer_type_list = getattr(decoder, "layer_type_list", None)
-        if layer_type_list is not None and Symbols.MAMBA in layer_type_list:
-            (mamba_conv_states_shape, mamba_ssm_states_shape) = (
+        layer_config_list = getattr(decoder, "layer_config_list", None)
+        if layer_config_list is not None and any(
+            type(layer_config) is MambaLayerConfig for layer_config in layer_config_list
+        ):
+            mamba_conv_states_shape, mamba_ssm_states_shape = (
                 decoder.mamba_state_shapes_per_request()
             )
             if conv_states_dtype is None:
@@ -73,12 +74,12 @@ def from_model(
             elif ssm_states_dtype is None:
                 ssm_states_dtype = model.config.params_dtype
             mamba_chunk_size = 128
-            for layer_type, layer in zip(decoder.layer_type_list, decoder.layers):
-                if layer_type == Symbols.MAMBA and hasattr(layer, 'mixer'):
+            for layer_config, layer in zip(layer_config_list, decoder.layers):
+                if type(layer_config) is MambaLayerConfig and hasattr(layer, 'mixer'):
                     mamba_chunk_size = layer.mixer.chunk_size
                     break
             return cls(
-                layer_type_list=layer_type_list,
+                layer_config_list=list(layer_config_list),
                 conv_states_shape=mamba_conv_states_shape,
                 ssm_states_shape=mamba_ssm_states_shape,
                 conv_states_dtype=conv_states_dtype,

megatron/core/inference/contexts/dynamic_context.py

@@ -2,7 +2,6 @@
 
 import logging
 import math
-import operator
 import warnings
 from contextlib import nullcontext
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
@@ -30,13 +29,16 @@
 )
 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 (
-    Symbols,
-    get_layer_maps_from_layer_type_list,
-)
 from megatron.core.package_info import __version__ as mcore_version
+from megatron.core.ssm.gdn_layer_config import GDNLayerConfig
+from megatron.core.ssm.mamba_layer_config import MambaLayerConfig
+from megatron.core.ssm.mlp_layer_config import MLPLayerConfig
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.core.transformer.attention_layer_config import AttentionLayerConfig
 from megatron.core.transformer.enums import InferenceCudaGraphScope
+from megatron.core.transformer.experimental_attention_variant.dsa_layer_config import DSALayerConfig
+from megatron.core.transformer.mla_layer_config import MLALayerConfig
+from megatron.core.transformer.moe.moe_layer_config import MoELayerConfig
 from megatron.core.transformer.moe.token_dispatcher_inference import (
     InferenceAllGatherDispatcherBase,
     NCCLAllGatherDispatcher,
@@ -438,17 +440,32 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
                     "boundaries are not rounded between decode chunks."
                 )
 
-            # For hybrid models, the layer map converts the global layer index to the
-            # corresponding attention layer index or Mamba layer index depending on the
-            # layer type.
-            attention_layer_map, dsa_layer_map, gdn_layer_map, mamba_layer_map = (
-                operator.itemgetter(
-                    Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.MAMBA
-                )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list))
-            )
-
-            if len(gdn_layer_map) > 0:
-                raise NotImplementedError("GDN layers are not supported for inference.")
+            # Convert each global layer index to the corresponding attention or
+            # Mamba cache index. Keep the attention variants in separate maps to
+            # preserve the legacy cache-indexing behavior.
+            attention_layer_map = {}
+            dsa_layer_map = {}
+            mamba_layer_map = {}
+            for global_layer_idx, layer_config in enumerate(
+                mamba_inference_state_config.layer_config_list
+            ):
+                if type(layer_config) is AttentionLayerConfig:
+                    attention_layer_map[global_layer_idx] = len(attention_layer_map)
+                elif type(layer_config) is DSALayerConfig:
+                    dsa_layer_map[global_layer_idx] = len(dsa_layer_map)
+                elif type(layer_config) is MambaLayerConfig:
+                    mamba_layer_map[global_layer_idx] = len(mamba_layer_map)
+                elif type(layer_config) is GDNLayerConfig:
+                    raise NotImplementedError("GDN layers are not supported for inference.")
+                elif type(layer_config) is MLALayerConfig:
+                    # Hybrid MLA cache mapping was not supported by the legacy path.
+                    continue
+                elif type(layer_config) in (MLPLayerConfig, MoELayerConfig):
+                    continue
+                else:
+                    raise ValueError(
+                        f"Unexpected hybrid layer config type: {type(layer_config).__name__}"
+                    )
 
             self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map)
             self.num_mamba_layers = len(mamba_layer_map)