Fix KV-cache scales dropped on Qwen Megatron-Core HF export - #2332
Conversation
_GPTModelExporter only emits KV-cache state for layers whose export mapping defines a `core_attention` rule. SelfAttentionScaling was wired for Llama and Nemotron but never for Qwen, so `_self_attention_scaling` never ran: `kv_cache_quant_algo` was left unset and a checkpoint calibrated with an FP8 or NVFP4 KV cache silently exported and served an unquantized one. Add the rule to `qwen3_causal_lm_export` and `qwen25_causal_lm_export`. That covers all six Qwen architectures, since `qwen3vl_causal_lm_export` and `qwen3_5_vl_causal_lm_export` derive from the former through `with_language_model_prefix`. GatedDeltaNet linear-attention layers have no `core_attention` submodule and are skipped by the existing hasattr guard. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.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 (3)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe change adds ChangesQwen KV-cache export
Estimated code review effort: 1 (Trivial) | ~5 minutes Merge Risk: ⚪ Minimal · up to Qwen Megatron-Core exports now retain KV-cache quantization metadata rather than silently serving calibrated caches unquantized. The mapping coverage validates the supported Qwen and derived vision-language architectures, with no remaining merge-blocking risk identified. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 2 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (claude-opus-5) — DM the bot to share feedback.
Small, well-scoped bug fix: adds the missing core_attention: SelfAttentionScaling(...) rule to qwen3_causal_lm_export and qwen25_causal_lm_export, mirroring exactly what mcore_llama.py already does. I verified the mechanism against the code:
_GPTModelExporter._get_transformer_layer_state_dictonly invokesself.rules["core_attention"]when the rule exists, so Qwen was indeed silently skipping_self_attention_scalingand leavingkv_cache_dtypeunset — the described root cause checks out.- The VLM claim holds:
with_language_model_prefixreconstructs eachCustomModuleMappingviatype(m)(...), soqwen3vl_causal_lm_export/qwen3_5_vl_causal_lm_exportinherit the rule with themodel.language_model.layers.{}.self_attn.prefix the test asserts. - No regression for non-KV-quantized Qwen exports:
get_kv_cache_scaling_factorreturns[None, None]when the bmm quantizers are absent/disabled (get_scaling_factorshort-circuits onis_enabled), andget_kv_cache_dtypereturnsQUANTIZATION_NONE, so no strayk_scale/v_scaletensors and nokv_cache_quant_algofield appear. GatedDeltaNet layers take thelinear_attnbranch and never reach the rule. - No import-side change is needed, and
_verify_exported_keysonly checks source-minus-exported, so the extra scale tensors are harmless.
Test is config-level (tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py), the same shape as the neighbouring test_moe_layout_choice.py, which is appropriate since the bug itself is config-level; it fails without the fix per the PR description. CHANGELOG updated. Licensing: only the new test file's header, which matches LICENSE_HEADER verbatim — no license concern.
One non-blocking observation for the owner: the new test pins NemotronForCausalLM → backbone.layers.{}.mixer., but every other rule in nemotron_causal_lm_export uses model.layers.{}.... (the backbone.../mixer. naming belongs to Nemotron-H). That pre-existing prefix looks like a copy-paste from the Nemotron-H mapping and the test now enshrines it; worth a separate look/fix rather than blocking this PR.
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2332 +/- ##
==========================================
- Coverage 79.31% 78.80% -0.51%
==========================================
Files 527 527
Lines 61482 61482
==========================================
- Hits 48765 48453 -312
- Misses 12717 13029 +312
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:
|
2 participants
What does this PR do?
Type of change: Bug fix
Megatron-Core → HuggingFace export silently dropped KV-cache quantization for every Qwen
architecture. A checkpoint calibrated with an FP8 (or NVFP4) KV cache exported with
kv_cache_quant_algounset, so the served model used an unquantized KV cache while the recipeand the Megatron checkpoint both said otherwise. Nothing warned.
_GPTModelExporteronly emits KV-cache state for layers whose architecture mapping defines acore_attentionrule:SelfAttentionScalingwas wired inmcore_llama.pyandmcore_nemotron.pybut never inmcore_qwen.py, so_self_attention_scalingnever ran for Qwen:self.kv_cache_dtypestayedunset and
_gather_kv_cache_dtype()returnedNone.Adding the rule to
qwen3_causal_lm_exportandqwen25_causal_lm_exportcovers all six Qwenarchitectures —
qwen3vl_causal_lm_exportandqwen3_5_vl_causal_lm_exportderive fromqwen3_causal_lm_exportthroughwith_language_model_prefix, which rewrites the prefix tomodel.language_model.layers.{}.self_attn.automatically. GatedDeltaNet linear-attention layershave no
core_attentionsubmodule, so the existinghasattrguard skips them.Known remaining gap, not addressed here:
deepseek_causal_lm_export,gptoss_causal_lm_exportandllama4_causal_lm_exportare missing the same rule. I could notvalidate those end to end, and DeepSeek's MLA uses different KV projection names, so they need
their own change rather than a copy of this one.
Usage
No API or flag change. The mapping now resolves for every Qwen architecture:
Testing
New
tests/gpu_megatron/torch/export/plugins/test_mcore_export_mappings.py(9 cases). It needsno GPU but imports
mcore_common, so it sits besidetest_moe_layout_choice.py, which is thesame shape. Confirmed the guard actually fires — reverting only
mcore_qwen.pygives6 failed / 3 passed (the Llama and Nemotron controls pass either way); with the fix,
9 passed.
End to end on a GB200 node in
nvcr.io/nvidia/nemo:26.08: quantizedQwen/Qwen3.5-0.8Bwith aW4A16-NVFP4 MLP / FP8-attention /
kv_fp8_castrecipe viaexamples/megatron_bridge/quantize.py, thenexport_quantized_megatron_to_hf.py.hf_quant_config.jsonnvidia/Qwen3.6-35B-A3B-NVFP4kv_cache_quant_algoNoneFP8FP8k_scale/v_scaletensorsThe absent scale tensors are correct for
kv_fp8_cast:use_constant_amaxpins the amax to theE4M3 maxbound, so
export_amax()returns nothing forget_scaling_factorand the runtime usesthe implicit 1.0 scale, while
get_kv_cache_dtypestill reports FP8 fromnum_bits. Thereleased checkpoint has exactly this shape, which is what the "after" column was checked
against.
The rest of the exported layer map is unchanged by this PR and was spot-checked against the
released checkpoint: NVFP4 W4A16 on the MLP projections, FP8 on
linear_attn.{in_proj_qkv,in_proj_z,out_proj}andself_attn.{q,k,v,o}_proj, within_proj_a/in_proj_b/conv1d/mtp.*excluded.pre-commit run --files ...passes on all changed files.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Found while reproducing the
nvidia/Qwen3.6-35B-A3B-NVFP4recipe throughexamples/megatron_bridge/rather thanexamples/hf_ptq/. Labeledcherry-pick-0.47.0.🤖 Generated with Claude Code
Summary by CodeRabbit