[model] support nemotron

src/mcore_bridge/config/model_config.py

@@ -196,6 +196,9 @@ class ModelConfig(TransformerConfig):
     attention_output_gate: bool = False
     linear_decoupled_in_proj: bool = False
 
+    # nemotron_h (hybrid mamba2 + attention + moe)
+    hybrid_layer_pattern: Optional[str] = None
+
     # dsa
     experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = None
     dsa_indexer_n_heads: Optional[int] = None

src/mcore_bridge/config/parser.py

@@ -1,6 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import torch.nn.functional as F
 from functools import partial
+from megatron.core.activations import squared_relu
 from transformers import PretrainedConfig
 from typing import Any, Dict
 
@@ -14,7 +15,7 @@
     'num_attention_heads': ['num_attention_heads'],
     'num_query_groups': ['num_key_value_heads'],
     'max_position_embeddings': ['max_position_embeddings'],
-    'layernorm_epsilon': ['rms_norm_eps'],
+    'layernorm_epsilon': ['rms_norm_eps', 'layer_norm_epsilon'],
     'rotary_base': ['rope_theta'],
     'padded_vocab_size': ['vocab_size'],
     'attention_dropout': ['attention_dropout'],
@@ -26,7 +27,7 @@
     'hf_model_type': ['model_type'],
     # moe
     'moe_ffn_hidden_size': ['moe_intermediate_size'],
-    'moe_shared_expert_intermediate_size': ['shared_expert_intermediate_size'],
+    'moe_shared_expert_intermediate_size': ['shared_expert_intermediate_size', 'moe_shared_expert_intermediate_size'],
     'moe_router_topk': ['num_experts_per_tok', 'moe_topk', 'moe_k', 'top_k_experts'],
     'moe_router_num_groups': ['n_group'],
     'moe_router_group_topk': ['topk_group'],
@@ -67,6 +68,14 @@
     'mhc_sinkhorn_iterations': ['hc_sinkhorn_iters'],
     'moe_n_hash_layers': ['mlp_layer_types'],
     'activation_func_clamp_value': ['swiglu_limit'],
+    # nemotron_h / mamba2
+    'mamba_num_heads': ['mamba_num_heads'],
+    'mamba_head_dim': ['mamba_head_dim'],
+    'mamba_state_dim': ['ssm_state_size', 'mamba_state_dim'],
+    'mamba_num_groups': ['n_groups', 'mamba_num_groups'],
+    'hybrid_layer_pattern': ['hybrid_override_pattern'],
+    'fp32_residual_connection': ['residual_in_fp32'],
+    'mtp_hybrid_override_pattern': ['mtp_hybrid_override_pattern'],
     # other
     'original_max_position_embeddings': ['original_max_position_embeddings'],
     'partial_rotary_factor': ['partial_rotary_factor'],
@@ -255,6 +264,19 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         res['add_qkv_bias'] = False
         res['moe_router_score_function'] = 'sigmoid'
         res['moe_router_load_balancing_type'] = 'seq_aux_loss'
+    elif llm_model_type == 'nemotron_h':
+        res['is_hybrid_model'] = True
+        res['position_embedding_type'] = 'none'
+        # relu^2 ("relu2") activation: non-gated, so fc1 is a single up_proj (no gate_proj).
+        res['swiglu'] = False
+        res['gated_linear_unit'] = False
+        res['activation_func'] = squared_relu
+        res['add_bias_linear'] = False
+        res['add_qkv_bias'] = False
+        res['qk_layernorm'] = False
+        res['moe_router_score_function'] = 'sigmoid'
+        res['moe_router_enable_expert_bias'] = True
+        res['moe_router_load_balancing_type'] = 'seq_aux_loss'
 
     if 'partial_rotary_factor' not in res and 'partial_rotary_factor' in rope_scaling:
         res['partial_rotary_factor'] = rope_scaling['partial_rotary_factor']

src/mcore_bridge/model/constant.py

@@ -12,6 +12,7 @@ class LLMModelType:
     bailing_hybrid = 'bailing_hybrid'
     deepseek_v4 = 'deepseek_v4'
     glm_moe_dsa = 'glm_moe_dsa'
+    nemotron_h = 'nemotron_h'
 
     qwen3_emb = 'qwen3_emb'
 

src/mcore_bridge/model/gpt_model.py

