[model] support nemotron-super-vl

src/mcore_bridge/bridge/gpt_bridge.py

@@ -450,7 +450,11 @@ def _broadcast_ep_pp(self, tensor, is_expert):
             dist.all_reduce(src_rank, group=pp_group)
             src_rank = dist.get_global_rank(pp_group, src_rank.item())
             meta_data = torch.zeros(10, dtype=torch.int64, device='cuda')
-            dtype_mapping = [torch.float64, torch.float32, torch.float16, torch.bfloat16, torch.uint8, torch.int32]
+            # Append NCCL-supported buffer dtypes without renumbering existing metadata.
+            dtype_mapping = [
+                torch.float64, torch.float32, torch.float16, torch.bfloat16, torch.uint8, torch.int32, torch.int64,
+                torch.bool, torch.int8
+            ]
             dtype_mapping_r = {v: k for k, v in enumerate(dtype_mapping)}
             if tensor is None:
                 dist.broadcast(meta_data, src=src_rank, group=pp_group)

src/mcore_bridge/config/parser.py

@@ -122,6 +122,7 @@
     'mamba_num_groups': ['n_groups', 'mamba_num_groups'],
     'hybrid_layer_pattern': ['hybrid_override_pattern'],
     'fp32_residual_connection': ['residual_in_fp32'],
+    'moe_latent_size': ['moe_latent_size'],
     'mtp_hybrid_override_pattern': ['mtp_hybrid_override_pattern'],
     # other
     'original_max_position_embeddings': ['original_max_position_embeddings'],
@@ -195,6 +196,8 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
     res['hf_config'] = hf_config
     hf_model_type = res.get('hf_model_type')
     llm_model_type = res.get('llm_model_type') or hf_model_type
+    if hf_model_type == 'nemotron_h_omni':
+        llm_model_type = 'nemotron_h'
     res['llm_model_type'] = llm_model_type
 
     first_k_dense_replace = res.pop('first_k_dense_replace', None)

src/mcore_bridge/model/constant.py

@@ -47,6 +47,7 @@ class MLLMModelType:
     muse_glimmer = 'muse_glimmer'
 
     deepseek_v41 = 'deepseek_v41'
+    nemotron_h_omni = 'nemotron_h_omni'
 
 
 class ModelType(LLMModelType, MLLMModelType):

src/mcore_bridge/model/gpts/nemotron_h.py

@@ -16,7 +16,7 @@
 
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.tuners import LoraParallelLinear
-from mcore_bridge.utils import get_logger
+from mcore_bridge.utils import deep_getattr, get_logger
 
 from ..constant import ModelType
 from ..register import ModelLoader, ModelMeta, register_model
@@ -226,18 +226,22 @@ def _set_mamba_in_proj(self, mg_mixer, hf_state_dict, to_mcore: bool):
         mg_param = None if mg_mixer is None else mg_mixer.in_proj.weight
         self._set_mamba_packed(mg_param, hf_state_dict, 'in_proj.weight', self._mamba_block_sizes(), to_mcore)
 
+    @staticmethod
+    def _mamba_conv1d_mg_key(mg_mixer, suffix):
+        if mg_mixer is not None and hasattr(mg_mixer, f'conv1d_{suffix}'):
+            return f'conv1d_{suffix}'
+        return f'conv1d.{suffix}'
+
     def _set_mamba_conv1d(self, mg_mixer, hf_state_dict, to_mcore: bool):
         """conv1d mirrors `in_proj` minus the dt block: [x, B, C] along dim 0."""
-        if self.tp_size == 1:
-            self._set_state_dict(mg_mixer, 'conv1d_weight', hf_state_dict, 'conv1d.weight', to_mcore)
-            self._set_state_dict(mg_mixer, 'conv1d_bias', hf_state_dict, 'conv1d.bias', to_mcore)
-            return
-        d_inner = self.config.mamba_num_heads * self.config.mamba_head_dim
-        bc = self.config.mamba_num_groups * self.config.mamba_state_dim
-        blocks = [d_inner, bc, bc]
-        for mg_name, hf_name in (('conv1d_weight', 'conv1d.weight'), ('conv1d_bias', 'conv1d.bias')):
-            mg_param = None if mg_mixer is None else getattr(mg_mixer, mg_name)
-            self._set_mamba_packed(mg_param, hf_state_dict, hf_name, blocks, to_mcore)
+        blocks = self._mamba_block_sizes()[1:-1]
+        for suffix in ('weight', 'bias'):
+            mg_key = self._mamba_conv1d_mg_key(mg_mixer, suffix)
+            mg_param = deep_getattr(mg_mixer, mg_key)
+            if self.tp_size == 1:
+                self._set_state_dict(mg_mixer, mg_key, hf_state_dict, f'conv1d.{suffix}', to_mcore)
+            else:
+                self._set_mamba_packed(mg_param, hf_state_dict, f'conv1d.{suffix}', blocks, to_mcore)
 
     def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool, is_mtp: bool = False):
         """Convert the channel-mixing half of a layer: MoE for 'E'."""
