Changing hybrid_override_pattern to hybrid_layer_pattern to mirror mcore change

src/megatron/bridge/models/mamba/mamba_builder.py

@@ -93,11 +93,12 @@ class MambaModelConfig(ModelConfig):
     on the embedded ``transformer`` config are accessible directly on this object
     via ``__getattr__``/``__setattr__`` proxying.
 
-    Supports hybrid SSM/attention architectures via ``hybrid_attention_ratio``,
-    ``hybrid_mlp_ratio``, and ``hybrid_override_pattern``.
+    Supports hybrid SSM/attention architectures via ``hybrid_layer_pattern``
 
     Note:
         ``vocab_size`` must be set before passing this config to ``MambaModelBuilder``.
+        ``hybrid_attention_ratio``,``hybrid_mlp_ratio``, and
+        ``hybrid_override_pattern`` are deprecated and will be removed in a future release.
     """
 
     builder: ClassVar[str] = "megatron.bridge.models.mamba.MambaModelBuilder"
@@ -108,6 +109,7 @@ class MambaModelConfig(ModelConfig):
     hybrid_attention_ratio: float = 0.0
     hybrid_mlp_ratio: float = 0.0
     hybrid_override_pattern: str | None = None
+    hybrid_layer_pattern: str | None = None
     seq_length: int = 8192
     # Mamba with no attention has no need for position embeddings, so none is default
     position_embedding_type: Literal["learned_absolute", "rope", "none"] = "none"
@@ -222,9 +224,7 @@ def build_model(
             mamba_stack_spec=mamba_stack_spec,
             vocab_size=padded_vocab_size,
             max_sequence_length=self._model_config.seq_length,
-            hybrid_attention_ratio=self._model_config.hybrid_attention_ratio,
-            hybrid_mlp_ratio=self._model_config.hybrid_mlp_ratio,
-            hybrid_override_pattern=self._model_config.hybrid_override_pattern,
+            hybrid_layer_pattern=self._model_config.hybrid_layer_pattern,
             fp16_lm_cross_entropy=self._model_config.fp16_lm_cross_entropy,
             parallel_output=self._model_config.parallel_output,
             share_embeddings_and_output_weights=self._model_config.share_embeddings_and_output_weights,

src/megatron/bridge/models/mamba/mamba_provider.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import warnings
 from dataclasses import dataclass
 from typing import Callable, Literal, Optional, Union
 
@@ -22,11 +23,13 @@
 from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
 from megatron.core.post_training.modelopt.mamba.model_specs import get_mamba_stack_modelopt_spec
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.ssm.mamba_hybrid_layer_allocation import get_hybrid_total_layer_count
 from megatron.core.transformer import ModuleSpec
 from megatron.core.transformer.enums import AttnBackend
 
 from megatron.bridge.models.model_provider import ModelProviderMixin
 from megatron.bridge.models.transformer_config import TransformerConfig
+from megatron.bridge.utils.common_utils import get_rank_safe
 from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 
 
@@ -97,6 +100,7 @@ class MambaModelProvider(TransformerConfig, ModelProviderMixin[MCoreMambaModel])
     hybrid_attention_ratio: float = 0.0
     hybrid_mlp_ratio: float = 0.0
     hybrid_override_pattern: Optional[str] = None
+    hybrid_layer_pattern: Optional[str] = None
     seq_length: int = 8192
     # Mamba with no attention has no need for position embeddings, so none is default
     position_embedding_type: Literal["learned_absolute", "rope", "none"] = "none"
@@ -127,6 +131,47 @@ class MambaModelProvider(TransformerConfig, ModelProviderMixin[MCoreMambaModel])
     # If True, restore the modelopt_state that contains quantization, sparsity, speculative decoding transformation state.
     restore_modelopt_state: bool = False
 
+    def finalize(self) -> None:
+        """Finalize the Mamba model provider.
+        Calculates the number of layers from the hybrid_layer_pattern.
+        Executes the deferred MCore post-init logic.
+        """
+        # Check if hybrid_override_pattern is specified and throw deprecation warning
+        used_hybrid_override_pattern = False
+        if self.hybrid_override_pattern is not None:
+            assert self.hybrid_layer_pattern is None, (
+                "hybrid_override_pattern and hybrid_layer_pattern cannot both be specified. "
+                "hybrid_override_pattern is deprecated; use hybrid_layer_pattern instead."
+            )
+            if get_rank_safe() == 0:
+                warnings.warn(
+                    "hybrid_override_pattern is deprecated. Use hybrid_layer_pattern instead.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            self.hybrid_layer_pattern = self.hybrid_override_pattern
+            used_hybrid_override_pattern = True
+
+        # Check if hybrid_layer_pattern is specified and derive num_layers from pattern
+        if self.hybrid_layer_pattern is not None:
+            # Derive num_layers from pattern
+            num_layers_in_pattern = get_hybrid_total_layer_count(self.hybrid_layer_pattern)
+            if self.num_layers is not None:
+                if used_hybrid_override_pattern:
+                    assert self.num_layers == num_layers_in_pattern, (
+                        f"num_layers ({self.num_layers}) does not match the number of layers "
+                        f"derived from hybrid_override_pattern ({num_layers_in_pattern}). "
+                        f"Please correct num_layers or the pattern."
+                    )
+                else:
+                    assert False, (
+                        "If hybrid_layer_pattern is specified, num_layers should not be specified. "
+                        "The number of layers is derived from the pattern."
+                    )
+            self.num_layers = num_layers_in_pattern
+
+        super().finalize()
+
     def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreMambaModel:
         """Configure and instantiate a Megatron Core Mamba model based on this configuration.
 
