From 303cc3856ab00130269ef461dd5327b16a8b8b8e Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Mon, 3 Aug 2026 21:05:14 +0800 Subject: [PATCH 1/3] wip --- src/mcore_bridge/config/model_config.py | 4 + src/mcore_bridge/config/parser.py | 34 +- src/mcore_bridge/model/constant.py | 1 + src/mcore_bridge/model/gpt_model.py | 2 + src/mcore_bridge/model/gpts/__init__.py | 4 +- src/mcore_bridge/model/gpts/nemotron_h.py | 378 ++++++++++++++++++++++ src/mcore_bridge/model/register.py | 6 +- tests/_nemotron_fwd.py | 163 ++++++++++ tests/_nemotron_layerdiff.py | 169 ++++++++++ tests/_nemotron_mixerdiff.py | 148 +++++++++ tests/_nemotron_moediff.py | 129 ++++++++ tests/_nemotron_rt.py | 173 ++++++++++ 12 files changed, 1206 insertions(+), 5 deletions(-) create mode 100644 src/mcore_bridge/model/gpts/nemotron_h.py create mode 100644 tests/_nemotron_fwd.py create mode 100644 tests/_nemotron_layerdiff.py create mode 100644 tests/_nemotron_mixerdiff.py create mode 100644 tests/_nemotron_moediff.py create mode 100644 tests/_nemotron_rt.py diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index e6d0294..fa2bccd 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -196,6 +196,10 @@ 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 + mtp_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..ee7ecab 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,13 @@ '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'], # other 'original_max_position_embeddings': ['original_max_position_embeddings'], 'partial_rotary_factor': ['partial_rotary_factor'], @@ -255,6 +263,28 @@ 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': + pattern = res.get('hybrid_layer_pattern') + 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' + moe_layer_freq = ['1' if ch == 'E' else '0' for ch in pattern] + res['moe_layer_freq'] = f"[{','.join(moe_layer_freq)}]" + if 'E' not in pattern: + res.pop('num_moe_experts', None) + # MTP: HF exposes num_nextn_predict_layers + its own mtp_hybrid_override_pattern. + mtp_pattern = getattr(hf_config, 'mtp_hybrid_override_pattern', None) + if mtp_pattern: + res['mtp_hybrid_layer_pattern'] = mtp_pattern 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..01080d0 --- /dev/null +++ b/src/mcore_bridge/model/gpts/nemotron_h.py @@ -0,0 +1,378 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from typing import Optional + +from mcore_bridge.bridge import GPTBridge + +from ..constant import ModelType +from ..register import ModelLoader, ModelMeta, register_model +from .bailing_moe import BailingMoeBridge + + +class NemotronHBridge(BailingMoeBridge): + """Bridge for Nemotron-3.5 Hybrid (Mamba2 + Attention + MoE) model. + + Handles weight conversion between HuggingFace and Megatron-Core formats + for three layer types determined by hybrid_override_pattern: + M = Mamba2 SSM layer + E = MoE expert layer (128 routed + 1 shared) + * = Attention layer (GQA) + """ + 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' + + # Nemotron attention uses separate q/k/v projections; restore the base + # implementation (BailingMoeBridge overrides it with fused query_key_value). + _set_qkv = GPTBridge._set_qkv + + def _get_layer_type(self, layer_idx): + """Parse hybrid_layer_pattern to get layer type. + + `layer_idx == -1` is used by the MTP path; MTP layers have their own pattern + (`mtp_hybrid_layer_pattern`) and must not index the backbone pattern. + """ + pattern = self.config.hybrid_layer_pattern + assert 0 <= layer_idx < len(pattern), ( + f'layer_idx {layer_idx} out of range for hybrid_layer_pattern of length {len(pattern)}. ' + 'MTP layers must be dispatched via mtp_hybrid_layer_pattern, not the backbone pattern.') + return {'M': 'mamba', 'E': 'moe', '*': 'attention', '-': 'mlp'}[pattern[layer_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): + # Not hf_grouped, not gate_up merged format. + return False, False + + def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool): + """Dispatch attention/mamba weight conversion based on layer type.""" + 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 attention/mixer part here; 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 exactly the same packed layout and shapes, so every tensor + maps 1:1 with no reordering. + """ + 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_state_dict(mg_mixer, 'in_proj.weight', hf_state_dict, 'in_proj.weight', to_mcore) + 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) + 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 _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool, is_mtp: bool = False): + """Dispatch MoE weight conversion for E-type layers.""" + layer_type = self._get_layer_type(layer_idx) + if layer_type == '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`. + """ + 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) + # `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. `linear_fc1_up` is an alias of `linear_fc1` that + # avoids the base `_get_weight` gated-fc1 reshape (which forces a [2, X, Y] + # view); the real module path is still linear_fc1. + if to_mcore: + self._set_weight(mg_mlp.linear_fc1.weight, hf_state_dict['up_proj.weight'].load(), + 'linear_fc1_up.weight') + else: + fc1 = None if mg_mlp is None else mg_mlp.linear_fc1.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 _build_mamba_layer_cls(): + """MambaLayer subclass that matches TransformerBlock's calling convention. + + Three mismatches have to be absorbed so a Mamba layer can live inside a plain + GPTModel/TransformerBlock instead of mcore's dedicated MambaStack: + + * `TransformerBlock.build_layer` always passes `vp_stage=`, while + `MambaLayer.__init__` only accepts `pp_layer_offset`. + * `TransformerLayer` adds the pipeline offset to `layer_number` internally, but + `MambaLayer` stores it verbatim and expects the caller to supply + `pp_layer_offset`. Without this, PP>1 ranks report local layer numbers (e.g. + [1, 4] instead of [3, 4]) and `GPTBridge._convert` indexes the wrong layer. + * `TransformerBlock.forward` calls layers with the full `TransformerLayer.forward` + keyword set (`context`, `attention_bias`, `rotary_pos_cos`, ...) and unpacks a + `(hidden_states, context)` pair; `MambaLayer.forward` accepts only a small subset + and returns just `hidden_states`. + """ + from megatron.core.ssm.mamba_layer import MambaLayer + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + + class _MambaLayerCompat(MambaLayer): + + def __init__(self, config, submodules, layer_number: int = 1, *args, vp_stage=None, **kwargs): + offset = get_transformer_layer_offset(config, vp_stage=vp_stage) + super().__init__( + config, + submodules, + layer_number=layer_number + offset, + *args, + pp_layer_offset=offset, + **kwargs) + self.vp_stage = vp_stage + + def forward( + self, + hidden_states, + attention_mask=None, + context=None, + context_mask=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + inference_context=None, + packed_seq_params=None, + sequence_len_offset=None, + padding_mask=None, + *, + inference_params=None, + ): + # Mamba has no cross-attention and no positional encoding; the extra + # TransformerLayer kwargs are inapplicable and intentionally dropped. + hidden_states = super().forward( + hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + inference_params=inference_params, + packed_seq_params=packed_seq_params, + ) + # TransformerBlock unpacks `(hidden_states, context)`. + return hidden_states, context + + return _MambaLayerCompat + + +class NemotronHLoader(ModelLoader): + """Loader for Nemotron-3.5 that builds dynamic layer specs. + + Uses hybrid_layer_pattern to assign different layer specs: + M -> MambaLayer (SSM) + E -> MoE layer (routed + shared experts) + * -> Standard attention (GQA) + - -> Dense MLP + + MTP is not supported yet: the HF checkpoint stores MTP layers under a separate + `mtp.layers.*` tree with its own `mtp_hybrid_override_pattern`, whose two-level + index flattening is not implemented here. + """ + + _mamba_layer_cls = None + + def __init__(self, config): + super().__init__(config) + if config.mtp_num_layers: + raise NotImplementedError( + 'nemotron_h MTP conversion is not implemented. The HF checkpoint keeps MTP under ' + '`mtp.layers.*` with its own mtp_hybrid_override_pattern, which requires dedicated ' + 'index-flattening mappings. Set mtp_num_layers=None to convert the backbone only.') + + def get_transformer_layer_spec(self, vp_stage=None): + """Build per-layer specs based on hybrid_layer_pattern. + + Each layer holds exactly ONE mixer, so the unused half of the standard + (attention + MLP) layer must be stripped, otherwise every 'E' layer would build an + unused attention block and every '*' layer an unused MLP -- ~600M phantom params + for this checkpoint, randomly initialized and picked up by the optimizer. + + `moe_layer_freq` (derived from the same pattern in parser.py) already decides which + layers get a MoE vs dense MLP; here we drop whichever submodule the layer type does + not use, and swap 'M' layers for MambaLayer entirely. + """ + from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + pattern = self.config.hybrid_layer_pattern + transformer_layer_spec = super().get_transformer_layer_spec(vp_stage=vp_stage) + # `super()` returns only this PP/VPP stage's layers, so local index 0 is not + # necessarily global layer 0. The pattern is indexed globally. + offset = get_transformer_layer_offset(self.config, vp_stage=vp_stage) + for i, layer_spec in enumerate(transformer_layer_spec.layer_specs): + ch = pattern[offset + i] + if ch == 'M': + # A fresh spec per layer: layer_specs entries must not alias each other, + # matching the base class `_deepcopy_layer_spec` contract. + transformer_layer_spec.layer_specs[i] = self._get_mamba_layer_spec() + continue + submodules = layer_spec.submodules + if ch == '*': + # Attention-only layer: no FFN. Its pre-norm is fused into linear_qkv. + # `mlp_bda` must go too: it unpacks its input as (output, bias), which an + # IdentityOp mlp does not produce. + submodules.mlp = IdentityOp + submodules.pre_mlp_layernorm = IdentityOp + submodules.mlp_bda = IdentityFuncOp + else: + # 'E'/'-': FFN-only layer. Its pre-norm is fused into the MLP's fc1 + # (or pre_mlp_layernorm for MoE), so drop attention and its norm. + submodules.self_attention = IdentityOp + submodules.input_layernorm = IdentityOp + submodules.self_attn_bda = IdentityFuncOp + return transformer_layer_spec + + def _get_mamba_layer_spec(self): + """Build a MambaLayer spec for Mamba2 SSM layers. + + The Bridge expects a standalone pre-norm (`norm`) and a mixer with + plain `in_proj`/`out_proj`, so TENorm + TEColumnParallelLinear are used + instead of the fused TELayerNormColumnParallelLinear variant. + """ + try: + 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.ssm.mamba_layer import MambaLayerSubmodules + from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules + from megatron.core.transformer.spec_utils import ModuleSpec + except ImportError as e: + raise ImportError('NemotronHLoader requires a megatron-core version with Mamba2 SSM support ' + '(megatron.core.ssm.mamba_layer / mamba_mixer).') from e + if self._mamba_layer_cls is None: + self._mamba_layer_cls = _build_mamba_layer_cls() + return ModuleSpec( + module=self._mamba_layer_cls, + submodules=MambaLayerSubmodules( + norm=TENorm, + mixer=ModuleSpec( + module=MambaMixer, + submodules=MambaMixerSubmodules( + in_proj=TEColumnParallelLinear, + out_proj=TERowParallelLinear, + ), + ), + mamba_bda=get_bias_dropout_add, + ), + ) + + +register_model( + ModelMeta( + ModelType.nemotron_h, + ['nemotron_h'], + bridge_cls=NemotronHBridge, + loader=NemotronHLoader, + )) 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/tests/_nemotron_fwd.py b/tests/_nemotron_fwd.py new file mode 100644 index 0000000..7680d26 --- /dev/null +++ b/tests/_nemotron_fwd.py @@ -0,0 +1,163 @@ +"""Forward-consistency check: HF NemotronHForCausalLM vs converted MCore GPTModel. + +The round-trip tests only prove the weight transport is reversible; they say nothing +about whether the MCore model *computes* the same thing. This script loads the real +checkpoint into both stacks and compares logits. + +Run (single H20, ~62GB model so keep TP=1 and expect high memory): + torchrun --nproc_per_node=1 tests/_nemotron_fwd.py +Optionally limit layers for a cheap smoke run: + NUM_LAYERS=8 torchrun --nproc_per_node=1 tests/_nemotron_fwd.py +""" +import os + +import torch +import torch.distributed as dist +from megatron.core import parallel_state as mpu +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + +MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' + 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') + + +def _restore_clobbered_weights(hf_model): + """Undo two HF `_init_weights` bugs that discard trained weights. + + This checkpoint's remote code overwrites weights AFTER they are loaded: + + * `dt_bias`: `module.dt_bias.copy_(inv_dt)` with a random `inv_dt`; the + `_no_reinit = True` marker is set afterwards and never checked. + * `out_proj.weight`: because `rescale_prenorm_residual=True`, it runs + `kaiming_uniform_` then divides by sqrt(num_layers) -- unconditionally. + + Without restoring both, HF is running partly random weights and is useless + as a numerical reference. + """ + import json + + from safetensors import safe_open + index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) + weight_map = index['weight_map'] + restored = 0 + for name, param in hf_model.named_parameters(): + key = name.replace('model.backbone', 'backbone') + if not (key.endswith('mixer.dt_bias') or key.endswith('mixer.out_proj.weight')): + continue + if key not in weight_map: + continue + with safe_open(os.path.join(MODEL_PATH, weight_map[key]), framework='pt') as f: + tensor = f.get_tensor(key) + with torch.no_grad(): + param.copy_(tensor.to(param.device, param.dtype)) + restored += 1 + print(f'RES restored {restored} clobbered weights (HF _init_weights bug)') + + +def main(): + seq_len = int(os.environ.get('SEQ_LEN', 16)) + num_layers_override = os.environ.get('NUM_LAYERS') + + dist.init_process_group('nccl') + torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) + mpu.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(1234) + + from transformers import AutoConfig, AutoModelForCausalLM + + import mcore_bridge.model.gpts # noqa: F401 + from mcore_bridge.config.model_config import ModelConfig + from mcore_bridge.config.parser import hf_to_mcore_config + from mcore_bridge.model.register import get_mcore_model + + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + if num_layers_override: + n = int(num_layers_override) + hf_config.num_hidden_layers = n + hf_config.hybrid_override_pattern = hf_config.hybrid_override_pattern[:n] + # MTP is not supported by the bridge; disable so both sides match. + hf_config.num_nextn_predict_layers = 0 + + torch.manual_seed(0) + input_ids = torch.randint(0, hf_config.vocab_size, (1, seq_len), device='cuda') + + # ---- HF reference ---- + hf_model = AutoModelForCausalLM.from_pretrained( + MODEL_PATH, config=hf_config, torch_dtype=torch.bfloat16, + trust_remote_code=True).cuda().eval() + # HF BUG WORKAROUND: this checkpoint's `_init_weights` unconditionally does + # `module.dt_bias.copy_(inv_dt)` with a *random* inv_dt, and only sets + # `_no_reinit = True` afterwards without ever checking it. So the trained + # dt_bias from safetensors is discarded and HF runs with random values. + # Restore it from the checkpoint so the reference is actually the trained model. + _restore_clobbered_weights(hf_model) + with torch.no_grad(): + hf_logits = hf_model(input_ids).logits.float() + del hf_model + torch.cuda.empty_cache() + print(f'RES hf_logits {tuple(hf_logits.shape)} ' + f'mean={hf_logits.mean():.5f} std={hf_logits.std():.5f}') + + # ---- MCore under test ---- + overrides = hf_to_mcore_config(hf_config) + overrides.update(params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) + # The cuDNN fused-attention backend fails to load its sublibrary in this container; + # flash is equivalent for correctness purposes here. + backend = os.environ.get('ATTN_BACKEND', 'flash') + if backend: + from megatron.core.transformer.enums import AttnBackend + overrides['attention_backend'] = getattr(AttnBackend, backend) + cfg = ModelConfig(**overrides) + models = get_mcore_model(cfg) + cfg.bridge.load_weights(models, MODEL_PATH) + mg_model = models[0].cuda().eval() + + position_ids = torch.arange(seq_len, device='cuda').unsqueeze(0) + attention_mask = torch.tril( + torch.ones((1, 1, seq_len, seq_len), device='cuda', dtype=torch.bool)).logical_not() + with torch.no_grad(): + mg_logits = mg_model( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + ).float() + if mg_logits.shape[0] == seq_len: # [s, b, h] -> [b, s, h] + mg_logits = mg_logits.transpose(0, 1) + mg_logits = mg_logits[..., :hf_logits.shape[-1]] + print(f'RES mg_logits {tuple(mg_logits.shape)} ' + f'mean={mg_logits.mean():.5f} std={mg_logits.std():.5f}') + + diff = (hf_logits - mg_logits).abs() + rel = diff.max() / hf_logits.abs().max() + hf_top = hf_logits.argmax(-1) + mg_top = mg_logits.argmax(-1) + agree = (hf_top == mg_top).float().mean() + print(f'RES max_abs_diff={diff.max():.6f} mean_abs_diff={diff.mean():.6f} ' + f'rel={rel:.6f} argmax_agree={agree:.4f}') + print(f'RES hf_top={hf_top.flatten()[:8].tolist()}') + print(f'RES mg_top={mg_top.flatten()[:8].tolist()}') + # Where argmax disagrees, check whether it's a near-tie (bf16 noise flipping the + # order of two nearly-equal logits) rather than a real behavioural difference. + mism = (hf_top != mg_top).nonzero() + for pos in mism[:5]: + b, t = pos.tolist() + hv, mv = hf_logits[b, t], mg_logits[b, t] + top2 = hv.topk(2).values + print(f'RES tie@t={t} hf_top1-top2_gap={float(top2[0] - top2[1]):.5f} ' + f'hf@hf_top={float(hv[hf_top[b, t]]):.5f} hf@mg_top={float(hv[mg_top[b, t]]):.5f} ' + f'delta={float(hv[hf_top[b, t]] - hv[mg_top[b, t]]):.5f}') + # Rank correlation is the robust check: argmax can flip on ties. + k = 20 + hf_set = hf_logits.topk(k, -1).indices + mg_set = mg_logits.topk(k, -1).indices + overlap = sum(len(set(a.tolist()) & set(b.tolist())) / k + for a, b in zip(hf_set.reshape(-1, k), mg_set.reshape(-1, k))) + overlap /= hf_set.reshape(-1, k).shape[0] + print(f'RES top{k}_overlap={overlap:.4f}') + if agree >= 0.9 and rel < 0.05 and overlap > 0.95: + print('RES FORWARD CONSISTENCY PASS (bf16-level)') + else: + print('RES FORWARD CONSISTENCY FAIL') + + +if __name__ == '__main__': + main() diff --git a/tests/_nemotron_layerdiff.py b/tests/_nemotron_layerdiff.py new file mode 100644 index 0000000..897ccca --- /dev/null +++ b/tests/_nemotron_layerdiff.py @@ -0,0 +1,169 @@ +"""Layer-by-layer hidden-state diff between HF and MCore (single process, one load each). + +Bisects where the two stacks diverge instead of only comparing final logits. +Both models are built with the same small layer count so this fits comfortably in memory. + + NUM_LAYERS=4 SEQ_LEN=8 torchrun --nproc_per_node=1 tests/_nemotron_layerdiff.py +""" +import os + +import torch +import torch.distributed as dist +from megatron.core import parallel_state as mpu +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + +MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' + 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') + + +def _restore_clobbered_weights(hf_model): + """Undo two HF `_init_weights` bugs that discard trained weights. + + This checkpoint's remote code overwrites weights AFTER they are loaded: + + * `dt_bias`: `module.dt_bias.copy_(inv_dt)` with a random `inv_dt`; the + `_no_reinit = True` marker is set afterwards and never checked. + * `out_proj.weight`: because `rescale_prenorm_residual=True`, it runs + `kaiming_uniform_` then divides by sqrt(num_layers) -- unconditionally. + + Without restoring both, HF is running partly random weights and is useless + as a numerical reference. + """ + import json + + from safetensors import safe_open + index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) + weight_map = index['weight_map'] + restored = 0 + for name, param in hf_model.named_parameters(): + key = name.replace('model.backbone', 'backbone') + if not (key.endswith('mixer.dt_bias') or key.endswith('mixer.out_proj.weight')): + continue + if key not in weight_map: + continue + with safe_open(os.path.join(MODEL_PATH, weight_map[key]), framework='pt') as f: + tensor = f.get_tensor(key) + with torch.no_grad(): + param.copy_(tensor.to(param.device, param.dtype)) + restored += 1 + print(f'RES restored {restored} clobbered weights (HF _init_weights bug)') + + +def main(): + seq_len = int(os.environ.get('SEQ_LEN', 8)) + n_layers = int(os.environ.get('NUM_LAYERS', 4)) + + dist.init_process_group('nccl') + torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) + mpu.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(1234) + + from transformers import AutoConfig, AutoModelForCausalLM + + import mcore_bridge.model.gpts # noqa: F401 + from mcore_bridge.config.model_config import ModelConfig + from mcore_bridge.config.parser import hf_to_mcore_config + from mcore_bridge.model.register import get_mcore_model + + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + hf_config.num_hidden_layers = n_layers + hf_config.hybrid_override_pattern = hf_config.hybrid_override_pattern[:n_layers] + hf_config.num_nextn_predict_layers = 0 + pattern = hf_config.hybrid_override_pattern + print(f'RES pattern={pattern}') + + torch.manual_seed(0) + input_ids = torch.randint(0, hf_config.vocab_size, (1, seq_len), device='cuda') + + # ---- HF: capture every backbone layer output ---- + hf_model = AutoModelForCausalLM.from_pretrained( + MODEL_PATH, config=hf_config, torch_dtype=torch.bfloat16, + trust_remote_code=True).cuda().eval() + _restore_clobbered_weights(hf_model) # HF _init_weights clobbers dt_bias with random inv_dt + hf_acts = {} + + hf_ins = {} + + def mk_hook(idx): + def hook(_mod, inp, out): + if inp: + hf_ins[idx] = inp[0].detach().float() + hf_acts[idx] = (out[0] if isinstance(out, tuple) else out).detach().float() + return hook + + for i, layer in enumerate(hf_model.backbone.layers): + layer.register_forward_hook(mk_hook(i)) + hf_emb = {} + hf_model.backbone.embeddings.register_forward_hook( + lambda m, i, o: hf_emb.__setitem__(0, o.detach().float())) + with torch.no_grad(): + hf_model(input_ids) + del hf_model + torch.cuda.empty_cache() + + # ---- MCore: same capture ---- + overrides = hf_to_mcore_config(hf_config) + overrides.update(params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) + from megatron.core.transformer.enums import AttnBackend + overrides['attention_backend'] = AttnBackend.flash + cfg = ModelConfig(**overrides) + models = get_mcore_model(cfg) + cfg.bridge.load_weights(models, MODEL_PATH) + mg_model = models[0].cuda().eval() + + mg_acts = {} + + mg_ins = {} + + def mk_hook_mg(idx): + def hook(_mod, inp, out): + if inp: + mg_ins[idx] = inp[0].detach().float() + t = out[0] if isinstance(out, tuple) else out + mg_acts[idx] = t.detach().float() + return hook + + for i, layer in enumerate(mg_model.decoder.layers): + layer.register_forward_hook(mk_hook_mg(i)) + mg_emb = {} + mg_model.embedding.register_forward_hook( + lambda m, i, o: mg_emb.__setitem__(0, o.detach().float())) + + position_ids = torch.arange(seq_len, device='cuda').unsqueeze(0) + attention_mask = torch.tril( + torch.ones((1, 1, seq_len, seq_len), device='cuda', dtype=torch.bool)).logical_not() + with torch.no_grad(): + mg_model(input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask) + + def norm(t): + """HF is [b, s, h]; MCore is [s, b, h]. Normalize to [s, h].""" + if t.dim() == 3: + if t.shape[0] == 1 and t.shape[1] == seq_len: + return t[0] + if t.shape[1] == 1 and t.shape[0] == seq_len: + return t[:, 0] + return t.reshape(seq_len, -1) + + if 0 in hf_emb and 0 in mg_emb: + a, b = norm(hf_emb[0]), norm(mg_emb[0]) + print(f'RES embedding max_abs_diff={(a - b).abs().max():.6f}') + for i in range(n_layers): + if i not in hf_acts or i not in mg_acts: + print(f'RES layer{i} MISSING hf={i in hf_acts} mg={i in mg_acts}') + continue + a, b = norm(hf_acts[i]), norm(mg_acts[i]) + d = (a - b).abs() + scale = a.abs().max().clamp(min=1e-6) + if i in hf_ins and i in mg_ins: + ia, ib = norm(hf_ins[i]), norm(mg_ins[i]) + di = (ia - ib).abs() + print(f'RES layer{i} IN max_abs_diff={di.max():.6f} ' + f'hf_std={ia.std():.5f} mg_std={ib.std():.5f}') + print(f'RES layer{i} type={pattern[i]!r} hf_shape={tuple(hf_acts[i].shape)} ' + f'mg_shape={tuple(mg_acts[i].shape)} ' + f'hf_std={a.std():.5f} mg_std={b.std():.5f} ' + f'max_abs_diff={d.max():.6f} mean={d.mean():.6f} rel={d.max() / scale:.6f}') + + +if __name__ == '__main__': + main() diff --git a/tests/_nemotron_mixerdiff.py b/tests/_nemotron_mixerdiff.py new file mode 100644 index 0000000..a3c36fe --- /dev/null +++ b/tests/_nemotron_mixerdiff.py @@ -0,0 +1,148 @@ +"""Isolate the Mamba2 mixer: run HF's and mcore's on identical weights + input. + +Loads only one Mamba layer's weights out of the checkpoint, so this is cheap. + torchrun --nproc_per_node=1 tests/_nemotron_mixerdiff.py +""" +import json +import os + +import torch +import torch.distributed as dist +from megatron.core import parallel_state as mpu +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from safetensors import safe_open + +MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' + 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') + + +def load_layer0_mixer(): + index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) + weight_map = index['weight_map'] + out = {} + for key, shard in weight_map.items(): + if key.startswith('backbone.layers.0.mixer.'): + with safe_open(os.path.join(MODEL_PATH, shard), framework='pt') as f: + out[key[len('backbone.layers.0.mixer.'):]] = f.get_tensor(key).cuda() + return out + + +def main(): + seq_len = int(os.environ.get('SEQ_LEN', 8)) + dist.init_process_group('nccl') + torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) + mpu.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(1234) + + from transformers import AutoConfig + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + w = load_layer0_mixer() + print('RES loaded mixer keys:', sorted(w)) + + torch.manual_seed(0) + # HF wants [b, s, h]; mcore wants [s, b, h]. + x_bsh = torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.bfloat16, device='cuda') + + # ---- HF mixer ---- + # Go through transformers' dynamic-module loader so the file's relative imports resolve. + from transformers.dynamic_module_utils import get_class_from_dynamic_module + hf_mixer_cls = get_class_from_dynamic_module( + 'modeling_nemotron_h.NemotronHMamba2Mixer', MODEL_PATH) + hf_mixer = hf_mixer_cls(hf_config, layer_idx=0).cuda().to(torch.bfloat16).eval() + missing, unexpected = hf_mixer.load_state_dict(w, strict=False) + print(f'RES hf_mixer missing={list(missing)} unexpected={list(unexpected)}') + with torch.no_grad(): + hf_out = hf_mixer(x_bsh).float() + print(f'RES hf_out {tuple(hf_out.shape)} mean={hf_out.mean():.6f} std={hf_out.std():.6f}') + + # ---- mcore mixer ---- + import mcore_bridge.model.gpts # noqa: F401 + from mcore_bridge.config.model_config import ModelConfig + from mcore_bridge.config.parser import hf_to_mcore_config + overrides = hf_to_mcore_config(hf_config) + overrides.update(num_layers=1, hybrid_layer_pattern='M', moe_layer_freq='[0]', + params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) + cfg = ModelConfig(**overrides) + + from megatron.core.extensions.transformer_engine import (TEColumnParallelLinear, + TERowParallelLinear) + from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules + mg_mixer = MambaMixer( + cfg, + MambaMixerSubmodules(in_proj=TEColumnParallelLinear, out_proj=TERowParallelLinear), + d_model=cfg.hidden_size, + layer_number=1, + pg_collection=_pg(), + ).cuda().to(torch.bfloat16).eval() + + mg_sd = { + 'in_proj.weight': w['in_proj.weight'], + 'conv1d_weight': w['conv1d.weight'], + 'conv1d_bias': w['conv1d.bias'], + 'A_log': w['A_log'], + 'D': w['D'], + 'dt_bias': w['dt_bias'], + 'norm.weight': w['norm.weight'], + 'out_proj.weight': w['out_proj.weight'], + } + incompat = mg_mixer.load_state_dict(mg_sd, strict=False) + print(f'RES mg_mixer missing={list(incompat.missing_keys)} ' + f'unexpected={list(incompat.unexpected_keys)}') + + x_sbh = x_bsh.transpose(0, 1).contiguous() + with torch.no_grad(): + mg_out, mg_bias = mg_mixer(x_sbh) + mg_out = mg_out.float().transpose(0, 1) # -> [b, s, h] + if mg_bias is not None: + mg_out = mg_out + mg_bias.float() + print(f'RES mg_out {tuple(mg_out.shape)} mean={mg_out.mean():.6f} std={mg_out.std():.6f}') + + d = (hf_out - mg_out).abs() + print(f'RES mixer max_abs_diff={d.max():.6f} mean_abs_diff={d.mean():.6f} ' + f'rel={d.max() / hf_out.abs().max():.6f}') + + # ---- Now the FULL block: norm + mixer + residual ---- + # Load layer0's outer norm too. + import json as _json + idx = _json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) + nk = 'backbone.layers.0.norm.weight' + with safe_open(os.path.join(MODEL_PATH, idx['weight_map'][nk]), framework='pt') as f: + outer_norm_w = f.get_tensor(nk).cuda() + + hf_block_cls = get_class_from_dynamic_module('modeling_nemotron_h.NemotronHBlock', MODEL_PATH) + hf_block = hf_block_cls(hf_config, layer_idx=0).cuda().to(torch.bfloat16).eval() + bsd = {f'mixer.{k}': v for k, v in w.items()} + bsd['norm.weight'] = outer_norm_w + inc = hf_block.load_state_dict(bsd, strict=False) + print(f'RES hf_block missing={list(inc.missing_keys)} unexpected={list(inc.unexpected_keys)}') + with torch.no_grad(): + hf_blk = hf_block(x_bsh) + hf_blk = (hf_blk[0] if isinstance(hf_blk, tuple) else hf_blk).float() + + from mcore_bridge.model.gpts.nemotron_h import _build_mamba_layer_cls, NemotronHLoader + loader = NemotronHLoader(cfg) + mspec = loader._get_mamba_layer_spec() + from megatron.core.transformer.spec_utils import build_module + mg_layer = build_module(mspec, config=cfg, layer_number=1, + pg_collection=_pg(), vp_stage=None).cuda().to(torch.bfloat16).eval() + lsd = dict(mg_sd) + lsd = {f'mixer.{k}': v for k, v in mg_sd.items()} + lsd['norm.weight'] = outer_norm_w + inc2 = mg_layer.load_state_dict(lsd, strict=False) + print(f'RES mg_layer missing={[k for k in inc2.missing_keys if "_extra_state" not in k]} ' + f'unexpected={list(inc2.unexpected_keys)}') + with torch.no_grad(): + out = mg_layer(x_sbh) + mg_blk = (out[0] if isinstance(out, tuple) else out).float().transpose(0, 1) + d2 = (hf_blk - mg_blk).abs() + print(f'RES BLOCK max_abs_diff={d2.max():.6f} mean_abs_diff={d2.mean():.6f} ' + f'rel={d2.max() / hf_blk.abs().max():.6f}') + + +def _pg(): + from megatron.core.process_groups_config import ProcessGroupCollection + return ProcessGroupCollection.use_mpu_process_groups() + + +if __name__ == '__main__': + main() diff --git a/tests/_nemotron_moediff.py b/tests/_nemotron_moediff.py new file mode 100644 index 0000000..9044b9d --- /dev/null +++ b/tests/_nemotron_moediff.py @@ -0,0 +1,129 @@ +"""Compare a single 'E' (MoE) block: HF NemotronHBlock vs mcore TransformerLayer. + +The Mamba block already matches in isolation (see _nemotron_mixerdiff.py), so this +checks the other layer type. Loads only layer 1's weights. + + torchrun --nproc_per_node=1 tests/_nemotron_moediff.py +""" +import json +import os + +import torch +import torch.distributed as dist +from megatron.core import parallel_state as mpu +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from safetensors import safe_open + +MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' + 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') +LAYER = int(os.environ.get('LAYER', 1)) # layer 1 is 'E' in MEMEM*... + + +def load_layer(prefix): + index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) + out = {} + for key, shard in index['weight_map'].items(): + if key.startswith(prefix): + with safe_open(os.path.join(MODEL_PATH, shard), framework='pt') as f: + out[key[len(prefix):]] = f.get_tensor(key).cuda() + return out + + +def _pg(): + from megatron.core.process_groups_config import ProcessGroupCollection + return ProcessGroupCollection.use_mpu_process_groups() + + +def main(): + seq_len = int(os.environ.get('SEQ_LEN', 8)) + dist.init_process_group('nccl') + torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) + mpu.initialize_model_parallel(1, 1) + model_parallel_cuda_manual_seed(1234) + + from transformers import AutoConfig + from transformers.dynamic_module_utils import get_class_from_dynamic_module + + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + # NEMOTRONH_ATTENTION_CLASSES is keyed by _attn_implementation; standalone block + # construction skips the usual model-level default, so set it explicitly. + hf_config._attn_implementation = os.environ.get('HF_ATTN', 'eager') + pattern = hf_config.hybrid_override_pattern + print(f'RES layer{LAYER} type={pattern[LAYER]!r}') + + w = load_layer(f'backbone.layers.{LAYER}.') + print(f'RES n_weights={len(w)}') + + torch.manual_seed(0) + if os.environ.get('REAL_EMB'): + # Use the checkpoint's real embedding rows: their magnitude is far from N(0,1), + # and Mamba's conv/SSM path is scale sensitive. + ek = 'backbone.embeddings.weight' + _idx = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) + with safe_open(os.path.join(MODEL_PATH, _idx['weight_map'][ek]), framework='pt') as f: + emb = f.get_tensor(ek).cuda() + ids = torch.randint(0, emb.shape[0], (1, seq_len), device='cuda') + x_bsh = emb[ids].to(torch.bfloat16) + print(f'RES using REAL embedding, std={x_bsh.float().std():.5f}') + else: + x_bsh = torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.bfloat16, device='cuda') + print(f'RES using randn input, std={x_bsh.float().std():.5f}') + + # ---- HF block ---- + hf_block_cls = get_class_from_dynamic_module('modeling_nemotron_h.NemotronHBlock', MODEL_PATH) + hf_block = hf_block_cls(hf_config, layer_idx=LAYER).cuda().to(torch.bfloat16).eval() + inc = hf_block.load_state_dict(w, strict=False) + print(f'RES hf missing={list(inc.missing_keys)[:5]} unexpected={list(inc.unexpected_keys)[:5]}') + with torch.no_grad(): + hf_out = hf_block(x_bsh) + hf_out = (hf_out[0] if isinstance(hf_out, tuple) else hf_out).float() + print(f'RES hf_out mean={hf_out.mean():.6f} std={hf_out.std():.6f}') + + # ---- MCore layer via the real Loader spec ---- + import mcore_bridge.model.gpts # noqa: F401 + from mcore_bridge.config.model_config import ModelConfig + from mcore_bridge.config.parser import hf_to_mcore_config + from mcore_bridge.model.gpts.nemotron_h import NemotronHLoader + + ch = pattern[LAYER] + overrides = hf_to_mcore_config(hf_config) + overrides.update(num_layers=1, hybrid_layer_pattern=ch, + moe_layer_freq='[1]' if ch == 'E' else '[0]', + params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) + from megatron.core.transformer.enums import AttnBackend + overrides['attention_backend'] = AttnBackend.flash + cfg = ModelConfig(**overrides) + + loader = NemotronHLoader(cfg) + spec = loader.get_transformer_layer_spec() + from megatron.core.transformer.spec_utils import build_module + mg_layer = build_module(spec.layer_specs[0], config=cfg, layer_number=1, + pg_collection=_pg(), vp_stage=None).cuda().to(torch.bfloat16).eval() + print('RES mg params:', sorted(n for n, _ in mg_layer.named_parameters())[:12]) + + # Drive the real Bridge so the mapping under test is exercised. + class Lazy: + def __init__(self, t): + self.t = t + + def load(self): + return self.t + + sd = {f'backbone.layers.0.{k}': Lazy(v) for k, v in w.items()} + cfg.bridge._set_layer_state(mg_layer, sd, 'backbone.layers.', 0, True) + + x_sbh = x_bsh.transpose(0, 1).contiguous() + attn_mask = torch.tril( + torch.ones((1, 1, seq_len, seq_len), device='cuda', dtype=torch.bool)).logical_not() + with torch.no_grad(): + out = mg_layer(x_sbh, attention_mask=attn_mask) + mg_out = (out[0] if isinstance(out, tuple) else out).float().transpose(0, 1) + print(f'RES mg_out mean={mg_out.mean():.6f} std={mg_out.std():.6f}') + + d = (hf_out - mg_out).abs() + print(f'RES BLOCK({ch}) max_abs_diff={d.max():.6f} mean_abs_diff={d.mean():.6f} ' + f'rel={d.max() / hf_out.abs().max():.6f}') + + +if __name__ == '__main__': + main() diff --git a/tests/_nemotron_rt.py b/tests/_nemotron_rt.py new file mode 100644 index 0000000..3d7115c --- /dev/null +++ b/tests/_nemotron_rt.py @@ -0,0 +1,173 @@ +"""Nemotron-H round-trip harness: HF -> MCore -> HF must be bit-exact. + +Parallel layout is taken from env vars so the same script covers TP/PP/EP: + TP, PP, EP, ETP (default 1), PATTERN (default 'ME*') + +Run e.g.: + EP=2 PATTERN='ME*' torchrun --nproc_per_node=2 tests/_nemotron_rt.py +""" +import os + +import torch +import torch.distributed as dist +from megatron.core import parallel_state as mpu +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + +MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' + 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') +NUM_EXPERTS = 4 + + +class Lazy: + """Mimics SafetensorLazyLoader's handle: Bridge calls `.load()`.""" + + def __init__(self, tensor): + self.tensor = tensor + + def load(self): + return self.tensor + + +def build_hf_state_dict(cfg, mixer, pattern, rand): + hidden = cfg.hidden_size + sd = { + 'backbone.embeddings.weight': Lazy(rand(256, hidden)), + 'lm_head.weight': Lazy(rand(256, hidden)), + 'backbone.norm_f.weight': Lazy(rand(hidden)), + } + for i, ch in enumerate(pattern): + p = f'backbone.layers.{i}.' + sd[f'{p}norm.weight'] = Lazy(rand(hidden)) + if ch == 'M': + conv_dim = mixer.d_inner + 2 * mixer.ngroups * mixer.d_state + in_dim = mixer.d_inner * 2 + 2 * mixer.ngroups * mixer.d_state + mixer.nheads + sd[f'{p}mixer.in_proj.weight'] = Lazy(rand(in_dim, hidden)) + sd[f'{p}mixer.conv1d.weight'] = Lazy(rand(conv_dim, 1, 4)) + sd[f'{p}mixer.conv1d.bias'] = Lazy(rand(conv_dim)) + # A_log/D are fp32 in mcore; dt_bias is bf16. + sd[f'{p}mixer.A_log'] = Lazy(rand(mixer.nheads, dtype=torch.float32)) + sd[f'{p}mixer.D'] = Lazy(rand(mixer.nheads, dtype=torch.float32)) + sd[f'{p}mixer.dt_bias'] = Lazy(rand(mixer.nheads)) + sd[f'{p}mixer.norm.weight'] = Lazy(rand(mixer.d_inner)) + sd[f'{p}mixer.out_proj.weight'] = Lazy(rand(hidden, mixer.d_inner)) + elif ch == 'E': + sd[f'{p}mixer.gate.weight'] = Lazy(rand(NUM_EXPERTS, hidden)) + sd[f'{p}mixer.gate.e_score_correction_bias'] = Lazy( + rand(NUM_EXPERTS, dtype=torch.float32)) + moe_ffn = cfg.moe_ffn_hidden_size + for e in range(NUM_EXPERTS): + sd[f'{p}mixer.experts.{e}.up_proj.weight'] = Lazy(rand(moe_ffn, hidden)) + sd[f'{p}mixer.experts.{e}.down_proj.weight'] = Lazy(rand(hidden, moe_ffn)) + shared = cfg.moe_shared_expert_intermediate_size + sd[f'{p}mixer.shared_experts.up_proj.weight'] = Lazy(rand(shared, hidden)) + sd[f'{p}mixer.shared_experts.down_proj.weight'] = Lazy(rand(hidden, shared)) + elif ch == '*': + head_dim = cfg.kv_channels + q = cfg.num_attention_heads * head_dim + kv = cfg.num_query_groups * head_dim + sd[f'{p}mixer.q_proj.weight'] = Lazy(rand(q, hidden)) + sd[f'{p}mixer.k_proj.weight'] = Lazy(rand(kv, hidden)) + sd[f'{p}mixer.v_proj.weight'] = Lazy(rand(kv, hidden)) + sd[f'{p}mixer.o_proj.weight'] = Lazy(rand(hidden, q)) + else: + raise ValueError(f'unsupported pattern char {ch!r}') + return sd + + +def main(): + tp = int(os.environ.get('TP', 1)) + pp = int(os.environ.get('PP', 1)) + ep = int(os.environ.get('EP', 1)) + etp = int(os.environ.get('ETP', tp)) + pattern = os.environ.get('PATTERN', 'ME*') + + dist.init_process_group('nccl') + local_rank = int(os.environ.get('LOCAL_RANK', 0)) + torch.cuda.set_device(local_rank) + mpu.initialize_model_parallel( + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + expert_model_parallel_size=ep, + expert_tensor_parallel_size=etp, + ) + model_parallel_cuda_manual_seed(1234) # MambaMixer uses get_cuda_rng_tracker().fork() + + from transformers import AutoConfig + + import mcore_bridge.model.gpts # noqa: F401 (triggers registration) + from mcore_bridge.config.model_config import ModelConfig + from mcore_bridge.config.parser import hf_to_mcore_config + from mcore_bridge.model.register import get_mcore_model + + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + overrides = hf_to_mcore_config(hf_config) + overrides.update( + num_layers=len(pattern), + hybrid_layer_pattern=pattern, + moe_layer_freq='[' + ','.join('1' if c == 'E' else '0' for c in pattern) + ']', + num_moe_experts=NUM_EXPERTS, + padded_vocab_size=256, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + expert_model_parallel_size=ep, + expert_tensor_parallel_size=etp, + params_dtype=torch.bfloat16, + bf16=True, + ) + cfg = ModelConfig(**overrides) + + models = get_mcore_model(cfg) + bridge = cfg.bridge + + # Any rank holding a Mamba layer can report the mixer dims; they are TP-invariant + # on the HF side, so derive them from config instead of the (possibly absent) module. + class _Dims: + nheads = cfg.mamba_num_heads + d_inner = cfg.mamba_num_heads * cfg.mamba_head_dim + ngroups = cfg.mamba_num_groups + d_state = cfg.mamba_state_dim + + gen = torch.Generator(device='cuda').manual_seed(4321) # identical data on every rank + + def rand(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, generator=gen, dtype=dtype, device='cuda') + + sd = build_hf_state_dict(cfg, _Dims, pattern, rand) + original = {k: v.load().clone() for k, v in sd.items()} + + for model in models: + list(bridge._convert([model], sd, '', True, 'Loading: ')) + + exported = dict(bridge.export_weights(models, target_device='cuda')) + + if local_rank != 0: + dist.barrier() + return + + tag = f'TP{tp}/PP{pp}/EP{ep} pattern={pattern}' + missing = sorted(set(original) - set(exported)) + extra = sorted(set(exported) - set(original)) + bad = [] + for key, want in original.items(): + got = exported.get(key) + if got is None: + continue + if tuple(got.shape) != tuple(want.shape): + bad.append((key, 'shape', tuple(want.shape), tuple(got.shape))) + elif not torch.equal(got.to(want.dtype).cpu(), want.cpu()): + delta = (got.to(want.dtype).cpu().float() - want.cpu().float()).abs().max() + bad.append((key, 'value', float(delta))) + + print(f'RES [{tag}] keys={len(exported)} missing={missing} extra={extra} ' + f'mismatches={len(bad)}') + for item in bad[:10]: + print(' ', item) + if not bad and not missing and not extra: + print(f'RES [{tag}] ROUNDTRIP EXACT PASS') + else: + print(f'RES [{tag}] ROUNDTRIP FAILED') + dist.barrier() + + +if __name__ == '__main__': + main() From 327200f58ad9df072511be65bd5bd83ac49fdf77 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 11 Aug 2026 21:25:55 +0800 Subject: [PATCH 2/3] fix --- src/mcore_bridge/config/model_config.py | 1 - src/mcore_bridge/config/parser.py | 10 +- src/mcore_bridge/model/gpts/nemotron_h.py | 439 ++++++++++++---------- src/mcore_bridge/model/hybrid_model.py | 86 +++++ src/mcore_bridge/patcher.py | 39 ++ tests/_nemotron_fwd.py | 163 -------- tests/_nemotron_layerdiff.py | 169 --------- tests/_nemotron_mixerdiff.py | 148 -------- tests/_nemotron_moediff.py | 129 ------- tests/_nemotron_rt.py | 173 --------- tests/test_llm.py | 7 +- 11 files changed, 371 insertions(+), 993 deletions(-) create mode 100644 src/mcore_bridge/model/hybrid_model.py delete mode 100644 tests/_nemotron_fwd.py delete mode 100644 tests/_nemotron_layerdiff.py delete mode 100644 tests/_nemotron_mixerdiff.py delete mode 100644 tests/_nemotron_moediff.py delete mode 100644 tests/_nemotron_rt.py diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index fa2bccd..a472003 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -198,7 +198,6 @@ class ModelConfig(TransformerConfig): # nemotron_h (hybrid mamba2 + attention + moe) hybrid_layer_pattern: Optional[str] = None - mtp_hybrid_layer_pattern: Optional[str] = None # dsa experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = None diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py index ee7ecab..a22ba6a 100644 --- a/src/mcore_bridge/config/parser.py +++ b/src/mcore_bridge/config/parser.py @@ -75,6 +75,7 @@ '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'], @@ -264,7 +265,6 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]: res['moe_router_score_function'] = 'sigmoid' res['moe_router_load_balancing_type'] = 'seq_aux_loss' elif llm_model_type == 'nemotron_h': - pattern = res.get('hybrid_layer_pattern') 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). @@ -277,14 +277,6 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]: res['moe_router_score_function'] = 'sigmoid' res['moe_router_enable_expert_bias'] = True res['moe_router_load_balancing_type'] = 'seq_aux_loss' - moe_layer_freq = ['1' if ch == 'E' else '0' for ch in pattern] - res['moe_layer_freq'] = f"[{','.join(moe_layer_freq)}]" - if 'E' not in pattern: - res.pop('num_moe_experts', None) - # MTP: HF exposes num_nextn_predict_layers + its own mtp_hybrid_override_pattern. - mtp_pattern = getattr(hf_config, 'mtp_hybrid_override_pattern', None) - if mtp_pattern: - res['mtp_hybrid_layer_pattern'] = mtp_pattern 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/gpts/nemotron_h.py b/src/mcore_bridge/model/gpts/nemotron_h.py index 01080d0..bf500f3 100644 --- a/src/mcore_bridge/model/gpts/nemotron_h.py +++ b/src/mcore_bridge/model/gpts/nemotron_h.py @@ -1,23 +1,42 @@ # 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 -from .bailing_moe import BailingMoeBridge + +logger = get_logger() -class NemotronHBridge(BailingMoeBridge): - """Bridge for Nemotron-3.5 Hybrid (Mamba2 + Attention + MoE) model. +class NemotronHBridge(GPTBridge): + """HuggingFace <-> Megatron-Core weight conversion for Nemotron-3.5. - Handles weight conversion between HuggingFace and Megatron-Core formats - for three layer types determined by hybrid_override_pattern: - M = Mamba2 SSM layer - E = MoE expert layer (128 routed + 1 shared) - * = Attention layer (GQA) + 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' @@ -31,22 +50,25 @@ class NemotronHBridge(BailingMoeBridge): 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' - # Nemotron attention uses separate q/k/v projections; restore the base - # implementation (BailingMoeBridge overrides it with fused query_key_value). - _set_qkv = GPTBridge._set_qkv + _LAYER_TYPES = {'M': 'mamba', 'E': 'moe', '*': 'attention', '-': 'mlp'} - def _get_layer_type(self, layer_idx): - """Parse hybrid_layer_pattern to get layer type. + def _get_layer_type(self, layer_idx: int): + """Resolve a layer's family from the pattern. - `layer_idx == -1` is used by the MTP path; MTP layers have their own pattern - (`mtp_hybrid_layer_pattern`) and must not index the backbone 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. """ - pattern = self.config.hybrid_layer_pattern - assert 0 <= layer_idx < len(pattern), ( - f'layer_idx {layer_idx} out of range for hybrid_layer_pattern of length {len(pattern)}. ' - 'MTP layers must be dispatched via mtp_hybrid_layer_pattern, not the backbone pattern.') - return {'M': 'mamba', 'E': 'moe', '*': 'attention', '-': 'mlp'}[pattern[layer_idx]] + 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 @@ -63,11 +85,17 @@ def _get_tp_split_dim(self, mg_key: Optional[str]) -> Optional[int]: return super()._get_tp_split_dim(mg_key) def _get_hf_experts_attr(self, is_mtp: bool = False): - # Not hf_grouped, not gate_up merged format. + # 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): - """Dispatch attention/mamba weight conversion based on layer type.""" + """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 @@ -80,7 +108,7 @@ def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: boo 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 attention/mixer part here; their norm is handled in _set_layer_mlp + # '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): @@ -88,8 +116,13 @@ def _set_mamba_state(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bo 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 exactly the same packed layout and shapes, so every tensor - maps 1:1 with no reordering. + 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: @@ -97,9 +130,8 @@ def _set_mamba_state(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bo else: hf_state_dict = {} mg_mixer = None if mg_layer is None else mg_layer.mixer - self._set_state_dict(mg_mixer, 'in_proj.weight', hf_state_dict, 'in_proj.weight', to_mcore) - 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) + 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) @@ -115,17 +147,97 @@ def _set_mamba_state(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bo 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` may be None on a PP rank that does not own this layer; the TP all-gather + still has to run collectively, so pass None straight through to `_all_gather_tp` + (which tolerates None) and only touch it when this rank actually holds the weight. + """ + 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) + if gathered is None: + return + merged = self._merge_packed_dim0(gathered, block_sizes) + # `_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): - """Dispatch MoE weight conversion for E-type layers.""" - layer_type = self._get_layer_type(layer_idx) - if layer_type == 'moe': + """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 + # mamba / attention layers have no MLP sub-module. return hf_state_dict def _set_mlp_state( @@ -143,7 +255,7 @@ def _set_mlp_state( 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`. + 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) @@ -156,16 +268,19 @@ def _set_mlp_state( 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) + 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, @@ -181,14 +296,14 @@ def _set_mlp_state( hf_state_dict[f'{start_idx + i}.{hf_name}.weight'] = weight[i].clone() del weight else: - # dense MLP / shared expert. `linear_fc1_up` is an alias of `linear_fc1` that - # avoids the base `_get_weight` gated-fc1 reshape (which forces a [2, X, Y] - # view); the real module path is still linear_fc1. + # 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(mg_mlp.linear_fc1.weight, hf_state_dict['up_proj.weight'].load(), - 'linear_fc1_up.weight') + self._set_weight(fc1_module.weight, hf_state_dict['up_proj.weight'].load(), 'linear_fc1_up.weight') else: - fc1 = None if mg_mlp is None else mg_mlp.linear_fc1.weight.data + 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() @@ -200,161 +315,68 @@ def _set_mlp_state( 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. -def _build_mamba_layer_cls(): - """MambaLayer subclass that matches TransformerBlock's calling convention. - - Three mismatches have to be absorbed so a Mamba layer can live inside a plain - GPTModel/TransformerBlock instead of mcore's dedicated MambaStack: - - * `TransformerBlock.build_layer` always passes `vp_stage=`, while - `MambaLayer.__init__` only accepts `pp_layer_offset`. - * `TransformerLayer` adds the pipeline offset to `layer_number` internally, but - `MambaLayer` stores it verbatim and expects the caller to supply - `pp_layer_offset`. Without this, PP>1 ranks report local layer numbers (e.g. - [1, 4] instead of [3, 4]) and `GPTBridge._convert` indexes the wrong layer. - * `TransformerBlock.forward` calls layers with the full `TransformerLayer.forward` - keyword set (`context`, `attention_bias`, `rotary_pos_cos`, ...) and unpacks a - `(hidden_states, context)` pair; `MambaLayer.forward` accepts only a small subset - and returns just `hidden_states`. - """ - from megatron.core.ssm.mamba_layer import MambaLayer - from megatron.core.transformer.transformer_layer import get_transformer_layer_offset - - class _MambaLayerCompat(MambaLayer): - - def __init__(self, config, submodules, layer_number: int = 1, *args, vp_stage=None, **kwargs): - offset = get_transformer_layer_offset(config, vp_stage=vp_stage) - super().__init__( - config, - submodules, - layer_number=layer_number + offset, - *args, - pp_layer_offset=offset, - **kwargs) - self.vp_stage = vp_stage - - def forward( - self, - hidden_states, - attention_mask=None, - context=None, - context_mask=None, - rotary_pos_emb=None, - rotary_pos_cos=None, - rotary_pos_sin=None, - rotary_pos_cos_sin=None, - attention_bias=None, - inference_context=None, - packed_seq_params=None, - sequence_len_offset=None, - padding_mask=None, - *, - inference_params=None, - ): - # Mamba has no cross-attention and no positional encoding; the extra - # TransformerLayer kwargs are inapplicable and intentionally dropped. - hidden_states = super().forward( - hidden_states, - attention_mask=attention_mask, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb, - inference_params=inference_params, - packed_seq_params=packed_seq_params, - ) - # TransformerBlock unpacks `(hidden_states, context)`. - return hidden_states, context - - return _MambaLayerCompat + 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): - """Loader for Nemotron-3.5 that builds dynamic layer specs. - - Uses hybrid_layer_pattern to assign different layer specs: - M -> MambaLayer (SSM) - E -> MoE layer (routed + shared experts) - * -> Standard attention (GQA) - - -> Dense MLP - - MTP is not supported yet: the HF checkpoint stores MTP layers under a separate - `mtp.layers.*` tree with its own `mtp_hybrid_override_pattern`, whose two-level - index flattening is not implemented here. - """ - - _mamba_layer_cls = None + model_cls = HybridModel - def __init__(self, config): - super().__init__(config) - if config.mtp_num_layers: - raise NotImplementedError( - 'nemotron_h MTP conversion is not implemented. The HF checkpoint keeps MTP under ' - '`mtp.layers.*` with its own mtp_hybrid_override_pattern, which requires dedicated ' - 'index-flattening mappings. Set mtp_num_layers=None to convert the backbone only.') + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + """Return a `HybridStack` spec with a standalone pre-norm on Mamba layers. - def get_transformer_layer_spec(self, vp_stage=None): - """Build per-layer specs based on hybrid_layer_pattern. - - Each layer holds exactly ONE mixer, so the unused half of the standard - (attention + MLP) layer must be stripped, otherwise every 'E' layer would build an - unused attention block and every '*' layer an unused MLP -- ~600M phantom params - for this checkpoint, randomly initialized and picked up by the optimizer. - - `moe_layer_freq` (derived from the same pattern in parser.py) already decides which - layers get a MoE vs dense MLP; here we drop whichever submodule the layer type does - not use, and swap 'M' layers for MambaLayer entirely. - """ - from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp - from megatron.core.transformer.transformer_layer import get_transformer_layer_offset - pattern = self.config.hybrid_layer_pattern - transformer_layer_spec = super().get_transformer_layer_spec(vp_stage=vp_stage) - # `super()` returns only this PP/VPP stage's layers, so local index 0 is not - # necessarily global layer 0. The pattern is indexed globally. - offset = get_transformer_layer_offset(self.config, vp_stage=vp_stage) - for i, layer_spec in enumerate(transformer_layer_spec.layer_specs): - ch = pattern[offset + i] - if ch == 'M': - # A fresh spec per layer: layer_specs entries must not alias each other, - # matching the base class `_deepcopy_layer_spec` contract. - transformer_layer_spec.layer_specs[i] = self._get_mamba_layer_spec() - continue - submodules = layer_spec.submodules - if ch == '*': - # Attention-only layer: no FFN. Its pre-norm is fused into linear_qkv. - # `mlp_bda` must go too: it unpacks its input as (output, bias), which an - # IdentityOp mlp does not produce. - submodules.mlp = IdentityOp - submodules.pre_mlp_layernorm = IdentityOp - submodules.mlp_bda = IdentityFuncOp - else: - # 'E'/'-': FFN-only layer. Its pre-norm is fused into the MLP's fc1 - # (or pre_mlp_layernorm for MoE), so drop attention and its norm. - submodules.self_attention = IdentityOp - submodules.input_layernorm = IdentityOp - submodules.self_attn_bda = IdentityFuncOp - return transformer_layer_spec - - def _get_mamba_layer_spec(self): - """Build a MambaLayer spec for Mamba2 SSM layers. - - The Bridge expects a standalone pre-norm (`norm`) and a mixer with - plain `in_proj`/`out_proj`, so TENorm + TEColumnParallelLinear are used - instead of the fused TELayerNormColumnParallelLinear variant. + 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. """ - try: - 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.ssm.mamba_layer import MambaLayerSubmodules - from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules - from megatron.core.transformer.spec_utils import ModuleSpec - except ImportError as e: - raise ImportError('NemotronHLoader requires a megatron-core version with Mamba2 SSM support ' - '(megatron.core.ssm.mamba_layer / mamba_mixer).') from e - if self._mamba_layer_cls is None: - self._mamba_layer_cls = _build_mamba_layer_cls() - return ModuleSpec( - module=self._mamba_layer_cls, + 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( @@ -367,12 +389,29 @@ def _get_mamba_layer_spec(self): 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, - )) +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/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/_nemotron_fwd.py b/tests/_nemotron_fwd.py deleted file mode 100644 index 7680d26..0000000 --- a/tests/_nemotron_fwd.py +++ /dev/null @@ -1,163 +0,0 @@ -"""Forward-consistency check: HF NemotronHForCausalLM vs converted MCore GPTModel. - -The round-trip tests only prove the weight transport is reversible; they say nothing -about whether the MCore model *computes* the same thing. This script loads the real -checkpoint into both stacks and compares logits. - -Run (single H20, ~62GB model so keep TP=1 and expect high memory): - torchrun --nproc_per_node=1 tests/_nemotron_fwd.py -Optionally limit layers for a cheap smoke run: - NUM_LAYERS=8 torchrun --nproc_per_node=1 tests/_nemotron_fwd.py -""" -import os - -import torch -import torch.distributed as dist -from megatron.core import parallel_state as mpu -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed - -MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' - 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') - - -def _restore_clobbered_weights(hf_model): - """Undo two HF `_init_weights` bugs that discard trained weights. - - This checkpoint's remote code overwrites weights AFTER they are loaded: - - * `dt_bias`: `module.dt_bias.copy_(inv_dt)` with a random `inv_dt`; the - `_no_reinit = True` marker is set afterwards and never checked. - * `out_proj.weight`: because `rescale_prenorm_residual=True`, it runs - `kaiming_uniform_` then divides by sqrt(num_layers) -- unconditionally. - - Without restoring both, HF is running partly random weights and is useless - as a numerical reference. - """ - import json - - from safetensors import safe_open - index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) - weight_map = index['weight_map'] - restored = 0 - for name, param in hf_model.named_parameters(): - key = name.replace('model.backbone', 'backbone') - if not (key.endswith('mixer.dt_bias') or key.endswith('mixer.out_proj.weight')): - continue - if key not in weight_map: - continue - with safe_open(os.path.join(MODEL_PATH, weight_map[key]), framework='pt') as f: - tensor = f.get_tensor(key) - with torch.no_grad(): - param.copy_(tensor.to(param.device, param.dtype)) - restored += 1 - print(f'RES restored {restored} clobbered weights (HF _init_weights bug)') - - -def main(): - seq_len = int(os.environ.get('SEQ_LEN', 16)) - num_layers_override = os.environ.get('NUM_LAYERS') - - dist.init_process_group('nccl') - torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) - mpu.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(1234) - - from transformers import AutoConfig, AutoModelForCausalLM - - import mcore_bridge.model.gpts # noqa: F401 - from mcore_bridge.config.model_config import ModelConfig - from mcore_bridge.config.parser import hf_to_mcore_config - from mcore_bridge.model.register import get_mcore_model - - hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) - if num_layers_override: - n = int(num_layers_override) - hf_config.num_hidden_layers = n - hf_config.hybrid_override_pattern = hf_config.hybrid_override_pattern[:n] - # MTP is not supported by the bridge; disable so both sides match. - hf_config.num_nextn_predict_layers = 0 - - torch.manual_seed(0) - input_ids = torch.randint(0, hf_config.vocab_size, (1, seq_len), device='cuda') - - # ---- HF reference ---- - hf_model = AutoModelForCausalLM.from_pretrained( - MODEL_PATH, config=hf_config, torch_dtype=torch.bfloat16, - trust_remote_code=True).cuda().eval() - # HF BUG WORKAROUND: this checkpoint's `_init_weights` unconditionally does - # `module.dt_bias.copy_(inv_dt)` with a *random* inv_dt, and only sets - # `_no_reinit = True` afterwards without ever checking it. So the trained - # dt_bias from safetensors is discarded and HF runs with random values. - # Restore it from the checkpoint so the reference is actually the trained model. - _restore_clobbered_weights(hf_model) - with torch.no_grad(): - hf_logits = hf_model(input_ids).logits.float() - del hf_model - torch.cuda.empty_cache() - print(f'RES hf_logits {tuple(hf_logits.shape)} ' - f'mean={hf_logits.mean():.5f} std={hf_logits.std():.5f}') - - # ---- MCore under test ---- - overrides = hf_to_mcore_config(hf_config) - overrides.update(params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) - # The cuDNN fused-attention backend fails to load its sublibrary in this container; - # flash is equivalent for correctness purposes here. - backend = os.environ.get('ATTN_BACKEND', 'flash') - if backend: - from megatron.core.transformer.enums import AttnBackend - overrides['attention_backend'] = getattr(AttnBackend, backend) - cfg = ModelConfig(**overrides) - models = get_mcore_model(cfg) - cfg.bridge.load_weights(models, MODEL_PATH) - mg_model = models[0].cuda().eval() - - position_ids = torch.arange(seq_len, device='cuda').unsqueeze(0) - attention_mask = torch.tril( - torch.ones((1, 1, seq_len, seq_len), device='cuda', dtype=torch.bool)).logical_not() - with torch.no_grad(): - mg_logits = mg_model( - input_ids=input_ids, - position_ids=position_ids, - attention_mask=attention_mask, - ).float() - if mg_logits.shape[0] == seq_len: # [s, b, h] -> [b, s, h] - mg_logits = mg_logits.transpose(0, 1) - mg_logits = mg_logits[..., :hf_logits.shape[-1]] - print(f'RES mg_logits {tuple(mg_logits.shape)} ' - f'mean={mg_logits.mean():.5f} std={mg_logits.std():.5f}') - - diff = (hf_logits - mg_logits).abs() - rel = diff.max() / hf_logits.abs().max() - hf_top = hf_logits.argmax(-1) - mg_top = mg_logits.argmax(-1) - agree = (hf_top == mg_top).float().mean() - print(f'RES max_abs_diff={diff.max():.6f} mean_abs_diff={diff.mean():.6f} ' - f'rel={rel:.6f} argmax_agree={agree:.4f}') - print(f'RES hf_top={hf_top.flatten()[:8].tolist()}') - print(f'RES mg_top={mg_top.flatten()[:8].tolist()}') - # Where argmax disagrees, check whether it's a near-tie (bf16 noise flipping the - # order of two nearly-equal logits) rather than a real behavioural difference. - mism = (hf_top != mg_top).nonzero() - for pos in mism[:5]: - b, t = pos.tolist() - hv, mv = hf_logits[b, t], mg_logits[b, t] - top2 = hv.topk(2).values - print(f'RES tie@t={t} hf_top1-top2_gap={float(top2[0] - top2[1]):.5f} ' - f'hf@hf_top={float(hv[hf_top[b, t]]):.5f} hf@mg_top={float(hv[mg_top[b, t]]):.5f} ' - f'delta={float(hv[hf_top[b, t]] - hv[mg_top[b, t]]):.5f}') - # Rank correlation is the robust check: argmax can flip on ties. - k = 20 - hf_set = hf_logits.topk(k, -1).indices - mg_set = mg_logits.topk(k, -1).indices - overlap = sum(len(set(a.tolist()) & set(b.tolist())) / k - for a, b in zip(hf_set.reshape(-1, k), mg_set.reshape(-1, k))) - overlap /= hf_set.reshape(-1, k).shape[0] - print(f'RES top{k}_overlap={overlap:.4f}') - if agree >= 0.9 and rel < 0.05 and overlap > 0.95: - print('RES FORWARD CONSISTENCY PASS (bf16-level)') - else: - print('RES FORWARD CONSISTENCY FAIL') - - -if __name__ == '__main__': - main() diff --git a/tests/_nemotron_layerdiff.py b/tests/_nemotron_layerdiff.py deleted file mode 100644 index 897ccca..0000000 --- a/tests/_nemotron_layerdiff.py +++ /dev/null @@ -1,169 +0,0 @@ -"""Layer-by-layer hidden-state diff between HF and MCore (single process, one load each). - -Bisects where the two stacks diverge instead of only comparing final logits. -Both models are built with the same small layer count so this fits comfortably in memory. - - NUM_LAYERS=4 SEQ_LEN=8 torchrun --nproc_per_node=1 tests/_nemotron_layerdiff.py -""" -import os - -import torch -import torch.distributed as dist -from megatron.core import parallel_state as mpu -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed - -MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' - 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') - - -def _restore_clobbered_weights(hf_model): - """Undo two HF `_init_weights` bugs that discard trained weights. - - This checkpoint's remote code overwrites weights AFTER they are loaded: - - * `dt_bias`: `module.dt_bias.copy_(inv_dt)` with a random `inv_dt`; the - `_no_reinit = True` marker is set afterwards and never checked. - * `out_proj.weight`: because `rescale_prenorm_residual=True`, it runs - `kaiming_uniform_` then divides by sqrt(num_layers) -- unconditionally. - - Without restoring both, HF is running partly random weights and is useless - as a numerical reference. - """ - import json - - from safetensors import safe_open - index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) - weight_map = index['weight_map'] - restored = 0 - for name, param in hf_model.named_parameters(): - key = name.replace('model.backbone', 'backbone') - if not (key.endswith('mixer.dt_bias') or key.endswith('mixer.out_proj.weight')): - continue - if key not in weight_map: - continue - with safe_open(os.path.join(MODEL_PATH, weight_map[key]), framework='pt') as f: - tensor = f.get_tensor(key) - with torch.no_grad(): - param.copy_(tensor.to(param.device, param.dtype)) - restored += 1 - print(f'RES restored {restored} clobbered weights (HF _init_weights bug)') - - -def main(): - seq_len = int(os.environ.get('SEQ_LEN', 8)) - n_layers = int(os.environ.get('NUM_LAYERS', 4)) - - dist.init_process_group('nccl') - torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) - mpu.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(1234) - - from transformers import AutoConfig, AutoModelForCausalLM - - import mcore_bridge.model.gpts # noqa: F401 - from mcore_bridge.config.model_config import ModelConfig - from mcore_bridge.config.parser import hf_to_mcore_config - from mcore_bridge.model.register import get_mcore_model - - hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) - hf_config.num_hidden_layers = n_layers - hf_config.hybrid_override_pattern = hf_config.hybrid_override_pattern[:n_layers] - hf_config.num_nextn_predict_layers = 0 - pattern = hf_config.hybrid_override_pattern - print(f'RES pattern={pattern}') - - torch.manual_seed(0) - input_ids = torch.randint(0, hf_config.vocab_size, (1, seq_len), device='cuda') - - # ---- HF: capture every backbone layer output ---- - hf_model = AutoModelForCausalLM.from_pretrained( - MODEL_PATH, config=hf_config, torch_dtype=torch.bfloat16, - trust_remote_code=True).cuda().eval() - _restore_clobbered_weights(hf_model) # HF _init_weights clobbers dt_bias with random inv_dt - hf_acts = {} - - hf_ins = {} - - def mk_hook(idx): - def hook(_mod, inp, out): - if inp: - hf_ins[idx] = inp[0].detach().float() - hf_acts[idx] = (out[0] if isinstance(out, tuple) else out).detach().float() - return hook - - for i, layer in enumerate(hf_model.backbone.layers): - layer.register_forward_hook(mk_hook(i)) - hf_emb = {} - hf_model.backbone.embeddings.register_forward_hook( - lambda m, i, o: hf_emb.__setitem__(0, o.detach().float())) - with torch.no_grad(): - hf_model(input_ids) - del hf_model - torch.cuda.empty_cache() - - # ---- MCore: same capture ---- - overrides = hf_to_mcore_config(hf_config) - overrides.update(params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) - from megatron.core.transformer.enums import AttnBackend - overrides['attention_backend'] = AttnBackend.flash - cfg = ModelConfig(**overrides) - models = get_mcore_model(cfg) - cfg.bridge.load_weights(models, MODEL_PATH) - mg_model = models[0].cuda().eval() - - mg_acts = {} - - mg_ins = {} - - def mk_hook_mg(idx): - def hook(_mod, inp, out): - if inp: - mg_ins[idx] = inp[0].detach().float() - t = out[0] if isinstance(out, tuple) else out - mg_acts[idx] = t.detach().float() - return hook - - for i, layer in enumerate(mg_model.decoder.layers): - layer.register_forward_hook(mk_hook_mg(i)) - mg_emb = {} - mg_model.embedding.register_forward_hook( - lambda m, i, o: mg_emb.__setitem__(0, o.detach().float())) - - position_ids = torch.arange(seq_len, device='cuda').unsqueeze(0) - attention_mask = torch.tril( - torch.ones((1, 1, seq_len, seq_len), device='cuda', dtype=torch.bool)).logical_not() - with torch.no_grad(): - mg_model(input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask) - - def norm(t): - """HF is [b, s, h]; MCore is [s, b, h]. Normalize to [s, h].""" - if t.dim() == 3: - if t.shape[0] == 1 and t.shape[1] == seq_len: - return t[0] - if t.shape[1] == 1 and t.shape[0] == seq_len: - return t[:, 0] - return t.reshape(seq_len, -1) - - if 0 in hf_emb and 0 in mg_emb: - a, b = norm(hf_emb[0]), norm(mg_emb[0]) - print(f'RES embedding max_abs_diff={(a - b).abs().max():.6f}') - for i in range(n_layers): - if i not in hf_acts or i not in mg_acts: - print(f'RES layer{i} MISSING hf={i in hf_acts} mg={i in mg_acts}') - continue - a, b = norm(hf_acts[i]), norm(mg_acts[i]) - d = (a - b).abs() - scale = a.abs().max().clamp(min=1e-6) - if i in hf_ins and i in mg_ins: - ia, ib = norm(hf_ins[i]), norm(mg_ins[i]) - di = (ia - ib).abs() - print(f'RES layer{i} IN max_abs_diff={di.max():.6f} ' - f'hf_std={ia.std():.5f} mg_std={ib.std():.5f}') - print(f'RES layer{i} type={pattern[i]!r} hf_shape={tuple(hf_acts[i].shape)} ' - f'mg_shape={tuple(mg_acts[i].shape)} ' - f'hf_std={a.std():.5f} mg_std={b.std():.5f} ' - f'max_abs_diff={d.max():.6f} mean={d.mean():.6f} rel={d.max() / scale:.6f}') - - -if __name__ == '__main__': - main() diff --git a/tests/_nemotron_mixerdiff.py b/tests/_nemotron_mixerdiff.py deleted file mode 100644 index a3c36fe..0000000 --- a/tests/_nemotron_mixerdiff.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Isolate the Mamba2 mixer: run HF's and mcore's on identical weights + input. - -Loads only one Mamba layer's weights out of the checkpoint, so this is cheap. - torchrun --nproc_per_node=1 tests/_nemotron_mixerdiff.py -""" -import json -import os - -import torch -import torch.distributed as dist -from megatron.core import parallel_state as mpu -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from safetensors import safe_open - -MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' - 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') - - -def load_layer0_mixer(): - index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) - weight_map = index['weight_map'] - out = {} - for key, shard in weight_map.items(): - if key.startswith('backbone.layers.0.mixer.'): - with safe_open(os.path.join(MODEL_PATH, shard), framework='pt') as f: - out[key[len('backbone.layers.0.mixer.'):]] = f.get_tensor(key).cuda() - return out - - -def main(): - seq_len = int(os.environ.get('SEQ_LEN', 8)) - dist.init_process_group('nccl') - torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) - mpu.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(1234) - - from transformers import AutoConfig - hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) - w = load_layer0_mixer() - print('RES loaded mixer keys:', sorted(w)) - - torch.manual_seed(0) - # HF wants [b, s, h]; mcore wants [s, b, h]. - x_bsh = torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.bfloat16, device='cuda') - - # ---- HF mixer ---- - # Go through transformers' dynamic-module loader so the file's relative imports resolve. - from transformers.dynamic_module_utils import get_class_from_dynamic_module - hf_mixer_cls = get_class_from_dynamic_module( - 'modeling_nemotron_h.NemotronHMamba2Mixer', MODEL_PATH) - hf_mixer = hf_mixer_cls(hf_config, layer_idx=0).cuda().to(torch.bfloat16).eval() - missing, unexpected = hf_mixer.load_state_dict(w, strict=False) - print(f'RES hf_mixer missing={list(missing)} unexpected={list(unexpected)}') - with torch.no_grad(): - hf_out = hf_mixer(x_bsh).float() - print(f'RES hf_out {tuple(hf_out.shape)} mean={hf_out.mean():.6f} std={hf_out.std():.6f}') - - # ---- mcore mixer ---- - import mcore_bridge.model.gpts # noqa: F401 - from mcore_bridge.config.model_config import ModelConfig - from mcore_bridge.config.parser import hf_to_mcore_config - overrides = hf_to_mcore_config(hf_config) - overrides.update(num_layers=1, hybrid_layer_pattern='M', moe_layer_freq='[0]', - params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) - cfg = ModelConfig(**overrides) - - from megatron.core.extensions.transformer_engine import (TEColumnParallelLinear, - TERowParallelLinear) - from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules - mg_mixer = MambaMixer( - cfg, - MambaMixerSubmodules(in_proj=TEColumnParallelLinear, out_proj=TERowParallelLinear), - d_model=cfg.hidden_size, - layer_number=1, - pg_collection=_pg(), - ).cuda().to(torch.bfloat16).eval() - - mg_sd = { - 'in_proj.weight': w['in_proj.weight'], - 'conv1d_weight': w['conv1d.weight'], - 'conv1d_bias': w['conv1d.bias'], - 'A_log': w['A_log'], - 'D': w['D'], - 'dt_bias': w['dt_bias'], - 'norm.weight': w['norm.weight'], - 'out_proj.weight': w['out_proj.weight'], - } - incompat = mg_mixer.load_state_dict(mg_sd, strict=False) - print(f'RES mg_mixer missing={list(incompat.missing_keys)} ' - f'unexpected={list(incompat.unexpected_keys)}') - - x_sbh = x_bsh.transpose(0, 1).contiguous() - with torch.no_grad(): - mg_out, mg_bias = mg_mixer(x_sbh) - mg_out = mg_out.float().transpose(0, 1) # -> [b, s, h] - if mg_bias is not None: - mg_out = mg_out + mg_bias.float() - print(f'RES mg_out {tuple(mg_out.shape)} mean={mg_out.mean():.6f} std={mg_out.std():.6f}') - - d = (hf_out - mg_out).abs() - print(f'RES mixer max_abs_diff={d.max():.6f} mean_abs_diff={d.mean():.6f} ' - f'rel={d.max() / hf_out.abs().max():.6f}') - - # ---- Now the FULL block: norm + mixer + residual ---- - # Load layer0's outer norm too. - import json as _json - idx = _json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) - nk = 'backbone.layers.0.norm.weight' - with safe_open(os.path.join(MODEL_PATH, idx['weight_map'][nk]), framework='pt') as f: - outer_norm_w = f.get_tensor(nk).cuda() - - hf_block_cls = get_class_from_dynamic_module('modeling_nemotron_h.NemotronHBlock', MODEL_PATH) - hf_block = hf_block_cls(hf_config, layer_idx=0).cuda().to(torch.bfloat16).eval() - bsd = {f'mixer.{k}': v for k, v in w.items()} - bsd['norm.weight'] = outer_norm_w - inc = hf_block.load_state_dict(bsd, strict=False) - print(f'RES hf_block missing={list(inc.missing_keys)} unexpected={list(inc.unexpected_keys)}') - with torch.no_grad(): - hf_blk = hf_block(x_bsh) - hf_blk = (hf_blk[0] if isinstance(hf_blk, tuple) else hf_blk).float() - - from mcore_bridge.model.gpts.nemotron_h import _build_mamba_layer_cls, NemotronHLoader - loader = NemotronHLoader(cfg) - mspec = loader._get_mamba_layer_spec() - from megatron.core.transformer.spec_utils import build_module - mg_layer = build_module(mspec, config=cfg, layer_number=1, - pg_collection=_pg(), vp_stage=None).cuda().to(torch.bfloat16).eval() - lsd = dict(mg_sd) - lsd = {f'mixer.{k}': v for k, v in mg_sd.items()} - lsd['norm.weight'] = outer_norm_w - inc2 = mg_layer.load_state_dict(lsd, strict=False) - print(f'RES mg_layer missing={[k for k in inc2.missing_keys if "_extra_state" not in k]} ' - f'unexpected={list(inc2.unexpected_keys)}') - with torch.no_grad(): - out = mg_layer(x_sbh) - mg_blk = (out[0] if isinstance(out, tuple) else out).float().transpose(0, 1) - d2 = (hf_blk - mg_blk).abs() - print(f'RES BLOCK max_abs_diff={d2.max():.6f} mean_abs_diff={d2.mean():.6f} ' - f'rel={d2.max() / hf_blk.abs().max():.6f}') - - -def _pg(): - from megatron.core.process_groups_config import ProcessGroupCollection - return ProcessGroupCollection.use_mpu_process_groups() - - -if __name__ == '__main__': - main() diff --git a/tests/_nemotron_moediff.py b/tests/_nemotron_moediff.py deleted file mode 100644 index 9044b9d..0000000 --- a/tests/_nemotron_moediff.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Compare a single 'E' (MoE) block: HF NemotronHBlock vs mcore TransformerLayer. - -The Mamba block already matches in isolation (see _nemotron_mixerdiff.py), so this -checks the other layer type. Loads only layer 1's weights. - - torchrun --nproc_per_node=1 tests/_nemotron_moediff.py -""" -import json -import os - -import torch -import torch.distributed as dist -from megatron.core import parallel_state as mpu -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from safetensors import safe_open - -MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' - 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') -LAYER = int(os.environ.get('LAYER', 1)) # layer 1 is 'E' in MEMEM*... - - -def load_layer(prefix): - index = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) - out = {} - for key, shard in index['weight_map'].items(): - if key.startswith(prefix): - with safe_open(os.path.join(MODEL_PATH, shard), framework='pt') as f: - out[key[len(prefix):]] = f.get_tensor(key).cuda() - return out - - -def _pg(): - from megatron.core.process_groups_config import ProcessGroupCollection - return ProcessGroupCollection.use_mpu_process_groups() - - -def main(): - seq_len = int(os.environ.get('SEQ_LEN', 8)) - dist.init_process_group('nccl') - torch.cuda.set_device(int(os.environ.get('LOCAL_RANK', 0))) - mpu.initialize_model_parallel(1, 1) - model_parallel_cuda_manual_seed(1234) - - from transformers import AutoConfig - from transformers.dynamic_module_utils import get_class_from_dynamic_module - - hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) - # NEMOTRONH_ATTENTION_CLASSES is keyed by _attn_implementation; standalone block - # construction skips the usual model-level default, so set it explicitly. - hf_config._attn_implementation = os.environ.get('HF_ATTN', 'eager') - pattern = hf_config.hybrid_override_pattern - print(f'RES layer{LAYER} type={pattern[LAYER]!r}') - - w = load_layer(f'backbone.layers.{LAYER}.') - print(f'RES n_weights={len(w)}') - - torch.manual_seed(0) - if os.environ.get('REAL_EMB'): - # Use the checkpoint's real embedding rows: their magnitude is far from N(0,1), - # and Mamba's conv/SSM path is scale sensitive. - ek = 'backbone.embeddings.weight' - _idx = json.load(open(os.path.join(MODEL_PATH, 'model.safetensors.index.json'))) - with safe_open(os.path.join(MODEL_PATH, _idx['weight_map'][ek]), framework='pt') as f: - emb = f.get_tensor(ek).cuda() - ids = torch.randint(0, emb.shape[0], (1, seq_len), device='cuda') - x_bsh = emb[ids].to(torch.bfloat16) - print(f'RES using REAL embedding, std={x_bsh.float().std():.5f}') - else: - x_bsh = torch.randn(1, seq_len, hf_config.hidden_size, dtype=torch.bfloat16, device='cuda') - print(f'RES using randn input, std={x_bsh.float().std():.5f}') - - # ---- HF block ---- - hf_block_cls = get_class_from_dynamic_module('modeling_nemotron_h.NemotronHBlock', MODEL_PATH) - hf_block = hf_block_cls(hf_config, layer_idx=LAYER).cuda().to(torch.bfloat16).eval() - inc = hf_block.load_state_dict(w, strict=False) - print(f'RES hf missing={list(inc.missing_keys)[:5]} unexpected={list(inc.unexpected_keys)[:5]}') - with torch.no_grad(): - hf_out = hf_block(x_bsh) - hf_out = (hf_out[0] if isinstance(hf_out, tuple) else hf_out).float() - print(f'RES hf_out mean={hf_out.mean():.6f} std={hf_out.std():.6f}') - - # ---- MCore layer via the real Loader spec ---- - import mcore_bridge.model.gpts # noqa: F401 - from mcore_bridge.config.model_config import ModelConfig - from mcore_bridge.config.parser import hf_to_mcore_config - from mcore_bridge.model.gpts.nemotron_h import NemotronHLoader - - ch = pattern[LAYER] - overrides = hf_to_mcore_config(hf_config) - overrides.update(num_layers=1, hybrid_layer_pattern=ch, - moe_layer_freq='[1]' if ch == 'E' else '[0]', - params_dtype=torch.bfloat16, bf16=True, mtp_num_layers=None) - from megatron.core.transformer.enums import AttnBackend - overrides['attention_backend'] = AttnBackend.flash - cfg = ModelConfig(**overrides) - - loader = NemotronHLoader(cfg) - spec = loader.get_transformer_layer_spec() - from megatron.core.transformer.spec_utils import build_module - mg_layer = build_module(spec.layer_specs[0], config=cfg, layer_number=1, - pg_collection=_pg(), vp_stage=None).cuda().to(torch.bfloat16).eval() - print('RES mg params:', sorted(n for n, _ in mg_layer.named_parameters())[:12]) - - # Drive the real Bridge so the mapping under test is exercised. - class Lazy: - def __init__(self, t): - self.t = t - - def load(self): - return self.t - - sd = {f'backbone.layers.0.{k}': Lazy(v) for k, v in w.items()} - cfg.bridge._set_layer_state(mg_layer, sd, 'backbone.layers.', 0, True) - - x_sbh = x_bsh.transpose(0, 1).contiguous() - attn_mask = torch.tril( - torch.ones((1, 1, seq_len, seq_len), device='cuda', dtype=torch.bool)).logical_not() - with torch.no_grad(): - out = mg_layer(x_sbh, attention_mask=attn_mask) - mg_out = (out[0] if isinstance(out, tuple) else out).float().transpose(0, 1) - print(f'RES mg_out mean={mg_out.mean():.6f} std={mg_out.std():.6f}') - - d = (hf_out - mg_out).abs() - print(f'RES BLOCK({ch}) max_abs_diff={d.max():.6f} mean_abs_diff={d.mean():.6f} ' - f'rel={d.max() / hf_out.abs().max():.6f}') - - -if __name__ == '__main__': - main() diff --git a/tests/_nemotron_rt.py b/tests/_nemotron_rt.py deleted file mode 100644 index 3d7115c..0000000 --- a/tests/_nemotron_rt.py +++ /dev/null @@ -1,173 +0,0 @@ -"""Nemotron-H round-trip harness: HF -> MCore -> HF must be bit-exact. - -Parallel layout is taken from env vars so the same script covers TP/PP/EP: - TP, PP, EP, ETP (default 1), PATTERN (default 'ME*') - -Run e.g.: - EP=2 PATTERN='ME*' torchrun --nproc_per_node=2 tests/_nemotron_rt.py -""" -import os - -import torch -import torch.distributed as dist -from megatron.core import parallel_state as mpu -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed - -MODEL_PATH = ('/root/.cache/modelscope/hub/models/nv-community/' - 'EA-NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16-07202026') -NUM_EXPERTS = 4 - - -class Lazy: - """Mimics SafetensorLazyLoader's handle: Bridge calls `.load()`.""" - - def __init__(self, tensor): - self.tensor = tensor - - def load(self): - return self.tensor - - -def build_hf_state_dict(cfg, mixer, pattern, rand): - hidden = cfg.hidden_size - sd = { - 'backbone.embeddings.weight': Lazy(rand(256, hidden)), - 'lm_head.weight': Lazy(rand(256, hidden)), - 'backbone.norm_f.weight': Lazy(rand(hidden)), - } - for i, ch in enumerate(pattern): - p = f'backbone.layers.{i}.' - sd[f'{p}norm.weight'] = Lazy(rand(hidden)) - if ch == 'M': - conv_dim = mixer.d_inner + 2 * mixer.ngroups * mixer.d_state - in_dim = mixer.d_inner * 2 + 2 * mixer.ngroups * mixer.d_state + mixer.nheads - sd[f'{p}mixer.in_proj.weight'] = Lazy(rand(in_dim, hidden)) - sd[f'{p}mixer.conv1d.weight'] = Lazy(rand(conv_dim, 1, 4)) - sd[f'{p}mixer.conv1d.bias'] = Lazy(rand(conv_dim)) - # A_log/D are fp32 in mcore; dt_bias is bf16. - sd[f'{p}mixer.A_log'] = Lazy(rand(mixer.nheads, dtype=torch.float32)) - sd[f'{p}mixer.D'] = Lazy(rand(mixer.nheads, dtype=torch.float32)) - sd[f'{p}mixer.dt_bias'] = Lazy(rand(mixer.nheads)) - sd[f'{p}mixer.norm.weight'] = Lazy(rand(mixer.d_inner)) - sd[f'{p}mixer.out_proj.weight'] = Lazy(rand(hidden, mixer.d_inner)) - elif ch == 'E': - sd[f'{p}mixer.gate.weight'] = Lazy(rand(NUM_EXPERTS, hidden)) - sd[f'{p}mixer.gate.e_score_correction_bias'] = Lazy( - rand(NUM_EXPERTS, dtype=torch.float32)) - moe_ffn = cfg.moe_ffn_hidden_size - for e in range(NUM_EXPERTS): - sd[f'{p}mixer.experts.{e}.up_proj.weight'] = Lazy(rand(moe_ffn, hidden)) - sd[f'{p}mixer.experts.{e}.down_proj.weight'] = Lazy(rand(hidden, moe_ffn)) - shared = cfg.moe_shared_expert_intermediate_size - sd[f'{p}mixer.shared_experts.up_proj.weight'] = Lazy(rand(shared, hidden)) - sd[f'{p}mixer.shared_experts.down_proj.weight'] = Lazy(rand(hidden, shared)) - elif ch == '*': - head_dim = cfg.kv_channels - q = cfg.num_attention_heads * head_dim - kv = cfg.num_query_groups * head_dim - sd[f'{p}mixer.q_proj.weight'] = Lazy(rand(q, hidden)) - sd[f'{p}mixer.k_proj.weight'] = Lazy(rand(kv, hidden)) - sd[f'{p}mixer.v_proj.weight'] = Lazy(rand(kv, hidden)) - sd[f'{p}mixer.o_proj.weight'] = Lazy(rand(hidden, q)) - else: - raise ValueError(f'unsupported pattern char {ch!r}') - return sd - - -def main(): - tp = int(os.environ.get('TP', 1)) - pp = int(os.environ.get('PP', 1)) - ep = int(os.environ.get('EP', 1)) - etp = int(os.environ.get('ETP', tp)) - pattern = os.environ.get('PATTERN', 'ME*') - - dist.init_process_group('nccl') - local_rank = int(os.environ.get('LOCAL_RANK', 0)) - torch.cuda.set_device(local_rank) - mpu.initialize_model_parallel( - tensor_model_parallel_size=tp, - pipeline_model_parallel_size=pp, - expert_model_parallel_size=ep, - expert_tensor_parallel_size=etp, - ) - model_parallel_cuda_manual_seed(1234) # MambaMixer uses get_cuda_rng_tracker().fork() - - from transformers import AutoConfig - - import mcore_bridge.model.gpts # noqa: F401 (triggers registration) - from mcore_bridge.config.model_config import ModelConfig - from mcore_bridge.config.parser import hf_to_mcore_config - from mcore_bridge.model.register import get_mcore_model - - hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) - overrides = hf_to_mcore_config(hf_config) - overrides.update( - num_layers=len(pattern), - hybrid_layer_pattern=pattern, - moe_layer_freq='[' + ','.join('1' if c == 'E' else '0' for c in pattern) + ']', - num_moe_experts=NUM_EXPERTS, - padded_vocab_size=256, - tensor_model_parallel_size=tp, - pipeline_model_parallel_size=pp, - expert_model_parallel_size=ep, - expert_tensor_parallel_size=etp, - params_dtype=torch.bfloat16, - bf16=True, - ) - cfg = ModelConfig(**overrides) - - models = get_mcore_model(cfg) - bridge = cfg.bridge - - # Any rank holding a Mamba layer can report the mixer dims; they are TP-invariant - # on the HF side, so derive them from config instead of the (possibly absent) module. - class _Dims: - nheads = cfg.mamba_num_heads - d_inner = cfg.mamba_num_heads * cfg.mamba_head_dim - ngroups = cfg.mamba_num_groups - d_state = cfg.mamba_state_dim - - gen = torch.Generator(device='cuda').manual_seed(4321) # identical data on every rank - - def rand(*shape, dtype=torch.bfloat16): - return torch.randn(*shape, generator=gen, dtype=dtype, device='cuda') - - sd = build_hf_state_dict(cfg, _Dims, pattern, rand) - original = {k: v.load().clone() for k, v in sd.items()} - - for model in models: - list(bridge._convert([model], sd, '', True, 'Loading: ')) - - exported = dict(bridge.export_weights(models, target_device='cuda')) - - if local_rank != 0: - dist.barrier() - return - - tag = f'TP{tp}/PP{pp}/EP{ep} pattern={pattern}' - missing = sorted(set(original) - set(exported)) - extra = sorted(set(exported) - set(original)) - bad = [] - for key, want in original.items(): - got = exported.get(key) - if got is None: - continue - if tuple(got.shape) != tuple(want.shape): - bad.append((key, 'shape', tuple(want.shape), tuple(got.shape))) - elif not torch.equal(got.to(want.dtype).cpu(), want.cpu()): - delta = (got.to(want.dtype).cpu().float() - want.cpu().float()).abs().max() - bad.append((key, 'value', float(delta))) - - print(f'RES [{tag}] keys={len(exported)} missing={missing} extra={extra} ' - f'mismatches={len(bad)}') - for item in bad[:10]: - print(' ', item) - if not bad and not missing and not extra: - print(f'RES [{tag}] ROUNDTRIP EXACT PASS') - else: - print(f'RES [{tag}] ROUNDTRIP FAILED') - dist.barrier() - - -if __name__ == '__main__': - main() 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() From 846dbf1642d0b692f65349bd5c8d52a5dda10006 Mon Sep 17 00:00:00 2001 From: hjh0119 Date: Tue, 11 Aug 2026 21:41:08 +0800 Subject: [PATCH 3/3] fix save --- src/mcore_bridge/model/gpts/nemotron_h.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/mcore_bridge/model/gpts/nemotron_h.py b/src/mcore_bridge/model/gpts/nemotron_h.py index bf500f3..744602e 100644 --- a/src/mcore_bridge/model/gpts/nemotron_h.py +++ b/src/mcore_bridge/model/gpts/nemotron_h.py @@ -185,9 +185,10 @@ def _merge_packed_dim0(self, gathered, block_sizes): 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` may be None on a PP rank that does not own this layer; the TP all-gather - still has to run collectively, so pass None straight through to `_all_gather_tp` - (which tolerates None) and only touch it when this rank actually holds the weight. + `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: @@ -198,9 +199,9 @@ def _set_mamba_packed(self, mg_param, hf_state_dict, hf_key, block_sizes, to_mco 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) - if gathered is None: - return - merged = self._merge_packed_dim0(gathered, block_sizes) + 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: