From 17111cb13c9c904e42fa7856fd387274ff7bc32d Mon Sep 17 00:00:00 2001 From: hazelduan Date: Fri, 21 Aug 2026 11:09:33 +0800 Subject: [PATCH 1/4] fix: make Nemotron-H registration optional --- src/mcore_bridge/model/gpts/__init__.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py index 6b4c034..2eff630 100644 --- a/src/mcore_bridge/model/gpts/__init__.py +++ b/src/mcore_bridge/model/gpts/__init__.py @@ -1,3 +1,8 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe, - qwen3_emb, qwen3_next) +from importlib.util import find_spec + +from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, olmoe, qwen3_emb, + qwen3_next) + +if find_spec('megatron.core.models.hybrid') is not None: + from . import nemotron_h From 04e21c54687fc8e2c3ad16f23870e511b12f81de Mon Sep 17 00:00:00 2001 From: hazelduan Date: Sat, 22 Aug 2026 15:00:30 +0800 Subject: [PATCH 2/4] fix: localize Nemotron-H dependency guard --- src/mcore_bridge/model/gpts/__init__.py | 9 ++----- src/mcore_bridge/model/gpts/nemotron_h.py | 30 ++++++++++++++++------- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py index 2eff630..6b4c034 100644 --- a/src/mcore_bridge/model/gpts/__init__.py +++ b/src/mcore_bridge/model/gpts/__init__.py @@ -1,8 +1,3 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -from importlib.util import find_spec - -from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, olmoe, qwen3_emb, - qwen3_next) - -if find_spec('megatron.core.models.hybrid') is not None: - from . import nemotron_h +from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe, + qwen3_emb, qwen3_next) diff --git a/src/mcore_bridge/model/gpts/nemotron_h.py b/src/mcore_bridge/model/gpts/nemotron_h.py index 744602e..176871a 100644 --- a/src/mcore_bridge/model/gpts/nemotron_h.py +++ b/src/mcore_bridge/model/gpts/nemotron_h.py @@ -9,8 +9,6 @@ import torch from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TENorm, TERowParallelLinear from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add -from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules -from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules from megatron.core.transformer.spec_utils import ModuleSpec @@ -21,9 +19,21 @@ from mcore_bridge.utils import get_logger from ..constant import ModelType -from ..hybrid_model import HybridModel from ..register import ModelLoader, ModelMeta, register_model +try: + from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_stack_spec + + from ..hybrid_model import HybridModel + + _HYBRID_MODEL_AVAILABLE = True +except ImportError as error: + if not (error.name or '').startswith('megatron.core.models.hybrid'): + raise + HybridModel = HybridStack = HybridStackSubmodules = hybrid_stack_spec = None + _HYBRID_MODEL_AVAILABLE = False + logger = get_logger() @@ -410,9 +420,11 @@ def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[in return model -register_model(ModelMeta( - ModelType.nemotron_h, - ['nemotron_h'], - bridge_cls=NemotronHBridge, - loader=NemotronHLoader, -)) +if _HYBRID_MODEL_AVAILABLE: + register_model( + ModelMeta( + ModelType.nemotron_h, + ['nemotron_h'], + bridge_cls=NemotronHBridge, + loader=NemotronHLoader, + )) From e7621d6295e308044d60feb048da7b27be9518d4 Mon Sep 17 00:00:00 2001 From: hazelduan Date: Sat, 29 Aug 2026 15:47:12 +0800 Subject: [PATCH 3/4] fix: normalize packed sequence length metadata --- src/mcore_bridge/model/gpts/qwen3_next.py | 5 +++-- src/mcore_bridge/model/mm_gpts/qwen3_5.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py index 0d4e290..ea73758 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next.py +++ b/src/mcore_bridge/model/gpts/qwen3_next.py @@ -471,10 +471,11 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): # Note: for packed inputs, we do not perform padding_free unpadding. # Doing so would allow different sequences to see each other; for efficiency we keep this implementation. if thd_format: + max_seqlen_q = int(packed_seq_params.max_seqlen_q) new_hidden_states = hidden_states.new_zeros( - (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item(), hidden_states.shape[-1])) + (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) attention_mask = hidden_states.new_zeros( - (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item()), dtype=torch.bool) + (packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) cu_seqlens_q = packed_seq_params.cu_seqlens_q for i in range(packed_seq_params.num_samples): start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1] diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py index 8ba5378..e26a0d6 100644 --- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py +++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py @@ -40,10 +40,11 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): # Note: for packed inputs, we do not perform padding_free unpadding. # Doing so would allow different sequences to see each other; for efficiency we keep this implementation. if thd_format: + max_seqlen_q = int(packed_seq_params.max_seqlen_q) new_hidden_states = hidden_states.new_zeros( - (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item(), hidden_states.shape[-1])) + (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) attention_mask = hidden_states.new_zeros( - (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item()), dtype=torch.bool) + (packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) cu_seqlens_q = packed_seq_params.cu_seqlens_q for i in range(packed_seq_params.num_samples): start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1] From 4b6115bae5268aa654387c0e13d675e454e75512 Mon Sep 17 00:00:00 2001 From: hazelduan Date: Sat, 29 Aug 2026 16:43:31 +0800 Subject: [PATCH 4/4] style: apply yapf formatting --- src/mcore_bridge/model/gpts/qwen3_next.py | 3 +-- src/mcore_bridge/model/mm_gpts/qwen3_5.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py index ea73758..f91f037 100644 --- a/src/mcore_bridge/model/gpts/qwen3_next.py +++ b/src/mcore_bridge/model/gpts/qwen3_next.py @@ -474,8 +474,7 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): max_seqlen_q = int(packed_seq_params.max_seqlen_q) new_hidden_states = hidden_states.new_zeros( (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) - attention_mask = hidden_states.new_zeros( - (packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) + attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) cu_seqlens_q = packed_seq_params.cu_seqlens_q for i in range(packed_seq_params.num_samples): start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1] diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py index e26a0d6..b673bf7 100644 --- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py +++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py @@ -43,8 +43,7 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): max_seqlen_q = int(packed_seq_params.max_seqlen_q) new_hidden_states = hidden_states.new_zeros( (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) - attention_mask = hidden_states.new_zeros( - (packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) + attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) cu_seqlens_q = packed_seq_params.cu_seqlens_q for i in range(packed_seq_params.num_samples): start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1]