@@ -166,9 +211,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreMa
             mamba_stack_spec=mamba_stack_spec,
             vocab_size=padded_vocab_size,
             max_sequence_length=self.seq_length,
-            hybrid_attention_ratio=self.hybrid_attention_ratio,
-            hybrid_mlp_ratio=self.hybrid_mlp_ratio,
-            hybrid_override_pattern=self.hybrid_override_pattern,
+            hybrid_layer_pattern=self.hybrid_layer_pattern,
             fp16_lm_cross_entropy=self.fp16_lm_cross_entropy,
             parallel_output=self.parallel_output,
             share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,

src/megatron/bridge/models/nemotron_vl/nemotron_vl_bridge.py

@@ -38,6 +38,15 @@
 class NemotronVLBridge(MegatronModelBridge):
     """Conversion utilities between HF Nemotron-VL and Megatron-Core format."""
 
+    # Extend CONFIG_MAPPING with Nemotron-VL specific fields
+    CONFIG_MAPPING = MegatronModelBridge.CONFIG_MAPPING + [
+        # Mamba-specific fields
+        ("hybrid_override_pattern", "hybrid_layer_pattern"),
+    ]
+
+    # Remove num_hidden_layers from CONFIG_MAPPING as it is derived from hybrid_layer_pattern
+    CONFIG_MAPPING.remove(("num_hidden_layers", "num_layers"))
+
     # ------------------------------------------------------------------
     # Provider translation
     # ------------------------------------------------------------------

src/megatron/bridge/models/nemotron_vl/nemotron_vl_provider.py

@@ -132,9 +132,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None):  # noqa:
             img_h=512,
             img_w=512,
             patch_dim=16,
-            hybrid_attention_ratio=0.0,
-            hybrid_mlp_ratio=0.0,
-            hybrid_override_pattern=self.hybrid_override_pattern,
+            hybrid_layer_pattern=self.hybrid_layer_pattern,
             image_token_index=131072,
             pixel_shuffle=True,
             max_num_tiles=12,

src/megatron/bridge/models/nemotronh/nemotron_h_bridge.py

@@ -61,14 +61,17 @@ class NemotronHBridge(MegatronModelBridge):
         ("mamba_num_heads", "mamba_num_heads"),
         ("n_groups", "mamba_num_groups"),
         ("ssm_state_size", "mamba_state_dim"),
-        ("hybrid_override_pattern", "hybrid_override_pattern"),
+        ("hybrid_override_pattern", "hybrid_layer_pattern"),
         ("residual_in_fp32", "fp32_residual_connection"),
         ("use_bias", "add_bias_linear"),
         ("layer_norm_epsilon", "layernorm_epsilon"),
         # MoE-specific fields (already in base but with different HF names)
         ("moe_shared_expert_intermediate_size", "moe_shared_expert_intermediate_size"),
     ]
 
