From 54eaceb6366c53d335e48c7674f68c9f78d9704e Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 12 Aug 2026 06:35:45 -0700 Subject: [PATCH 1/2] Export quantized/co-trained MTP weights instead of copying BF16 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_get_mtp_state_dict` previously copied the MTP (multi-token prediction) head verbatim from the BF16 pretrained model (`# TODO Implement MTP export for quantized MTP`), because `_get_state_dict` only walks `model.decoder.layers` and never `model.mtp` — so the `key not in self._state_dict` guard was always true. Any quantization or co-training applied to the MTP head during QAD was silently discarded at export; the draft head shipped as the original BF16 weights. This walks the live MCore `model.mtp` module and applies the same quantization rules used for the base decoder, mirroring `_get_eagle_module_state_dict`. The MTP inner attention/MoE layers are structurally identical to backbone hybrid layers, so the base layer walker is reused with a restricted set of `mtp.*` naming rules aliased onto the standard rule keys (emitting `mtp.layers.{}.` HF keys). The old BF16-copy path is kept as `_copy_mtp_state_dict_from_pretrained`, used only when the live model has no `mtp` module. Adds the missing `mtp.*` inner-layer export rules to `nemotron_h_causal_lm_export` (attention qkv/o_proj/norm, MoE router/experts/shared_experts); the predictor projection rules (`mtp.enorm/hnorm/eh_proj/final_layernorm`) already existed. Round-trips with the `is_mtp` keys in `nemotron_h_causal_lm_import`. Co-Authored-By: Claude Opus 4.8 --- .../torch/export/plugins/mcore_nemotron.py | 19 ++++- .../torch/export/unified_export_megatron.py | 80 ++++++++++++++++++- 2 files changed, 94 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/plugins/mcore_nemotron.py b/modelopt/torch/export/plugins/mcore_nemotron.py index 24bd8144055..b657c8b54c1 100644 --- a/modelopt/torch/export/plugins/mcore_nemotron.py +++ b/modelopt/torch/export/plugins/mcore_nemotron.py @@ -155,9 +155,26 @@ # Grouped local experts (TEGroupedMLP: fused per-expert weights) "experts.linear_fc1": GroupedMLPSlicing("backbone.layers.{}.mixer.experts.{{}}.up_proj"), "experts.linear_fc2": GroupedMLPSlicing("backbone.layers.{}.mixer.experts.{{}}.down_proj"), - # MTP + # MTP predictor projections (outer MultiTokenPredictionLayer) "mtp.enorm": NameRemapping("mtp.layers.{}.enorm."), "mtp.hnorm": NameRemapping("mtp.layers.{}.hnorm."), "mtp.eh_proj": NameRemapping("mtp.layers.{}.eh_proj."), "mtp.final_layernorm": NameRemapping("mtp.layers.{}.final_layernorm."), + # MTP inner attention / MoE layers. Structurally identical to the backbone hybrid + # layers, so these mirror the base rules above with the `mtp.layers.{}` prefix; the + # `mtp.` namespace is aliased onto the standard rule keys by _get_mtp_state_dict. + "mtp.input_layernorm": NameRemapping("mtp.layers.{}.norm."), + "mtp.fused_norm": NameRemapping("mtp.layers.{}.norm.weight"), + "mtp.linear_qkv": QKVSlicing("mtp.layers.{}.mixer."), + "mtp.linear_proj": NameRemapping("mtp.layers.{}.mixer.o_proj."), + "mtp.pre_mlp_layernorm": NameRemapping("mtp.layers.{}.norm."), + "mtp.router": NameRemapping( + "mtp.layers.{}.mixer.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}} + ), + "mtp.shared_experts.linear_fc1": NameRemapping("mtp.layers.{}.mixer.shared_experts.up_proj."), + "mtp.shared_experts.linear_fc2": NameRemapping("mtp.layers.{}.mixer.shared_experts.down_proj."), + "mtp.local_experts.linear_fc1": NameRemapping("mtp.layers.{}.mixer.experts.{}.up_proj."), + "mtp.local_experts.linear_fc2": NameRemapping("mtp.layers.{}.mixer.experts.{}.down_proj."), + "mtp.experts.linear_fc1": GroupedMLPSlicing("mtp.layers.{}.mixer.experts.{{}}.up_proj"), + "mtp.experts.linear_fc2": GroupedMLPSlicing("mtp.layers.{}.mixer.experts.{{}}.down_proj"), } diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index ace6c2825df..81e4090d430 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -597,12 +597,84 @@ def _get_transformer_layer_state_dict(self, layer, layer_id): self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id) def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]: - """Export the MTP module. + """Export the MTP (Multi-Token Prediction) module. - Currently, we copy the BF16 MTP weights from the pretrained model if the pretrained model has MTP layers. + Walks the live MCore ``model.mtp`` module and applies the same quantization + rules used for the base decoder, so the exported draft head reflects the + actual (quantized / co-trained) weights. Falls back to copying the BF16 MTP + weights from the pretrained model only when the live model has no ``mtp`` + module (e.g. exporting a base-only checkpoint that grafts a pretrained head). + """ + mtp = getattr(self.model, "mtp", None) + if mtp is None or not hasattr(mtp, "layers") or len(mtp.layers) == 0: + return self._copy_mtp_state_dict_from_pretrained() + + # The MTP inner attention / MoE layers are structurally identical to the base + # decoder layers, so we reuse the base layer walker. We alias the ``mtp.*`` + # naming rules onto the standard rule keys for the duration of the walk so the + # walker emits ``mtp.layers.{}.`` HF keys instead of ``backbone.layers.{}.``. + # A *restricted* alias set is used on purpose: any base rule key the walker + # references but that has no ``mtp.`` variant is simply absent (and its call is + # guarded), rather than silently emitting a wrong ``backbone.`` prefix. + mtp_rules = { + key[len("mtp.") :]: rule for key, rule in self.rules.items() if key.startswith("mtp.") + } + + saved_rules = self.rules + saved_state_dict = self._state_dict + self.rules = mtp_rules + self._state_dict = OrderedDict() + try: + for mtp_layer in mtp.layers: + inner_layers = mtp_layer.mtp_model_layer.layers + first_id = inner_layers[0].layer_number - 1 + last_id = inner_layers[-1].layer_number - 1 + + # Outer predictor projections attach to the first inner HF index. + if "enorm" in self.rules: + self.rules["enorm"](mtp_layer.enorm, first_id) + if "hnorm" in self.rules: + self.rules["hnorm"](mtp_layer.hnorm, first_id) + if "eh_proj" in self.rules: + self.rules["eh_proj"](mtp_layer.eh_proj, first_id) + + # Inner hybrid stack (e.g. [attention, MoE] for the ``*E`` pattern) + # reuses the base decoder walker with the aliased mtp rules. + for inner in inner_layers: + hf_layer_id = inner.layer_number - 1 + if isinstance(inner, MambaLayer): + self._get_mamba_layer_state_dict(inner, hf_layer_id) + elif isinstance(inner, TransformerLayer): + self._get_transformer_layer_state_dict(inner, hf_layer_id) + else: + raise ValueError( + "Only TransformerLayer or MambaLayer are supported in the MTP block." + ) + + # The MTP block's own final layernorm attaches to the last inner HF index. + final_layernorm = getattr(mtp_layer, "final_layernorm", None) + if ( + "final_layernorm" in self.rules + and final_layernorm is not None + and not isinstance(final_layernorm, IdentityOp) + ): + self.rules["final_layernorm"](final_layernorm, last_id) + + mtp_state_dict = self._state_dict + finally: + self.rules = saved_rules + self._state_dict = saved_state_dict + + if len(mtp_state_dict) > 0: + print(f"Exported {len(mtp_state_dict)} MTP tensors from the live model") + return mtp_state_dict + + def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]: + """Fallback: copy the BF16 MTP weights from the pretrained model. + + Used only when the live model has no ``mtp`` module. This does not reflect any + quantization or co-training applied to the MTP head. """ - # TODO Implement MTP export for quantized MTP - # Hacky version for now: copy MTP weights from pretrained model mtp_state_dict = {} if not self._hf_pretrained_model_name: return mtp_state_dict From 89ef5962bdc68a14d94017964bd0e909d6ca3918 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 12 Aug 2026 20:17:37 -0700 Subject: [PATCH 2/2] Address review: reuse base layer rules via is_mtp instead of duplicate mtp rules Per review (@jenchen13): the MTP inner attention/MoE layers are structurally identical to the backbone hybrid layers, so instead of adding duplicate `mtp.*` inner-layer rules, thread an `is_mtp` flag through the base layer walker (`_get_transformer_layer_state_dict` / `_get_mamba_layer_state_dict`) and the remapping helpers. When set, the helper rewrites the target root (`backbone`/`model` -> `mtp`), exactly mirroring the importer. Only the predictor-specific keys (enorm/hnorm/eh_proj/final_layernorm) remain dedicated `mtp.*` rules. This keeps the import and export rule books symmetric and avoids rule duplication. Co-Authored-By: Claude Opus 4.8 --- .../torch/export/plugins/mcore_nemotron.py | 23 +-- .../torch/export/unified_export_megatron.py | 187 +++++++++++------- 2 files changed, 121 insertions(+), 89 deletions(-) diff --git a/modelopt/torch/export/plugins/mcore_nemotron.py b/modelopt/torch/export/plugins/mcore_nemotron.py index b657c8b54c1..ecd694e414b 100644 --- a/modelopt/torch/export/plugins/mcore_nemotron.py +++ b/modelopt/torch/export/plugins/mcore_nemotron.py @@ -155,26 +155,13 @@ # Grouped local experts (TEGroupedMLP: fused per-expert weights) "experts.linear_fc1": GroupedMLPSlicing("backbone.layers.{}.mixer.experts.{{}}.up_proj"), "experts.linear_fc2": GroupedMLPSlicing("backbone.layers.{}.mixer.experts.{{}}.down_proj"), - # MTP predictor projections (outer MultiTokenPredictionLayer) + # MTP predictor projections (outer MultiTokenPredictionLayer). The MTP inner + # attention/MoE layers are structurally identical to the backbone hybrid layers, so + # they reuse the base rules above via is_mtp=True (which retargets the backbone/model + # root to mtp, mirroring the importer) — only these predictor-specific keys are + # dedicated. "mtp.enorm": NameRemapping("mtp.layers.{}.enorm."), "mtp.hnorm": NameRemapping("mtp.layers.{}.hnorm."), "mtp.eh_proj": NameRemapping("mtp.layers.{}.eh_proj."), "mtp.final_layernorm": NameRemapping("mtp.layers.{}.final_layernorm."), - # MTP inner attention / MoE layers. Structurally identical to the backbone hybrid - # layers, so these mirror the base rules above with the `mtp.layers.{}` prefix; the - # `mtp.` namespace is aliased onto the standard rule keys by _get_mtp_state_dict. - "mtp.input_layernorm": NameRemapping("mtp.layers.{}.norm."), - "mtp.fused_norm": NameRemapping("mtp.layers.{}.norm.weight"), - "mtp.linear_qkv": QKVSlicing("mtp.layers.{}.mixer."), - "mtp.linear_proj": NameRemapping("mtp.layers.{}.mixer.o_proj."), - "mtp.pre_mlp_layernorm": NameRemapping("mtp.layers.{}.norm."), - "mtp.router": NameRemapping( - "mtp.layers.{}.mixer.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}} - ), - "mtp.shared_experts.linear_fc1": NameRemapping("mtp.layers.{}.mixer.shared_experts.up_proj."), - "mtp.shared_experts.linear_fc2": NameRemapping("mtp.layers.{}.mixer.shared_experts.down_proj."), - "mtp.local_experts.linear_fc1": NameRemapping("mtp.layers.{}.mixer.experts.{}.up_proj."), - "mtp.local_experts.linear_fc2": NameRemapping("mtp.layers.{}.mixer.experts.{}.down_proj."), - "mtp.experts.linear_fc1": GroupedMLPSlicing("mtp.layers.{}.mixer.experts.{{}}.up_proj"), - "mtp.experts.linear_fc2": GroupedMLPSlicing("mtp.layers.{}.mixer.experts.{{}}.down_proj"), } diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 81e4090d430..fdf10fb60d5 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -498,103 +498,129 @@ def _get_fused_norm_weight(self, module, primary_key: str = "fused_norm"): return None, None return fused_key, weight - def _get_transformer_layer_state_dict(self, layer, layer_id): + def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.input_layernorm, IdentityOp): - self.rules["input_layernorm"](layer.input_layernorm, layer_id) + self.rules["input_layernorm"](layer.input_layernorm, layer_id, is_mtp=is_mtp) else: fused_key, norm_weight = self._get_fused_norm_weight( getattr(layer.self_attention, "linear_qkv", None), primary_key="fused_input_layernorm", ) if norm_weight is not None: - self.rules[fused_key](norm_weight, layer_id) + self.rules[fused_key](norm_weight, layer_id, is_mtp=is_mtp) if not isinstance(layer.self_attention, IdentityOp): if "MLASelfAttention" in str(type(layer.self_attention)): if hasattr(layer.self_attention, "linear_q_proj"): - self.rules["linear_q_proj"](layer.self_attention.linear_q_proj, layer_id) + self.rules["linear_q_proj"]( + layer.self_attention.linear_q_proj, layer_id, is_mtp=is_mtp + ) else: self.rules["linear_q_down_proj"]( - layer.self_attention.linear_q_down_proj, layer_id + layer.self_attention.linear_q_down_proj, layer_id, is_mtp=is_mtp + ) + self.rules["linear_q_layernorm"]( + layer.self_attention.q_layernorm, layer_id, is_mtp=is_mtp + ) + self.rules["linear_q_up_proj"]( + layer.self_attention.linear_q_up_proj, layer_id, is_mtp=is_mtp ) - self.rules["linear_q_layernorm"](layer.self_attention.q_layernorm, layer_id) - self.rules["linear_q_up_proj"](layer.self_attention.linear_q_up_proj, layer_id) self.rules["linear_kv_down_proj"]( - layer.self_attention.linear_kv_down_proj, layer_id + layer.self_attention.linear_kv_down_proj, layer_id, is_mtp=is_mtp + ) + self.rules["linear_kv_layernorm"]( + layer.self_attention.kv_layernorm, layer_id, is_mtp=is_mtp + ) + self.rules["linear_kv_up_proj"]( + layer.self_attention.linear_kv_up_proj, layer_id, is_mtp=is_mtp ) - self.rules["linear_kv_layernorm"](layer.self_attention.kv_layernorm, layer_id) - self.rules["linear_kv_up_proj"](layer.self_attention.linear_kv_up_proj, layer_id) - self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id) + self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id, is_mtp=is_mtp) else: if layer.self_attention.q_layernorm is not None and not isinstance( layer.self_attention.q_layernorm, (IdentityOp, L2Norm) ): - self.rules["q_layernorm"](layer.self_attention.q_layernorm, layer_id) - self.rules["k_layernorm"](layer.self_attention.k_layernorm, layer_id) - self.rules["linear_qkv"](layer.self_attention.linear_qkv, layer_id) + self.rules["q_layernorm"]( + layer.self_attention.q_layernorm, layer_id, is_mtp=is_mtp + ) + self.rules["k_layernorm"]( + layer.self_attention.k_layernorm, layer_id, is_mtp=is_mtp + ) + self.rules["linear_qkv"](layer.self_attention.linear_qkv, layer_id, is_mtp=is_mtp) if ( hasattr(layer.self_attention, "core_attention") and "core_attention" in self.rules ): # KV cache quant export - self.rules["core_attention"](layer.self_attention.core_attention, layer_id) - self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id) + self.rules["core_attention"]( + layer.self_attention.core_attention, layer_id, is_mtp=is_mtp + ) + self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id, is_mtp=is_mtp) if getattr(layer.self_attention.core_attention, "softmax_offset", None) is not None: self.rules["softmax_offset"]( - layer.self_attention.core_attention.softmax_offset, layer_id + layer.self_attention.core_attention.softmax_offset, layer_id, is_mtp=is_mtp ) if not isinstance(layer.pre_mlp_layernorm, IdentityOp): - self.rules["pre_mlp_layernorm"](layer.pre_mlp_layernorm, layer_id) + self.rules["pre_mlp_layernorm"](layer.pre_mlp_layernorm, layer_id, is_mtp=is_mtp) elif not isinstance(layer.mlp, IdentityOp) and "MoE" not in str(type(layer.mlp)): fused_key, norm_weight = self._get_fused_norm_weight( getattr(layer.mlp, "linear_fc1", None), primary_key="fused_pre_mlp_layernorm", ) if norm_weight is not None: - self.rules[fused_key](norm_weight, layer_id) + self.rules[fused_key](norm_weight, layer_id, is_mtp=is_mtp) if not isinstance(layer.mlp, IdentityOp): if "MoE" in str(type(layer.mlp)): - self.rules["router"](layer.mlp.router, layer_id, dtype=self.moe_router_dtype) + self.rules["router"]( + layer.mlp.router, layer_id, dtype=self.moe_router_dtype, is_mtp=is_mtp + ) if hasattr(layer.mlp, "fc1_latent_proj") and layer.mlp.fc1_latent_proj is not None: - self.rules["fc1_latent_proj"](layer.mlp.fc1_latent_proj, layer_id) + self.rules["fc1_latent_proj"]( + layer.mlp.fc1_latent_proj, layer_id, is_mtp=is_mtp + ) if hasattr(layer.mlp, "fc2_latent_proj") and layer.mlp.fc2_latent_proj is not None: - self.rules["fc2_latent_proj"](layer.mlp.fc2_latent_proj, layer_id) + self.rules["fc2_latent_proj"]( + layer.mlp.fc2_latent_proj, layer_id, is_mtp=is_mtp + ) if hasattr(layer.mlp, "shared_experts") and layer.mlp.shared_experts is not None: self.rules["shared_experts.linear_fc1"]( - layer.mlp.shared_experts.linear_fc1, layer_id + layer.mlp.shared_experts.linear_fc1, layer_id, is_mtp=is_mtp ) self.rules["shared_experts.linear_fc2"]( - layer.mlp.shared_experts.linear_fc2, layer_id + layer.mlp.shared_experts.linear_fc2, layer_id, is_mtp=is_mtp ) if hasattr(layer.mlp.experts, "local_experts"): if not self.rules.get("use_packed_local_experts", False): for expert_id, expert in enumerate(layer.mlp.experts.local_experts): self.rules["local_experts.linear_fc1"]( - expert.linear_fc1, layer_id, expert_id + expert.linear_fc1, layer_id, expert_id, is_mtp=is_mtp ) self.rules["local_experts.linear_fc2"]( - expert.linear_fc2, layer_id, expert_id + expert.linear_fc2, layer_id, expert_id, is_mtp=is_mtp ) else: # For llama 4, in hf unified checkpoint, all local experts share one scale self.rules["local_experts.linear_fc1"]( - layer.mlp.experts.local_experts, layer_id + layer.mlp.experts.local_experts, layer_id, is_mtp=is_mtp ) self.rules["local_experts.linear_fc2"]( - layer.mlp.experts.local_experts, layer_id + layer.mlp.experts.local_experts, layer_id, is_mtp=is_mtp ) elif "experts.linear_fc1" in self.rules: # TEGroupedMLP: experts use fused grouped GEMM with a single # linear_fc1/linear_fc2 for all experts (no local_experts attribute). # Uses "experts.linear_fc1" rule (GroupedMLPMerging) instead of # "local_experts.linear_fc1" which expects per-expert iteration. - self.rules["experts.linear_fc1"](layer.mlp.experts.linear_fc1, layer_id) - self.rules["experts.linear_fc2"](layer.mlp.experts.linear_fc2, layer_id) + self.rules["experts.linear_fc1"]( + layer.mlp.experts.linear_fc1, layer_id, is_mtp=is_mtp + ) + self.rules["experts.linear_fc2"]( + layer.mlp.experts.linear_fc2, layer_id, is_mtp=is_mtp + ) else: - self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id) - self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id) + self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id, is_mtp=is_mtp) + self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id, is_mtp=is_mtp) def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]: """Export the MTP (Multi-Token Prediction) module. @@ -610,19 +636,11 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]: return self._copy_mtp_state_dict_from_pretrained() # The MTP inner attention / MoE layers are structurally identical to the base - # decoder layers, so we reuse the base layer walker. We alias the ``mtp.*`` - # naming rules onto the standard rule keys for the duration of the walk so the - # walker emits ``mtp.layers.{}.`` HF keys instead of ``backbone.layers.{}.``. - # A *restricted* alias set is used on purpose: any base rule key the walker - # references but that has no ``mtp.`` variant is simply absent (and its call is - # guarded), rather than silently emitting a wrong ``backbone.`` prefix. - mtp_rules = { - key[len("mtp.") :]: rule for key, rule in self.rules.items() if key.startswith("mtp.") - } - - saved_rules = self.rules + # decoder layers, so the same layer walker + rules are reused with is_mtp=True, + # which swaps the ``backbone``/``model`` target root for ``mtp`` (mirroring the + # importer). Only the predictor-specific projections (enorm/hnorm/eh_proj) and the + # MTP block's own final_layernorm use dedicated ``mtp.*`` rules. saved_state_dict = self._state_dict - self.rules = mtp_rules self._state_dict = OrderedDict() try: for mtp_layer in mtp.layers: @@ -631,21 +649,21 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]: last_id = inner_layers[-1].layer_number - 1 # Outer predictor projections attach to the first inner HF index. - if "enorm" in self.rules: - self.rules["enorm"](mtp_layer.enorm, first_id) - if "hnorm" in self.rules: - self.rules["hnorm"](mtp_layer.hnorm, first_id) - if "eh_proj" in self.rules: - self.rules["eh_proj"](mtp_layer.eh_proj, first_id) + if "mtp.enorm" in self.rules: + self.rules["mtp.enorm"](mtp_layer.enorm, first_id) + if "mtp.hnorm" in self.rules: + self.rules["mtp.hnorm"](mtp_layer.hnorm, first_id) + if "mtp.eh_proj" in self.rules: + self.rules["mtp.eh_proj"](mtp_layer.eh_proj, first_id) # Inner hybrid stack (e.g. [attention, MoE] for the ``*E`` pattern) - # reuses the base decoder walker with the aliased mtp rules. + # reuses the base decoder walker; is_mtp=True retargets to mtp.layers.{}. for inner in inner_layers: hf_layer_id = inner.layer_number - 1 if isinstance(inner, MambaLayer): - self._get_mamba_layer_state_dict(inner, hf_layer_id) + self._get_mamba_layer_state_dict(inner, hf_layer_id, is_mtp=True) elif isinstance(inner, TransformerLayer): - self._get_transformer_layer_state_dict(inner, hf_layer_id) + self._get_transformer_layer_state_dict(inner, hf_layer_id, is_mtp=True) else: raise ValueError( "Only TransformerLayer or MambaLayer are supported in the MTP block." @@ -654,15 +672,14 @@ def _get_mtp_state_dict(self) -> dict[str, torch.Tensor]: # The MTP block's own final layernorm attaches to the last inner HF index. final_layernorm = getattr(mtp_layer, "final_layernorm", None) if ( - "final_layernorm" in self.rules + "mtp.final_layernorm" in self.rules and final_layernorm is not None and not isinstance(final_layernorm, IdentityOp) ): - self.rules["final_layernorm"](final_layernorm, last_id) + self.rules["mtp.final_layernorm"](final_layernorm, last_id) mtp_state_dict = self._state_dict finally: - self.rules = saved_rules self._state_dict = saved_state_dict if len(mtp_state_dict) > 0: @@ -731,24 +748,24 @@ def _copy_mtp_state_dict_from_pretrained(self) -> dict[str, torch.Tensor]: self.exclude_modules.append("mtp*") return mtp_state_dict - def _get_mamba_layer_state_dict(self, layer, layer_id): + def _get_mamba_layer_state_dict(self, layer, layer_id, is_mtp=False): if not isinstance(layer.norm, IdentityOp): - self.rules["norm"](layer.norm, layer_id) + self.rules["norm"](layer.norm, layer_id, is_mtp=is_mtp) else: # TE spec: norm is fused into in_proj (QuantTELayerNormColumnParallelLinear). # Mamba uses the legacy single-key `fused_norm` rule (Nemotron-H style). fused_key, norm_weight = self._get_fused_norm_weight(layer.mixer.in_proj) if norm_weight is not None: - self.rules[fused_key](norm_weight, layer_id) + self.rules[fused_key](norm_weight, layer_id, is_mtp=is_mtp) - self.rules["mixer_norm"](layer.mixer.norm, layer_id) - self.rules["A_log"](layer.mixer.A_log, layer_id) - self.rules["D"](layer.mixer.D, layer_id) - self.rules["dt_bias"](layer.mixer.dt_bias, layer_id) + self.rules["mixer_norm"](layer.mixer.norm, layer_id, is_mtp=is_mtp) + self.rules["A_log"](layer.mixer.A_log, layer_id, is_mtp=is_mtp) + self.rules["D"](layer.mixer.D, layer_id, is_mtp=is_mtp) + self.rules["dt_bias"](layer.mixer.dt_bias, layer_id, is_mtp=is_mtp) - self.rules["conv1d"](_get_mamba_conv1d(layer.mixer), layer_id) - self.rules["in_proj"](layer.mixer.in_proj, layer_id) - self.rules["out_proj"](layer.mixer.out_proj, layer_id) + self.rules["conv1d"](_get_mamba_conv1d(layer.mixer), layer_id, is_mtp=is_mtp) + self.rules["in_proj"](layer.mixer.in_proj, layer_id, is_mtp=is_mtp) + self.rules["out_proj"](layer.mixer.out_proj, layer_id, is_mtp=is_mtp) def _get_medusa_heads_state_dict(self): medusa_heads = getattr(self.model, "medusa_heads", None) @@ -1009,6 +1026,18 @@ def _record_excluded_module(self, prefix: str): if layer_name not in self.exclude_modules: self.exclude_modules.append(layer_name) + @staticmethod + def _mtp_prefix(prefix: str) -> str: + """Rewrite a base-model target prefix to its MTP counterpart. + + Mirrors the importer so import/export naming stays symmetric: the MTP inner + layers reuse the base decoder rules, only the ``backbone``/``model`` root is + swapped for ``mtp``. + """ + if "backbone" in prefix: + return prefix.replace("backbone", "mtp") + return prefix.replace("model", "mtp") + def _name_remapping( self, module: torch.nn.Module | torch.Tensor, @@ -1016,7 +1045,10 @@ def _name_remapping( skip_output_scale: bool = True, mapping={}, dtype: torch.dtype | None = None, + is_mtp: bool = False, ): + if is_mtp: + prefix = self._mtp_prefix(prefix) if dtype is None: dtype = self.dtype @@ -1055,8 +1087,10 @@ def _name_remapping( self._state_dict[prefix + source_key] = val def _gated_mlp_slicing( - self, module, prefix, gate_proj_name="gate_proj", up_proj_name="up_proj" + self, module, prefix, gate_proj_name="gate_proj", up_proj_name="up_proj", is_mtp=False ): + if is_mtp: + prefix = self._mtp_prefix(prefix) name_to_value, qformat, block_size = self._get_quantized_state( module, self.dtype, prefix=prefix ) @@ -1116,7 +1150,7 @@ def _gated_mlp_slicing( self._state_dict[gate_proj_key] = val.detach().clone() self._state_dict[up_proj_key] = val.detach().clone() - def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): + def _grouped_mlp_slicing(self, module, prefix, parallel_config=None, is_mtp=False): """Export TEGroupedMLP weight0..weight{N-1} as one HF-style entry per expert. At EP>1, local ids are mapped to global via ``module.local_expert_indices`` @@ -1125,6 +1159,8 @@ def _grouped_mlp_slicing(self, module, prefix, parallel_config=None): Reverse of _grouped_mlp_merging in the importer. """ + if is_mtp: + prefix = self._mtp_prefix(prefix) num_experts = module.num_gemms state_dict = module.state_dict() @@ -1299,7 +1335,10 @@ def _qkv_slicing( q_proj_name="q_proj", k_proj_name="k_proj", v_proj_name="v_proj", + is_mtp=False, ): + if is_mtp: + prefix = self._mtp_prefix(prefix) name_to_value, qformat, block_size = self._get_quantized_state( module, self.dtype, prefix=prefix ) @@ -1428,9 +1467,11 @@ def _qkv_slicing( self._state_dict[v_proj_key] = val.detach().clone() def _self_attention_scaling( - self, module, prefix, k_scale_name="k_scale", v_scale_name="v_scale" + self, module, prefix, k_scale_name="k_scale", v_scale_name="v_scale", is_mtp=False ): """KV cache scaling for CoreAttention module.""" + if is_mtp: + prefix = self._mtp_prefix(prefix) k_scale_key = prefix + k_scale_name v_scale_key = prefix + v_scale_name if hasattr(module, "k_bmm_quantizer") and hasattr(module, "v_bmm_quantizer"): @@ -1444,8 +1485,10 @@ def _self_attention_scaling( # FP8 KV Cache is supported in VLLM; NVFP4 supported in TRTLLM self.kv_cache_dtype = kv_cache_dtype - def _pack_name_remapping(self, module, prefix, layer_type=None): + def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False): """Pack name remapping into one tensor.""" + if is_mtp: + prefix = self._mtp_prefix(prefix) weight_list = [] weight_scale_list = [] weight_scale_2_list = [] @@ -1510,8 +1553,10 @@ def _pack_name_remapping(self, module, prefix, layer_type=None): if merged_input_scale is not None: self._state_dict[prefix + "_input_scale"] = merged_input_scale - def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None): + def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=False): """Pack name remapping into one tensor.""" + if is_mtp: + prefix = self._mtp_prefix(prefix) weight_list = [] weight_scale_list = [] weight_scale_2_list = []