[6466805] Add Dynamo ONNX export support for quantized models

CHANGELOG.rst

@@ -40,6 +40,7 @@ Changelog
 
 *Quantization*
 
+- Add opt-in Dynamo ONNX export through ``get_onnx_bytes_and_metadata(..., dynamo_export=True, onnx_opset=21)`` for FP8, INT8, INT4 AWQ, MXFP8, NVFP4, and mixed AutoQuant PyTorch models with canonical Linear, MatMul, or Gemm weight paths. The legacy exporter remains the default.
 - ONNX quantization with Autotune now benchmarks placements in the requested runtime precision and retains calibrated INT8/FP8 Q/DQ only when it meets the configured TensorRT speedup threshold (1.02x by default); otherwise it saves the high-precision no-Q/DQ model.
 - 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.

docs/source/guides/_pytorch_quantization.rst

@@ -70,11 +70,25 @@ 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, ModelOpt's helper supplies the operator translations and weight
+postprocessing required by the ``torch.export``-based ONNX workflow.
 
 .. code-block:: python
 
-    torch.onnx.export(model, sample_input, onnx_file)
+    from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata
+
+    onnx_bytes, _ = get_onnx_bytes_and_metadata(
+        model,
+        (sample_input,),
+        dynamo_export=True,
+        onnx_opset=21,
+    )
+    OnnxBytes.from_bytes(onnx_bytes).write_to_disk("onnx_model")
+
+Dynamo export supports FP8, INT8, INT4 AWQ, MXFP8, NVFP4, and mixed AutoQuant on
+canonical Linear, MatMul, or Gemm weight paths with opset 21 or newer. Direct
+PyTorch ONNX export is unsupported. The legacy helper path remains the default
+and can be selected explicitly with ``dynamo_export=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)
 

examples/torch_onnx/README.md

@@ -61,7 +61,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 uses opset 20 by default. Dynamo export requires opset 21 or newer.*
 
 ### Usage
 
@@ -76,6 +76,18 @@ Without `--recipe`, `--qformat` selects a quantization preset. Pass a built-in r
 path to `--recipe` to use a PTQ or AutoQuantize recipe instead. The recipe is authoritative when
 provided, so `--qformat` is ignored.
 
+Opt in to the `torch.export`-based ONNX exporter with opset 21 or newer:
+
+```bash
+python torch_quant_to_onnx.py \
+    --timm_model_name=vit_base_patch16_224 \
+    --qformat=fp8 --onnx_save_path=vit_base_patch16_224.fp8.onnx \
+    --dynamo_export --onnx_opset=24
+```
+
+Dynamo export is limited to canonical Linear, MatMul, or Gemm weight paths. Models containing INT4
+AWQ weights, including AutoQuant selections, must omit `--trt_build`.
+
 Convolutional architectures such as ResNet support only FP8 and INT8 quantization. MXFP8, NVFP4,
 INT4_AWQ, and AutoQuantize are not supported for these models because TensorRT does not provide
 the required convolution kernels.

examples/torch_onnx/torch_quant_to_onnx.py

@@ -36,6 +36,7 @@
 import modelopt.torch.quantization as mtq
 from modelopt.recipe import ModelOptAutoQuantizeRecipe, ModelOptPTQRecipe, load_recipe
 from modelopt.recipe.presets import QUANT_CFG_CHOICES
+from modelopt.torch._deploy.utils.torch_onnx import is_int4_quantized
 from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.plugins.custom import CUSTOM_POST_CONVERSION_PLUGINS
 
@@ -536,6 +537,17 @@ def main():
         help="The save path to save the ONNX model.",
         type=str,
     )
