diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index e6d0294..a472003 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/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 diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py index 338fa5f..a22ba6a 100644 --- a/src/mcore_bridge/config/parser.py +++ b/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'] diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py index f1ba4ad..108c349 100644 --- a/src/mcore_bridge/model/constant.py +++ b/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' diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py index 2158b90..8ed9994 100644 --- a/src/mcore_bridge/model/gpt_model.py +++ b/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) diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py index e79d01a..6b4c034 100644 --- a/src/mcore_bridge/model/gpts/__init__.py +++ b/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) diff --git a/src/mcore_bridge/model/gpts/nemotron_h.py b/src/mcore_bridge/model/gpts/nemotron_h.py new file mode 100644 index 0000000..744602e --- /dev/null +++ b/src/mcore_bridge/model/gpts/nemotron_h.py @@ -0,0 +1,418 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Nemotron-3.5 (hybrid Mamba2 + Attention + MoE) on megatron-core's HybridModel. + +Upstream deprecated `GPTModel` in favour of `HybridModel` (Megatron-LM #5911). On +`HybridModel` one pattern symbol *is* one layer, so the `IdentityOp` stripping and the +`MambaLayer` compat shim that a `GPTModel` build would need are both unnecessary here, +and MTP can span several heterogeneous inner layers. +""" +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 +from typing import Optional + +from mcore_bridge.bridge import GPTBridge +from mcore_bridge.tuners import LoraParallelLinear +from mcore_bridge.utils import get_logger + +from ..constant import ModelType +from ..hybrid_model import HybridModel +from ..register import ModelLoader, ModelMeta, register_model + +logger = get_logger() + + +class NemotronHBridge(GPTBridge): + """HuggingFace <-> Megatron-Core weight conversion for Nemotron-3.5. + + Layer families come from `hybrid_layer_pattern`, one symbol per layer: + M = Mamba2 SSM, E = MoE (routed + shared), * = attention (GQA), - = dense MLP + + All three families sit under the same HF `mixer.` prefix, so dispatch is driven by + the pattern rather than by the key name. + """ + + hf_embed_key = 'backbone.embeddings.weight' + hf_layers_prefix = 'backbone.layers' + hf_final_layernorm_key = 'backbone.norm_f.weight' + hf_lm_head_key = 'lm_head.weight' + hf_attn_prefix = 'mixer' + hf_mlp_prefix = 'mixer' + hf_input_layernorm_key = 'norm.weight' + hf_o_proj_key = 'o_proj' + hf_q_norm_key = 'q_norm.weight' + hf_k_norm_key = 'k_norm.weight' + hf_gate_key = 'gate.weight' + hf_expert_bias_key = 'gate.e_score_correction_bias' + hf_shared_expert_key = 'shared_experts' + hf_mtp_prefix = 'mtp.layers' + hf_mtp_final_layernorm_key = 'final_layernorm.weight' + + _LAYER_TYPES = {'M': 'mamba', 'E': 'moe', '*': 'attention', '-': 'mlp'} + + def _get_layer_type(self, layer_idx: int): + """Resolve a layer's family from the pattern. + + A negative index means "MTP inner layer i" (see `_convert_mtp_layer`), which is + described by `mtp_hybrid_override_pattern` rather than the backbone pattern. + """ + if layer_idx < 0: + pattern = self.config.mtp_hybrid_override_pattern + idx = -layer_idx - 1 + else: + pattern = self.config.hybrid_layer_pattern + idx = layer_idx + assert 0 <= idx < len(pattern), f'layer index {idx} out of range for pattern {pattern!r}' + return self._LAYER_TYPES[pattern[idx]] + + def _get_tp_split_dim(self, mg_key: Optional[str]) -> Optional[int]: + # `D` and `conv1d_{weight,bias}` are flat nn.Parameters on MambaMixer (no dot in the + # relative key for `D`), so the base class keyword lookup cannot classify them. + if mg_key in {'D', 'conv1d_weight', 'conv1d_bias'}: + return 0 + if mg_key == 'mixer.norm.weight': + # Inner gated RMSNorm of the Mamba mixer is sharded over d_inner. + return 0 + if mg_key is not None and mg_key.split('.', 1)[0] in {'linear_fc1', 'linear_fc1_up'}: + # relu^2 is non-gated, so linear_fc1 is a plain [ffn, hidden] column-parallel + # weight. The base class returns 1 because it assumes the gated [2, X, Y] layout. + return 0 + return super()._get_tp_split_dim(mg_key) + + def _get_hf_experts_attr(self, is_mtp: bool = False): + # Experts are stored one module per expert, with separate up/down projections. + return False, False + + def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore): + # `HybridStack` names the trailing norm `final_norm` (`TransformerBlock` uses + # `final_layernorm`). + self._set_state_dict(lm_model, 'decoder.final_norm.weight', hf_state_dict, self.hf_final_layernorm_key, + to_mcore) + + def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool): + """Convert the sequence-mixing half of a layer: attention or Mamba.""" + layer_type = self._get_layer_type(layer_idx) + if layer_type == 'attention': + mg_attn = None if mg_layer is None else mg_layer.self_attention + hf_state_dict.update( + self._set_attn_state(mg_attn, hf_state_dict, f'{self.hf_attn_prefix}.', layer_idx, to_mcore)) + # Pre-norm is fused into linear_qkv (TELayerNormColumnParallelLinear). + self._set_state_dict(mg_layer, 'self_attention.linear_qkv.layer_norm_weight', hf_state_dict, + self.hf_input_layernorm_key, to_mcore) + elif layer_type == 'mamba': + hf_state_dict.update(self._set_mamba_state(mg_layer, hf_state_dict, layer_idx, to_mcore)) + # MambaLayer keeps a standalone pre-norm (`norm`, not `input_layernorm`). + self._set_state_dict(mg_layer, 'norm.weight', hf_state_dict, self.hf_input_layernorm_key, to_mcore) + # 'moe'/'-' layers have no sequence mixer; their norm is handled in _set_layer_mlp. + return hf_state_dict + + def _set_mamba_state(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool): + """Convert Mamba2 SSM weights under the `mixer.` prefix. + + MambaMixer keeps `conv1d_weight` / `conv1d_bias` as flat nn.Parameters (not a + `conv1d` submodule), and `in_proj` is a single packed [z, x, B, C, dt] projection. + The HF checkpoint uses the same packed layout, so at TP=1 every tensor maps 1:1. + + Under TP the packed tensors need per-block slicing rather than one contiguous cut: + upstream sizes each block by its own local width (`d_inner_local_tp`, + `ngroups_local_tp * d_state`, `nheads_local_tp`), so a rank owns a slice of *every* + block. A naive split of the concatenation would land inside one block and silently + hand a rank the wrong projections -- see `_split_packed_dim0`. + """ + hf_prefix = f'{self.hf_attn_prefix}.' + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + else: + hf_state_dict = {} + mg_mixer = None if mg_layer is None else mg_layer.mixer + self._set_mamba_in_proj(mg_mixer, hf_state_dict, to_mcore) + self._set_mamba_conv1d(mg_mixer, hf_state_dict, to_mcore) + self._set_state_dict(mg_mixer, 'A_log', hf_state_dict, 'A_log', to_mcore) + self._set_state_dict(mg_mixer, 'D', hf_state_dict, 'D', to_mcore) + self._set_state_dict(mg_mixer, 'dt_bias', hf_state_dict, 'dt_bias', to_mcore) + self._set_state_dict(mg_mixer, 'out_proj.weight', hf_state_dict, 'out_proj.weight', to_mcore) + # Inner gated RMSNorm, present when the mixer uses rmsnorm. + has_inner_norm = False if mg_mixer is None else getattr(mg_mixer, 'norm', None) is not None + has_inner_norm = self._reduce_tensor_pp_group(has_inner_norm, to_mcore) + if has_inner_norm: + self._set_state_dict(mg_layer, 'mixer.norm.weight', hf_state_dict, 'norm.weight', to_mcore) + if to_mcore: + hf_state_dict = {} + else: + hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) + return hf_state_dict + + def _mamba_block_sizes(self): + """Global dim-0 sizes of the [z, x, B, C, dt] blocks packed into `in_proj`. + + Derived from `self.config` rather than the mixer instance, so it stays valid on PP + ranks that do not hold this Mamba layer (`mg_mixer is None`). + """ + d_inner = self.config.mamba_num_heads * self.config.mamba_head_dim + bc = self.config.mamba_num_groups * self.config.mamba_state_dim + return [d_inner, d_inner, bc, bc, self.config.mamba_num_heads] + + def _split_packed_dim0(self, tensor, block_sizes): + """Take this TP rank's slice out of each packed block along dim 0.""" + out, offset = [], 0 + for size in block_sizes: + local = size // self.tp_size + start = offset + self.tp_rank * local + out.append(tensor[start:start + local]) + offset += size + return torch.cat(out, dim=0) + + def _merge_packed_dim0(self, gathered, block_sizes): + """Inverse of `_split_packed_dim0`. + + `_all_gather_tp` concatenates the per-rank shards along dim 0, so `gathered` reads + [rank0 blocks..., rank1 blocks..., ...]. Regroup it back into whole global blocks. + """ + local_total = sum(size // self.tp_size for size in block_sizes) + shards = [gathered[i * local_total:(i + 1) * local_total] for i in range(self.tp_size)] + blocks, offset = [], 0 + for size in block_sizes: + local = size // self.tp_size + blocks.append(torch.cat([s[offset:offset + local] for s in shards], dim=0)) + offset += local + return torch.cat(blocks, dim=0) + + def _set_mamba_packed(self, mg_param, hf_state_dict, hf_key, block_sizes, to_mcore: bool): + """Load/export a packed Mamba tensor whose dim-0 blocks are each TP-sharded. + + `mg_param` is None on a PP rank that does not own this layer. Both collectives still + have to run on every rank: the TP all-gather tolerates None, and the PP broadcast is + what actually hands the merged tensor to the non-owning ranks -- returning early + instead would drop this layer from the export entirely. + """ + if to_mcore: + if mg_param is None: + return + weight = hf_state_dict[hf_key].load() + # `_set_weight` would split by `_get_tp_split_dim`, i.e. one contiguous cut, so + # slice per block here and hand it the already-local shard (tp_dim None). + self._set_weight(mg_param, self._split_packed_dim0(weight, block_sizes), None) + else: + gathered = self._all_gather_tp(None if mg_param is None else mg_param.data, 0, False) + merged = None if gathered is None else self._merge_packed_dim0(gathered, block_sizes) + # Non-owning PP ranks receive the merged tensor here; owning ranks send it. + merged = self._broadcast_ep_pp(merged, False) + # `_all_gather_tp` leaves the result on cuda; the generic export path applies + # `_target_device` when it writes into hf_state_dict, so do the same here. + if self._target_device is not None: + merged = merged.to(self._target_device) + hf_state_dict[hf_key] = merged + + def _set_mamba_in_proj(self, mg_mixer, hf_state_dict, to_mcore: bool): + """`in_proj` packs [z, x, B, C, dt]; each block is TP-sharded on its own.""" + if self.tp_size == 1: + self._set_state_dict(mg_mixer, 'in_proj.weight', hf_state_dict, 'in_proj.weight', to_mcore) + return + 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) + + 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) + + 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'.""" + if self._get_layer_type(layer_idx) == 'moe': + mg_mlp = None if mg_layer is None else mg_layer.mlp + 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)) + 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. + return hf_state_dict + + def _set_mlp_state( + self, + mg_mlp, + hf_state_dict, + hf_prefix: str, + layer_idx: int, + to_mcore: bool, + ep_rank: Optional[int] = None, + is_mtp: bool = False, + ): + """Map linear_fc1 <-> up_proj 1:1. + + Nemotron uses relu^2, which is non-gated: there is no gate_proj, so linear_fc1 is a + plain [ffn, hidden] tensor rather than the merged [gate_proj; up_proj] layout the + base class assumes. That single assumption is why the base implementation cannot be + reused here; everything else still goes through `_set_state_dict` / `_set_weight`. + """ + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + else: + hf_state_dict = {} + is_expert = ep_rank is not None + if not self._peft_format: + if is_expert: + num_local_experts = self.config.num_moe_experts // self.ep_size + start_idx = ep_rank * num_local_experts + for mg_name, hf_name in [('linear_fc1', 'up_proj'), ('linear_fc2', 'down_proj')]: + mg_linear = None if mg_mlp is None else getattr(mg_mlp, mg_name) + # Under LoRA the expert linear is wrapped, and the per-expert `weight{i}` + # live on the wrapped module. This branch exports merged base weights + # (`_peft_format` is False), so unwrap before indexing. + if isinstance(mg_linear, LoraParallelLinear): + mg_linear = mg_linear.base_layer + # `linear_fc1_up` aliases linear_fc1 to bypass the base gated-fc1 + # reshape; TP dim is registered for both names in _get_tp_split_dim. + tp_key = 'linear_fc1_up.weight' if mg_name == 'linear_fc1' else 'linear_fc2.weight' + if to_mcore: + weight = torch.concat([ + hf_state_dict[f'{start_idx + i}.{hf_name}.weight'].load() for i in range(num_local_experts) + ], + dim=0) + self._set_weight([getattr(mg_linear, f'weight{i}') for i in range(num_local_experts)], + weight, + tp_key, + is_expert=True) + else: + mg_weight = None if mg_linear is None else [ + getattr(mg_linear, f'weight{i}').data for i in range(num_local_experts) + ] + # `_get_weight` reshapes to [num_local_experts, ffn, hidden]. + weight, _ = self._get_weight(mg_weight, tp_key, is_expert=True) + if weight is not None: + for i in range(num_local_experts): + hf_state_dict[f'{start_idx + i}.{hf_name}.weight'] = weight[i].clone() + del weight + else: + # dense MLP / shared expert, same non-gated fc1 handling as above. + fc1_module = None if mg_mlp is None else mg_mlp.linear_fc1 + if isinstance(fc1_module, LoraParallelLinear): + fc1_module = fc1_module.base_layer + if to_mcore: + self._set_weight(fc1_module.weight, hf_state_dict['up_proj.weight'].load(), 'linear_fc1_up.weight') + else: + fc1 = None if fc1_module is None else fc1_module.weight.data + weight, _ = self._get_weight(fc1, 'linear_fc1_up.weight') + if weight is not None: + hf_state_dict['up_proj.weight'] = weight.clone() + del weight + self._set_state_dict(mg_mlp, 'linear_fc2.weight', hf_state_dict, 'down_proj.weight', to_mcore) + if to_mcore: + hf_state_dict = {} + else: + hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) + return hf_state_dict + + def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): + """Map one MTP depth, whose inner layers span several HF indices. + + With `mtp_hybrid_override_pattern='*E'` a depth holds two inner layers, and HF stores + them as two `mtp.layers.{0,1}` entries: index 0 carries `enorm`/`hnorm`/`eh_proj` + plus the attention mixer, index 1 the MoE mixer plus `final_layernorm`. mcore keeps + both under `mtp.layers[depth].mtp_model_layer.layers[i]`, so the base class + assumption of a single `mtp_layer.transformer_layer` does not hold. + """ + pattern = self.config.mtp_hybrid_override_pattern + mtp_layer = lm_model.mtp.layers[layer_idx] if hasattr(lm_model, 'mtp') else None + n_inner = len(pattern) + exported = {} + for inner_idx in range(n_inner): + inner_prefix = f'{hf_prefix}{layer_idx * n_inner + inner_idx}.' + if to_mcore: + inner_sd = self._remove_prefix(hf_state_dict, inner_prefix) + if not inner_sd: + logger.info(f'MTP inner layer {inner_prefix} safetensors weights not found, ' + 'this part will be randomly initialized.') + continue + else: + inner_sd = {} + inner_layer = None if mtp_layer is None else mtp_layer.mtp_model_layer.layers[inner_idx] + # enorm/hnorm/eh_proj live on the MTP layer itself and only exist on inner 0. + if inner_idx == 0: + for key in ['enorm.weight', 'hnorm.weight', 'eh_proj.weight']: + self._set_state_dict(mtp_layer, key, inner_sd, key, to_mcore) + self._fp8_skip_modules.update({'eh_proj'}) + if inner_idx == n_inner - 1: + self._set_state_dict(mtp_layer, 'final_layernorm.weight', inner_sd, self.hf_mtp_final_layernorm_key, + to_mcore) + # Negative index selects `mtp_hybrid_override_pattern` in `_get_layer_type`. + mtp_layer_idx = -(inner_idx + 1) + inner_sd.update(self._set_layer_attn(inner_layer, inner_sd, mtp_layer_idx, to_mcore)) + inner_sd.update(self._set_layer_mlp(inner_layer, inner_sd, mtp_layer_idx, to_mcore, is_mtp=True)) + if not to_mcore: + exported.update(self._add_prefix(inner_sd, inner_prefix)) + return {} if to_mcore else exported + + +class NemotronHLoader(ModelLoader): + model_cls = HybridModel + + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + """Return a `HybridStack` spec with a standalone pre-norm on Mamba layers. + + Upstream's default mamba spec fuses the pre-norm into `in_proj` + (`TELayerNormColumnParallelLinear`), which renames the weight to + `mixer.in_proj.layer_norm_weight`. This checkpoint stores it as a separate + `norm.weight`, so `TENorm` keeps the Bridge mapping one-to-one. + """ + submodules = HybridStackSubmodules( + **{ + field: getattr(hybrid_stack_spec.submodules, field) + for field in hybrid_stack_spec.submodules.__dataclass_fields__ + }) + # Separate norm from in_proj: the fused TELayerNormColumnParallelLinear would rename the + # weight to `mixer.in_proj.layer_norm_weight`, while this checkpoint stores a standalone + # `norm.weight`. Keeping them separate makes the Bridge mapping one-to-one. + submodules.mamba_layer = ModuleSpec( + module=MambaLayer, + submodules=MambaLayerSubmodules( + norm=TENorm, + mixer=ModuleSpec( + module=MambaMixer, + submodules=MambaMixerSubmodules( + in_proj=TEColumnParallelLinear, + out_proj=TERowParallelLinear, + ), + ), + mamba_bda=get_bias_dropout_add, + ), + ) + return ModuleSpec(module=HybridStack, submodules=submodules) + + def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): + """Build via `HybridModel`, skipping the base class's layer_specs post-processing. + + `ModelLoader.build_model` rewrites `spec.layer_specs` (MLA / router / TransformerLayer + substitution); a `HybridStack` spec exposes per-layer-family submodules instead, and + this model needs none of those substitutions. + """ + model = self.model_cls( + config=self.config, + transformer_layer_spec=self.get_transformer_layer_spec(vp_stage=vp_stage), + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + self._set_linear_is_expert(model) + return model + + +register_model(ModelMeta( + ModelType.nemotron_h, + ['nemotron_h'], + bridge_cls=NemotronHBridge, + loader=NemotronHLoader, +)) diff --git a/src/mcore_bridge/model/hybrid_model.py b/src/mcore_bridge/model/hybrid_model.py new file mode 100644 index 0000000..9e6424d --- /dev/null +++ b/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 diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py index e0d6541..847c9b0 100644 --- a/src/mcore_bridge/model/register.py +++ b/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': diff --git a/src/mcore_bridge/patcher.py b/src/mcore_bridge/patcher.py index abf8ad0..7e189b9 100644 --- a/src/mcore_bridge/patcher.py +++ b/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 diff --git a/tests/test_llm.py b/tests/test_llm.py index 92531e9..1a5a664 100644 --- a/tests/test_llm.py +++ b/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()