@@ -203,6 +203,8 @@ def _preprocess(
         return decoder_input, mtp_decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset
 
     def _set_inv_freq(self):
+        if getattr(self, 'rotary_pos_emb', None) is None:
+            return
         new_inv_freq, self.config.attention_scaling = get_rope_inv_freq(self.config)
         self.rotary_pos_emb.inv_freq = new_inv_freq.to(self.rotary_pos_emb.inv_freq.device)
 

src/mcore_bridge/model/gpts/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, olmoe, qwen3_emb,
-               qwen3_next)
+from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe,
+               qwen3_emb, qwen3_next)

src/mcore_bridge/model/hybrid_model.py

@@ -0,0 +1,86 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import math
+import torch
+from megatron.core import mpu
+from megatron.core.models.hybrid.hybrid_model import HybridModel as McoreHybridModel
+from megatron.core.transformer.spec_utils import ModuleSpec
+from typing import Optional
+
+from mcore_bridge.config import ModelConfig
+from mcore_bridge.utils import split_cp_inputs
+
+
+class HybridModel(McoreHybridModel):
+    """Thin adapter over megatron-core's HybridModel.
+
+    Upstream `HybridModel` already covers embedding, the hybrid layer stack, MTP
+    (via `process_mtp_loss`) and the loss/logits tail, so only two things are added
+    here:
+
+    1. Translate `ModelConfig` into the upstream constructor signature.
+    2. Build `padding_mask`, which upstream takes as a forward argument but never
+       computes. Deriving it needs the CP size, the TP size and the current TP rank,
+       so it stays on this side rather than leaking into the caller.
+    """
+
+    config: ModelConfig
+
+    def __init__(
+        self,
+        config: ModelConfig,
+        transformer_layer_spec: ModuleSpec,
+        pre_process: bool = True,
+        post_process: bool = True,
+        vp_stage: Optional[int] = None,
+    ):
+        # `ModelLoader.build_model` passes the stack spec positionally as
+        # `transformer_layer_spec`; upstream names the same argument `hybrid_stack_spec`.
+        # MTP needs no separate spec here: upstream derives it from the `/` suffix of
+        # `hybrid_layer_pattern`.
+        vocab_size = math.ceil(
+            config.padded_vocab_size / config.tensor_model_parallel_size) * config.tensor_model_parallel_size
+        super().__init__(
+            config,
+            transformer_layer_spec,
+            vocab_size,
+            config.max_position_embeddings,
+            hybrid_layer_pattern=config.hybrid_layer_pattern,
+            pre_process=pre_process,
+            post_process=post_process,
+            share_embeddings_and_output_weights=not config.untie_embeddings_and_output_weights,
+            position_embedding_type=config.position_embedding_type,
+            rotary_base=config.rotary_base,
+            vp_stage=vp_stage,
+        )
+
+    def _get_padding_mask(self, attention_mask) -> Optional[torch.Tensor]:
+        """Mark fully-padded sequence positions, sharded to match the hidden states."""
+        if isinstance(attention_mask, dict):
+            attention_mask = attention_mask['full_attention']
+        if attention_mask is None:
+            return None
+        padding_mask = ~((~attention_mask).sum(dim=(1, 2)) > 0)
+        if self.config.context_parallel_size > 1:
+            padding_mask = split_cp_inputs(padding_mask, None, 1)
+        tp_size = self.config.tensor_model_parallel_size
+        if self.config.sequence_parallel and tp_size > 1:
+            assert padding_mask.shape[1] % tp_size == 0, f'padding_mask.shape: {padding_mask.shape}'
+            padding_mask = torch.chunk(padding_mask, tp_size, dim=1)[mpu.get_tensor_model_parallel_rank()]
+        return padding_mask.contiguous()
+
+    def forward(self, input_ids, position_ids, attention_mask=None, *args, packed_seq_params=None, **kwargs):
+        padding_mask = None
+        if packed_seq_params is None:
+            padding_mask = self._get_padding_mask(attention_mask)
+        return super().forward(
+            input_ids,
+            position_ids,
+            attention_mask,
+            *args,
+            packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
+            **kwargs,
+        )
+
+    def get_input_tensor(self):
+        return self.decoder.input_tensor

src/mcore_bridge/model/register.py

@@ -160,7 +160,11 @@ def _set_transformer_layer(self, transformer_layer_spec):
 
     def _replace_mla_attention(self, transformer_layer_spec):
         for layer_spec in transformer_layer_spec.layer_specs:
