mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Noeyy/fix bug 6701777 (#2402)
### 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 -->
This commit is contained in:
@@ -35,3 +35,6 @@ quant_cfg:
|
||||
- quantizer_name: '*softmax_quantizer'
|
||||
cfg:
|
||||
$import: fp8
|
||||
- parent_class: 'nn.LayerNorm'
|
||||
quantizer_name: '*'
|
||||
enable: false
|
||||
|
||||
@@ -39,3 +39,6 @@ quant_cfg:
|
||||
- quantizer_name: '*bmm2_output_quantizer'
|
||||
cfg:
|
||||
$import: fp8
|
||||
- parent_class: 'nn.LayerNorm'
|
||||
quantizer_name: '*'
|
||||
enable: false
|
||||
|
||||
Reference in New Issue
Block a user