Repository navigation
Keep nn.LayerNorm inputs unquantized in the default PTQ presets - #2814
kevalmorabia97 merged 1 commit into
Conversation
nn.LayerNorm became a quant module in 0.44 so that recipes can opt into its output quantizer for FP8 MHA. As a side effect, the bare "*input_quantizer" wildcard also enables its input quantizer, so every preset or recipe that imports default_disabled_quantizers and enables "*input_quantizer" quantizes the residual stream entering each LayerNorm. Disable the LayerNorm quantizers in default_disabled_quantizers, as for BatchNorm and Embedding; a recipe that wants them enables them after importing the unit. Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
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 2 functions across 1 files. (2 skipped: 2 … | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: excluding nn.LayerNorm inputs from quantization in the default PTQ presets. Linked Issues check No active directly linked issue remains. Issue #1289is closed and supplies historical context only. Therefore, no linked-issue coding requirement applies to this pull request.Out of Scope Changes check The reported changes stay within the stated objective. The default_disabled_quantizers.yamlrule prevents defaultnn.LayerNorminput quantization, the unit test coversINT8_DEFAULT_CFG, `FP8_DEF…Security Anti-Patterns No stated security anti-pattern is introduced. The authoritative diff changes only a changelog, a YAML quantizer rule, and one test file; it adds no modelopt or examples Python code. The test addition…
Full details: Docstring Coverage
-
Explanation
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 2 functions across 1 files. (2 skipped: 2 unsupported.)
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
/ok to test 294e8e4 |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the exclusion correctly follows preset enables and preserves the FP8 MHA override; human sign-off is needed for the justified test-file extension.
Needs action:
- Confirm the additive regression coverage in
tests/unit/torch/quantization/test_quant_layernorm.py; the existing identity test and assertions remain unchanged.
No action needed:
- The three preset regressions check that Linear inputs remain quantized while LayerNorm inputs are disabled.
- Tests were inspected, not executed.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2814 +/- ##
==========================================
+ Coverage 69.52% 78.32% +8.79%
==========================================
Files 649 649
Lines 72081 72286 +205
==========================================
+ Hits 50114 56616 +6502
+ Misses 21967 15670 -6297
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:
|
3 participants
What does this PR do?
Type of change: Bug fix
#1289 (in 0.44.0 through #1350) registered
nn.LayerNorminQuantModuleRegistryso that FP8 MHA recipes can quantize its output. As a side effect, presets and recipes that importconfigs/ptq/units/default_disabled_quantizersand enable a bare*input_quantizer(INT8_DEFAULT_CFG,FP8_DEFAULT_CFG,NVFP4_DEFAULT_CFG, ...) also quantize every LayerNorm input, i.e. the residual stream (49 extra NVFP4 quantizers in XLM-R large). In ONNX export this becomes a dynamic quantize in front ofLayerNormalizationthat no GEMM consumes. Weight-only and layer-scoped presets are not affected.This adds a
parent_class: 'nn.LayerNorm'exclusion todefault_disabled_quantizers, in the same form as the existingnn.BatchNorm*,nn.LeakyReLUandnn.Embeddingentries. LayerNorm quantization stays opt-in: the FP8 MHA override inexamples/torch_onnx/torch_quant_to_onnx.py, appended after the preset, enables the LayerNorm output quantizer and already disables the input quantizer, so its result is unchanged. As with the BatchNorm entries, annn.LayerNormsubclass registered as its own quant class is not covered.#2402 fixed the same regression for the diffusers NVFP4 presets with the same entry; those presets do not import this unit. Out of scope:
NVFP4_FP8_MHA_CONFIG,model_type/vit/ptq/fp8.yamland the diffusersfp8,int8andnvfp4_fp8_convpresets also skip this unit and still quantize LayerNorm inputs.Effect on
BAAI/bge-reranker-v2-m3(XLM-R large), NVFP4 W4A4 fake quant (ModelOpt 0.47.0, same as main here), 60 queries × 100 candidates each:Usage
No API change. A config that wants LayerNorm input quantization adds the entry after the default exclusions:
Testing
test_default_presets_keep_layernorm_inputs_unquantized(INT8_DEFAULT_CFG,FP8_DEFAULT_CFG,NVFP4_DEFAULT_CFG): fails on main, passes here.pytest tests/unit/torch/quantization tests/unit/recipe: 1912 passed, 1 skipped.pytest tests/examples/torch_onnx/test_torch_quant_to_onnx.py -k "vit_tiny or swin"(RTX PRO 6000, TensorRT 11.3): the 12 FP8, INT8, MXFP8 and auto cases pass. The 3 NVFP4 cases fail the same way on main, in the FP32 NVFP4 export that [6721556] Fix mixed-dtype NVFP4 ONNX precision conversion #2415 fixes, and pass with a local fix for it.pre-commit run --fileson the changed files: pass.bge-reranker-v2-m3withNVFP4_DEFAULT_CFGhas 0 LayerNorm input quantizers enabled, and its TensorRT engine matches fake quant (Recall@10 89.8% for both) with the [6721556] Fix mixed-dtype NVFP4 ONNX precision conversion #2415 export fix applied locally.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).mtq.quantizecalls with these presets no longer add LayerNorm input quantizers, as before 0.44. Saved checkpoints restore unchanged, because restore replays their stored config. Usage shows how to re-enable them.CONTRIBUTING.md: N/AAdditional Information
Found while quantizing
BAAI/bge-reranker-v2-m3to NVFP4; its TensorRT engine also needs the FP32 NVFP4 export fix in #2415.Summary by CodeRabbit
LayerNormmodules, avoiding added error in residual streams.LayerNorminput quantization can enable it after importing the relevant preset.