Noeyy/fix bug 6701777

[noeyy-mino]

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 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.).

• 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?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

N/A

Summary by CodeRabbit

• Bug Fixes
  • Improved Diffusers NVFP4 and NVFP4/FP8 MHA quantization presets by excluding LayerNorm modules from quantization.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f774d80b-daa7-41c3-8917-cd18e8f851a3

📥 Commits

Reviewing files that changed from the base of the PR and between 59d93af and 42c250f.

📒 Files selected for processing (2)

• modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml
• modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

The Diffusers NVFP4 presets now disable all quantizers under nn.LayerNorm parent modules.

Changes

Diffusers PTQ preset updates

Layer / File(s)	Summary
LayerNorm quantizer exclusions modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml, modelopt_recipes/configs/ptq/presets/diffusers/nvfp4_fp8_mha.yaml	Both presets add rules that disable quantizers within nn.LayerNorm modules.

Priority: ⬇️ Low

Estimated code review effort: 1 (Trivial) | ~5 minutes

Change: Bug fix

Suggested reviewers: shengliangxu

Merge Risk: ⚪ Minimal · up to 6b69d

The preset-only change is ready to merge; no actionable current-head risk remains.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Title check	❓ Inconclusive	The title indicates that the pull request fixes bug 6701777, but it does not describe the LayerNorm exclusions in the Diffusers NVFP4 presets. The "Noeyy/" prefix also reduces clarity.	Rename the pull request with a concise, descriptive title such as "Exclude LayerNorm from Diffusers NVFP4 presets".

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	PASS. The review-scoped diff changes only two YAML quantization presets: modelopt_recipes/configs/ptq/presets/diffusers/nvfp4.yaml and nvfp4_fp8_mha.yaml. The additions only disable quantizers und…

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.77%. Comparing base (59d93af) to head (42c250f).

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     

Flag	Coverage Δ
examples-diffusers	20.64% <ø> (-0.01%)	⬇️
examples-gpt-oss	13.19% <ø> (+0.02%)	⬆️
examples-hf_ptq	22.19% <ø> (+0.57%)	⬆️
examples-llm_distill	13.26% <ø> (+0.01%)	⬆️
examples-llm_eval	17.10% <ø> (+0.06%)	⬆️
examples-llm_qat	17.44% <ø> (+0.04%)	⬆️
examples-llm_sparsity	15.74% <ø> (-0.01%)	⬇️
examples-megatron_bridge	26.05% <ø> (-0.20%)	⬇️
examples-specdec_bench	12.94% <ø> (+0.02%)	⬆️
examples-speculative_decoding	17.51% <ø> (-0.14%)	⬇️
examples-torch_onnx	21.62% <ø> (-0.03%)	⬇️
examples-torch_trt	14.95% <ø> (+<0.01%)	⬆️
gpu	58.18% <ø> (+25.95%)	⬆️
unit	57.51% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ajrasane]

I think this might break the TensorRT layer fusion. Could you check what is the latency of the model before and after this change?