@@ -246,6 +250,11 @@ def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool
             hf_state_dict.update(
                 self._set_moe_state(
                     mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore, is_mtp=is_mtp))
+            # nemotron3.5-super-vl
+            if self.config.moe_latent_size:
+                for name in ('fc1_latent_proj', 'fc2_latent_proj'):
+                    self._set_state_dict(mg_mlp, f'{name}.weight', hf_state_dict, f'{self.hf_mlp_prefix}.{name}.weight',
+                                         to_mcore)
             self._set_state_dict(mg_layer, 'pre_mlp_layernorm.weight', hf_state_dict, self.hf_input_layernorm_key,
                                  to_mcore)
         # mamba / attention layers have no MLP sub-module.

src/mcore_bridge/model/hybrid_model.py

@@ -1,4 +1,5 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import inspect
 import math
 import torch
 from megatron.core import mpu
@@ -24,6 +25,11 @@ class HybridModel(McoreHybridModel):
     """
 
     config: ModelConfig
+    # These kwargs would otherwise be cleared by the multimodal embedding hook / PP stage.
+    extra_forward_keys = ('loss_mask', 'padding_mask')
+    _upstream_forward_keys = set(inspect.signature(McoreHybridModel.forward).parameters) - {
+        'self', 'input_ids', 'position_ids', 'attention_mask', 'packed_seq_params', 'padding_mask'
+    }
 
     def __init__(
         self,
@@ -65,7 +71,7 @@ def _get_padding_mask(self, attention_mask) -> Optional[torch.Tensor]:
             attention_mask = attention_mask['full_attention']
         if attention_mask is None:
             return None
-        padding_mask = attention_mask.all(dim=(1, 2))
+        padding_mask = ~attention_mask.bool() if attention_mask.dim() == 2 else attention_mask.all(dim=(1, 2))
         if self.config.context_parallel_size > 1:
             padding_mask = split_cp_inputs(padding_mask, None, 1)
         tp_size = self.config.tensor_model_parallel_size
@@ -81,7 +87,16 @@ def forward(self,
                 *args,
                 packed_seq_params=None,
                 padding_mask=None,
+                extra_block_kwargs=None,
                 **kwargs):
+        # MMGPT separates media kwargs from forward controls; Hybrid has no extra_block_kwargs.
+        extra_block_kwargs = extra_block_kwargs or {}
+        if padding_mask is None:
+            padding_mask = extra_block_kwargs.get('padding_mask')
+        for key in self.extra_forward_keys:
+            if key != 'padding_mask' and key in extra_block_kwargs:
+                kwargs.setdefault(key, extra_block_kwargs[key])
+        kwargs = {key: value for key, value in kwargs.items() if key in self._upstream_forward_keys}
         if padding_mask is None and packed_seq_params is None:
             padding_mask = self._get_padding_mask(attention_mask)
         return super().forward(

src/mcore_bridge/model/mm_gpt_model.py

@@ -102,7 +102,7 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
         inference_context = kwargs.pop('inference_context', None)
-        extra_kwargs = {k: kwargs[k] for k in self.language_model.extra_forward_keys}
+        extra_kwargs = {k: kwargs[k] for k in self.language_model.extra_forward_keys if k in kwargs}
         # Compatible with legacy mcore-bridge behavior.
         cp_size = self.config.context_parallel_size
         cp_partition_mode = getattr(self.config, 'cp_partition_mode', 'zigzag')

src/mcore_bridge/model/mm_gpts/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from . import (deepseek_v41, gemma4, glm, glm5_next, internvl, kimi_vl, llama4, llava, minicpmv4_6, muse_glimmer, qwen,
-               qwen3_5, qwen3_5_gdn, qwen3_asr, qwen3_omni, qwen3_vl, qwen4_exp)
+from . import (deepseek_v41, gemma4, glm, glm5_next, internvl, kimi_vl, llama4, llava, minicpmv4_6, muse_glimmer,
+               nemotron_h_omni, qwen, qwen3_5, qwen3_5_gdn, qwen3_asr, qwen3_omni, qwen3_vl, qwen4_exp)