[NVBug 5702186] Fix awq model export for Gemma3 (#793)

## What does this PR do?

**Type of change:** Bug fix <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**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
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

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

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

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

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

Author: meenchen

modelopt/torch/export/quant_utils.py

@@ -1101,6 +1101,16 @@ def fuse_prequant_to_linear(model: torch.nn.Module, fuse_grouped_heads=False):
                     setattr(linear_pqs_from, "fused_with_prequant", True)
 
 
+def _layernorm_uses_weight_plus_one(module: torch.nn.Module) -> bool:
+    if any(
+        name in type(module).__name__
+        for name in ["LayerNorm1P", "GemmaRMSNorm", "Gemma2RMSNorm", "Gemma3RMSNorm"]
+    ):
+        return True
+
+    return bool(hasattr(module, "zero_centered_gamma") and module.zero_centered_gamma)
+
+
 def fuse_prequant_layernorm(
     layernorm_module: torch.nn.Module,
     modules: list[torch.Tensor],
@@ -1116,13 +1126,17 @@ def fuse_prequant_layernorm(
         fused_bias = bias * avg_pre_quant_scale
         layernorm_output_scaled = (normalization(input) * fused_weight) + fused_bias
     """
-    layernorm_module.weight = torch.nn.Parameter(
-        layernorm_module.weight * getattr(modules[0].input_quantizer, "_pre_quant_scale")
+    pre_quant_scale = getattr(modules[0].input_quantizer, "_pre_quant_scale").to(
+        layernorm_module.weight.device
     )
+    if _layernorm_uses_weight_plus_one(layernorm_module):
+        # For norms that use (1 + weight) in forward, fold pre_quant_scale into the effective weight.
+        fused_weight = (layernorm_module.weight + 1.0) * pre_quant_scale - 1.0
+    else:
+        fused_weight = layernorm_module.weight * pre_quant_scale
+    layernorm_module.weight = torch.nn.Parameter(fused_weight.to(layernorm_module.weight.dtype))
     if hasattr(layernorm_module, "bias") and layernorm_module.bias is not None:
-        layernorm_module.bias = torch.nn.Parameter(
-            layernorm_module.bias * getattr(modules[0].input_quantizer, "_pre_quant_scale")
-        )
+        layernorm_module.bias = torch.nn.Parameter(layernorm_module.bias * pre_quant_scale)
     # Pre_quant_scales of modules must not be exported, since they have been fused with layernorm
     for module in modules:
         delattr(module.input_quantizer, "_pre_quant_scale")