mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Ajinkya RasaneandCodex
cf1f48fa0f
### What does this PR do?
Type of change: Bug fix
Adds a compact SDXL and SDXL-Turbo mixed-precision FP4 recipe:
- block-16 NVFP4 for non-QKV Linear/GEMM layers;
- FP8 for Conv2d layers;
- high-precision Q/K/V projection Linears to preserve TensorRT
horizontal fusion;
- optional FP8 MHA quantization.
For SDXL FP4 export, Conv2d quantizers export directly through the
shared FP8 custom-op path. The previous `generate_fp8_scales` plus
`convert_zp_fp8` INT8 zero-point workaround is removed. The graph then
uses the existing FP8 Q/DQ normalization and `NVFP4QuantExporter`
lowering, with opset 23 for FLOAT4 support. Flux FP8 export also saves
the graph returned by its RoPE weight conversion.
This PR also changes shared exporter behavior:
- `_fp8_quantize` refreshes ONNX shape/type inference after applying the
custom FP8 operator's uint8 output metadata, affecting all FP8 ONNX
exports through this symbolic.
- `_quantized_sdpa` derives `disable_fp8_mha` from the live Q/K/V
quantizer state instead of a restored private module flag.
Other model recipe configurations remain unchanged.
### Usage
```bash
python quantize.py \
--model sdxl-1.0 \
--model-dtype Half \
--trt-high-precision-dtype Half \
--format fp4 \
--block-size 16 \
--batch-size 2 \
--calib-size 128 \
--n-steps 20 \
--quantized-torch-ckpt-save-path ./sdxl-fp4 \
--onnx-dir ./onnx-sdxl-fp4
```
### Testing
- CPU-only focused and generic NVFP4 exporter tests: 44 passed in 4.35
seconds.
- Focused Flux returned-graph save test: 1 passed.
- Required Linux unit CI at `034fe23ec` passed with the `all` dependency
set, including `tests/unit/examples/test_diffusers_fp4.py`.
- Latest changed-file pre-commit checks: all passed.
- TensorRT 10.14 on a B200 GPU:
- 302 native block-scaled NVFP4 GEMM tactics;
- 38 native FP8 Conv tactics;
- no FP4 Q/K/V projections;
- all 11 FP16 Q/K/V projection-fusion groups preserved;
- three alternating batch-2 profiles measured 18.614 ms FP4 versus
20.028 ms FP16 median UNet latency, a 7.06% reduction.
- FP8 SDXL/SD3 ONNX-to-TensorRT end-to-end runs were not executed
because they require explicit approval. The existing end-to-end test
matrix now includes SD3 FP8 alongside SDXL FP8.
### 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?: ✅ — no public API or CLI flags
change; the shared changes preserve the intended FP8 export and
attention behavior.
- 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)?:
✅ — the shared NVFP4 opset, FP8 shape-inference, and Diffusers
attention-policy changes are recorded under bug fixes.
- Did you get Claude approval on this PR?: N/A
### Additional Information
Tracking: [5565357]
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added SDXL support for mixed NVFP4/FP8 quantization, including
convolution and softmax handling.
- Added an SDXL quantization preset for streamlined post-training
quantization workflows.
- Expanded FP4 ONNX export support to Flux and SDXL, with improved
FP4/FP8 graph processing and export reliability.
- Added automatic quantization policy and format restoration from
checkpoints.
- **Documentation**
- Documented SDXL layer behavior, optional FP8 attention quantization,
and Blackwell/TensorRT requirements for FP4 and FP8 deployment.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
> 🤖 _Generated by Codex (AI agent)._
---------
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Co-authored-by: Codex <codex@openai.com>
95 lines
3.9 KiB
Python
95 lines
3.9 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import torch.nn as nn
|
|
from calib.plugin_calib import PercentileCalibrator
|
|
|
|
from modelopt.torch.opt.config_loader import load_config
|
|
from modelopt.torch.quantization.config import QuantizeConfig
|
|
|
|
FP8_DEFAULT_CONFIG = load_config(
|
|
"configs/ptq/presets/diffusers/fp8", schema_type=QuantizeConfig
|
|
).model_dump(exclude_unset=True)
|
|
INT8_DEFAULT_CONFIG = load_config(
|
|
"configs/ptq/presets/diffusers/int8", schema_type=QuantizeConfig
|
|
).model_dump(exclude_unset=True)
|
|
NVFP4_DEFAULT_CONFIG = load_config(
|
|
"configs/ptq/presets/diffusers/nvfp4", schema_type=QuantizeConfig
|
|
).model_dump(exclude_unset=True)
|
|
NVFP4_FP8_MHA_CONFIG = load_config(
|
|
"configs/ptq/presets/diffusers/nvfp4_fp8_mha", schema_type=QuantizeConfig
|
|
).model_dump(exclude_unset=True)
|
|
NVFP4_FP8_CONV_CONFIG = load_config(
|
|
"configs/ptq/presets/diffusers/nvfp4_fp8_conv", schema_type=QuantizeConfig
|
|
).model_dump(exclude_unset=True)
|
|
|
|
|
|
def set_quant_config_attr(quant_config, trt_high_precision_dtype, quant_algo, **kwargs):
|
|
algo_cfg = {"method": quant_algo}
|
|
|
|
if quant_algo == "smoothquant" and "alpha" in kwargs:
|
|
algo_cfg["alpha"] = kwargs["alpha"]
|
|
elif quant_algo == "svdquant":
|
|
if "lowrank" in kwargs:
|
|
algo_cfg["lowrank"] = kwargs["lowrank"]
|
|
# Layers excluded from the SVDQuant algorithm (no AWQ smoothing, no
|
|
# low-rank branch); they stay quantized with plain max calibration.
|
|
if kwargs.get("skip_layers"):
|
|
algo_cfg["skip_layers"] = kwargs["skip_layers"]
|
|
quant_config["algorithm"] = algo_cfg
|
|
|
|
for entry in quant_config["quant_cfg"]:
|
|
p = entry.get("cfg", {})
|
|
if isinstance(p, dict) and "num_bits" in p and "trt_high_precision_dtype" not in p:
|
|
p["trt_high_precision_dtype"] = trt_high_precision_dtype
|
|
|
|
|
|
def reset_set_int8_config(quant_config, percentile, n_steps, collect_method, backbone):
|
|
"""Add PercentileCalibrator to Conv2d input quantizers.
|
|
|
|
Linear layers are left unchanged — their axis settings come from the base
|
|
quant_config (e.g. INT8_SMOOTHQUANT_CFG or INT8_DEFAULT_CONFIG).
|
|
|
|
Args:
|
|
quant_config: The quantization configuration dictionary
|
|
percentile: Percentile value for calibration
|
|
n_steps: Number of calibration steps
|
|
collect_method: Method for collecting calibration statistics
|
|
backbone: The model backbone to analyze layer types
|
|
"""
|
|
for name, module in backbone.named_modules():
|
|
if isinstance(module, nn.Conv2d):
|
|
aq_name = f"*{name}*input_quantizer*"
|
|
quant_config["quant_cfg"].append(
|
|
{
|
|
"quantizer_name": aq_name,
|
|
"cfg": {
|
|
"num_bits": 8,
|
|
"axis": None,
|
|
"calibrator": (
|
|
PercentileCalibrator,
|
|
(),
|
|
{
|
|
"num_bits": 8,
|
|
"axis": None,
|
|
"percentile": percentile,
|
|
"total_step": n_steps,
|
|
"collect_method": collect_method,
|
|
},
|
|
),
|
|
},
|
|
}
|
|
)
|