-            self_attention = layer_spec.submodules.self_attention
+            # Hybrid models (e.g. nemotron_h) may have layers with no attention submodule
+            # at all (MambaLayer) or with it replaced by IdentityOp (FFN-only layers).
+            self_attention = getattr(layer_spec.submodules, 'self_attention', None)
+            if not hasattr(self_attention, 'module'):
+                continue
             if self_attention.module is McoreMLASelfAttention:
                 self_attention.module = MLASelfAttention
             elif getattr(self_attention.module, '__name__', None) == 'AbsorbedMLASelfAttention':

src/mcore_bridge/patcher.py

@@ -8,6 +8,7 @@
 from megatron.core.models.common.embeddings import rope_utils
 from megatron.core.models.common.embeddings.rotary_pos_embedding import MultimodalRotaryEmbedding
 from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.moe.router import TopKRouter
 from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionBlock, get_mtp_layer_offset
 from packaging import version
 from peft.tuners.tuners_utils import BaseTuner
@@ -237,9 +238,25 @@ def apply_rotary_pos_emb(
 
 
 def _patch_mtp():
+    """Unroll the MTP block over `mtp_unroll_steps` for the GPTModel build.
+
+    This rewrite drives the layer with `decoder_input` / `layer_number`, which only the
+    GPTModel-path MTP layer accepts. `HybridStack`-based MTP layers take neither, and
+    upstream's own `MultiTokenPredictionBlock.forward` already unrolls them, so hybrid
+    models must keep the upstream implementation.
+    """
+    origin_forward = MultiTokenPredictionBlock.forward
 
     def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, hidden_states: torch.Tensor,
                 attention_mask: torch.Tensor, **kwargs) -> torch.Tensor:
+        if getattr(self.config, 'is_hybrid_model', False):
+            return origin_forward(
+                self,
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                **kwargs)
         # get hidden states from previous mtp stages
         get_offset_kwargs = {} if self.vp_stage is None else {'vp_stage': self.vp_stage}
         mtp_decoder_input = decoder_input = kwargs.pop('decoder_input', None)
@@ -283,6 +300,27 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, hidden_st
     MultiTokenPredictionBlock.forward = forward
 
 
+def _patch_moe_expert_bias_padding_mask():
+    """Align the padding mask with `routing_map` in `TopKRouter._apply_expert_bias`.
+
+    `TopKRouter.routing` flattens `padding_mask` to `[num_tokens]`, but
+    `_apply_expert_bias` then computes `routing_map & (~padding_mask)` against a
+    `[num_tokens, num_experts]` map, so the 1-D mask broadcasts over the expert dim and
+    raises a size mismatch. Restore the trailing dim so it broadcasts over experts instead.
+
+    Only reachable with `moe_router_enable_expert_bias` and a non-None `padding_mask`,
+    i.e. non-packed batches -- packed runs pass `padding_mask=None` and never hit it.
+    """
+    origin_apply_expert_bias = TopKRouter._apply_expert_bias
+
+    def _apply_expert_bias(self, routing_map, padding_mask=None):
+        if padding_mask is not None and padding_mask.dim() == routing_map.dim() - 1:
+            padding_mask = padding_mask.unsqueeze(-1)
+        return origin_apply_expert_bias(self, routing_map, padding_mask=padding_mask)
+
+    TopKRouter._apply_expert_bias = _apply_expert_bias
+
+
 def apply_patch():
     _patch_flash_attn()
     _patch_transformer_engine()
@@ -297,4 +335,5 @@ def apply_patch():
     _patch_TELinear()
     _patch_mrope()
     _patch_mtp()
+    _patch_moe_expert_bias_padding_mask()
     from mcore_bridge import tuners  # apply patch

tests/test_llm.py

@@ -162,6 +162,10 @@ def test_bailing():
     _test_model('inclusionAI/Ling-mini-2.0')
 
 
+def test_nemotron_h():
+    _test_model('nv-community/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16')
+
+
 if __name__ == '__main__':
     # test_qwen2()
     # test_llama2()
@@ -195,4 +199,5 @@ def test_bailing():
     # test_minimax_m2()
     # test_glm4_moe_lite()
     # test_olmoe()
-    test_bailing()
+    # test_bailing()
+    test_nemotron_h()