mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
[OMNIML-2244] Create the nvfp4 quant exporter (#636)
## What does this PR do? **Type of change:** New feature **Overview:** - Implemented the NVFP4QuantExporter - Deprecated fp4qdq_to_2dq - Updated tests ## Usage ```python python torch_quant_to_onnx.py --quantize_mode=nvfp4 \ --onnx_save_path=vit_base_patch16_224.nvfp4.onnx \ --calibration_data_size 64 \ --batch_size 128 ``` ## Testing <!-- Mention how have you tested your change if applicable. --> ``` python evaluate.py --onnx_path=vit_base_patch16_224.nvfp4.onnx \ --model_name=vit_base_patch16_224 \ --results_path=./results.txt \ --batch_size 128 ``` Results: ``` The top1 accuracy of the model is 84.39% The top5 accuracy of the model is 97.312% Inference latency of the model is 7.22412 ms ``` ## 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/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: No - Deprecated fp4qdq_to_2dq - **Did you write any new necessary tests?**: No - **Did you add or update any necessary documentation?**: No - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
This commit is contained in:
@@ -47,7 +47,7 @@ from diffusers.models.transformers.transformer_ltx import LTXVideoTransformer3DM
|
||||
from diffusers.models.unets import UNet2DConditionModel
|
||||
from torch.onnx import export as onnx_export
|
||||
|
||||
from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq
|
||||
from modelopt.onnx.export import NVFP4QuantExporter
|
||||
from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers
|
||||
from modelopt.torch.utils import torch_to
|
||||
|
||||
@@ -547,6 +547,6 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision):
|
||||
else:
|
||||
flux_convert_rope_weight_type(onnx_model)
|
||||
if precision == "fp4":
|
||||
onnx_model = fp4qdq_to_2dq(onnx_model)
|
||||
onnx_model = NVFP4QuantExporter.process_model(onnx_model)
|
||||
save_onnx(onnx_model, q_output)
|
||||
shutil.rmtree(tmp_subfolder, ignore_errors=True)
|
||||
|
||||
Reference in New Issue
Block a user