+    # Remove num_hidden_layers from CONFIG_MAPPING as it is derived from hybrid_layer_pattern
+    CONFIG_MAPPING.remove(("num_hidden_layers", "num_layers"))
+
     # Additional files to copy during HF export (reasoning parser utilities)
     ADDITIONAL_FILE_PATTERNS = ["*reasoning_parser.py"]
 

src/megatron/bridge/models/nemotronh/nemotron_h_provider.py

@@ -60,8 +60,7 @@ class NemotronHModelProvider(MambaModelProvider):
 class NemotronHModelProvider4B(NemotronHModelProvider):
     """Configuration for a 4B parameter Nemotron-H model."""
 
-    hybrid_override_pattern: str = "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"
-    num_layers: int = 52
+    hybrid_layer_pattern: str = "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"
     hidden_size: int = 3072
     mamba_num_heads: int = 112
     kv_channels: int = 128
@@ -75,8 +74,7 @@ class NemotronHModelProvider4B(NemotronHModelProvider):
 class NemotronHModelProvider8B(NemotronHModelProvider):
     """Configuration for a 8B parameter Nemotron-H model."""
 
-    hybrid_override_pattern: str = "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"
-    num_layers: int = 52
+    hybrid_layer_pattern: str = "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"
     hidden_size: int = 4096
     mamba_state_dim: int = 128
     mamba_num_heads: int = 128
@@ -88,10 +86,9 @@ class NemotronHModelProvider8B(NemotronHModelProvider):
 class NemotronHModelProvider47B(NemotronHModelProvider):
     """Configuration for a 47B parameter Nemotron-H model."""
 
