[6466805] Add Dynamo ONNX export support for quantized models

CHANGELOG.rst

@@ -8,6 +8,8 @@ Changelog
 
 *Quantization*
 
+- Add Dynamo ONNX export support for INT8, INT4 AWQ, FP8, MXFP8, NVFP4, and mixed AutoQuant models. Enable it with
+  ``dynamo_export=True`` and ONNX opset 21 or newer; the legacy exporter remains the default.
 - Add a Muse Glimmer AutoQuantize recipe that searches language-model MLP projections, self-attention projections, and ``lm_head`` over W4A16 NVFP4 Four-Over-Six, FP8, and BF16 fallback at 5.5 effective bits while leaving the vision tower unquantized.
 - Add ``examples/alpamayo/qad.py``, which runs quantization-aware distillation on the quantized Alpamayo checkpoint produced by ``examples/alpamayo/quantize.py``. It distills the quantized VLM against the original FP16 VLM with ``QADTrainer``, supports FSDP2 for multi-GPU runs, and ``--export`` reassembles the trained VLM into a full AlpamayoR1 checkpoint that ``AlpamayoR1.from_pretrained`` can reload.
 - Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path.

docs/source/guides/_pytorch_quantization.rst

@@ -70,11 +70,40 @@ To verify that the quantizer nodes are placed correctly in the model, let's prin
     mtq.print_quant_summary(model)
 
 
-After PTQ, the model can be exported to ONNX with the normal PyTorch ONNX export flow.
+After PTQ, models whose non-strict Dynamo capture succeeds can be exported directly
+to ONNX without a translation table using ONNX opset 21 or newer.
 
 .. code-block:: python
 
-    torch.onnx.export(model, sample_input, onnx_file)
+    torch.onnx.export(
+        model,
+        (sample_input,),
+        onnx_file,
+        dynamo=True,
+        opset_version=24,
+    )
+
+For strict capture, provide ModelOpt's custom translations.
+
+.. code-block:: python
+
+    from modelopt.torch.quantization.export_onnx import get_dynamo_onnx_translation_table
+
+    exported_program = torch.export.export(model, (sample_input,), strict=True)
+    torch.onnx.export(
+        exported_program,
+        (),
+        onnx_file,
+        dynamo=True,
+        opset_version=24,
+        custom_translation_table=get_dynamo_onnx_translation_table(),
+    )
+
+Use the ``get_onnx_bytes_and_metadata(..., dynamo_export=True, onnx_opset=24)``
+helper for block-quantized formats because it adds this table automatically and
+performs the required weight postprocessing. The ``examples/torch_onnx`` workflow
+demonstrates this complete path. The legacy TorchScript exporter remains available
+with ``dynamo=False``.
 
 ModelOpt also supports direct export of Huggingface or Megatron-Bridge/Megatron-LM LLM models to TensorRT-LLM for deployment.
 Please see :doc:`TensorRT-LLM Deployment <../deployment/1_tensorrt_llm>` for more details.

examples/onnx_ptq/download_example_onnx.py

@@ -22,7 +22,15 @@
 from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata
 
 
-def export_to_onnx(model, input_shape, onnx_save_path, device, weights_dtype="fp32"):
+def export_to_onnx(
+    model,
+    input_shape,
+    onnx_save_path,
+    device,
+    weights_dtype="fp32",
+    dynamo_export=False,
+    onnx_opset=20,
+):
     """Export the torch model to ONNX format."""
     # Create input tensor with same precision as model's first parameter
     input_dtype = model.parameters().__next__().dtype
@@ -34,6 +42,8 @@ def export_to_onnx(model, input_shape, onnx_save_path, device, weights_dtype="fp
         dummy_input=(input_tensor,),
         weights_dtype=weights_dtype,
         model_name=model_name,
+        dynamo_export=dynamo_export,
+        onnx_opset=onnx_opset,
     )
     onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes)
 
@@ -64,8 +74,22 @@ def export_to_onnx(model, input_shape, onnx_save_path, device, weights_dtype="fp
         action="store_true",
         help="Whether to export the ONNX model in FP16.",
     )
+    parser.add_argument(
+        "--dynamo_export",
+        action="store_true",
+        help="Use the torch.export-based ONNX exporter.",
+    )
+    parser.add_argument(
+        "--onnx_opset",
+        type=int,
+        default=20,
+        help="ONNX opset version. Dynamo quantization export requires opset 21 or newer.",
+    )
     args = parser.parse_args()
 
+    if args.dynamo_export and args.onnx_opset < 21:
+        parser.error("--dynamo_export requires --onnx_opset=21 or newer.")
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = timm.create_model(args.timm_model_name, pretrained=True, num_classes=1000).to(device)
     data_config = timm.data.resolve_model_data_config(model)
