Repository navigation
Implement canonical ModelOpt quantized HF export through Megatron Bridge #4804
Description
Parent PoR
Part of #5773.
Goal
Wire Megatron-Bridge topology reconstruction and HF naming to ModelOpt functional export so quantized Megatron models can be emitted in the canonical ModelOpt HF schema for native vLLM deployment and refit.
Megatron-Bridge scope
- Reconstruct logical HF weights across TP, ETP, EP, and PP.
- Apply fused QKV, gate/up, and expert topology transforms while carrying opaque ModelOpt export state.
- Call ModelOpt-owned packing and config APIs and expose the resulting canonical tensors and sidecars.
- Validate offline/functional parity, distributed topology transforms, and native vLLM load/refit behavior.
- Reject unsupported schemas, topology transforms, and runtime fallbacks explicitly.
Out of scope for Megatron-Bridge
- Quantization format tables, scale formulas, packing rules, sidecar semantics, or HF quantization-config construction.
- vLLM kernel or custom quantization-method ownership.
- Generic handling of arbitrary vendor quantized checkpoint schemas.
Vendor-specific source checkpoint import remains an explicit model adapter concern as described in #5773.
Timeline
TBD. Implementation steps and cross-repository dependencies will be linked as the APIs are finalized.
Activity
[-]Natively support quantized megatron to HF export for vLLM deployment (retires ModelOpt custom logic)[/-][+]Implement canonical ModelOpt quantized HF export through Megatron Bridge[/+]kevalmorabia97 commentedon Sep 30, 2026
Re-evaluated this against current
mainof both repos (2026-09-30) and the open PRs, with the goal of fully deprecating ModelOpt's custom Megatron↔HF logic (unified_export_megatron.py,plugins/mcore_*.py,plugins/megatron_importer.py) in favor of the bridge. Below is the current state and the proposed split of remaining work.Where things stand
MBridge
mainis essentially unchanged since July.export_hf_weights_modelopt:quantization_configand has no save-to-dir wrapper;w13_weight/w2_weight.It has no callers outside unit tests.
examples/quantization/export.pystill calls ModelOpt'smtex.export_mcore_gpt_to_hf.feat(export): add functional quantized weight export NVIDIA/Model-Optimizer#2251 + refactor(modelopt): delegate real-quant export to ModelOpt #5812 implement the ownership split in this issue:
quantization_config.feat(export): add functional quantized weight export NVIDIA/Model-Optimizer#2251 covers FP8, FP8_PC_PT, FP8_PB_WO, MXFP8, NVFP4 (dynamic and static), W4A16/W4A8 NVFP4, MXFP4 and mixed precision. That is a superset of what the custom Megatron exporter writes a quant config for, except IQ. The custom exporter already rejects AWQ/SmoothQuant, so this is no regression. Both PRs have been drafts since 2026-09-10.
The custom exporter is TP=1 only; it shards through PP. The bridge gathers across TP / EP / ETP / PP.
The custom path keeps growing: ~59 commits since June, and GLM-5.x alone is +1.1k LoC ([3/3] Add GLM-5.3-Flash support to Megatron-Bridge PTQ and HF export NVIDIA/Model-Optimizer#2539). So the migration is worth prioritizing.
Decisions
config.jsonquantization_configis needed; drophf_quant_config.json.quantized_layers(config.py#L921-L933, weight_utils.py#L224-L265).fused_qkv_a_proj) and TRT-LLM (fused_a) fuseq_a_proj+kv_a_proj_with_mqainto one GEMM.weight_scale_2differ.input_scale(modeling_deepseekv3.py#L471).mla_down_proj_fusiongives the pair a single quantizer natively.import_mcore_gpt_from_hf/megatron_importer.py). BF16 HF→Megatron import is already covered by AutoBridge. The remaining users are Megatron-LM'sexamples/post_training/modelopt(convert_model.py,quantize.py,prune.py;export.pyuses the exporter). Those examples get migrated to AutoBridge or retired together with the exporter.Test plan (Megatron-Bridge CI, with every PR)
Parallelism parity against the current ModelOpt export.
Quantize a small model per architecture once, then export the same Megatron checkpoint two ways:
export_mcore_gpt_to_hf(TP=1, PP=1/2);Assert that everything vLLM reads is identical:
config.json, includingquantization_config(quant_algo,quantized_layers, exclude list, KV scheme);How tensors are split across safetensors files may differ. Any intentional difference (e.g. MLA scale unification) is listed explicitly and treated as having no baseline.
Also export a checkpoint quantized at TP/EP > 1, so amax resharding on load is covered.
Once the ModelOpt exporter is removed, the bridge TP1/PP1 export becomes the reference for parallelism invariance, plus small checked-in golden hashes.
Format × architecture matrix: NVFP4, W4A16 NVFP4, FP8 (per-tensor / per-channel / 128×128 block), MXFP8, mixed precision, KV FP8.
Structural checks on every export:
weight_scale_2/input_scale;exclude_modulescovers exactly the unquantized linears;quantization_configor*_scale_invleaks through.vLLM runs only where there is no identical baseline. If the bridge export is identical to a vLLM-validated checkpoint from the ModelOpt exporter, a vLLM run adds no information. vLLM (Blackwell for NVFP4) is needed only for:
Even there it is a one-time load plus known-answer prompts, not a full eval.
Pilot: Qwen3.6-35B-A3B W4A4 NVFP4 from the ModelOpt tutorial (
examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B). One vLLM load-and-generate, only to confirm the parity harness compares everything vLLM reads.Remaining work
ModelOpt
k_scale/v_scalefromk_bmm/v_bmmquantizers, pluskv_cache_quant_algo. It is hard-set to NONE today.SharedWeightGlobalAmaxStatepatterns; dynamic NVFP4 / FP8 per-tensor via a post-calibration step.unified_export_megatron.pybefore deleting it:export_mcore_gpt_to_hf_vllm_fq, used by Megatron-LMexport.py): port to the bridge's BF16 export of folded weights plus a quantizer-state sidecar;export_extra_modules): decide separately.examples/megatron_bridgeand the Megatron-LMpost_training/modeloptexamples to the bridge. Deprecateexport_mcore_gpt_to_hf/import_mcore_gpt_from_hffor one release, then remove them together with themcore_*rule tables.MBridge, generic (details in #5812)
with_megatron_names.save_hf_pretrained-style) that consumes the export plan and writes the packed tensors plus the inlinequantization_config. It should also strip any source-checkpointquantization_configand*_scale_invcompanions.transpose_on_exportexperts.exclude_modulesderived from disabled quantizers, plus a per-HF-output "export unquantized" override for HF weights that MCore fuses into one quantized linear (below).weight_scale_2/input_scaleacross fused sibling groups.MBridge, per bridge
MambaInProjMapping(Nemotron-H, Falcon-H1), Nemotron-H MTP flattening, QKVG / ConcatenatedQKV / asymmetric QKV (Step, MiMo, Bailing, …)in_proj_a/in_proj_b(inside the fusedin_proj)shared_expert_gate, routers, Gemma4 per-layer-embedding projections, vision towers, MTP (Qwen3.5)cc @cuichenx @mxinO
kevalmorabia97 commentedon Sep 30, 2026
Execution plan: pilot, then remaining PRs
This breaks the checklist above into 8 PRs. Every PR lands with its CI tests from the test plan: bitwise parity against ModelOpt's PP-only export at TP2 / PP2 / EP2 / ETP2, plus structural checks. Implementation is expected to be quick. Most of the time goes into tests, so each PR can be merged with confidence.
Pilot: Qwen3.6-35B-A3B W4A4 NVFP4 (PRs 1–4, target ~2 weeks)
The reference is the checkpoint from ModelOpt's tutorial (
examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B). It is mixed precision:lm_head;conv1d, GDNin_proj_a/band the vision tower.That makes it a good first target: it covers grouped experts, mixed precision, the attention output gate, GDN and a VLM wrapper.
k_scale/v_scalefromk_bmm/v_bmmquantizers, pluskv_cache_quant_algo) and MLAlinear_q_down_proj/linear_kv_down_projamax unification after calibration (static NVFP4 viaSharedWeightGlobalAmaxStatepatterns; dynamic NVFP4 / FP8 per-tensor via a post-calibration step)main(incl. #6074with_megatron_names), plus:- offline save API (packed tensors plus the inline
- parity harness and fused-sibling scale check;
- switch
Parity tests on small Llama / Qwen3 / Qwen3-MoE models (NVFP4, FP8, mixed)quantization_config; drops source quant metadata);exclude_modulesfrom disabled quantizers;examples/quantization/export.pyto the bridge.GDNLinearMappingSeparate), per-HF-output unquantized override forin_proj_a/b, KV-cache config. Parity tests on a small Qwen3.5/3.6-MoE modelPilot exit criteria:
config.json/quantization_config, the tokenizer and the processor config.Remaining PRs
megatron_to_hf:MambaInProjMapping, Nemotron-H MTP flattening, QKVG / ConcatenatedQKV / asymmetric QKV (Nemotron-H +MTP);transpose_on_exportexperts (Qwen3-VL-MoE);examples/megatron_bridgeand its tutorials to the bridge;unified_export_megatron.py: vLLM fake-quant export via the bridge's BF16 export (folded weights plus quantizer-state sidecar), IQ formats via the functional API;export_mcore_gpt_to_hf/import_mcore_gpt_from_hf.examples/post_training/modelopt(convert_model.py,quantize.py,prune.py,export.py) to AutoBridge, or retire themunified_export_megatron.py, theplugins/mcore_*.pyrule tables andmegatron_importer.pyone release after the deprecationvLLM testing: needed only where no identical ModelOpt-exported baseline exists. That means architectures and formats the old exporter doesn't cover (e.g. Qwen3-VL-MoE, Gemma3-VL, MXFP8, FP8 per-channel) and intentional differences (e.g. MLA scale unification). Even there it is a one-time load plus known-answer prompts.
Out of scope for this deprecation:
export_extra_modules): decided separately;main.Rough timeline: pilot ~2 weeks; everything through PR 7 ~5–6 engineer-weeks. The calendar depends mostly on reviews and releases:
cc @cuichenx @mxinO
kevalmorabia97 commentedon Oct 14, 2026
Status 2026-10-14
feat(export): add functional quantized weight export NVIDIA/Model-Optimizer#2251: approved, CI green; waiting on a second code-owner approval. KV-cache export and MLA unification (PR 2) not started.
refactor(modelopt): delegate real-quant export to ModelOpt #5812: ready for review. Still needs:
importorskipfix for CI;It lands together with feat(modelopt): add format-general real-quant vLLM rollout RL#3945.
vLLM fake-quant (moved ahead of PR 6): [2/5] Export Megatron fakequant quantizer state and YAML NVIDIA/Model-Optimizer#2627 through [5/5] Preserve source config in HF fakequant export NVIDIA/Model-Optimizer#2630 set the sidecar format to
quantizer_state.pth+quant_recipe.yaml, which vLLM detects automatically. The bridge side needs two changes:modelopt_weights.pt;ENABLE_BRIDGE_QUANT_MAPPINGshould be on by default, or unmapped quantizer buffers skipped. Without it, saving a quantized model crashes.