Reject unsupported heterogeneous hybrid inference

Signed-off-by: Philip Petrakian <ppetrakian@nvidia.com>

Author: Phlip79

docs/user-guide/hybrid-model-migration.md

@@ -366,11 +366,6 @@ remove `--decoder-first-pipeline-num-layers` and
 `--decoder-last-pipeline-num-layers`. Express virtual-pipeline segmentation
 with additional pipe-delimited segments instead.
 
-The declarative `HybridModelBuilder` currently rejects virtual pipeline
-parallelism. Pipe-defined virtual stages are supported by the
-`pretrain_hybrid.py` CLI builder, but custom builder users must avoid VPP or use
-a path that explicitly supports it.
-
 ### Update custom providers and conversion mappings
 
 Custom providers and conversion mappings also need to account for these API and
@@ -398,3 +393,12 @@ Before starting a long run:
   parameter counts by layer.
 - Save and reload one new checkpoint to confirm that the new optimizer and RNG
   state resume correctly.
+
+Direct layer specs are currently training-first. Dynamic inference rejects a
+direct architecture when occurrence configurations are incompatible with its
+model-global cache or runtime buffers, whether the occurrences differ from one
+another or a family-uniform override differs from the base `TransformerConfig`.
+Examples include differing Mamba state dimensions and attention KV dimensions
+or MoE top-k values that differ from the base config. Legacy
+`hybrid_layer_pattern` models retain their existing inference path and do not
+use this direct-spec validator.

megatron/core/inference/apis/_llm_base.py

@@ -17,7 +17,7 @@
 
 import torch.distributed as dist
 
-from megatron.core.inference.config import InferenceConfig
+from megatron.core.inference.config import InferenceConfig, validate_dynamic_inference_model
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine, EngineState
 from megatron.core.inference.inference_request import DynamicInferenceRequest
@@ -279,6 +279,7 @@ def __init__(
             inference_config = InferenceConfig()
 
         # Build the engine pipeline. Mirrors examples/inference/gpt/gpt_dynamic_inference.py.
+        validate_dynamic_inference_model(model)
         context = DynamicInferenceContext(model.config, inference_config)
         wrapper = GPTInferenceWrapper(model, context)
         controller = TextGenerationController(inference_wrapped_model=wrapper, tokenizer=tokenizer)

megatron/core/inference/config.py

@@ -12,6 +12,29 @@
 from megatron.core.utils import get_attr_wrapped_model
 
 
+def validate_dynamic_inference_model(model: MegatronModule) -> None:
+    """Reject incompatible direct specs before model-global inference buffers are allocated."""
+
+    from megatron.core.models.hybrid.hybrid_architecture import ResolvedHybridArchitecture
+
+    try:
+        architecture = get_attr_wrapped_model(model, "resolved_hybrid_architecture")
+    except RuntimeError:
+        return
+    if not isinstance(architecture, ResolvedHybridArchitecture):
+        return
+    if architecture.source != "direct":
+        return
+
+    model_config = get_attr_wrapped_model(model, "config", allow_none=False)
+    if architecture.has_incompatible_dynamic_inference_shapes(model_config):
+        raise NotImplementedError(
+            "Direct HybridModel occurrence configurations are incompatible with dynamic "
+            "inference's model-global cache or runtime buffers (Mamba state/cache dimensions, "
+            "attention KV dimensions, or MoE router top-k values)."
+        )
+
+
 @dataclass
 class MambaInferenceStateConfig:
     """
@@ -49,14 +72,41 @@ def from_model(
         model: MegatronModule,
         conv_states_dtype: Optional[torch.dtype] = None,
         ssm_states_dtype: Optional[torch.dtype] = None,
+        *,
+        validate_dynamic_inference: bool = True,
     ) -> Optional["MambaInferenceStateConfig"]:
-        """Returns Mamba inference state config from the model if it is a hybrid model."""
+        """Returns Mamba inference state config from the model if it is a hybrid model.
+
+        Args:
+            validate_dynamic_inference: Validate direct specs against the model-global buffers
+                used by dynamic inference. The explicit legacy static engine disables this
+                because its Mamba, attention, and MoE state is allocated by each layer.
+        """
         from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols
 
+        if validate_dynamic_inference:
+            validate_dynamic_inference_model(model)
         decoder = get_attr_wrapped_model(model, "decoder")
         layer_type_list = getattr(decoder, "layer_type_list", None)
+        # HybridStack's first-class API exposes stable semantic names, while
+        # dynamic inference's layer maps intentionally retain legacy symbols.
+        semantic_to_symbol = {
+            "mamba": Symbols.MAMBA,
+            "gdn": Symbols.GDN,
+            "attention": Symbols.ATTENTION,
+            "dsa": Symbols.DS_ATTENTION,
+            "mla": Symbols.MLA,
+            "mlp": Symbols.MLP,
+            "moe": Symbols.MOE,
+        }
+        if layer_type_list is not None and any(
+            layer_type in semantic_to_symbol for layer_type in layer_type_list
+        ):
+            layer_type_list = [
+                semantic_to_symbol.get(layer_type, layer_type) for layer_type in layer_type_list
+            ]
         if layer_type_list is not None and Symbols.MAMBA in layer_type_list:
-            (mamba_conv_states_shape, mamba_ssm_states_shape) = (
+            mamba_conv_states_shape, mamba_ssm_states_shape = (
                 decoder.mamba_state_shapes_per_request()
             )
             if conv_states_dtype is None:
@@ -73,7 +123,7 @@ 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):
+            for layer_type, layer in zip(layer_type_list, decoder.layers):
                 if layer_type == Symbols.MAMBA and hasattr(layer, 'mixer'):
                     mamba_chunk_size = layer.mixer.chunk_size
                     break

megatron/core/inference/engines/static_engine.py

@@ -94,7 +94,7 @@ def __init__(
         self.scheduler = Scheduler(max_batch_size=max_batch_size)
 
         mamba_inference_state_config = MambaInferenceStateConfig.from_model(
-            self.inference_wrapped_model.model
+            self.inference_wrapped_model.model, validate_dynamic_inference=not legacy
         )
 
         try:

megatron/core/models/hybrid/hybrid_model.py

@@ -641,6 +641,21 @@ def forward(
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
         in_inference_mode = InferenceMode.is_active()
+        resolved_architecture = getattr(self, "resolved_hybrid_architecture", None)
+
+        if (
+            in_inference_mode
+            and inference_context is not None
+            and inference_context.is_dynamic_batching()
+            and resolved_architecture is not None
+            and resolved_architecture.source == "direct"
+            and resolved_architecture.has_incompatible_dynamic_inference_shapes(self.config)
+        ):
+            raise NotImplementedError(
+                "Direct HybridModel occurrence configurations are incompatible with dynamic "
+                "inference's model-global cache or runtime buffers (Mamba state/cache "
+                "dimensions, attention KV dimensions, or MoE router top-k values)."
+            )
 
         if in_inference_mode:
             assert runtime_gather_output, "Inference must always gather TP logits"