+    parser.add_argument(
+        "--dynamo_export",
+        action="store_true",
+        help="Use the torch.export-based ONNX exporter. Requires ONNX opset 21 or newer.",
+    )
+    parser.add_argument(
+        "--onnx_opset",
+        type=int,
+        default=20,
+        help="ONNX opset version.",
+    )
     parser.add_argument(
         "--calibration_data_size",
         type=int,
@@ -608,6 +620,9 @@ def main():
 
     args = parser.parse_args()
 
+    if args.dynamo_export and args.onnx_opset < 21:
+        parser.error("--dynamo_export requires --onnx_opset=21 or newer.")
+
     recipe = load_recipe(args.recipe) if args.recipe is not None else None
     if recipe is not None and not isinstance(
         recipe, (ModelOptPTQRecipe, ModelOptAutoQuantizeRecipe)
@@ -702,6 +717,12 @@ def main():
     # Blackwell has no tactic for an FP8 Q→Conv fusion on the first RGB layer.
     _disable_low_channel_fp8_conv_input_quantizers(quantized_model)
 
+    if args.trt_build and is_int4_quantized(quantized_model):
+        parser.error(
+            "--trt_build is not supported when the exported graph contains INT4 AWQ weights; "
+            "export ONNX without --trt_build."
+        )
+
     # Print quantization summary
     print("\nQuantization Summary:")
     mtq.print_quant_summary(quantized_model)
@@ -726,6 +747,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/base_exporter.py

@@ -20,6 +20,72 @@
 import onnx
 
 
+def _sync_initializer_metadata(graph: onnx.GraphProto, tensor: onnx.TensorProto) -> None:
+    """Synchronize declarations for an initializer after changing its type or shape."""
+    tensor_type = onnx.helper.make_tensor_value_info(
+        tensor.name, tensor.data_type, tensor.dims
+    ).type
+    for value_info in (*graph.input, *graph.value_info, *graph.output):
+        if value_info.name == tensor.name:
+            value_info.type.CopyFrom(tensor_type)
+
+
+def _replace_initializer(graph: onnx.GraphProto, tensor: onnx.TensorProto) -> None:
+    """Replace an initializer and synchronize any existing type declarations."""
+    existing = next((item for item in graph.initializer if item.name == tensor.name), None)
+    if existing is None:
+        graph.initializer.append(tensor)
+    else:
+        existing.CopyFrom(tensor)
+
+    _sync_initializer_metadata(graph, tensor)
+
+
+def _materialize_initializer_input(
+    graph: onnx.GraphProto,
+    node: onnx.NodeProto,
+    input_index: int,
+    tensor: onnx.TensorProto,
+) -> None:
+    """Materialize a constant input, splitting shared values per quantized weight."""
+    input_name = node.input[input_index]
+    consumers = [candidate for candidate in graph.node if input_name in candidate.input]
+    producer = next((candidate for candidate in graph.node if input_name in candidate.output), None)
+
+    if len(consumers) > 1:
+        tensor.name = f"{tensor.name}_{node.output[0]}"
+    elif producer is not None and producer.op_type == "Constant":
+        graph.node.remove(producer)
+
+    node.input[input_index] = tensor.name
+    _replace_initializer(graph, tensor)
+
+
+def _single_consumer(consumers, name: str, weight_name: str) -> onnx.NodeProto:
+    matches = consumers.get(name, [])
+    if len(matches) != 1:
+        raise NotImplementedError(
+            f"Unsupported Dynamo quantized weight topology for '{weight_name}': "
+            f"expected one consumer; found {len(matches)}."
+        )
+    return matches[0]
+
+
+def _validate_linear_weight_path(consumers, marker: onnx.NodeProto) -> None:
+    """Require a straight marker-to-MatMul/Gemm weight path."""
+    _single_consumer(consumers, marker.input[0], marker.input[0])
+    value_name = marker.output[0]
+    consumer = _single_consumer(consumers, value_name, marker.input[0])
+    while consumer.op_type in {"Cast", "Transpose"}:
+        value_name = consumer.output[0]
+        consumer = _single_consumer(consumers, value_name, marker.input[0])
+    if consumer.op_type not in {"MatMul", "Gemm"} or consumer.input[1] != value_name:
+        raise NotImplementedError(
+            f"Unsupported Dynamo quantized weight topology for '{marker.input[0]}': "
+            "expected terminal MatMul/Gemm at input 1."
+        )
+
+
 class ONNXQuantExporter(ABC):
     """Base class for ONNX quantizer exporters."""
 

modelopt/onnx/export/fp8_exporter.py

@@ -144,6 +144,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)
 
@@ -479,6 +480,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
@@ -497,6 +502,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"
                 )