[NVBug 5702186] Fix awq model export for Gemma3

[meenchen]

What does this PR do?

Type of change: Bug fix

Overview: norms laers in Gemma that use (1 + weight) in forward, we will fold pre_quant_scale into the effective weight. That is to find folded w' subject to: 1 + w' = (1 + w) * s => w' = (1 + w) * s -1

Usage

# Add a code snippet demonstrating how to use this

Testing

./scripts/huggingface_example.sh --model google/gemma-3-1b-it --quant int4_awq

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

Summary by CodeRabbit

• Improvements
  • Enhanced quantization utilities to better handle various LayerNorm variants and normalization patterns, including support for weight-offset variants and zero-centered gamma configurations.
  • Optimized pre-quantization layer normalization fusion to apply conditional weight scaling strategies based on normalization type.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

A single file in the quantization export module is refactored to introduce conditional weight-folding logic for LayerNorm fusion. A new helper function detects LayerNorm variants using weight-plus-one patterns, and the pre-quantization scale folding is now conditionally applied either as direct weight multiplication or via the detected weight-plus-one mechanism.

Changes

Cohort / File(s)	Summary
LayerNorm quantization fusion logic modelopt/torch/export/quant_utils.py	Added _layernorm_uses_weight_plus_one() helper to detect LayerNorm variants (e.g., LayerNorm1P, Gemma RMSNorm) with zero-centered gamma. Refactored fuse_prequant_layernorm() to conditionally fold pre_quant_scale: applies (weight + 1) pattern for detected variants, otherwise simple multiplication. Bias folding updated to use pre_quant_scale directly.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the specific bug being fixed (awq model export for Gemma3) and references the NVBug ticket, directly matching the core change in the PR.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.22%. Comparing base (db76b1e) to head (707140c).
⚠️ Report is 7 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #793      +/-   ##
==========================================
- Coverage   74.23%   74.22%   -0.01%     
==========================================
  Files         192      192              
  Lines       19033    19035       +2     
==========================================
  Hits        14129    14129              
- Misses       4904     4906       +2     

☔ View full report in Codecov by Sentry.
📢 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.

[Edwardf0t1]

LGTM

[Edwardf0t1]

Does LayerNorm1P appears in Gemma3 only?

[meenchen]

No, but nemotron uses it, so adding it in the change for future proof: https://github.com/huggingface/transformers/blob/24807bfcf4a21286fa2a7e728f381ddaaca7bbc7/src/transformers/models/nemotron/modeling_nemotron.py#L88