Repository navigation
[https://nvbugs/6701777] Explicitly exclude nn.LayerNorm from the diffusers NVFP4 presets. - #2402
Merged
Merged
Conversation
Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com>
Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com>
Contributor
|
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: 📥 CommitsReviewing files that changed from the base of the PR and between 59d93af and 42c250f72d219f94da5350d1f6c2b6cee5f4373d. 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe Diffusers NVFP4 presets now disable all quantizers under ChangesDiffusers PTQ preset updates
Priority: ⬇️ Low Estimated code review effort: 1 (Trivial) | ~5 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to The preset-only change is ready to merge; no actionable current-head risk remains. 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2402 +/- ##
==========================================
+ Coverage 71.15% 78.77% +7.62%
==========================================
Files 543 543
Lines 64346 64346
==========================================
+ Hits 45785 50690 +4905
+ Misses 18561 13656 -4905
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:
|
noeyy-mino
force-pushed
the
noeyy/fix_bug_6701777
branch
from
September 11, 2026 10:53
42c250f to
6b69da9
Compare
Contributor
|
I think this might break the TensorRT layer fusion. Could you check what is the latency of the model before and after this change? |
Contributor
Author
|
Hi @ajrasane, Thanks for pointing this out. I checked the effective configs: both diffusers NVFP4 presets already disable LayerNorm output quantization before this PR. The effective change is therefore disabling the newly enabled NVFP4 input quantizer on LayerNorm; the existing Linear activation/weight QDQ rules remain unchanged. Test on B100 ModelOpt with this fix, TRT: 11.1.0 2026-09-14 10:34:52 [09/14/2026-02:34:24] [I] Latency: min = 67.0042 ms, max = 72.2706 ms, mean = 68.703 ms, median = 68.0374 ms, percentile(90%) = 70.9759 ms, percentile(95%) = 71.3879 ms, percentile(99%) = 72.2706 ms The regression came from the newly registered LayerNorm input quantizers being matched by the existing NVFP4 input wildcard. The intent is to restore the pre-regression quantization scope, not to claim that diffusion models cannot support LayerNorm output quantization. |
jingyu-ml
approved these changes
Sep 15, 2026
kevalmorabia97
enabled auto-merge (squash)
September 16, 2026 20:20
noeyy-mino
changed the title
shengliangxu
approved these changes
Sep 22, 2026
chadvoegele
added
the
cherry-pick-done
chadvoegele
pushed a commit
that referenced
this pull request
Sep 22, 2026
### What does this PR do? Type of change: Bug fix: 6701777 Regression source: "[OMNIML-3349] Add FP8 MHA quantization support for HuggingFace ViT" (#1289), merged into 0.44.0rc3 via the batch cherry-pick #1350. This PR: 1. Registers nn.LayerNorm as a QuantModule for the first time (modelopt/torch/quantization/nn/modules/quant_layernorm.py), intended to let FP8_DEFAULT_CFG's BMM input / LayerNorm output quantizer rules apply to ViT. 2. Removes the prior forced Cast-alignment logic in export_onnx.py that used to normalize Q/DQ node dtypes to trt_high_precision_dtype. Root Cause: Once nn.LayerNorm became a registered QuantModule, these wildcards started unintentionally matching norm1.norm inside FLUX's AdaLayerNormZero block — an elementwise_affine=False LayerNorm with no learnable weight/bias. Its input got routed through NVFP4 Q/DQ (emitted as Float32) while its synthesized affine scale remained native BFloat16, producing the dtype mismatch. Chosen fix: Explicitly exclude nn.LayerNorm from the diffusers NVFP4 presets rather than touching the global QuantModuleRegistry (which ViT FP8 MHA still needs). Add, in both modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml and nvfp4_fp8_mha.yaml, after the existing weight/input wildcard rules (list order matters — later entries override earlier ones): - parent_class: 'nn.LayerNorm' quantizer_name: '*' enable: false ### Usage ``` python examples/diffusers/quantization/quantize.py --model flux-dev --format fp4 --batch-size 2 --percentile 1.0 --alpha 0.8 --quant-algo max --n-steps 20 --quantized-torch-ckpt-save-path /tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4.pt --onnx-dir /tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4 --collect-method default --calib-size 128 --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 trtexec --onnx=/tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4/model.onnx --builderOptimizationLevel=4 --saveEngine=/tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4/model.plan --stronglyTyped --minShapes=hidden_states:1x1024x64,img_ids:1024x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 --optShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 --maxShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 ``` ### Testing The above test commands. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved Diffusers NVFP4 and NVFP4/FP8 MHA quantization presets by excluding LayerNorm modules from quantization. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
chadvoegele
pushed a commit
that referenced
this pull request
Sep 22, 2026
### What does this PR do? Type of change: Bug fix: 6701777 Regression source: "[OMNIML-3349] Add FP8 MHA quantization support for HuggingFace ViT" (#1289), merged into 0.44.0rc3 via the batch cherry-pick #1350. This PR: 1. Registers nn.LayerNorm as a QuantModule for the first time (modelopt/torch/quantization/nn/modules/quant_layernorm.py), intended to let FP8_DEFAULT_CFG's BMM input / LayerNorm output quantizer rules apply to ViT. 2. Removes the prior forced Cast-alignment logic in export_onnx.py that used to normalize Q/DQ node dtypes to trt_high_precision_dtype. Root Cause: Once nn.LayerNorm became a registered QuantModule, these wildcards started unintentionally matching norm1.norm inside FLUX's AdaLayerNormZero block — an elementwise_affine=False LayerNorm with no learnable weight/bias. Its input got routed through NVFP4 Q/DQ (emitted as Float32) while its synthesized affine scale remained native BFloat16, producing the dtype mismatch. Chosen fix: Explicitly exclude nn.LayerNorm from the diffusers NVFP4 presets rather than touching the global QuantModuleRegistry (which ViT FP8 MHA still needs). Add, in both modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml and nvfp4_fp8_mha.yaml, after the existing weight/input wildcard rules (list order matters — later entries override earlier ones): - parent_class: 'nn.LayerNorm' quantizer_name: '*' enable: false ### Usage ``` python examples/diffusers/quantization/quantize.py --model flux-dev --format fp4 --batch-size 2 --percentile 1.0 --alpha 0.8 --quant-algo max --n-steps 20 --quantized-torch-ckpt-save-path /tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4.pt --onnx-dir /tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4 --collect-method default --calib-size 128 --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 trtexec --onnx=/tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4/model.onnx --builderOptimizationLevel=4 --saveEngine=/tmp/pytest-of-root/pytest-0/test_diffusers_quant_export_on0/flux-dev-fp4/model.plan --stronglyTyped --minShapes=hidden_states:1x1024x64,img_ids:1024x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 --optShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 --maxShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 ``` ### Testing The above test commands. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved Diffusers NVFP4 and NVFP4/FP8 MHA quantization presets by excluding LayerNorm modules from quantization. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
chadvoegele
added a commit
that referenced
this pull request
Sep 22, 2026
#2438 #2402 #2434 (#2504) ## Cherry-picked PRs - #2202 - #2317 - #2314 - #2403 - #2203 - #2413 - #2436 - #2451 - #2336 - #2438 - #2402 - #2434 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added Step-3.7 PTQ recipes for routed experts and MLPs using NVFP4, with FP8 KV-cache support. * Added SDXL FP4 quantization with FP8 convolution support. * Expanded FP4 ONNX export guidance to include Flux and SDXL. * Added BF16 compatibility for FP8 ONNX export. * **Bug Fixes** * Improved ONNX handling for large models, external data, metadata, opsets, and output types. * Improved attention quantization and invalid dynamic-scale handling. * Quantization now reports unmatched weight configurations, with an override for valid pipeline-parallel workflows. * **Documentation** * Updated calibration, ONNX opset, CLI, Step-3.7, and TensorRT-LLM deployment guidance. * Updated deployment guidance to use unified Hugging Face exports with the TensorRT-LLM PyTorch backend. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Signed-off-by: Chad Voegele <cvoegele@nvidia.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: Ajinkya Rasane <ajinkyaashwin@gmail.com> Signed-off-by: Yue <yueshen@nvidia.com> Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Co-authored-by: Zhiyu <zhiyuc@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Co-authored-by: Ajinkya Rasane <ajinkyaashwin@gmail.com> Co-authored-by: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Co-authored-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com>
kevalmorabia97
pushed a commit
that referenced
this pull request
Oct 15, 2026
### What does this PR do? Type of change: Bug fix #1289 (in 0.44.0 through #1350) registered `nn.LayerNorm` in `QuantModuleRegistry` so that FP8 MHA recipes can quantize its output. As a side effect, presets and recipes that import `configs/ptq/units/default_disabled_quantizers` and 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 of `LayerNormalization` that no GEMM consumes. Weight-only and layer-scoped presets are not affected. This adds a `parent_class: 'nn.LayerNorm'` exclusion to `default_disabled_quantizers`, in the same form as the existing `nn.BatchNorm*`, `nn.LeakyReLU` and `nn.Embedding` entries. LayerNorm quantization stays opt-in: the FP8 MHA override in `examples/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, an `nn.LayerNorm` subclass 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.yaml` and the diffusers `fp8`, `int8` and `nvfp4_fp8_conv` presets 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: | NVFP4_DEFAULT_CFG (classifier excluded) | Ko-StrategyQA Recall@10 | same top-1 as FP32 | mean \|logit − FP32\| | MIRACL-ko Recall@10 | same top-1 | mean \|logit − FP32\| | |---|---|---|---|---|---|---| | as is (LayerNorm inputs quantized) | 87.6% | 80% | 0.94 | 91.1% | 80% | 1.05 | | LayerNorm quantizers disabled (this PR) | 90.6% | 90% | 0.61 | 94.5% | 87% | 0.66 | | FP32 | 88.9% | | | 93.4% | | | ### Usage No API change. A config that wants LayerNorm input quantization adds the entry after the default exclusions: ```python cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) cfg["quant_cfg"].append( {"parent_class": "nn.LayerNorm", "quantizer_name": "*input_quantizer", "cfg": {"num_bits": (4, 3), "axis": None}} ) ``` ### Testing - New `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 #2415 fixes, and pass with a local fix for it. - `pre-commit run --files` on the changed files: pass. - End to end: `bge-reranker-v2-m3` with `NVFP4_DEFAULT_CFG` has 0 LayerNorm input quantizers enabled, and its TensorRT engine matches fake quant (Recall@10 89.8% for both) with the #2415 export fix applied locally. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ❌ New `mtq.quantize` calls 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. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ Under 0.49.0; I can move the entry if this should go into 0.48.0. - Did you get Claude approval on this PR?: N/A ### Additional Information Found while quantizing `BAAI/bge-reranker-v2-m3` to NVFP4; its TensorRT engine also needs the FP32 NVFP4 export fix in #2415. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Default quantization presets no longer quantize inputs to `LayerNorm` modules, avoiding added error in residual streams. * Re-quantize models affected by this change. Recipes that need `LayerNorm` input quantization can enable it after importing the relevant preset. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix: 6701777
Regression source: "[OMNIML-3349] Add FP8 MHA
quantization support for HuggingFace ViT" (#1289), merged into
0.44.0rc3 via the batch cherry-pick #1350. This PR:
(modelopt/torch/quantization/nn/modules/quant_layernorm.py),
intended to let FP8_DEFAULT_CFG's BMM input / LayerNorm output
quantizer rules apply to ViT.
that used to normalize Q/DQ node dtypes to trt_high_precision_dtype.
Root Cause:
Once nn.LayerNorm became a registered QuantModule, these wildcards
started unintentionally matching norm1.norm inside FLUX's
AdaLayerNormZero block — an elementwise_affine=False LayerNorm with
no learnable weight/bias. Its input got routed through NVFP4 Q/DQ
(emitted as Float32) while its synthesized affine scale remained
native BFloat16, producing the dtype mismatch.
Chosen fix:
Explicitly exclude nn.LayerNorm from the diffusers NVFP4 presets
rather than touching the global QuantModuleRegistry (which ViT FP8
MHA still needs). Add, in both
modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml and
nvfp4_fp8_mha.yaml, after the existing weight/input wildcard rules
(list order matters — later entries override earlier ones):
quantizer_name: '*'
enable: false
Usage
Testing
The above test commands.
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.).CONTRIBUTING.md: N/AAdditional Information
N/A
Summary by CodeRabbit