@@ -79,5 +103,7 @@ def export_to_onnx(model, input_shape, onnx_save_path, device, weights_dtype="fp
         save_path,
         device,
         weights_dtype=weights_dtype,
+        dynamo_export=args.dynamo_export,
+        onnx_opset=args.onnx_opset,
     )
     print(f"{args.timm_model_name} model exported to {save_path}")

examples/torch_onnx/README.md

@@ -58,7 +58,7 @@ The `torch_quant_to_onnx.py` script quantizes [timm](https://github.com/huggingf
 - Postprocesses the ONNX model to be compatible with TensorRT.
 - Saves the final ONNX model.
 
-> *Opset 20 is used to export the torch models to ONNX.*
+> *The legacy exporter and opset 20 remain the defaults. Dynamo export requires opset 21 or newer.*
 
 ### Usage
 
@@ -69,6 +69,40 @@ python torch_quant_to_onnx.py \
     --onnx_save_path=<path to save the exported ONNX model>
 ```
 
+Use the torch.export-based ONNX exporter by selecting a compatible opset:
+
+```bash
+python torch_quant_to_onnx.py \
+    --timm_model_name=vit_base_patch16_224 \
+    --quantize_mode=fp8 \
+    --onnx_save_path=vit_base_patch16_224.fp8.onnx \
+    --dynamo_export \
+    --onnx_opset=24
+```
+
+ModelOpt's export helper supplies the custom translations and performs the weight
+postprocessing required by block-quantized formats. Use the helper for these formats.
+Advanced users calling `torch.onnx.export` directly can use the same translations
+for strict export:
+
+```python
+from modelopt.torch.quantization.export_onnx import get_dynamo_onnx_translation_table
+
+exported_program = torch.export.export(model, (sample_input,), strict=True)
+torch.onnx.export(
+    exported_program,
+    (),
+    "model.onnx",
+    dynamo=True,
+    opset_version=24,
+    custom_translation_table=get_dynamo_onnx_translation_table(),
+)
+```
+
+For models whose non-strict Dynamo capture succeeds and which do not require block-weight
+postprocessing, direct `torch.onnx.export(..., dynamo=True, opset_version=24)` is also
+supported. Pass `dynamo=False` to select the legacy TorchScript exporter explicitly.
+
 Quantization configs are loaded from the YAML preset recipes under
 `modelopt_recipes/configs/ptq/presets/model/`, selected by `--quantize_mode`. Pass
 `--recipe=<preset basename or path to a QuantizeConfig YAML>` to use a different

examples/torch_onnx/torch_quant_to_onnx.py

@@ -527,6 +527,17 @@ def main():
         action="store_true",
         help="Build a TensorRT engine from the exported ONNX model using trtexec.",
     )
+    parser.add_argument(
+        "--dynamo_export",
+        action="store_true",
+        help="Use the torch.export-based ONNX exporter.",
+    )
+    parser.add_argument(
+        "--onnx_opset",
+        type=int,
+        default=20,
+        help="ONNX opset version. Dynamo quantization export requires opset 21 or newer.",
+    )
     parser.add_argument(
         "--no_pretrained",
         action="store_true",
@@ -547,6 +558,9 @@ def main():
             "use --auto_quantization_formats instead."
         )
 
+    if args.dynamo_export and args.onnx_opset < 21:
+        parser.error("--dynamo_export requires --onnx_opset=21 or newer.")
+
     # Create model and move to appropriate device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model_kwargs = json.loads(args.model_kwargs) if args.model_kwargs else {}
@@ -650,6 +664,8 @@ def main():
         args.onnx_save_path,
         device,
         weights_dtype="fp16",
+        dynamo_export=args.dynamo_export,
+        onnx_opset=args.onnx_opset,
     )
 
     print(f"Quantized ONNX model is saved to {args.onnx_save_path}")

modelopt/onnx/export/fp8_exporter.py

@@ -136,6 +136,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
                 # Convert TRT DQ to native ONNX DequantizeLinear with FP8 weights
                 dq_op.inputs[0] = onnx_weights_fp8
                 dq_op.op = "DequantizeLinear"
+                dq_op.domain = ""
                 dq_op.outputs[0].dtype = dq_op.inputs[1].dtype
                 dq_op.outputs[0].shape = list(numpy_weights.shape)
 
@@ -457,6 +458,10 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         for node in graph.nodes:
             if node.op == "TRT_FP8QuantizeLinear":
                 node.op = "QuantizeLinear"
+                node.domain = ""
+                node.outputs[0].dtype = onnx.helper.tensor_dtype_to_np_dtype(
+                    onnx.TensorProto.FLOAT8E4M3FN
+                )
                 # Add FP8 zero_point if not present
                 if len(node.inputs) == 2:
                     # Create FP8 zero point constant
@@ -475,6 +480,7 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto:
         for node in graph.nodes:
             if node.op == "TRT_FP8DequantizeLinear":
                 node.op = "DequantizeLinear"
+                node.domain = ""
                 logger.debug(
                     f"Converted {node.name} from TRT_FP8DequantizeLinear to DequantizeLinear"
                 )