Export quantized/co-trained MTP weights instead of copying BF16 - #2174
Export quantized/co-trained MTP weights instead of copying BF16#2174yeyu-nvidia wants to merge 2 commits into
Conversation
`_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 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe PR adds live MTP traversal to Megatron export. It propagates MTP-specific prefixes through shared layer and projection rules, preserves the pretrained BF16 fallback, and restores exporter state after export. ChangesMTP export
Estimated code review effort: 4 (Complex) | ~45 minutes Mergeability Score: ⚪ Minimal · up to The PR changes MTP export to preserve live quantized or co-trained weights while retaining a pretrained fallback when no live MTP module exists; no actionable merge-blocking risk remains after normal checks and review. Sequence Diagram(s)sequenceDiagram
participant MegatronModel
participant MTPExporter
participant ExportRules
participant ExportState
MegatronModel->>MTPExporter: provide live MTP layers
MTPExporter->>ExportState: enable MTP export mode
MTPExporter->>ExportRules: traverse layers with is_mtp=true
ExportRules->>MTPExporter: return mtp-prefixed tensor entries
MTPExporter->>ExportState: restore prior exporter state
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/export/plugins/mcore_nemotron.py`:
- Around line 163-179: Add MTP-prefixed entries to the MTP rule mapping for
every Mamba walker key accessed by _get_mamba_layer_state_dict: norm,
mixer_norm, A_log, D, dt_bias, conv1d, in_proj, and out_proj. Map them to the
corresponding mtp.layers.{}.mixer.* paths, preserving the existing NameRemapping
or slicing behavior used by the base Mamba rules so _get_mtp_state_dict can
export Mamba-based MTP layers without KeyError.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b9290d4b-2650-402d-ba4e-c438024d580d
📒 Files selected for processing (2)
modelopt/torch/export/plugins/mcore_nemotron.pymodelopt/torch/export/unified_export_megatron.py
| # 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"), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Add MTP mappings for all Mamba walker rules.
_get_mtp_state_dict aliases only mtp.* rules. If an MTP inner layer is a MambaLayer, _get_mamba_layer_state_dict accesses norm, mixer_norm, A_log, D, dt_bias, conv1d, in_proj, and out_proj without guards. This mapping adds none of their mtp.* variants. Export then raises KeyError instead of exporting live MTP weights.
Add the corresponding mtp.* mappings with the mtp.layers.{}.mixer. prefix.
Proposed mapping additions
+ "mtp.norm": NameRemapping("mtp.layers.{}.norm."),
+ "mtp.mixer_norm": NameRemapping("mtp.layers.{}.mixer.norm."),
+ "mtp.A_log": NameRemapping("mtp.layers.{}.mixer.A_log"),
+ "mtp.D": NameRemapping("mtp.layers.{}.mixer.D"),
+ "mtp.dt_bias": NameRemapping("mtp.layers.{}.mixer.dt_bias"),
+ "mtp.conv1d": NameRemapping("mtp.layers.{}.mixer.conv1d."),
+ "mtp.in_proj": NameRemapping("mtp.layers.{}.mixer.in_proj."),
+ "mtp.out_proj": NameRemapping("mtp.layers.{}.mixer.out_proj."),This conflicts with the PR objective to reuse Mamba layer walkers.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # 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"), | |
| # 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"), | |
| "mtp.norm": NameRemapping("mtp.layers.{}.norm."), | |
| "mtp.mixer_norm": NameRemapping("mtp.layers.{}.mixer.norm."), | |
| "mtp.A_log": NameRemapping("mtp.layers.{}.mixer.A_log"), | |
| "mtp.D": NameRemapping("mtp.layers.{}.mixer.D"), | |
| "mtp.dt_bias": NameRemapping("mtp.layers.{}.mixer.dt_bias"), | |
| "mtp.conv1d": NameRemapping("mtp.layers.{}.mixer.conv1d."), | |
| "mtp.in_proj": NameRemapping("mtp.layers.{}.mixer.in_proj."), | |
| "mtp.out_proj": NameRemapping("mtp.layers.{}.mixer.out_proj."), |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/export/plugins/mcore_nemotron.py` around lines 163 - 179, Add
MTP-prefixed entries to the MTP rule mapping for every Mamba walker key accessed
by _get_mamba_layer_state_dict: norm, mixer_norm, A_log, D, dt_bias, conv1d,
in_proj, and out_proj. Map them to the corresponding mtp.layers.{}.mixer.*
paths, preserving the existing NameRemapping or slicing behavior used by the
base Mamba rules so _get_mtp_state_dict can export Mamba-based MTP layers
without KeyError.
There was a problem hiding this comment.
if the MTP layers are structurally identical to the base model layers, then maybe we don't need to add mtp.* rules that are duplicates of the base model rules? It should be possible to just use the base model rules. Adding these mtp rules also makes the import & export rules diverge.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="modelopt/torch/export/unified_export_megatron.py"
ast-grep outline "$target" --items all --type function --match '_get_mtp_state_dict|_get_mamba_layer_state_dict|_get_transformer_layer_state_dict' --view expanded
rg -n -A100 -B15 \
'def _get_mtp_state_dict\b|def _get_mamba_layer_state_dict\b|def _get_transformer_layer_state_dict\b' \
"$target"
rg -n -A5 -B5 \
'mtp\.\*|is_mtp|_get_mtp_state_dict\s*\(|_get_mamba_layer_state_dict\s*\(' \
modelopt/torch/exportLength of output: 56329
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -A80 -B15 \
'self\.rules\s*=|CustomModuleMapping|target_name_or_prefix|func_kwargs|name_remapping' \
modelopt/torch/export/unified_export_megatron.py \
modelopt/torch/export/plugins/mcore_custom.py \
modelopt/torch/export/plugins
rg -n -A80 -B15 \
'def _name_remapping\b|def _qkv_slicing\b|def _grouped_mlp_slicing\b|def .*remapp' \
modelopt/torch/export
rg -n -A5 -B5 \
'replace\("backbone", "mtp"\)|replace\("model", "mtp"\)|is_mtp' \
modelopt/torch/exportLength of output: 201035
@jenchen13, agreed. The current mtp.* aliases are needed by the export control flow, but they should not be manually duplicated.
The existing code creates self.rules from only mtp.* entries in _get_mtp_state_dict. Therefore, the Mamba walker cannot currently reuse the base rules. It raises KeyError for an MTP MambaLayer.
A better fix is to derive the MTP inner-layer export mappings from the base layer mappings. A helper can clone each layer-scoped CustomModuleMapping, replace the backbone.layers. prefix with mtp.layers., and add the cloned rule under mtp.<base-rule-key>. This follows the existing prefix-rewrite pattern in modelopt/torch/export/plugins/mcore_qwen3vl.py.
This change gives the MTP walker all base-layer rules, including Mamba rules, without a second hand-maintained mapping list. It also prevents future base Mamba or Transformer mapping changes from diverging from MTP export mappings.
I will withdraw the proposed manual mapping additions. The underlying export failure remains open until the aliases are generated from the base mappings or the exporter uses an equivalent prefix-rewrite mechanism.
🐇
🧬 Code Graph Analysis Results
modelopt/torch/export/plugins/mcore_custom.py
Lines 70–79 — CustomModuleMapping
class CustomModuleMapping:
"""A custom module mapping from Megatron Core to its HF counter part."""
def __init__(
self, func_name: str = "", target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}
):
"""Create a custom module mapping."""
self.func_name = func_name
self.target_name_or_prefix = target_name_or_prefix
self.func_kwargs = func_kwargsStores the mapping function name, target name/prefix, and optional keyword arguments such as is_mtp.
Lines 82–91 — NameRemapping
class NameRemapping(CustomModuleMapping):
"""A custom module mapping that renames of the modules."""
def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
"""Create a custom module mapping that renames of the modules."""
super().__init__(
func_name="name_remapping",
target_name_or_prefix=target_name_or_prefix,
func_kwargs=func_kwargs,
)Lines 142–151 — QKVSlicing
class QKVSlicing(CustomModuleMapping):
"""A custom module mapping that slices Q, K, V."""
def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
"""Create a custom module mapping that slices Q, K, V."""
super().__init__(
func_name="qkv_slicing",
target_name_or_prefix=target_name_or_prefix,
func_kwargs=func_kwargs,
)Lines 118–127 — GroupedMLPSlicing
class GroupedMLPSlicing(CustomModuleMapping):
"""A custom module mapping that slices fused grouped MLP weights into per-expert weights."""
def __init__(self, target_name_or_prefix: str = "", func_kwargs: dict[str, Any] = {}):
"""Create a custom module mapping that slices grouped MLP weights for export."""
super().__init__(
func_name="grouped_mlp_slicing",
target_name_or_prefix=target_name_or_prefix,
func_kwargs=func_kwargs,
)You are interacting with an AI system.
There was a problem hiding this comment.
Good call — done in 89ef596. I removed the duplicate mtp.* inner-layer rules and instead thread an is_mtp flag through the base layer walker + remapping helpers, which rewrite the target root backbone/model → mtp — mirroring exactly what the importer does. Only enorm/hnorm/eh_proj/final_layernorm stay as dedicated mtp.* rules now, so the import/export rule books stay symmetric.
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2174 +/- ##
===========================================
- Coverage 78.74% 66.16% -12.59%
===========================================
Files 522 522
Lines 60368 61519 +1151
===========================================
- Hits 47538 40705 -6833
- Misses 12830 20814 +7984
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| # 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"), |
There was a problem hiding this comment.
if the MTP layers are structurally identical to the base model layers, then maybe we don't need to add mtp.* rules that are duplicates of the base model rules? It should be possible to just use the base model rules. Adding these mtp rules also makes the import & export rules diverge.
| # 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 = { |
There was a problem hiding this comment.
you can just reuse the base layer rules + add the mtp specific ones for enorm, hnorm, eh_proj, final_layernorm
There was a problem hiding this comment.
Good call — done in 89ef596. I removed the duplicate mtp.* inner-layer rules and instead thread an is_mtp flag through the base layer walker + remapping helpers, which rewrite the target root backbone/model → mtp — mirroring exactly what the importer does. Only enorm/hnorm/eh_proj/final_layernorm stay as dedicated mtp.* rules now, so the import/export rule books stay symmetric.
End-to-end validation (NemotronH,
|
| tensor | fixed export vs BF16 pretrained source |
|---|---|
mtp.layers.0.eh_proj.weight |
differs |
mtp.layers.0.enorm / hnorm / norm |
differs |
mtp.layers.0.mixer.{q,k,v,o}_proj.weight |
differs |
Before this PR every mtp.* tensor was byte-identical to the pretrained model regardless of training (the copy path), so any QAD co-training of the MTP head was discarded at export. After the fix the exported head reflects the trained weights. Full layout (270 mtp.* tensors, layers.0=attention, layers.1=MoE) matches the nemotron_h_causal_lm_import round-trip.
Note on precision: in this particular checkpoint the MTP head was left unquantized (BF16, like the lm_head), so the exported MTP is BF16 with no weight_scale. The walker runs the same quantization rules as the base decoder, so it will emit NVFP4 weights + scales whenever the MTP module is quantized in the checkpoint — this validation just didn't exercise that path. @jenchen13 flagging in case the MTP head is expected to be quantized by the recipe.
…e 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 <noreply@anthropic.com>
2 participants
What
GPTModelExporter._get_mtp_state_dictcopied the MTP (multi-token prediction) head verbatim from the BF16 pretrained model instead of exporting the live model's MTP weights._get_state_dictonly walksmodel.decoder.layersand nevermodel.mtp, soself._state_dictnever contains anymtp.*keys and thekey not in self._state_dictguard was always true — every MTP tensor came from the pretrained safetensors. There was a standing# TODO Implement MTP export for quantized MTP.Consequence: any quantization or co-training applied to the MTP head during QAD was silently discarded at export; the exported draft head was always the original BF16 weights. This makes it impossible to evaluate MTP quantization or MTP co-training downstream.
Fix
_get_mtp_state_dictto walk the live MCoremodel.mtpmodule and apply 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 (_get_transformer_layer_state_dict/_get_mamba_layer_state_dict) is reused with a restricted set ofmtp.*naming rules aliased onto the standard rule keys — emittingmtp.layers.{}.HF keys. Restricted on purpose: any base rule key the walker references but that has nomtp.variant is simply absent (guarded), rather than silently emitting a wrongbackbone.prefix._copy_mtp_state_dict_from_pretrained, used only when the live model has nomtpmodule (e.g. exporting a base-only checkpoint that grafts a pretrained head).nemotron_h_causal_lm_export(attentionqkv/o_proj/norm, MoErouter/experts/shared_experts). The predictor projection rules (mtp.enorm/hnorm/eh_proj/final_layernorm) already existed.Validation
NemotronHForCausalLMmodel (num_nextn_predict_layers=1, hybrid*E), the walker reproduces exactly the expected HF key layout —mtp.layers.0= attention (enorm/hnorm/eh_proj/norm/mixer.{q,k,v,o}_proj),mtp.layers.1= MoE (norm/final_layernorm/mixer.gate/shared_experts/experts.{e}) — matching what the existingnemotron_h_causal_lm_import(is_mtpkeys) reads back.Note for reviewer
@jenchen13 — this is the MTP-export bug you flagged (the L599 BF16 copy). Would appreciate your review, especially on the assumption that the MTP inner layers can be driven through the base layer walker with only a prefix swap, and whether any non-
*EMTP configurations need additional inner-layer rules.Summary by CodeRabbit
New Features
Bug Fixes