[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:
Ajinkya Rasane
2025-12-04 22:56:07 +00:00
committed by GitHub
parent 097037d4dd
commit c6c9905929
7 changed files with 394 additions and 287 deletions
@@ -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)