-    hybrid_override_pattern: str = (
+    hybrid_layer_pattern: str = (
         "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---MM---M-M*-M-M-M-M-M-"
     )
-    num_layers: int = 98
     hidden_size: int = 8192
     mamba_state_dim: int = 256
     mamba_num_heads: int = 256
@@ -103,11 +100,10 @@ class NemotronHModelProvider47B(NemotronHModelProvider):
 class NemotronHModelProvider56B(NemotronHModelProvider):
     """Configuration for a 56B parameter Nemotron-H model."""
 
-    hybrid_override_pattern: str = (
+    hybrid_layer_pattern: str = (
         "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-"
     )
-    num_layers: int = 118
     hidden_size: int = 8192
     mamba_state_dim: int = 256
     mamba_num_heads: int = 256
@@ -121,8 +117,7 @@ class NemotronHModelProvider56B(NemotronHModelProvider):
 class NemotronNanoModelProvider9Bv2(NemotronHModelProvider):
     """Configuration for a 9B parameter Nemotron Nano v2 model."""
 
-    hybrid_override_pattern: str = "M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M-"
-    num_layers: int = 56
+    hybrid_layer_pattern: str = "M-M-M-MM-M-M-M*-M-M-M*-M-M-M-M*-M-M-M-M*-M-MM-M-M-M-M-M-"
     hidden_size: int = 4480
     mamba_num_heads: int = 128
     kv_channels: int = 128
@@ -137,8 +132,7 @@ class NemotronNanoModelProvider9Bv2(NemotronHModelProvider):
 class NemotronNanoModelProvider12Bv2(NemotronHModelProvider):
     """Configuration for the Nemotron Nano v2 12B model."""
 
-    hybrid_override_pattern: str = "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-"
-    num_layers: int = 62
+    hybrid_layer_pattern: str = "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-"
     hidden_size: int = 5120
     mamba_num_heads: int = 128
     kv_channels: int = 128
@@ -155,8 +149,7 @@ class Nemotron3NanoProvider(NemotronHModelProvider):
 
     seq_length: int = 262144
     num_query_groups: int = 2
-    hybrid_override_pattern: str = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"
-    num_layers: int = 52
+    hybrid_layer_pattern: str = "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"
     hidden_size: int = 2688
     mamba_num_heads: int = 64
     kv_channels: int = 128

src/megatron/bridge/training/mlm_compat/model.py

@@ -167,6 +167,7 @@ def _mamba_provider(
         hybrid_attention_ratio=args.hybrid_attention_ratio,
         hybrid_mlp_ratio=args.hybrid_mlp_ratio,
         hybrid_override_pattern=args.hybrid_override_pattern,
+        hybrid_layer_pattern=args.hybrid_layer_pattern,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         parallel_output=True,

src/megatron/bridge/training/utils/flop_utils.py

@@ -36,13 +36,13 @@ def num_floating_point_operations(cfg: ConfigContainer, batch_size: int = 1):
 
     def calculate_layer_counts():
         """Calculate the number of attention, Mamba, MLP, and MoE layers."""
-        if hasattr(cfg.model, "hybrid_override_pattern") and cfg.model.hybrid_override_pattern:
+        if hasattr(cfg.model, "hybrid_layer_pattern") and cfg.model.hybrid_layer_pattern:
             counts = {"M": 0, "*": 0, "-": 0, "E": 0}
             try:
                 parse_hybrid_pattern = importlib.import_module(
                     "megatron.core.ssm.mamba_hybrid_layer_allocation"
                 ).parse_hybrid_pattern
-                parsed = parse_hybrid_pattern(cfg.model.hybrid_override_pattern)
+                parsed = parse_hybrid_pattern(cfg.model.hybrid_layer_pattern)
                 if parsed.main_pattern:
                     for layer_type in parsed.main_pattern:
                         if layer_type in counts:
@@ -52,7 +52,7 @@ def calculate_layer_counts():
                         if layer_type in counts:
                             counts[layer_type] += parsed.mtp_num_depths
             except (ImportError, ModuleNotFoundError):
-                for layer_type in cfg.model.hybrid_override_pattern:
+                for layer_type in cfg.model.hybrid_layer_pattern:
                     if layer_type in counts:
                         counts[layer_type] += 1
             return counts["*"], counts["M"], counts["-"], counts["E"]
@@ -431,7 +431,7 @@ def transformer_flops():
         mtp_num_layers = getattr(cfg.model, "mtp_num_layers", None)
         if mtp_num_layers is None:
             # When using unified hybrid patterns, infer MTP depth count from the pattern.
-            hybrid_pattern = getattr(cfg.model, "hybrid_override_pattern", None)
+            hybrid_pattern = getattr(cfg.model, "hybrid_layer_pattern", None)
             if hybrid_pattern:
                 try:
                     parse_hybrid_pattern = importlib.import_module(

src/megatron/bridge/training/utils/train_utils.py

@@ -624,7 +624,7 @@ def training_log(
             track_names.append("z_loss")
 
         if config.model.is_hybrid_model:
-            layers = config.model.hybrid_override_pattern.count("E")
+            layers = config.model.hybrid_layer_pattern.count("E")
         else:
             layers = config.model.num_layers
 

tests/functional_tests/models/nemotron_vl/test_nemotron_vl_conversion.py

@@ -198,6 +198,24 @@ def test_nemotron_vl_conversion_parallelism(self, nemotron_vl_toy_model_path, tm
         test_output_dir = tmp_path / f"nemotron_vl_{test_name}"
         test_output_dir.mkdir(exist_ok=True)
 
+        # Modify config.json to add | separator for hybrid_override_pattern to be able to run PP > 1
+        config_file = Path(nemotron_vl_toy_model_path) / "config.json"
+        assert config_file.exists(), f"config.json not found at {config_file}"
+        with open(config_file) as f:
+            config_data = json.load(f)
+
+        if pp > 1:
+            config_data["hybrid_override_pattern"] = (
+                HF_NEMOTRON_VL_TOY_MODEL_OVERRIDES["hybrid_override_pattern"][:2]
+                + '|' 
+                + HF_NEMOTRON_VL_TOY_MODEL_OVERRIDES["hybrid_override_pattern"][2:]
+            )
+        else:
+            config_data["hybrid_override_pattern"] = HF_NEMOTRON_VL_TOY_MODEL_OVERRIDES["hybrid_override_pattern"]
+
+        with open(config_file, "w") as f:
+            json.dump(config_data, f, indent=2)
+
         # Run hf_megatron_roundtrip_multi_gpu.py with specified parallelism configuration on our toy model
         cmd = [
             "python",

tests/functional_tests/models/nemotronh/test_nemotron_h_conversion.py

@@ -209,6 +209,24 @@ def test_nemotronh_conversion_parallelism(self, nemotronh_toy_model_path, tmp_pa
         test_output_dir = tmp_path / f"nemotronh_{test_name}"
         test_output_dir.mkdir(exist_ok=True)
 
+        # Modify config.json to add | separator for hybrid_override_pattern to be able to run PP > 1
+        config_file = Path(nemotronh_toy_model_path) / "config.json"
+        assert config_file.exists(), f"config.json not found at {config_file}"
+        with open(config_file) as f:
+            config_data = json.load(f)
+
+        if pp > 1:
+            config_data["hybrid_override_pattern"] = (
+                HF_NEMOTRONH_TOY_MODEL_OVERRIDES["hybrid_override_pattern"][:2]
+                + '|' 
+                + HF_NEMOTRONH_TOY_MODEL_OVERRIDES["hybrid_override_pattern"][2:]
+            )
+        else:
+            config_data["hybrid_override_pattern"] = HF_NEMOTRONH_TOY_MODEL_OVERRIDES["hybrid_override_pattern"]
+
+        with open(config_file, "w") as f:
+            json.dump(config_data, f, indent=2)
+
         # Run hf_megatron_roundtrip_multi_gpu.py with specified parallelism configuration on our toy model
         cmd = [
             "python",

tests/functional_tests/recipes/test_nemotronh_recipes_finetune.py

@@ -195,7 +195,7 @@ def _finetune_wrapper_full(self, checkpoint_dir, **kwargs):
                 "nemotron_nano_9b_v2_lora",
                 {
                     "num_layers": 4,  # Match toy model
-                    "hybrid_override_pattern": "M*M-",  # Match toy model
+                    "hybrid_layer_pattern": "M*M-",  # Match toy model
                     "hidden_size": 640,  # Match toy model
                     "ffn_hidden_size": 2240,  # Match toy model
                     "num_attention_heads": 8,  # Match toy model
@@ -214,7 +214,7 @@ def _finetune_wrapper_full(self, checkpoint_dir, **kwargs):
                 "nemotron_nano_9b_v2_full",
                 {
                     "num_layers": 4,  # Match toy model
-                    "hybrid_override_pattern": "M*M-",  # Match toy model
+                    "hybrid_layer_pattern": "M*M-",  # Match toy model
                     "hidden_size": 640,  # Match toy model
                     "ffn_hidden_size": 2240,  # Match toy model
                     "num_attention_heads": 8,  # Match toy model
@@ -320,7 +320,7 @@ def test_nemotron_nano_v2_finetune_recipes(
 
 MEGATRON_NEMOTRON_3_NANO_OVERRIDES = {
     "num_layers": HF_NEMOTRON_3_NANO_TOY_MODEL_OVERRIDES["num_hidden_layers"],
-    "hybrid_override_pattern": HF_NEMOTRON_3_NANO_TOY_MODEL_OVERRIDES["hybrid_override_pattern"],
+    "hybrid_layer_pattern": HF_NEMOTRON_3_NANO_TOY_MODEL_OVERRIDES["hybrid_override_pattern"],
     "hidden_size": HF_NEMOTRON_3_NANO_TOY_MODEL_OVERRIDES["hidden_size"],
     "num_moe_experts": HF_NEMOTRON_3_NANO_TOY_MODEL_OVERRIDES["n_routed_experts"],
     "tensor_model_parallel_size": 1,

tests/functional_tests/recipes/test_nemotronh_recipes_pretrain.py

@@ -30,7 +30,7 @@
         nemotronh_4b_pretrain_config,
         "nemotronh_4b",
         {"tensor_model_parallel_size": 1, "pipeline_model_parallel_size": 1},
-        {"num_layers": 3, "hybrid_override_pattern": "M*-"},
+        {"num_layers": 3, "hybrid_layer_pattern": "M*-"},
     ),
 ]
 
@@ -41,7 +41,7 @@
         nemotron_nano_9b_v2_pretrain_config,
         "nemotron_nano_9b_v2",
         {"tensor_model_parallel_size": 1, "pipeline_model_parallel_size": 1},
-        {"num_layers": 3, "hybrid_override_pattern": "M*-", "sequence_parallel": False},
+        {"num_layers": 3, "hybrid_layer_pattern": "M*-", "sequence_parallel": False},
     ),
 ]
 
@@ -55,7 +55,7 @@
         {
             "hidden_size": 672,
             "num_layers": 3,
-            "hybrid_override_pattern": "M*E",
+            "hybrid_layer_pattern": "M*E",
             "num_moe_experts": 16,
             "moe_token_dispatcher_type": "alltoall",
             "moe_shared_expert_overlap": True,