[kevalmorabia97]

What does this PR do ?

save_hf_pretrained_modelopt / export_hf_weights_modelopt (#6657, #5812) now export a quantized GatedDeltaNet in_proj (Qwen3.5 / Qwen3.6) instead of raising NotImplementedError for GDNLinearMappingSeparate. The output matches ModelOpt's export_mcore_gpt_to_hf bit for bit.

Part of #4804 (PR 4 of its plan, the Qwen3.6-35B-A3B pilot); tracked as OMNIML-5977.

Changelog

• in_proj_qkv and in_proj_z are exported packed with their scales. Their rows come from the mapping's own split (split_gdn_linear_weights + _split_gdn_grouped_to_separate) applied to row indices, so the scales follow the same TP layout as the weights.
• in_proj_b and in_proj_a (the beta / alpha gates) are written in BF16 and listed in quantization_config.ignore. Megatron fuses them into the one quantized in_proj, so a recipe can't exclude them with a quantizer pattern; ModelOpt's exporter keeps them in BF16 the same way.
• The ModelOpt export hook passes outputs whose export spec is explicitly None through unchanged, instead of failing.

Tests

• Unit: the GDN state selection recovers every HF row of a Megatron in_proj built with merge_gdn_linear_weights at TP1 and TP2, and leaves the gates unquantized; the export hook passes unquantized outputs through. Both fail without this change.
• Functional (L0, 2 GPUs): test_modelopt_hf_export.py adds a toy Qwen3.5 (one GatedDeltaNet and one full-attention layer) in NVFP4: TP2 and PP2 write the TP1 checkpoint bit for bit, and ignore lists exactly the embedding, LM head, conv1d, in_proj_a and in_proj_b. The export script keeps conv1d unquantized, as ModelOpt's recipes do.

Run in nvcr.io/nvidian/nemo:26.10.rc9 with this branch overlaid:

• With nvidia-modelopt 0.48.0rc1 (this repo's pin): tests/unit_tests/models/test_modelopt_utils.py (78 passed) and test_modelopt_hf_export.py (4 passed, including the new Qwen3.5 case).
• With ModelOpt main: ModelOpt's tests/gpu_megatron/torch/export/test_quantized_mbridge_export.py (Test Megatron-Bridge quantized HF export against the Megatron exporter NVIDIA/Model-Optimizer#3066) extended to a tiny Qwen3.5-VL-MoE with every linear layer quantized, NVFP4 and FP8: the bridge export at TP2 matches export_mcore_gpt_to_hf at TP1 bit for bit (weights, scales, config.json, quantization_config; FP8 input_scale within 2% at TP2). Once this merges, I'll bump ModelOpt CI's Megatron-Bridge commit and switch that test's Qwen3.5-VL case from NVFP4 MLP-only to full NVFP4.

Not covered here: KV-cache export (waits for the nvidia-modelopt 0.49.0rc0 pin), Qwen3-Next's GDNLinearMapping, and Mamba (MambaInProjMapping, OMNIML-5978).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? N/A
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

Additional Information

• Related to Implement canonical ModelOpt quantized HF export through Megatron Bridge #4804

🤖 Generated with Claude Code

[github-actions]

Automatic Claude reviews have been retired. To request a pull-request review, post a comment containing:

/review

Add model=claude to use a Claude reviewer (the default is model=codex). mode=light|strict selects the review depth; for example, /review model=claude mode=strict. Comment /review help for all options.

[kevalmorabia97]

/review

[svcnvidia-nemo-ci]

Reviewed all four changed files and traced the affected GDN mappings, quantization-spec planning, export-hook lifecycle, and export callers. No verified findings at light review depth.

Static inspection only; tests were not executed.

Suggested test cases

• test_gdn_state_follows_the_in_proj_layout_and_keeps_gates_unquantized[1]: verify exact QKV/Z row recovery and explicit unquantized B/A states at TP1.
• test_gdn_state_follows_the_in_proj_layout_and_keeps_gates_unquantized[2]: retain the same assertions for the TP2-packed layout.
• test_export_hook_passes_unquantized_outputs_through: retain the original output name and tensor.
• test_quantized_export_is_independent_of_parallelism[qwen3_5-nvfp4-parallelisms3]: verify the exact ignore set, packed QKV/scales, and TP2/PP2 equality against TP1.
• test_quantized_export_is_independent_of_parallelism[qwen3-nvfp4-parallelisms0]: preserve TP2/PP2 checkpoint equality.
• test_quantized_export_is_independent_of_parallelism[qwen3-fp8-parallelisms1]: preserve FP8 output dtype and TP2 checkpoint equality.
• test_quantized_export_is_independent_of_parallelism[qwen3_moe-mixed-parallelisms2]: preserve router exclusions and EP2/ETP2/PP2 checkpoint equality.
• No perf tests impacted.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[kevalmorabia97]

/ok to test cf6f1cd

[mxinO]

Why do we leave the alpha/beta projections in BF16 here? ModelOpt fake-quantizes the entire fused Megatron in_proj, including those rows and their input, but export writes the raw gate weights and marks them unquantized. That seems to introduce a fake-quant/inference mismatch. I see the existing ModelOpt exporter does the same, but could we clarify the reason?

[kevalmorabia97]

Good catch. The gates stay in BF16 so the exported checkpoint matches what ModelOpt produces for HF PTQ and what serving expects:

• ModelOpt's recipes disable *linear_attn.in_proj_a* / *linear_attn.in_proj_b* (default_disabled_quantizers.yaml, Adds AutoQuant support for VLM / Qwen3.5-Qwen3.6 style models NVIDIA/Model-Optimizer#1381 and [OMNIML-2938] PTQ support of Mamba Model (Nemotron-H-8B-Base-8K) #1687), and the published nvidia/Qwen3.6-35B-A3B-NVFP4-QAD lists them in exclude_modules. vLLM fuses them into one in_proj_ba with only 2 × linear_num_value_heads rows (64 for Qwen3.6-35B-A3B, fewer per TP rank). That's too narrow for the FP4 kernel tiles ([OMNIML-2938] PTQ support of Mamba Model (Nemotron-H-8B-Base-8K) #1687 hit size_n ... not divisible by tile_n_size = 64), and it's 0.5% of in_proj, so quantizing it gains nothing.
• Megatron fuses all six sections behind one weight_quantizer, so a quantizer pattern can't exclude the gates there. Both export_mcore_gpt_to_hf and this PR drop them at export instead.

You're right that this leaves a mismatch: Megatron fake-quantizes those rows, and the export doesn't. Its effect is limited:

• Calibration is unaffected. Max calibration collects amax with quantization off, and the in_proj input amax is shared, so in_proj_qkv / in_proj_z get the same input_scale either way.
• The deployed gates are more precise than what was simulated. So a Megatron-side fake-quant eval is slightly pessimistic, and QAT trains these rows under noise that the export then removes. Neither should cost accuracy.

The proper fix is for ModelOpt's Megatron GDN support to leave those rows unquantized during fake quant, so simulation matches export. That's a ModelOpt change, so I'll track it there as a follow-up rather than in this PR.