diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 01a0fbe1b47..a712591a409 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/docs/source/guides/_pytorch_quantization.rst b/docs/source/guides/_pytorch_quantization.rst index f8f12b068ba..287b617d184 100644 --- a/docs/source/guides/_pytorch_quantization.rst +++ b/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. diff --git a/examples/onnx_ptq/download_example_onnx.py b/examples/onnx_ptq/download_example_onnx.py index e78ff2ecbf2..d4c61015de9 100644 --- a/examples/onnx_ptq/download_example_onnx.py +++ b/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) diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index bf50cfd2c69..078f1c1e6db 100644 --- a/examples/torch_onnx/README.md +++ b/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. diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 7450f124274..46044fc8424 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/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}") diff --git a/modelopt/onnx/export/base_exporter.py b/modelopt/onnx/export/base_exporter.py index 41d80c0e7ec..d919ff65739 100644 --- a/modelopt/onnx/export/base_exporter.py +++ b/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.""" diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 8e94f40f332..e13dcce4f32 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/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" ) diff --git a/modelopt/onnx/export/int4_exporter.py b/modelopt/onnx/export/int4_exporter.py index 0da217ae76f..2cc0be3f667 100644 --- a/modelopt/onnx/export/int4_exporter.py +++ b/modelopt/onnx/export/int4_exporter.py @@ -19,11 +19,110 @@ from onnx import numpy_helper from modelopt.onnx.logging_config import logger -from modelopt.onnx.quantization.graph_utils import get_tensor_producer_nodes +from modelopt.onnx.quantization.graph_utils import ( + get_tensor_consumer_nodes, + get_tensor_producer_nodes, +) from modelopt.onnx.quantization.qdq_utils import cast_initializer_to_dtype from modelopt.onnx.quantization.quant_utils import pack_weights_to_int4 -from .base_exporter import ONNXQuantExporter +from .base_exporter import ( + ONNXQuantExporter, + _materialize_initializer_input, + _replace_initializer, + _single_consumer, +) + + +def _optional_attribute(node: onnx.NodeProto, name: str): + attr = next((attr for attr in node.attribute if attr.name == name), None) + return onnx.helper.get_attribute_value(attr) if attr else None + + +def _constant_array(producers, name: str): + producer = producers.get(name) + if isinstance(producer, onnx.TensorProto): + return numpy_helper.to_array(producer) + value = ( + _optional_attribute(producer, "value") + if producer and producer.op_type == "Constant" + else None + ) + if value is None: + raise NotImplementedError("Dynamo ONNX export does not support dynamic weight shapes.") + return numpy_helper.to_array(value) + + +def _weight_dq_nodes(graph: onnx.GraphProto) -> list[onnx.NodeProto]: + initializers = {initializer.name for initializer in graph.initializer} + producers = get_tensor_producer_nodes(graph) + consumers = get_tensor_consumer_nodes(graph) + result = [] + for node in graph.node: + if node.op_type != "DequantizeLinear" or len(node.input) != 2: + continue + producer = producers.get(node.input[0]) + initializer_backed = node.input[0] in initializers or ( + producer is not None + and producer.op_type == "Reshape" + and producer.input[0] in initializers + ) + marker = _optional_attribute(node, "block_size") is not None + legacy = node.domain == "" and ( + _optional_attribute(node, "_target_shape") is not None + or any(consumer.op_type == "Reshape" for consumer in consumers.get(node.output[0], [])) + ) + if initializer_backed and (marker or legacy): + result.append(node) + elif marker and (producer is None or producer.op_type != "TRT_FP4DynamicQuantize"): + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{node.input[0]}': " + "data input must be a static initializer." + ) + return result + + +def _normalize_weight_paths(graph: onnx.GraphProto) -> list[onnx.NodeProto]: + """Normalize the supported Dynamo weight prefix to the legacy INT4 shape.""" + initializers = {initializer.name: initializer for initializer in graph.initializer} + producers = get_tensor_producer_nodes(graph, get_initializer_producers=True) + consumers = get_tensor_consumer_nodes(graph) + removed_outputs = set() + + for node in _weight_dq_nodes(graph): + weight_name = node.input[0] + producer = producers.get(weight_name) + if isinstance(producer, onnx.NodeProto): + source_name = producer.input[0] + _single_consumer(consumers, source_name, source_name) + _single_consumer(consumers, weight_name, source_name) + blocked_shape = [int(dim) for dim in _constant_array(producers, producer.input[1])] + weight = numpy_helper.to_array(initializers[source_name]).reshape(blocked_shape) + _replace_initializer(graph, numpy_helper.from_array(weight, source_name)) + node.input[0] = source_name + weight_name = source_name + removed_outputs.update(producer.output) + else: + _single_consumer(consumers, weight_name, weight_name) + if _optional_attribute(node, "block_size") is None: + node.attribute.append( + onnx.helper.make_attribute("block_size", initializers[weight_name].dims[-1]) + ) + + scale_name = node.input[1] + scale_producer = producers.get(scale_name) + if scale_name not in initializers and ( + not isinstance(scale_producer, onnx.NodeProto) or scale_producer.op_type != "Constant" + ): + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight_name}': scale must be constant." + ) + + if removed_outputs: + retained = [node for node in graph.node if removed_outputs.isdisjoint(node.output)] + del graph.node[:] + graph.node.extend(retained) + return _weight_dq_nodes(graph) class INT4QuantExporter(ONNXQuantExporter): @@ -34,87 +133,109 @@ def pre_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Pre-processes the ONNX model for INT4 quantization.""" graph = onnx_model.graph value_info_map = {value_info.name: value_info for value_info in graph.value_info} - weight_dq_nodes = [node for node in graph.node if node.op_type == "DequantizeLinear"] + weight_dq_nodes = _normalize_weight_paths(graph) tensor_producer_map = get_tensor_producer_nodes(graph, get_initializer_producers=True) + tensor_consumers = get_tensor_consumer_nodes(graph) - nodes_to_remove = [] + outputs_to_remove = set() for node in weight_dq_nodes: weight_name = node.input[0] logger.debug(f"Restructuring graph for weight {weight_name}") - ## Convert DequantizeLinear -> Reshape -> Transpose -> MatMul/Gemm to DequantizeLinear -> Matmul/Gemm - dq_child_nodes = [n for n in graph.node if node.output[0] in n.input] - reshape_node = dq_child_nodes[0] - nodes_to_remove.append(reshape_node.name) - assert reshape_node.op_type == "Reshape", f"Expected Reshape node for {node.name}" - reshape_node_output = reshape_node.output[0] - - # Remove constant node from reshape node - shape_constant_name = next(input for input in reshape_node.input if "Constant" in input) - nodes_to_remove.append(tensor_producer_map[shape_constant_name].name) - - # Get the shape of the output of the reshape node - store for compute_scales - reshape_output_value_info = value_info_map.get(reshape_node_output) - if reshape_output_value_info is not None: - weight_shape = [ - dim.dim_value for dim in reshape_output_value_info.type.tensor_type.shape.dim - ] + next_node = _single_consumer(tensor_consumers, node.output[0], weight_name) + weight_output = node.output[0] + path_input = node.output[0] + cast_node = None + reshape_node = None + target_shape_attr = next( + (attr for attr in node.attribute if attr.name == "_target_shape"), None + ) + weight_shape = ( + list(target_shape_attr.ints) + if target_shape_attr is not None + else list(next(item for item in graph.initializer if item.name == weight_name).dims) + ) + for _ in range(2): + if next_node.op_type == "Cast" and cast_node is None: + cast_node = next_node + elif next_node.op_type == "Reshape" and reshape_node is None: + reshape_node = next_node + outputs_to_remove.update(reshape_node.output) + reshape_output = value_info_map.get(reshape_node.output[0]) + weight_shape = ( + [dim.dim_value for dim in reshape_output.type.tensor_type.shape.dim] + if reshape_output is not None + else [ + int(dim) + for dim in _constant_array(tensor_producer_map, reshape_node.input[1]) + ] + ) + shape_producer = tensor_producer_map[reshape_node.input[1]] + if ( + isinstance(shape_producer, onnx.NodeProto) + and len(tensor_consumers.get(reshape_node.input[1], [])) == 1 + ): + outputs_to_remove.update(shape_producer.output) + else: + break + path_input = next_node.output[0] + next_node = _single_consumer(tensor_consumers, path_input, weight_name) + + target_shape = onnx.helper.make_attribute("_target_shape", weight_shape) + if target_shape_attr is None: + node.attribute.append(target_shape) else: - raise ValueError(f"Unable to determine shape of weight tensor {weight_name}") - - # Store target shape as attribute on DequantizeLinear node - target_shape_attr = node.attribute.add() - target_shape_attr.name = "_target_shape" - target_shape_attr.ints.extend(weight_shape) - - reshape_child_nodes = [n for n in graph.node if reshape_node.output[0] in n.input] - assert len(reshape_child_nodes) == 1, f"Expected exactly one child node for {node.name}" - - # Check if there's an optional Cast node between Reshape and Transpose/MatMul/Gemm - next_node = reshape_child_nodes[0] - if next_node.op_type == "Cast": - # Remove unnecessary Cast node - cast_node = next_node - nodes_to_remove.append(cast_node.name) - cast_child_nodes = [n for n in graph.node if cast_node.output[0] in n.input] - next_node = cast_child_nodes[0] - - # Store transpose permutation if present + target_shape_attr.CopyFrom(target_shape) + + if cast_node is not None: + value_info = value_info_map.get(node.output[0]) + source_dtype = ( + value_info.type.tensor_type.elem_type + if value_info is not None + else onnx.TensorProto.FLOAT + ) + if _optional_attribute(cast_node, "to") == source_dtype: + outputs_to_remove.update(cast_node.output) + else: + cast_node.input[0] = node.output[0] + weight_output = cast_node.output[0] + if next_node.op_type == "Transpose": transpose_node = next_node - nodes_to_remove.append(transpose_node.name) - assert transpose_node.op_type == "Transpose", ( - f"Expected Transpose node for {node.name}" - ) - perm = None - for attr in transpose_node.attribute: - if attr.name == "perm": - perm = [x for x in attr.ints] # noqa: C416 + outputs_to_remove.update(transpose_node.output) + path_input = transpose_node.output[0] + perm = _optional_attribute(transpose_node, "perm") assert perm is not None, f"Permutation not found for {node.name}" - # Store permutation as attribute on DequantizeLinear node - perm_attr = node.attribute.add() - perm_attr.name = "_transpose_perm" - perm_attr.ints.extend(perm) + node.attribute.append(onnx.helper.make_attribute("_transpose_perm", perm)) - transpose_child_nodes = [ - n for n in graph.node if transpose_node.output[0] in n.input - ] - assert len(transpose_child_nodes) == 1, ( - f"Expected exactly one matmul node for {node.name}" + matmul_node = _single_consumer( + tensor_consumers, transpose_node.output[0], weight_name ) - matmul_node = transpose_child_nodes[0] else: + perm = None matmul_node = next_node - assert matmul_node.op_type in ["MatMul", "Gemm"], ( - f"Expected MatMul or Gemm node for {node.name}" - ) - # Rewire MatMul to use DequantizeLinear output directly - matmul_node.input[1] = node.output[0] + if ( + matmul_node.op_type not in ["MatMul", "Gemm"] + or len(matmul_node.input) < 2 + or matmul_node.input[1] != path_input + ): + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight topology for '{weight_name}': " + "expected terminal MatMul/Gemm at input 1." + ) + axis = len(weight_shape) - 1 + axis = perm.index(axis) if perm is not None else axis + axis_attr = next((attr for attr in node.attribute if attr.name == "axis"), None) + if axis_attr is None: + node.attribute.append(onnx.helper.make_attribute("axis", axis)) + else: + axis_attr.i = axis + matmul_node.input[1] = weight_output # Remove transpose, reshape, and constant nodes - new_nodes = [node for node in graph.node if node.name not in nodes_to_remove] + new_nodes = [node for node in graph.node if outputs_to_remove.isdisjoint(node.output)] del graph.node[:] graph.node.extend(new_nodes) @@ -125,7 +246,7 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Computes the scales for the weights in the ONNX model for INT4 quantization.""" graph = onnx_model.graph initializer_map = {initializer.name: initializer for initializer in graph.initializer} - weight_dq_nodes = [node for node in graph.node if node.op_type == "DequantizeLinear"] + weight_dq_nodes = _weight_dq_nodes(graph) tensor_producer_map = get_tensor_producer_nodes(graph, get_initializer_producers=True) for node in weight_dq_nodes: @@ -135,28 +256,14 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: # Load weight and scale tensors weight = numpy_helper.to_array(initializer_map[weight_name]) - if scale_name in initializer_map: - scale = numpy_helper.to_array(initializer_map[scale_name]) - else: - scale_constant_node = tensor_producer_map[scale_name] - for attr in scale_constant_node.attribute: - if attr.name == "value": - tensor = attr.t - scale = numpy_helper.to_array(tensor) + scale = _constant_array(tensor_producer_map, scale_name) # Dequantize weight weight = weight / scale - block_size = weight.shape[-1] - - # Get target shape from metadata stored in pre_process - target_shape = None - transpose_perm = None - for attr in node.attribute: - if attr.name == "_target_shape": - target_shape = list(attr.ints) - elif attr.name == "_transpose_perm": - transpose_perm = list(attr.ints) + block_size = _optional_attribute(node, "block_size") or weight.shape[-1] + target_shape = _optional_attribute(node, "_target_shape") + transpose_perm = _optional_attribute(node, "_transpose_perm") assert target_shape is not None, f"Target shape not found for {node.name}" # Reshape weights and scales @@ -172,25 +279,12 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weight = weight.transpose(transpose_perm) scale = scale.transpose(transpose_perm) - # Handle scale tensor creation/update if scale_name not in initializer_map: - # Remove scale producer if it's a Constant node - scale_producer = tensor_producer_map[scale_name] - if scale_producer.op_type == "Constant": - graph.node.remove(scale_producer) - - # Create a new scale tensor scale_name = scale_name.replace("Constant_output_0", "scale") - scale_tensor = onnx.numpy_helper.from_array(scale, scale_name) - graph.initializer.append(scale_tensor) - node.input[1] = scale_name - else: - scale_tensor = onnx.numpy_helper.from_array(scale, scale_name) - initializer_map[scale_name].CopyFrom(scale_tensor) - - # Update weight tensor - weight_tensor = numpy_helper.from_array(weight, weight_name) - initializer_map[weight_name].CopyFrom(weight_tensor) + _materialize_initializer_input( + graph, node, 1, onnx.numpy_helper.from_array(scale, scale_name) + ) + _replace_initializer(graph, numpy_helper.from_array(weight, weight_name)) logger.debug(f"Computed scales for weight {weight_name} for INT4 quantization") @@ -207,7 +301,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Compresses the weights in the ONNX model for INT4 quantization.""" graph = onnx_model.graph initializer_map = {initializer.name: initializer for initializer in graph.initializer} - weight_dq_nodes = [node for node in graph.node if node.op_type == "DequantizeLinear"] + weight_dq_nodes = _weight_dq_nodes(graph) for node in weight_dq_nodes: weight_name = node.input[0] @@ -217,7 +311,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights_int4_onnx = onnx.numpy_helper.from_array(weights_int4_np, weight_name) weights_int4_onnx.data_type = onnx.TensorProto.INT4 weights_int4_onnx.dims[0] = weight_shape[0] - initializer_map[weight_name].CopyFrom(weights_int4_onnx) + _replace_initializer(graph, weights_int4_onnx) logger.debug(f"Converted {weight_name} to INT4 precision") return onnx_model @@ -255,10 +349,10 @@ def is_fp32_cast(node: onnx.NodeProto) -> bool: pqs_child_nodes = [n for n in graph.node if node.output[0] in n.input] assert len(pqs_child_nodes) == 1, f"Expected exactly one child node for {node.name}" cast_node = pqs_child_nodes[0] - assert cast_node.op_type == "Cast", f"Expected Cast node for {node.name}" - node.output.clear() - node.output.extend(cast_node.output) - nodes_to_remove.append(cast_node.name) + if cast_node.op_type == "Cast": + node.output.clear() + node.output.extend(cast_node.output) + nodes_to_remove.append(cast_node.name) # Remove unnecessary casts new_nodes = [node for node in graph.node if node.name not in nodes_to_remove] diff --git a/modelopt/onnx/export/mxfp8_exporter.py b/modelopt/onnx/export/mxfp8_exporter.py index 8c1e1f4df4f..cecb3fb6930 100644 --- a/modelopt/onnx/export/mxfp8_exporter.py +++ b/modelopt/onnx/export/mxfp8_exporter.py @@ -20,12 +20,20 @@ from onnx import numpy_helper from modelopt.onnx.logging_config import logger -from modelopt.onnx.quantization.graph_utils import get_tensor_producer_nodes +from modelopt.onnx.quantization.graph_utils import ( + get_tensor_consumer_nodes, + get_tensor_producer_nodes, +) from modelopt.onnx.quantization.qdq_utils import _cast_fp8, onnx_dtype_map from modelopt.onnx.quantization.quant_utils import compute_e8m0, get_amax from modelopt.onnx.utils import get_attribute, has_attribute -from .base_exporter import ONNXQuantExporter +from .base_exporter import ( + ONNXQuantExporter, + _materialize_initializer_input, + _replace_initializer, + _validate_linear_weight_path, +) E8_M0_BIAS = 127 DEFAULT_BLOCK_SIZE = 32 @@ -34,11 +42,11 @@ def _get_weight_dq_nodes(graph: onnx.GraphProto) -> list[onnx.NodeProto]: """Get weight DequantizeLinear nodes from the graph.""" + initializer_names = {initializer.name for initializer in graph.initializer} return [ node for node in graph.node - if node.op_type == "TRT_MXFP8DequantizeLinear" - and any(".weight" in inp for inp in node.input) + if node.op_type == "TRT_MXFP8DequantizeLinear" and node.input[0] in initializer_names ] @@ -70,6 +78,10 @@ class MXFP8QuantExporter(ONNXQuantExporter): @staticmethod def pre_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Pre-processes the ONNX model for MXFP8 quantization.""" + graph = onnx_model.graph + consumer_map = get_tensor_consumer_nodes(graph) + for node in _get_weight_dq_nodes(graph): + _validate_linear_weight_path(consumer_map, node) return onnx_model @staticmethod @@ -92,17 +104,16 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: se8m0_fp32 = compute_e8m0(amax, weight.shape, quant_axis, block_size) se8m0 = se8m0_fp32.astype(np.uint8) - # Remove scale producer if it's a Constant node scale_name = node.input[1] - scale_producer = tensor_producer_map[scale_name] - if scale_producer.op_type == "Constant": - graph.node.remove(scale_producer) - - # Create and add new scale tensor - scale_name_new = scale_name.replace("Constant_output_0", "scale") - scale_tensor = onnx.numpy_helper.from_array(se8m0, scale_name_new) - graph.initializer.append(scale_tensor) - node.input[1] = scale_name_new + if scale_name not in initializer_map: + scale_producer = tensor_producer_map.get(scale_name) + if scale_producer is None or scale_producer.op_type != "Constant": + raise NotImplementedError( + f"Unsupported Dynamo MXFP8 weight '{weight_name}': scale must be constant." + ) + scale_name = scale_name.replace("Constant_output_0", "scale") + scale_tensor = onnx.numpy_helper.from_array(se8m0, scale_name) + _materialize_initializer_input(graph, node, 1, scale_tensor) return onnx_model @@ -137,7 +148,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: vals=_cast_fp8(scaled_weight).tobytes(), raw=True, ) - initializer_map[weight_name].CopyFrom(weights_e4m3) + _replace_initializer(graph, weights_e4m3) logger.debug(f"Converted {weight_name} to MXFP8") return onnx_model diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 338e2725b14..7cb0370462e 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -31,7 +31,7 @@ ) from modelopt.torch.quantization.qtensor import NVFP4QTensor -from .base_exporter import ONNXQuantExporter +from .base_exporter import ONNXQuantExporter, _validate_linear_weight_path def _cast_fp4(array: np.ndarray) -> np.ndarray: @@ -138,7 +138,6 @@ def _add_input_value_info(graph, tensor_proto): sw_f32_per_tensor_proto = onnx.numpy_helper.from_array( sw_f32_per_tensor, sw_f32_per_tensor_name ) - sw_f8_per_block_proto = onnx.numpy_helper.from_array(sw_f8_per_block, sw_f8_per_block_name) sw_f8_per_block_proto = onnx.helper.make_tensor( name=sw_f8_per_block_name, data_type=onnx_dtype_map["Float8"], @@ -175,6 +174,7 @@ def _add_input_value_info(graph, tensor_proto): name=weight_name + "_DequantizeLinear_1", axis=-1, block_size=block_size, + domain="trt", ) # Add value_info for sw_f32 @@ -205,10 +205,14 @@ class NVFP4QuantExporter(ONNXQuantExporter): @staticmethod def pre_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Pre-processes the ONNX model for NVFP4 quantization. + """Rejects ambiguous weight and marker-output fanout.""" + tensor_consumers = get_tensor_consumer_nodes(onnx_model.graph) + for node in onnx_model.graph.node: + if node.op_type != "TRT_FP4QDQ": + continue + + _validate_linear_weight_path(tensor_consumers, node) - This is a no-op for NVFP4 quantization as no pre-processing is needed. - """ return onnx_model @staticmethod @@ -351,8 +355,10 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): assert maybe_matmul.op_type == "MatMul" node = maybe_matmul - # Create Cast nodes for each input of the target node except bias - for i, input_name in enumerate(node.input[:2]): + # Create Cast nodes for each input of the target node + for i, input_name in enumerate(node.input): + if not input_name: + continue cast_output_name = cast_output_cache.get((input_name, precision_dtype)) if cast_output_name is None: cast_output_suffix = "bf16" if precision_dtype == "BFloat16" else "f16" @@ -428,6 +434,10 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): graph.initializer.extend(new_initializers) logger.info(f"Removed {len(initializers_to_delete)} initializers") + if fp4_qdq_nodes and not any(opset.domain == "trt" for opset in onnx_model.opset_import): + onnx_model.opset_import.append(onnx.helper.make_opsetid("trt", 1)) + logger.info("Added TensorRT opset import") + utils.topologically_sort_graph_nodes(graph) return onnx_model diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index b217f1188b3..ef718803be5 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -43,6 +43,7 @@ NVFP4QuantExporter, ONNXQuantExporter, ) +from modelopt.onnx.export.base_exporter import _sync_initializer_metadata from modelopt.onnx.quantization.qdq_utils import qdq_to_dq, replace_zero_scale_with_smallest_nonzero from modelopt.onnx.utils import ( change_casts_to_fp16, @@ -517,11 +518,17 @@ def get_onnx_bytes_and_metadata( ModelMetadata: The model's meta data. Raises: - ValueError: If model is not an nn.Module or the requested precision conversion is unsupported. + ValueError: If model is not an nn.Module or the requested export configuration is unsupported. + NotImplementedError: If Dynamo export is requested with dynamic axes. """ if not isinstance(model, nn.Module): raise ValueError("Only PyTorch model compilation is supported.") + if dynamo_export and onnx_opset < 21: + raise ValueError("Dynamo ONNX export requires opset 21 or newer.") + if dynamo_export and dynamic_axes: + raise NotImplementedError("Dynamo ONNX export does not support dynamic_axes yet.") + assert weights_dtype in ["fp32", "fp16", "bf16"], ( "weights_dtype must be one of fp32, fp16, or bf16" ) @@ -546,11 +553,9 @@ def get_onnx_bytes_and_metadata( and not (uses_fp4 or uses_other_unsupported_quantizer) ) - # Standardize model args and also tensorize them so they also appear in the onnx graph! - # Floats/ints are tensorized when they are provided, but not tensorized when they are not - # provided which is somewhat inconsistent (we always tensorize them!) named_args, _ = standardize_named_model_args(model, dummy_input) - named_args = {k: _to_expected_onnx_type(v) for k, v in named_args.items()} + if not dynamo_export: + named_args = {name: _to_expected_onnx_type(value) for name, value in named_args.items()} # Also standardize dummy_input again so we can use it dummy_input = tuple(named_args.values()) @@ -622,7 +627,13 @@ def get_onnx_bytes_and_metadata( conv_wq_context = _disable_fp8_conv_weight_quantizers(model) if uses_fp8 else nullcontext() with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: additional_kwargs = {} - if not dynamo_export: + if dynamo_export: + from modelopt.torch.quantization._dynamo_onnx import _get_dynamo_onnx_translation_table + + additional_kwargs["custom_translation_table"] = _get_dynamo_onnx_translation_table() + if "fallback" in inspect.signature(torch.onnx.export).parameters: + additional_kwargs["fallback"] = False + else: additional_kwargs["dynamic_axes"] = dynamic_axes torch.onnx.export( model, @@ -661,11 +672,12 @@ def get_onnx_bytes_and_metadata( if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop: - if uses_other_unsupported_quantizer or uses_fp8: + preserve_block_io_types = dynamo_export and (uses_fp4 or uses_mxfp8) and weights_dtype == "fp32" + if (weights_dtype in ["fp16", "bf16"] or preserve_block_io_types) and not is_bf16_fp8_noop: + if (dynamo_export and uses_fp4) or uses_other_unsupported_quantizer or uses_fp8: onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, - keep_io_types=False, + keep_io_types=preserve_block_io_types, disable_shape_infer=True, check_fp16_ready=False, op_block_list=["QuantizeLinear", "DequantizeLinear", "Div"], @@ -678,6 +690,9 @@ def get_onnx_bytes_and_metadata( # Remove Cast(FP16->FP32) feeding Q/DQ scales so DQ stays FP16 for downstream # MatMul/Add layers under strongly-typed TRT parsing. onnx_opt_graph = fold_qdq_scale_fp16_to_fp32_casts(onnx_opt_graph) + if preserve_block_io_types: + for initializer in onnx_opt_graph.graph.initializer: + _sync_initializer_metadata(onnx_opt_graph.graph, initializer) else: onnx_opt_graph = convert_to_f16( onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False diff --git a/modelopt/torch/quantization/_dynamo_onnx.py b/modelopt/torch/quantization/_dynamo_onnx.py new file mode 100644 index 00000000000..cd24e88842a --- /dev/null +++ b/modelopt/torch/quantization/_dynamo_onnx.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""Private ONNXScript translations for quantized Dynamo export.""" + +import onnx +import onnxscript +import torch +from onnxscript.function_libs.torch_lib.tensor_typing import TFloat +from onnxscript.onnx_types import FLOAT4E2M1, FLOAT8E4M3FN + +from .export_onnx import onnx_dtype_map + +_OPSET = onnxscript.opset21 +_TRT_OPSET = onnxscript.values.Opset(domain="trt", version=1) + + +@onnxscript.script(_TRT_OPSET) +def _fp8_qdq(inputs: TFloat, scale: TFloat) -> TFloat: + quantized = _TRT_OPSET.TRT_FP8QuantizeLinear(inputs, scale) + return _TRT_OPSET.TRT_FP8DequantizeLinear(quantized, scale) + + +@onnxscript.script(_TRT_OPSET) +def _int4_dq(inputs: TFloat, scale: TFloat, axis: int, block_size: int) -> TFloat: + return _TRT_OPSET.DequantizeLinear(inputs, scale, axis=axis, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _fp4_qdq(inputs: TFloat, block_size: int) -> TFloat: + return _TRT_OPSET.TRT_FP4QDQ(inputs, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _fp4_dynamic_quantize( + inputs: TFloat, scale: TFloat, block_size: int +) -> tuple[FLOAT4E2M1[...], FLOAT8E4M3FN[...]]: + quantized, dynamic_scale = _TRT_OPSET.TRT_FP4DynamicQuantize( + inputs, scale, axis=-1, block_size=block_size, scale_type=17 + ) + return quantized, dynamic_scale + + +@onnxscript.script(_TRT_OPSET) +def _fp4_dq(inputs: FLOAT4E2M1[...], scale: TFloat, block_size: int) -> TFloat: + return _TRT_OPSET.DequantizeLinear(inputs, scale, axis=-1, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _mxfp8_dynamic_qdq(inputs: TFloat, block_size: int, output_dtype: int) -> TFloat: + quantized, scale = _TRT_OPSET.TRT_MXFP8DynamicQuantize( + inputs, axis=-1, block_size=block_size, output_dtype=17 + ) + return _TRT_OPSET.TRT_MXFP8DequantizeLinear( + quantized, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + +@onnxscript.script(_TRT_OPSET) +def _mxfp8_static_dq(inputs: TFloat, scale: TFloat, block_size: int, output_dtype: int) -> TFloat: + return _TRT_OPSET.TRT_MXFP8DequantizeLinear( + inputs, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + +def _cast(inputs, dtype: int): + return inputs if int(inputs.dtype) == dtype else _OPSET.Cast(inputs, to=dtype) + + +def _resolve_dtype(inputs, high_precision_dtype: str | None) -> int: + return ( + int(inputs.dtype) if high_precision_dtype is None else onnx_dtype_map[high_precision_dtype] + ) + + +def _static_shape(value) -> list[int]: + try: + return [int(dim) for dim in value.shape] + except (AttributeError, TypeError, ValueError): + raise NotImplementedError("Dynamo ONNX export does not support dynamic shapes.") from None + + +def _block_activation_shape( + inputs, block_size: int, quantizer_type: str | None +) -> list[int] | None: + if quantizer_type != "dynamic": + return None + shape = _static_shape(inputs) + if len(shape) not in (2, 3): + raise NotImplementedError("Dynamo ONNX block activation export supports rank 2 or 3 only.") + if shape[-1] % block_size: + raise NotImplementedError( + f"Dynamo ONNX block size {block_size} must divide the last dimension {shape[-1]}." + ) + return shape + + +def _translate_quantize_op( + inputs, + amax, + num_bits: int, + exponent_bits: int, + unsigned: bool, + narrow_range: bool, + high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, +): + if num_bits == 8 and exponent_bits == 4: + if amax is None: + scale = _OPSET.CastLike(1.0, inputs) + elif any(dim != 1 for dim in _static_shape(amax)): + raise AssertionError( + "E4M3 supports ONNX export only for per-tensor quantization with scalar amax." + ) + else: + scale = _OPSET.CastLike(_OPSET.Div(amax, 448.0), inputs) + return _fp8_qdq(inputs, scale) + + output_dtype = _resolve_dtype(inputs, high_precision_dtype) + if num_bits == 8 and exponent_bits == 0: + source_dtype = int(inputs.dtype) + if not unsigned: + assert not narrow_range, "ONNX does not support signed narrow-range INT8." + assert output_dtype in ( + source_dtype, + onnx.TensorProto.FLOAT, + onnx.TensorProto.BFLOAT16, + ), "TensorRT strongly typed mode requires Q/DQ in the input dtype, FP32, or BF16." + inputs = _cast(inputs, output_dtype) + quantized_axes = [index for index, dim in enumerate(_static_shape(amax)) if dim != 1] + if len(quantized_axes) > 1: + raise AssertionError("ONNX does not support multi-axis quantization.") + quantized_axis = quantized_axes[0] if quantized_axes else None + amax = _OPSET.Squeeze(_OPSET.Cast(amax, to=output_dtype)) + scale = _OPSET.Div(amax, float((1 << (7 + int(unsigned))) - 1)) + scale = _OPSET.Where(_OPSET.Equal(scale, 0.0), _OPSET.CastLike(1.0, scale), scale) + zero_point_dtype = onnx.TensorProto.UINT8 if unsigned else onnx.TensorProto.INT8 + zero_point = _OPSET.Cast(_OPSET.Mul(amax, 0.0), to=zero_point_dtype) + if quantized_axis is None: + quantized = _OPSET.QuantizeLinear(inputs, scale, zero_point) + output = _OPSET.DequantizeLinear(quantized, scale, zero_point) + else: + quantized = _OPSET.QuantizeLinear(inputs, scale, zero_point, axis=quantized_axis) + output = _OPSET.DequantizeLinear(quantized, scale, zero_point, axis=quantized_axis) + return output if output_dtype == source_dtype else _OPSET.Cast(output, to=source_dtype) + + if num_bits == 4 and exponent_bits == 0: + if unsigned: + raise NotImplementedError("Dynamo ONNX export supports signed INT4 only.") + if block_size is None or axis is None: + raise ValueError("INT4 ONNX export requires block_size and axis.") + scale = _OPSET.Div(_OPSET.Cast(amax, to=output_dtype), 7.0) + output = _int4_dq(inputs, scale, axis, block_size) + return output if int(inputs.dtype) == output_dtype else _OPSET.Cast(output, to=output_dtype) + + raise NotImplementedError( + f"Unsupported num_bits={num_bits}, exponent_bits={exponent_bits} for ONNX export." + ) + + +def _translate_dynamic_block_quantize_op( + inputs, + block_size: int, + amax, + num_bits: int, + exponent_bits: int, + scale_num_bits: int, + scale_exponent_bits: int, + high_precision_dtype: str | None = None, + quantizer_type: str | None = None, +): + activation_shape = _block_activation_shape(inputs, block_size, quantizer_type) + format_bits = (num_bits, exponent_bits, scale_num_bits, scale_exponent_bits) + + if format_bits == (4, 2, 8, 4): + if quantizer_type != "dynamic": + return _OPSET.Identity(_fp4_qdq(inputs, block_size)) + assert activation_shape is not None + output_dtype = _resolve_dtype(inputs, high_precision_dtype) + inputs = _cast(inputs, output_dtype) + if amax is None: + scale = _OPSET.Constant(value_float=1.0) + else: + scale = _OPSET.Div(_OPSET.Cast(amax, to=onnx.TensorProto.FLOAT), 2688.0) + scale = _OPSET.Where(_OPSET.Equal(scale, 0.0), _OPSET.CastLike(1.0, scale), scale) + quantized, dynamic_scale = _fp4_dynamic_quantize(inputs, scale, block_size) + quantized.dtype = onnxscript.ir.DataType.FLOAT4E2M1 + dynamic_scale.dtype = onnxscript.ir.DataType.FLOAT8E4M3FN + quantized.shape = onnxscript.ir.Shape(activation_shape) + scale_shape = [*activation_shape[:-1], activation_shape[-1] // block_size] + dynamic_scale.shape = onnxscript.ir.Shape(scale_shape) + dequantized_scale = _OPSET.DequantizeLinear(dynamic_scale, scale) + output = _fp4_dq(quantized, dequantized_scale, block_size) + return ( + output + if output_dtype == onnx.TensorProto.FLOAT + else _OPSET.Cast(output, to=output_dtype) + ) + + if format_bits == (8, 4, 9, 8): + output_dtype = int(inputs.dtype) + if quantizer_type == "dynamic": + output = _mxfp8_dynamic_qdq(inputs, block_size, output_dtype) + else: + output = _mxfp8_static_dq( + inputs, _OPSET.CastLike(1.0, inputs), block_size, output_dtype + ) + return _OPSET.Identity(output) + + raise NotImplementedError( + "Unsupported block format " + f"({exponent_bits}, {num_bits - exponent_bits - 1}) with scale format " + f"({scale_exponent_bits}, {scale_num_bits - scale_exponent_bits - 1})." + ) + + +def _get_dynamo_onnx_translation_table() -> dict: + dynamic_op = torch.ops.tensorrt.dynamic_block_quantize_op + return { + torch.ops.tensorrt.quantize_op.default: _translate_quantize_op, + dynamic_op.default: _translate_dynamic_block_quantize_op, + dynamic_op.overload: _translate_dynamic_block_quantize_op, + } diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 18b97ac2774..626dec54f6b 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -747,8 +747,9 @@ def _get_amax(self, inputs): reduce_axis = quant_utils.convert_quantization_axis_to_reduce_axis(inputs, self._axis) amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach() - amax = amax.detach() if is_torch_export_mode() else amax.data - return amax + if is_torch_export_mode() or torch.compiler.is_exporting(): + return amax.detach() + return amax.data def validate_attr( self, attr_value=None, attr_name="amax", raise_error=False, warn_error=False, name="" diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..84517ce8d99 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -119,11 +119,15 @@ def _quantize_impl( exponent_bits: int = 0, unsigned: bool = False, narrow_range: bool = True, + trt_high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, ): if num_bits == 8 and exponent_bits == 4: return scaled_e4m3_impl(inputs=inputs, amax=amax) elif isinstance(num_bits, int): - return fake_quant_impl( + quantize_impl = fake_quant_impl if inputs.is_cuda else _tensor_quant + return quantize_impl( inputs=inputs, amax=amax, num_bits=num_bits, @@ -143,6 +147,9 @@ def _quantize_impl_abstract( exponent_bits: int = 0, unsigned: bool = False, narrow_range: bool = True, + trt_high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, ) -> torch.Tensor: """Register an abstract implementation for quantizing tensor. @@ -162,6 +169,8 @@ def _dynamic_block_quantize_impl( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): scale_bits = (scale_exponent_bits, scale_num_bits - scale_exponent_bits - 1) if exponent_bits != 0: @@ -203,6 +212,8 @@ def _dynamic_block_quantize_impl_abstract( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): """Register an abstract implementation for dynamic block quantization. @@ -223,17 +234,20 @@ def _dynamic_block_quantize_impl_abstract( torch.library.define( "tensorrt::quantize_op", "(Tensor input, Tensor amax, int num_bits, int exponent_bits, " - "bool unsigned, bool narrow_range) -> Tensor", + "bool unsigned, bool narrow_range, str? trt_high_precision_dtype=None, " + "int? block_size=None, int? axis=None) -> Tensor", ) torch.library.define( "tensorrt::dynamic_block_quantize_op", "(Tensor input, int block_size, Tensor amax, int num_bits, int exponent_bits, " - "int scale_num_bits, int scale_exponent_bits) -> Tensor", + "int scale_num_bits, int scale_exponent_bits, str? trt_high_precision_dtype=None, " + "str? onnx_quantizer_type=None) -> Tensor", ) torch.library.define( "tensorrt::dynamic_block_quantize_op.overload", "(Tensor input, int block_size, None amax, int num_bits, int exponent_bits, " - "int scale_num_bits, int scale_exponent_bits) -> Tensor", + "int scale_num_bits, int scale_exponent_bits, str? trt_high_precision_dtype=None, " + "str? onnx_quantizer_type=None) -> Tensor", ) # Implement the None amax case @@ -245,6 +259,8 @@ def _dynamic_block_quantize_impl_none_amax( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): return torch.empty_like(inputs) @@ -371,7 +387,7 @@ def legacy_quant_func(): outputs = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) return outputs - if not inputs.is_cuda: + if not inputs.is_cuda and not torch.compiler.is_exporting(): outputs = legacy_quant_func() else: try: @@ -382,6 +398,9 @@ def legacy_quant_func(): exponent_bits=0, unsigned=unsigned, narrow_range=narrow_range, + trt_high_precision_dtype=trt_high_precision_dtype, + block_size=block_size, + axis=axis, ) except (AttributeError, ValueError): # AttributeError: cuda_ext is not imported, possibly due to CPU only installation @@ -447,6 +466,7 @@ def forward( exponent_bits=4, unsigned=False, narrow_range=False, + trt_high_precision_dtype=trt_high_precision_dtype, ) if bias is not None: @@ -461,7 +481,6 @@ def backward(ctx, grad_outputs): def _dynamic_block_quantize_forward( - ctx, inputs, block_size, amax, @@ -469,7 +488,6 @@ def _dynamic_block_quantize_forward( scale_bits, trt_high_precision_dtype=None, onnx_quantizer_type="dynamic", - pass_through_bwd=True, ): """Forward method.""" if isinstance(num_bits, int): @@ -490,6 +508,8 @@ def _dynamic_block_quantize_forward( exponent_bits, scale_num_bits, scale_exponent_bits, + trt_high_precision_dtype, + onnx_quantizer_type, ) return outputs @@ -551,7 +571,6 @@ def forward( """Forward method.""" _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax) return _dynamic_block_quantize_forward( - ctx, inputs, block_size, amax, @@ -559,7 +578,6 @@ def forward( scale_bits, trt_high_precision_dtype, onnx_quantizer_type, - pass_through_bwd, ) @staticmethod diff --git a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py index 2fd1a1c8504..8bb334ae638 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -67,6 +67,11 @@ def forward(self, inputs): return super().forward(inputs) + self.fp32_buffer +class _PythonScalarArgs(nn.Module): + def forward(self, x, flag: bool = True, count: int = 2, scale: float = 1.5): + return x * scale + count if flag else x / scale - count + + def _make_fp8_model(source_dtype, kind="fp8"): if kind == "format": model = nn.Sequential(*(nn.Linear(128, 128, bias=False) for _ in range(2))) @@ -130,7 +135,9 @@ def test_onnx_dynamo_export(skip_on_windows, model: BaseDeployModel): args = model.get_args() with pytest.raises(AssertionError) if model.compile_fail else nullcontext(): - onnx_bytes, _ = get_onnx_bytes_and_metadata(model, args, dynamo_export=True) + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, args, dynamo_export=True, onnx_opset=21 + ) onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) model_bytes = onnx_bytes_obj.get_onnx_model_file_bytes() @@ -141,6 +148,25 @@ def test_onnx_dynamo_export(skip_on_windows, model: BaseDeployModel): assert onnx.load_model_from_string(model_bytes) +def test_onnx_export_explicit_python_scalars(): + payload, _ = get_onnx_bytes_and_metadata( + _PythonScalarArgs().eval(), + (torch.ones(1), False, 3, 2.5), + dynamo_export=True, + onnx_opset=21, + ) + exported = onnx.load_model_from_string( + OnnxBytes.from_bytes(payload).get_onnx_model_file_bytes() + ) + + onnx.checker.check_model(exported) + assert [value.name for value in exported.graph.input] == ["x"] + assert [node.op_type for node in exported.graph.node if node.op_type in {"Div", "Sub"}] == [ + "Div", + "Sub", + ] + + @pytest.mark.parametrize("model", deploy_benchmark_all.values(), ids=deploy_benchmark_all.keys()) def test_onnx_export_and_inputs(model: BaseDeployModel): # try it for all potential numeric types diff --git a/tests/unit/torch/quantization/test_dynamo_onnx_export.py b/tests/unit/torch/quantization/test_dynamo_onnx_export.py new file mode 100644 index 00000000000..7eb424175d0 --- /dev/null +++ b/tests/unit/torch/quantization/test_dynamo_onnx_export.py @@ -0,0 +1,378 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""CPU contract tests for the supported Dynamo ONNX export path.""" + +import copy +import inspect + +import pytest +import torch +from _test_utils.torch.quantization.models import SimpleLinear +from torch import nn + +onnx = pytest.importorskip("onnx") +pytest.importorskip("onnxscript") + +import modelopt.torch.quantization as mtq +from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata +from modelopt.torch.quantization import tensor_quant +from modelopt.torch.quantization._dynamo_onnx import _get_dynamo_onnx_translation_table + +_NVFP4_ARGS = (4, 2, 8, 4, "Float") +_MXFP8_ARGS = (8, 4, 9, 8, "Float") + + +class _StrictQuantOp(nn.Module): + def __init__(self, quant_format): + super().__init__() + self.quant_format = quant_format + + def forward(self, x, amax): + if self.quant_format.startswith("fp8"): + return torch.ops.tensorrt.quantize_op.default( + x, + None if self.quant_format == "fp8_none" else amax, + 8, + 4, + False, + False, + "Float", + ) + if self.quant_format == "int8": + return torch.ops.tensorrt.quantize_op.default( + x, amax, 8, 0, False, False, "Float", None, 0 + ) + if self.quant_format.startswith("nvfp4"): + return torch.ops.tensorrt.dynamic_block_quantize_op.default( + x, 16, amax, *_NVFP4_ARGS, self.quant_format.removeprefix("nvfp4_") + ) + return torch.ops.tensorrt.dynamic_block_quantize_op.overload( + x, 32, None, *_MXFP8_ARGS, self.quant_format.removeprefix("mxfp8_") + ) + + +class _INT4ScaleFanout(nn.Module): + def __init__(self): + super().__init__() + self.weight0 = nn.Parameter(torch.randn(64, 128)) + self.weight1 = nn.Parameter(torch.randn(64, 128)) + self.register_buffer("amax0", torch.ones(64, 1)) + self.register_buffer("amax1", torch.ones(64, 1)) + + def forward(self, x): + quantized = [ + torch.ops.tensorrt.quantize_op.default(weight, amax, 4, 0, False, True, "Float", 128, 1) + for weight, amax in ((self.weight0, self.amax0), (self.weight1, self.amax1)) + ] + return x @ quantized[0].T + x @ quantized[1].T + + +def _assert_valid_without_custom_ops(model): + onnx.checker.check_model(model, full_check=True) + assert not model.functions + assert not any( + node.domain == "tensorrt" or node.op_type in {"quantize_op", "dynamic_block_quantize_op"} + for node in model.graph.node + ) + return model + + +def _helper_export(model, inputs, name, opset, tmp_path): + payload, _ = get_onnx_bytes_and_metadata( + model, + inputs, + model_name=name, + dynamo_export=True, + onnx_opset=opset, + ) + package = OnnxBytes.from_bytes(payload) + directory = tmp_path / name + package.write_to_disk(str(directory)) + path = directory / f"{package.model_name}.onnx" + return _assert_valid_without_custom_ops(onnx.load(path, load_external_data=True)) + + +def _raw_export(model, inputs, path, opset=21): + compatibility_options = ( + {"fallback": False} if "fallback" in inspect.signature(torch.onnx.export).parameters else {} + ) + torch.onnx.export( + model, + inputs, + path, + dynamo=True, + opset_version=opset, + custom_translation_table=_get_dynamo_onnx_translation_table(), + **compatibility_options, + ) + return _assert_valid_without_custom_ops(onnx.load(path)) + + +def _auto_quantize(model, sample_input, formats, effective_bits): + return mtq.auto_quantize( + model, + constraints={"effective_bits": effective_bits}, + quantization_formats=[copy.deepcopy(config) for config in formats], + data_loader=[sample_input], + forward_step=lambda candidate, batch: candidate(batch), + loss_func=lambda output, _batch: output.float().square().mean(), + num_calib_steps=1, + num_score_steps=1, + )[0] + + +def _linears(features, count=2): + return nn.Sequential(*(nn.Linear(features, features, bias=False) for _ in range(count))) + + +def _quantize(model, inputs, config): + return mtq.quantize( + model, copy.deepcopy(config), forward_loop=lambda candidate: candidate(*inputs) + ) + + +@pytest.fixture +def block_quant_export_only(monkeypatch): + real_op = tensor_quant.dynamic_block_quantize_op + + def cpu_stub(inputs, *args): + return real_op(inputs, *args) if torch.compiler.is_exporting() else inputs + + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_stub) + + +def _weight_dqs(model, op_type): + initializers = {item.name for item in model.graph.initializer} + return [ + node + for node in model.graph.node + if node.op_type == op_type and node.input[0] in initializers + ] + + +def _assert_fp16_gemm_with_float_io(model): + graph = model.graph + tensor_types = { + value.name: value.type.tensor_type.elem_type + for value in (*graph.input, *graph.value_info, *graph.output) + } + tensor_types.update( + {initializer.name: initializer.data_type for initializer in graph.initializer} + ) + gemms = [node for node in graph.node if node.op_type == "Gemm"] + assert gemms and all( + tensor_types[input_name] == onnx.TensorProto.FLOAT16 + for node in gemms + for input_name in node.input + ) + producers = {output: node for node in graph.node for output in node.output} + output_cast = producers[graph.output[0].name] + assert output_cast.op_type == "Cast" and output_cast.attribute[0].i == onnx.TensorProto.FLOAT + assert producers[output_cast.input[0]].op_type == "Gemm" + + +_STRICT_CASES = [ + ("fp8_none", 21, ("trt", "TRT_FP8QuantizeLinear")), + ("int8", 21, ("", "QuantizeLinear")), + ("nvfp4_dynamic", 21, ("trt", "TRT_FP4DynamicQuantize")), + ("nvfp4_static", 21, ("trt", "TRT_FP4QDQ")), + ("mxfp8_dynamic", 21, ("trt", "TRT_MXFP8DynamicQuantize")), + ("mxfp8_static", 24, ("trt", "TRT_MXFP8DequantizeLinear")), +] + + +@pytest.mark.parametrize(("quant_format", "opset", "expected_node"), _STRICT_CASES) +def test_private_table_strict_formats(tmp_path, quant_format, opset, expected_node): + sample_input = torch.randn(4, 32) + amax = torch.ones(4, 1) if quant_format == "int8" else torch.tensor(1.0) + exported_program = torch.export.export( + _StrictQuantOp(quant_format), (sample_input, amax), strict=True + ) + exported = _raw_export(exported_program, (), tmp_path / f"{quant_format}_{opset}.onnx", opset) + assert expected_node in {(node.domain, node.op_type) for node in exported.graph.node} + assert {item.domain: item.version for item in exported.opset_import}[""] == opset + + +def test_int8_fp16_bias_preserves_source_dtype(tmp_path): + inputs = (torch.randn(2, 32, dtype=torch.float16),) + model = _quantize(nn.Linear(32, 32, dtype=torch.float16).eval(), inputs, mtq.INT8_DEFAULT_CFG) + exported = _helper_export(model, inputs, "int8_fp16", 21, tmp_path) + assert exported.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT16 + + +@pytest.mark.parametrize( + ("fmt", "inputs", "amax", "message"), + [ + ("fp8", torch.randn(2, 4), torch.ones(2, 1), "scalar amax"), + ("int8", torch.randn(2, 3, 4), torch.ones(2, 1, 4), "multi-axis"), + ("nvfp4_dynamic", torch.randn(1, 2, 3, 32), torch.tensor(1.0), "rank 2 or 3"), + ], +) +def test_private_table_rejects_shapes(tmp_path, fmt, inputs, amax, message): + exported_program = torch.export.export(_StrictQuantOp(fmt), (inputs, amax), strict=True) + with pytest.raises(torch.onnx.OnnxExporterError, match=message): + _raw_export(exported_program, (), tmp_path / f"{fmt}.onnx") + + +@pytest.mark.parametrize("quant_format", ["int4", "mxfp8"]) +def test_postprocessors_split_shared_scales(tmp_path, quant_format): + if quant_format == "int4": + inputs = (torch.randn(2, 128),) + model = _INT4ScaleFanout().eval() + exporter, op_type = INT4QuantExporter, "DequantizeLinear" + else: + inputs = (torch.randn(2, 64),) + model = _quantize(_linears(64, 3).eval(), inputs, mtq.MXFP8_DEFAULT_CFG) + exporter, op_type = MXFP8QuantExporter, "TRT_MXFP8DequantizeLinear" + + exported = _raw_export(model, inputs, tmp_path / f"{quant_format}_shared_scale.onnx") + weight_dqs = _weight_dqs(exported, op_type) + assert len(weight_dqs) >= 2 and len({node.input[1] for node in weight_dqs}) == 1 + + exported = _assert_valid_without_custom_ops(exporter.process_model(exported)) + weight_dqs = _weight_dqs(exported, op_type) + assert len({node.input[1] for node in weight_dqs}) == len(weight_dqs) + + +def test_helper_exports_biased_nvfp4_with_default_fp32(tmp_path, block_quant_export_only): + inputs = (SimpleLinear.get_input(),) + model = _quantize(SimpleLinear().eval(), inputs, mtq.NVFP4_DEFAULT_CFG) + + exported = _helper_export(model, inputs, "biased_nvfp4", 21, tmp_path) + + _assert_fp16_gemm_with_float_io(exported) + + +def test_helper_exports_mxfp8_at_opset21(tmp_path, block_quant_export_only): + inputs = (torch.randn(2, 64),) + model = _quantize(nn.Linear(64, 64).eval(), inputs, mtq.MXFP8_DEFAULT_CFG) + exported = _helper_export(model, inputs, "mxfp8", 21, tmp_path) + assert any(node.op_type == "TRT_MXFP8DequantizeLinear" for node in exported.graph.node) + assert {onnx.TensorProto.FLOAT8E4M3FN, onnx.TensorProto.UINT8} <= { + item.data_type for item in exported.graph.initializer + } + _assert_fp16_gemm_with_float_io(exported) + + +@pytest.mark.parametrize("opset", [21, 24]) +def test_helper_exports_fp8_autoquant(tmp_path, opset): + sample_input = torch.randn(2, 16) + model = _auto_quantize( + _linears(16).eval(), + sample_input, + [mtq.FP8_DEFAULT_CFG], + effective_bits=8.0, + ) + exported = _helper_export(model, (sample_input,), f"two_linear_fp8_{opset}", opset, tmp_path) + assert {"QuantizeLinear", "DequantizeLinear"} <= {node.op_type for node in exported.graph.node} + fp8_weights = { + item.name + for item in exported.graph.initializer + if item.data_type == onnx.TensorProto.FLOAT8E4M3FN and list(item.dims) == [16, 16] + } + assert len(fp8_weights) == 2 + + +def test_helper_exports_fp16_int4_awq_at_opset21(tmp_path): + inputs = (torch.randn(2, 256, dtype=torch.float16),) + model = _quantize( + nn.Linear(256, 64, bias=False, dtype=torch.float16).eval(), + inputs, + mtq.INT4_AWQ_CFG, + ) + + exported = _helper_export(model, inputs, "linear_int4", 21, tmp_path) + + dq = next( + node + for node in exported.graph.node + if node.domain == "trt" and node.op_type == "DequantizeLinear" + ) + initializers = {initializer.name: initializer for initializer in exported.graph.initializer} + weight, scale = initializers[dq.input[0]], initializers[dq.input[1]] + assert (weight.data_type, list(weight.dims), list(scale.dims)) == ( + onnx.TensorProto.INT4, + [64, 256], + [64, 2], + ) + + +def test_helper_exports_mixed_autoquant_without_markers(tmp_path): + sample_input = torch.randn(2, 128) + model = _auto_quantize( + _linears(128).eval(), + sample_input, + [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.FP8_DEFAULT_CFG], + effective_bits=6.0, + ) + + exported = _helper_export(model, (sample_input,), "two_linear_mixed", 24, tmp_path) + + initializer_dtypes = {initializer.data_type for initializer in exported.graph.initializer} + assert {onnx.TensorProto.INT4, onnx.TensorProto.FLOAT8E4M3FN} <= initializer_dtypes + assert not any(node.op_type.startswith("TRT_") for node in exported.graph.node) + + +def test_custom_op_schemas_keep_legacy_positional_calls(): + inputs = torch.randn(4, 32) + amax = torch.tensor(1.0) + fp8 = torch.ops.tensorrt.quantize_op(inputs, amax, 8, 4, False, False) + mxfp8 = torch.ops.tensorrt.dynamic_block_quantize_op.overload(inputs, 32, None, 8, 4, 9, 8) + + assert (fp8.shape, fp8.dtype) == (inputs.shape, inputs.dtype) + assert (mxfp8.shape, mxfp8.dtype) == (inputs.shape, inputs.dtype) + + +def test_nvfp4_rejects_marker_output_fanout(): + weight = onnx.numpy_helper.from_array(torch.ones(4, 32).numpy(), "weight") + nodes = [ + onnx.helper.make_node("TRT_FP4QDQ", ["weight"], ["weight_dq"], block_size=16), + onnx.helper.make_node("Identity", ["weight_dq"], ["output0"]), + onnx.helper.make_node("Identity", ["weight_dq"], ["output1"]), + ] + graph = onnx.helper.make_graph(nodes, "fanout", [], [], [weight]) + with pytest.raises(NotImplementedError, match="expected one consumer"): + NVFP4QuantExporter.pre_process(onnx.helper.make_model(graph)) + + +def test_helper_rejects_unsupported_dynamo_options(): + model, inputs = nn.Linear(128, 64, bias=False), (torch.randn(1, 128),) + with pytest.raises(ValueError, match="opset 21 or newer"): + get_onnx_bytes_and_metadata(model, inputs, dynamo_export=True, onnx_opset=20) + with pytest.raises(NotImplementedError, match="dynamic_axes"): + get_onnx_bytes_and_metadata( + model, + inputs, + dynamo_export=True, + onnx_opset=21, + dynamic_axes={"input": {0: "batch"}}, + ) + + +def test_helper_disables_legacy_fallback(monkeypatch): + if "fallback" not in inspect.signature(torch.onnx.export).parameters: + pytest.skip("This PyTorch version has no legacy fallback option.") + + def assert_disabled(*args, fallback=True, **kwargs): + assert fallback is False + raise RuntimeError("checked") + + monkeypatch.setattr(torch.onnx, "export", assert_disabled) + with pytest.raises(RuntimeError, match="checked"): + get_onnx_bytes_and_metadata( + nn.Identity(), (torch.ones(1),), dynamo_export=True, onnx_opset=21 + )