Skip to content

Export quantized/co-trained MTP weights instead of copying BF16 #2174

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yeyu-nvidia wants to merge 2 commits into
base: main
Choose a base branch
from
+173 −52
Open

Export quantized/co-trained MTP weights instead of copying BF16 #2174

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
54eaceb
Export quantized/co-trained MTP weights instead of copying BF16
yeyu-nvidia Aug 12, 2026
89ef596
Address review: reuse base layer rules via is_mtp instead of duplicat…
yeyu-nvidia Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion modelopt/torch/export/plugins/mcore_nemotron.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,11 @@
# 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). 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."),
Expand Down
Loading
Loading