diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 68927078e83..a1f36ab59b9 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/docs/source/guides/_pytorch_quantization.rst b/docs/source/guides/_pytorch_quantization.rst index f8f12b068ba..dff2006f387 100644 --- a/docs/source/guides/_pytorch_quantization.rst +++ b/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. diff --git a/examples/onnx_ptq/download_example_onnx.py b/examples/onnx_ptq/download_example_onnx.py index e78ff2ecbf2..2f35aa5e4fb 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) @@ -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}") diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index f479bab3ae1..c8efa902d66 100644 --- a/examples/torch_onnx/README.md +++ b/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= ``` +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=` to use a different diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index e0ffc75a294..2ee261da794 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/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}") diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 427a7791f3b..745ce06e800 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/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" ) diff --git a/modelopt/onnx/export/int4_exporter.py b/modelopt/onnx/export/int4_exporter.py index 0da217ae76f..6c3f6405e74 100644 --- a/modelopt/onnx/export/int4_exporter.py +++ b/modelopt/onnx/export/int4_exporter.py @@ -15,17 +15,46 @@ """INT4 quantization exporter.""" +import math + import onnx 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 +def _get_weight_dq_nodes(graph: onnx.GraphProto) -> list[onnx.NodeProto]: + initializer_names = {initializer.name for initializer in graph.initializer} + tensor_producer_map = get_tensor_producer_nodes(graph) + + def _has_initializer_source(tensor_name: str) -> bool: + if tensor_name in initializer_names: + return True + producer = tensor_producer_map.get(tensor_name) + return ( + producer is not None + and producer.op_type == "Reshape" + and producer.input[0] in initializer_names + ) + + return [ + node + for node in graph.node + if node.op_type == "DequantizeLinear" + and node.domain == "trt" + and any(attr.name == "block_size" for attr in node.attribute) + and _has_initializer_source(node.input[0]) + ] + + class INT4QuantExporter(ONNXQuantExporter): """Exporter for INT4 quantization.""" @@ -33,59 +62,198 @@ class INT4QuantExporter(ONNXQuantExporter): 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"] - tensor_producer_map = get_tensor_producer_nodes(graph, get_initializer_producers=True) - - nodes_to_remove = [] - for node in weight_dq_nodes: + value_info_map = { + value_info.name: value_info + for value_info in (*graph.input, *graph.value_info, *graph.output) + } + weight_dq_nodes = _get_weight_dq_nodes(graph) + tensor_producer_map = get_tensor_producer_nodes(graph) + initializer_map = {initializer.name: initializer for initializer in graph.initializer} + tensor_consumer_map = get_tensor_consumer_nodes(graph) + node_outputs_to_remove = set() + constant_outputs_to_remove = set() + initializer_candidates_to_remove = set() + tensor_names = { + name for node in graph.node for name in (*node.input, *node.output) if name + } | {initializer.name for initializer in graph.initializer} + + def _get_value_info_shape(tensor_name: str) -> list[int] | None: + value_info = value_info_map.get(tensor_name) + if value_info is None: + return None + dims = value_info.type.tensor_type.shape.dim + if not all(dim.HasField("dim_value") for dim in dims): + return None + return [dim.dim_value for dim in dims] + + def _get_value_info_dtype(tensor_name: str) -> int | None: + value_info = value_info_map.get(tensor_name) + return None if value_info is None else value_info.type.tensor_type.elem_type + + def _get_constant_values(tensor_name: str) -> list[int] | None: + if tensor_name in initializer_map: + return numpy_helper.to_array(initializer_map[tensor_name]).reshape(-1).tolist() + producer = tensor_producer_map.get(tensor_name) + if producer is None or producer.op_type != "Constant": + return None + value = next((attr.t for attr in producer.attribute if attr.name == "value"), None) + return None if value is None else numpy_helper.to_array(value).reshape(-1).tolist() + + def _get_reshape_output_shape( + reshape_node: onnx.NodeProto, input_shape: list[int] + ) -> list[int]: + output_shape = _get_value_info_shape(reshape_node.output[0]) + if output_shape is not None: + return output_shape + + requested_shape = _get_constant_values(reshape_node.input[1]) + if requested_shape is None: + raise ValueError(f"Unable to determine shape for Reshape node {reshape_node.name}") + + allowzero = next( + (attr.i for attr in reshape_node.attribute if attr.name == "allowzero"), 0 + ) + output_shape = [] + inferred_axis = None + for axis, dim in enumerate(requested_shape): + if dim == 0 and not allowzero: + dim = input_shape[axis] + elif dim == -1: + if inferred_axis is not None: + raise ValueError(f"Multiple inferred dimensions in {reshape_node.name}") + inferred_axis = axis + dim = 1 + elif dim < 0: + raise ValueError(f"Invalid dimension {dim} in {reshape_node.name}") + output_shape.append(dim) + + input_size = math.prod(input_shape) + known_output_size = math.prod(output_shape) + if inferred_axis is not None: + if known_output_size == 0 or input_size % known_output_size: + raise ValueError(f"Invalid shape for Reshape node {reshape_node.name}") + output_shape[inferred_axis] = input_size // known_output_size + elif input_size != known_output_size: + raise ValueError(f"Invalid shape for Reshape node {reshape_node.name}") + return output_shape + + def _mark_shape_input_for_removal(reshape_node: onnx.NodeProto): + shape_name = reshape_node.input[1] + initializer_candidates_to_remove.add(shape_name) + producer = tensor_producer_map.get(shape_name) + if producer is not None and producer.op_type == "Constant": + constant_outputs_to_remove.update(producer.output) + + def _get_only_child(tensor_name: str, parent_name: str) -> onnx.NodeProto: + child_nodes = tensor_consumer_map.get(tensor_name, []) + assert len(child_nodes) == 1, f"Expected exactly one child node for {parent_name}" + return child_nodes[0] + + def _clone_shared_initializer_input( + node: onnx.NodeProto, input_index: int, path_index: int + ) -> str: + tensor_name = node.input[input_index] + if len(tensor_consumer_map[tensor_name]) <= 1: + return tensor_name + + tensor = initializer_map.get(tensor_name) + if tensor is None: + producer = tensor_producer_map.get(tensor_name) + if producer is None or producer.op_type != "Constant": + raise ValueError(f"Expected a constant shared input for {node.name}") + tensor = next((attr.t for attr in producer.attribute if attr.name == "value"), None) + if tensor is None: + raise ValueError(f"Expected a tensor value for {producer.name}") + constant_outputs_to_remove.update(producer.output) + else: + initializer_candidates_to_remove.add(tensor_name) + + base_name = f"{tensor_name}_int4_{path_index}" + unique_name = base_name + suffix = 0 + while unique_name in tensor_names: + suffix += 1 + unique_name = f"{base_name}_{suffix}" + tensor_names.add(unique_name) + + cloned_tensor = onnx.TensorProto() + cloned_tensor.CopyFrom(tensor) + cloned_tensor.name = unique_name + graph.initializer.append(cloned_tensor) + initializer_map[unique_name] = cloned_tensor + node.input[input_index] = unique_name + return unique_name + + for path_index, node in enumerate(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 - ] + if weight_name in initializer_map: + weight_shape = list(initializer_map[weight_name].dims) else: - raise ValueError(f"Unable to determine shape of weight tensor {weight_name}") + pre_reshape = tensor_producer_map.get(weight_name) + if ( + pre_reshape is None + or pre_reshape.op_type != "Reshape" + or pre_reshape.input[0] not in initializer_map + ): + raise ValueError( + f"Expected an initializer or constant Reshape input for {node.name}" + ) + + source_name = pre_reshape.input[0] + source_initializer = initializer_map[source_name] + weight_shape = list(source_initializer.dims) + blocked_shape = _get_reshape_output_shape(pre_reshape, weight_shape) + if math.prod(weight_shape) != math.prod(blocked_shape): + raise ValueError(f"Invalid blocked weight shape for {node.name}") + + blocked_initializer = onnx.TensorProto() + blocked_initializer.CopyFrom(source_initializer) + blocked_initializer.name = weight_name + del blocked_initializer.dims[:] + blocked_initializer.dims.extend(blocked_shape) + graph.initializer.append(blocked_initializer) + initializer_map[weight_name] = blocked_initializer + weight_shape = blocked_shape + + node_outputs_to_remove.update(pre_reshape.output) + initializer_candidates_to_remove.add(source_name) + _mark_shape_input_for_removal(pre_reshape) + + weight_name = _clone_shared_initializer_input(node, 0, path_index) + _clone_shared_initializer_input(node, 1, path_index) + + next_node = _get_only_child(node.output[0], node.name) + path_dtype = _get_value_info_dtype(node.output[0]) + preserved_cast = None + while next_node.op_type in {"Cast", "Reshape"}: + if next_node.op_type == "Reshape": + weight_shape = _get_reshape_output_shape(next_node, weight_shape) + if path_dtype is None: + path_dtype = _get_value_info_dtype(next_node.output[0]) + node_outputs_to_remove.update(next_node.output) + _mark_shape_input_for_removal(next_node) + else: + cast_dtype = next(attr.i for attr in next_node.attribute if attr.name == "to") + if path_dtype == cast_dtype: + node_outputs_to_remove.update(next_node.output) + else: + assert preserved_cast is None, ( + f"Expected at most one precision Cast node for {node.name}" + ) + preserved_cast = next_node + path_dtype = cast_dtype + next_node = _get_only_child(next_node.output[0], node.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 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}" - ) + node_outputs_to_remove.update(transpose_node.output) perm = None for attr in transpose_node.attribute: if attr.name == "perm": @@ -97,26 +265,73 @@ def pre_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: perm_attr.name = "_transpose_perm" perm_attr.ints.extend(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 = transpose_child_nodes[0] + matmul_node = _get_only_child(transpose_node.output[0], node.name) + quant_axis = perm.index(len(weight_shape) - 1) + output_shape = [weight_shape[axis] for axis in perm] else: matmul_node = next_node + quant_axis = len(weight_shape) - 1 + output_shape = weight_shape 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] - - # Remove transpose, reshape, and constant nodes - new_nodes = [node for node in graph.node if node.name not in nodes_to_remove] + weight_output = node.output[0] + if preserved_cast is not None: + preserved_cast.input[0] = node.output[0] + weight_output = preserved_cast.output[0] + axis_attr = next((attr for attr in node.attribute if attr.name == "axis"), None) + if axis_attr is None: + axis_attr = node.attribute.add() + axis_attr.name = "axis" + axis_attr.i = quant_axis + output_value_info = value_info_map.get(node.output[0]) + if output_value_info is None: + output_value_info = onnx.helper.make_tensor_value_info( + node.output[0], initializer_map[weight_name].data_type, output_shape + ) + graph.value_info.append(output_value_info) + value_info_map[node.output[0]] = output_value_info + output_dims = output_value_info.type.tensor_type.shape.dim + del output_dims[:] + for dim_value in output_shape: + output_dims.add().dim_value = dim_value + cast_output_value_info = value_info_map.get(weight_output) + if cast_output_value_info is not None: + output_dims = cast_output_value_info.type.tensor_type.shape.dim + del output_dims[:] + for dim_value in output_shape: + output_dims.add().dim_value = dim_value + # Rewire MatMul to use the normalized weight output. + matmul_node.input[1] = weight_output + + new_nodes = [ + node + for node in graph.node + if not any(output in node_outputs_to_remove for output in node.output) + ] + used_tensors = {input_name for node in new_nodes for input_name in node.input} + new_nodes = [ + node + for node in new_nodes + if not ( + node.op_type == "Constant" + and any(output in constant_outputs_to_remove for output in node.output) + and not any(output in used_tensors for output in node.output) + ) + ] + used_tensors = {input_name for node in new_nodes for input_name in node.input} + protected_tensors = used_tensors | {value.name for value in (*graph.input, *graph.output)} + new_initializers = [ + initializer + for initializer in graph.initializer + if initializer.name not in initializer_candidates_to_remove + or initializer.name in protected_tensors + ] del graph.node[:] graph.node.extend(new_nodes) + del graph.initializer[:] + graph.initializer.extend(new_initializers) return onnx_model @@ -125,7 +340,11 @@ 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"] + value_info_map = { + value_info.name: value_info + for value_info in (*graph.input, *graph.value_info, *graph.output) + } + weight_dq_nodes = _get_weight_dq_nodes(graph) tensor_producer_map = get_tensor_producer_nodes(graph, get_initializer_producers=True) for node in weight_dq_nodes: @@ -188,6 +407,14 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: scale_tensor = onnx.numpy_helper.from_array(scale, scale_name) initializer_map[scale_name].CopyFrom(scale_tensor) + scale_value_info = value_info_map.get(scale_name) + if scale_value_info is not None: + tensor_type = scale_value_info.type.tensor_type + tensor_type.elem_type = scale_tensor.data_type + del tensor_type.shape.dim[:] + for dim_value in scale_tensor.dims: + tensor_type.shape.dim.add().dim_value = dim_value + # Update weight tensor weight_tensor = numpy_helper.from_array(weight, weight_name) initializer_map[weight_name].CopyFrom(weight_tensor) @@ -207,7 +434,11 @@ 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"] + value_info_map = { + value_info.name: value_info + for value_info in (*graph.input, *graph.value_info, *graph.output) + } + weight_dq_nodes = _get_weight_dq_nodes(graph) for node in weight_dq_nodes: weight_name = node.input[0] @@ -218,6 +449,12 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights_int4_onnx.data_type = onnx.TensorProto.INT4 weights_int4_onnx.dims[0] = weight_shape[0] initializer_map[weight_name].CopyFrom(weights_int4_onnx) + if weight_name in value_info_map: + tensor_type = value_info_map[weight_name].type.tensor_type + tensor_type.elem_type = onnx.TensorProto.INT4 + del tensor_type.shape.dim[:] + for dim_value in weight_shape: + tensor_type.shape.dim.add().dim_value = dim_value logger.debug(f"Converted {weight_name} to INT4 precision") return onnx_model @@ -253,12 +490,11 @@ def is_fp32_cast(node: onnx.NodeProto) -> bool: for node in graph.node: if is_pre_quant_scale_node(node): 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 len(pqs_child_nodes) == 1 and pqs_child_nodes[0].op_type == "Cast": + cast_node = pqs_child_nodes[0] + 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..a6c6f2c2d52 100644 --- a/modelopt/onnx/export/mxfp8_exporter.py +++ b/modelopt/onnx/export/mxfp8_exporter.py @@ -20,7 +20,10 @@ 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 @@ -32,13 +35,25 @@ DEFAULT_QUANT_AXIS = -1 +def _sync_initializer_metadata(graph: onnx.GraphProto, initializer: onnx.TensorProto) -> None: + """Synchronize existing type and shape declarations for an initializer.""" + for value_info in (*graph.input, *graph.value_info, *graph.output): + if value_info.name != initializer.name: + continue + tensor_type = value_info.type.tensor_type + tensor_type.elem_type = initializer.data_type + del tensor_type.shape.dim[:] + for dim_value in initializer.dims: + tensor_type.shape.dim.add().dim_value = dim_value + + 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 +85,74 @@ 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 + weight_dq_nodes = _get_weight_dq_nodes(graph) + tensor_consumer_map = get_tensor_consumer_nodes(graph) + tensor_producer_map = get_tensor_producer_nodes(graph) + initializer_map = {initializer.name: initializer for initializer in graph.initializer} + tensor_names = { + name for node in graph.node for name in (*node.input, *node.output) if name + } | set(initializer_map) + initializer_candidates_to_remove = set() + constant_outputs_to_remove = set() + + def _clone_shared_input(node: onnx.NodeProto, input_index: int, path_index: int): + tensor_name = node.input[input_index] + if len(tensor_consumer_map[tensor_name]) <= 1: + return + + tensor = initializer_map.get(tensor_name) + if tensor is None: + producer = tensor_producer_map.get(tensor_name) + if producer is None or producer.op_type != "Constant": + raise ValueError(f"Expected a constant shared input for {node.name}") + tensor = next((attr.t for attr in producer.attribute if attr.name == "value"), None) + if tensor is None: + raise ValueError(f"Expected a tensor value for {producer.name}") + constant_outputs_to_remove.update(producer.output) + else: + initializer_candidates_to_remove.add(tensor_name) + + base_name = f"{tensor_name}_mxfp8_{path_index}" + unique_name = base_name + suffix = 0 + while unique_name in tensor_names: + suffix += 1 + unique_name = f"{base_name}_{suffix}" + tensor_names.add(unique_name) + + cloned_tensor = onnx.TensorProto() + cloned_tensor.CopyFrom(tensor) + cloned_tensor.name = unique_name + graph.initializer.append(cloned_tensor) + initializer_map[unique_name] = cloned_tensor + node.input[input_index] = unique_name + + for path_index, node in enumerate(weight_dq_nodes): + _clone_shared_input(node, 0, path_index) + _clone_shared_input(node, 1, path_index) + + used_tensors = {input_name for node in graph.node for input_name in node.input} + protected_tensors = used_tensors | {value.name for value in (*graph.input, *graph.output)} + new_initializers = [ + initializer + for initializer in graph.initializer + if initializer.name not in initializer_candidates_to_remove + or initializer.name in protected_tensors + ] + new_nodes = [ + node + for node in graph.node + if not ( + node.op_type == "Constant" + and any(output in constant_outputs_to_remove for output in node.output) + and not any(output in used_tensors for output in node.output) + ) + ] + del graph.initializer[:] + graph.initializer.extend(new_initializers) + del graph.node[:] + graph.node.extend(new_nodes) return onnx_model @staticmethod @@ -92,17 +175,21 @@ 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 in initializer_map: + scale_tensor = onnx.numpy_helper.from_array(se8m0, scale_name) + initializer_map[scale_name].CopyFrom(scale_tensor) + else: + scale_producer = tensor_producer_map[scale_name] + if scale_producer.op_type == "Constant": + graph.node.remove(scale_producer) + + 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) + initializer_map[scale_name_new] = scale_tensor + node.input[1] = scale_name_new + _sync_initializer_metadata(graph, scale_tensor) return onnx_model @@ -138,6 +225,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: raw=True, ) initializer_map[weight_name].CopyFrom(weights_e4m3) + _sync_initializer_metadata(graph, initializer_map[weight_name]) 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..b5a3a8db7c2 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -175,6 +175,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 @@ -428,6 +429,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 01fb754bbae..44f77beca0b 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -59,7 +59,10 @@ remove_node_training_mode, remove_redundant_casts, ) -from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers +from modelopt.torch.quantization.export_onnx import ( + configure_linear_module_onnx_quantizers, + get_dynamo_onnx_translation_table, +) from modelopt.torch.utils import flatten_tree, standardize_named_model_args from modelopt.torch.utils._pytree import TreeSpec @@ -527,11 +530,16 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module - # 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()} + # Standardize model args and pre-convert numeric inputs to ONNX tensor inputs. + # Dynamo must preserve signature defaults as Python constants for control-flow specialization. + # Legacy export retains its existing tensorization behavior. + named_args, args_with_default = standardize_named_model_args(model, dummy_input) + named_args = { + name: value + if dynamo_export and name in args_with_default + else _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()) @@ -595,7 +603,9 @@ def get_onnx_bytes_and_metadata( ) with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: additional_kwargs = {} - if not dynamo_export: + if dynamo_export: + additional_kwargs["custom_translation_table"] = get_dynamo_onnx_translation_table() + else: additional_kwargs["dynamic_axes"] = dynamic_axes torch.onnx.export( model, @@ -636,7 +646,8 @@ def get_onnx_bytes_and_metadata( if weights_dtype in ["fp16", "bf16"]: if ( - is_int4_quantized(model) + (dynamo_export and is_fp4_quantized(model)) + or is_int4_quantized(model) or is_mxfp8_quantized(model) or is_fp8_quantized(model) or is_int8_quantized(model) diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index e5778c3c96b..33a32e35171 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -103,6 +103,7 @@ """Utility to export a quantized torch model to quantized ONNX.""" import contextlib +from collections.abc import Callable from typing import TYPE_CHECKING import onnx @@ -111,6 +112,12 @@ from torch.onnx import symbolic_helper as sym_help if TYPE_CHECKING: + from onnxscript.function_libs.torch_lib.tensor_typing import ( # noqa: TC004 + TFloat as _DYNAMO_T_FLOAT, # noqa: N814 + ) + from onnxscript.onnx_types import FLOAT4E2M1 as _DYNAMO_FLOAT4_E2M1 # noqa: TC004 + from onnxscript.onnx_types import FLOAT8E4M3FN as _DYNAMO_FLOAT8_E4M3FN # noqa: TC004 + if hasattr(torch.onnx._internal, "jit_utils"): from torch.onnx._internal.jit_utils import GraphContext else: # torch >= 2.9 @@ -126,7 +133,492 @@ } mha_valid_precisions = {"Half", "BFloat16"} -torch_dtype_map = {"Float": torch.float32, "Half": torch.float16, "BFloat16": torch.bfloat16} +torch_dtype_map = { + "Float": torch.float32, + "Half": torch.float16, + "BFloat16": torch.bfloat16, +} +torch_onnx_dtype_map = { + torch.bfloat16: onnx.TensorProto.BFLOAT16, + torch.float16: onnx.TensorProto.FLOAT16, + torch.float32: onnx.TensorProto.FLOAT, + torch.int8: onnx.TensorProto.INT8, + torch.uint8: onnx.TensorProto.UINT8, +} + + +def _dynamo_identity(inputs: torch.Tensor) -> torch.Tensor: + """Keep a default-domain opset import in custom-op-only Dynamo graphs.""" + return torch.onnx.ops.symbolic( + "ai.onnx::Identity", + (inputs,), + dtype=inputs.dtype, + shape=inputs.shape, + version=21, + ) + + +def _dynamo_cast(inputs: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return torch.onnx.ops.symbolic( + "ai.onnx::Cast", + (inputs,), + {"to": torch_onnx_dtype_map[dtype]}, + dtype=dtype, + shape=inputs.shape, + version=21, + ) + + +def _dynamo_shape_with_axis_divided( + inputs: torch.Tensor, divisor: int, axis: int = -1 +) -> list[int | torch.SymInt]: + shape = list(inputs.shape) + shape[axis] = shape[axis] // divisor + return shape + + +def export_int8_dynamo( + inputs: torch.Tensor, + amax: torch.Tensor, + num_bits: int, + unsigned: bool, + narrow_range: bool, + trt_high_precision_dtype: str | None, +) -> torch.Tensor: + """Export INT8 Q/DQ with the Dynamo ONNX exporter.""" + assert num_bits == 8, "Number of bits must be 8 for INT8 ONNX export." + maxbound = (1 << (num_bits - 1 + int(unsigned))) - 1 + output_dtype = ( + inputs.dtype + if trt_high_precision_dtype is None + else torch_dtype_map[trt_high_precision_dtype] + ) + + if amax.numel() == 1: + axis = None + zero_point = torch.zeros((), dtype=output_dtype, device=amax.device) + else: + amax_init_shape = amax.shape + amax = amax.squeeze().detach() + assert len(amax.shape) == 1, "ONNX does not support multi-axis quantization." + zero_point = torch.zeros_like(amax, dtype=output_dtype) + axis = list(amax_init_shape).index(next(iter(amax.shape))) + + if not unsigned: + assert not narrow_range, "ONNX does not support unsigned narrow range INT8." + zero_point_dtype = torch.uint8 if unsigned else torch.int8 + zero_point = _dynamo_cast(zero_point, zero_point_dtype) + + scale = amax.to(output_dtype) / maxbound + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + assert output_dtype in (inputs.dtype, torch.float32, torch.bfloat16), ( + "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." + ) + + if output_dtype != inputs.dtype: + inputs = _dynamo_cast(inputs, output_dtype) + attrs = {} if axis is None else {"axis": axis} + quantized = torch.onnx.ops.symbolic( + "ai.onnx::QuantizeLinear", + (inputs, scale, zero_point), + attrs, + dtype=zero_point_dtype, + shape=inputs.shape, + version=21, + ) + return torch.onnx.ops.symbolic( + "ai.onnx::DequantizeLinear", + (quantized, scale, zero_point), + attrs, + dtype=output_dtype, + shape=inputs.shape, + version=21, + ) + + +def export_int4_dynamo( + inputs: torch.Tensor, + amax: torch.Tensor, + num_bits: int, + trt_high_precision_dtype: str | None, + block_size: int, + axis: int, +) -> torch.Tensor: + """Export INT4 DQ with the Dynamo ONNX exporter.""" + assert num_bits == 4, "Number of bits must be 4 for INT4 ONNX export." + output_dtype = ( + inputs.dtype + if trt_high_precision_dtype is None + else torch_dtype_map[trt_high_precision_dtype] + ) + scale = amax / 7.0 + return torch.onnx.ops.symbolic( + "trt::DequantizeLinear", + (inputs, scale), + {"axis": axis, "block_size": block_size}, + dtype=output_dtype, + shape=inputs.shape, + version=1, + ) + + +def export_fp8_dynamo( + inputs: torch.Tensor, + amax: torch.Tensor, + trt_high_precision_dtype: str | None, +) -> torch.Tensor: + """Export FP8 Q/DQ with the Dynamo ONNX exporter.""" + del trt_high_precision_dtype + scale = (amax / 448.0).to(inputs.dtype) + quantized = torch.onnx.ops.symbolic( + "trt::TRT_FP8QuantizeLinear", + (inputs, scale), + dtype=torch.uint8, + shape=inputs.shape, + version=1, + ) + return torch.onnx.ops.symbolic( + "trt::TRT_FP8DequantizeLinear", + (quantized, scale), + dtype=inputs.dtype, + shape=inputs.shape, + version=1, + ) + + +def export_fp4_dynamo( + inputs: torch.Tensor, + block_size: int, + amax: torch.Tensor | None, + num_bits: tuple[int, int], + trt_high_precision_dtype: str | None, + onnx_quantizer_type: str, +) -> torch.Tensor: + """Export NVFP4 quantization with the Dynamo ONNX exporter.""" + if onnx_quantizer_type != "dynamic": + output = torch.onnx.ops.symbolic( + "trt::TRT_FP4QDQ", + (inputs,), + {"block_size": block_size}, + dtype=inputs.dtype, + shape=inputs.shape, + version=1, + ) + return _dynamo_identity(output) + + assert num_bits == (2, 1) + output_dtype = ( + inputs.dtype + if trt_high_precision_dtype is None + else torch_dtype_map[trt_high_precision_dtype] + ) + if output_dtype != inputs.dtype: + inputs = _dynamo_cast(inputs, output_dtype) + if amax is None: + scale = torch.ones((), dtype=torch.float32, device=inputs.device) + else: + scale = amax.to(torch.float32) / (6.0 * 448.0) + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + + quantized_shape = _dynamo_shape_with_axis_divided(inputs, 2) + scale_shape = _dynamo_shape_with_axis_divided(inputs, block_size) + dynamic_quantize_args = ( + "trt::TRT_FP4DynamicQuantize", + (inputs, scale), + { + "axis": -1, + "block_size": block_size, + "scale_type": onnx.TensorProto.FLOAT8E4M3FN, + }, + ) + try: + quantized, dynamic_scale = torch.onnx.ops.symbolic_multi_out( + *dynamic_quantize_args, + dtypes=(torch.float4_e2m1fn_x2, torch.float8_e4m3fn), + shapes=(quantized_shape, scale_shape), + version=1, + ) + except RuntimeError as e: + if "Unsupported dtype: torch.float4_e2m1fn_x2" not in str(e): + raise + quantized, dynamic_scale = torch.onnx.ops.symbolic_multi_out( + *dynamic_quantize_args, + dtypes=(onnx.TensorProto.FLOAT4E2M1, torch.float8_e4m3fn), + shapes=(inputs.shape, scale_shape), + version=1, + ) + dequantized_scale = torch.onnx.ops.symbolic( + "ai.onnx::DequantizeLinear", + (dynamic_scale, scale), + dtype=torch.float32, + shape=scale_shape, + version=21, + ) + dequantized = torch.onnx.ops.symbolic( + "trt::DequantizeLinear", + (quantized, dequantized_scale), + {"axis": -1, "block_size": block_size}, + dtype=torch.float32, + shape=inputs.shape, + version=1, + ) + return dequantized if output_dtype == torch.float32 else _dynamo_cast(dequantized, output_dtype) + + +def export_mxfp8_dynamo( + inputs: torch.Tensor, + onnx_quantizer_type: str, + block_size: int, +) -> torch.Tensor: + """Export MXFP8 quantization with the Dynamo ONNX exporter.""" + if onnx_quantizer_type == "dynamic": + scale_shape = _dynamo_shape_with_axis_divided(inputs, block_size) + quantized, scale = torch.onnx.ops.symbolic_multi_out( + "trt::TRT_MXFP8DynamicQuantize", + (inputs,), + { + "axis": -1, + "block_size": block_size, + "output_dtype": onnx.TensorProto.FLOAT8E4M3FN, + }, + dtypes=(torch.float8_e4m3fn, torch.uint8), + shapes=(inputs.shape, scale_shape), + version=1, + ) + else: + quantized = inputs + scale = torch.ones((), dtype=inputs.dtype, device=inputs.device) + + output = torch.onnx.ops.symbolic( + "trt::TRT_MXFP8DequantizeLinear", + (quantized, scale), + { + "axis": -1, + "block_size": block_size, + "output_dtype": torch_onnx_dtype_map[inputs.dtype], + }, + dtype=inputs.dtype, + shape=inputs.shape, + version=1, + ) + return _dynamo_identity(output) + + +def get_dynamo_onnx_translation_table() -> dict[Callable, Callable]: + """Return ModelOpt custom-op translations for the Dynamo ONNX exporter.""" + import onnxscript + from onnxscript.function_libs.torch_lib.tensor_typing import TFloat + from onnxscript.onnx_types import FLOAT4E2M1, FLOAT8E4M3FN + + # ONNXScript resolves annotations from module globals while compiling nested functions. + globals().update( + { + "_DYNAMO_FLOAT4_E2M1": FLOAT4E2M1[...], + "_DYNAMO_FLOAT8_E4M3FN": FLOAT8E4M3FN[...], + "_DYNAMO_T_FLOAT": TFloat, + } + ) + op = onnxscript.opset21 + trt = onnxscript.values.Opset(domain="trt", version=1) + + @onnxscript.script(trt) + def _fp8_qdq(inputs: _DYNAMO_T_FLOAT, scale: _DYNAMO_T_FLOAT) -> _DYNAMO_T_FLOAT: + quantized = trt.TRT_FP8QuantizeLinear(inputs, scale) + return trt.TRT_FP8DequantizeLinear(quantized, scale) + + @onnxscript.script(trt) + def _int4_dq( + inputs: _DYNAMO_T_FLOAT, scale: _DYNAMO_T_FLOAT, axis: int, block_size: int + ) -> _DYNAMO_T_FLOAT: + return trt.DequantizeLinear(inputs, scale, axis=axis, block_size=block_size) + + @onnxscript.script(trt) + def _fp4_qdq(inputs: _DYNAMO_T_FLOAT, block_size: int) -> _DYNAMO_T_FLOAT: + return trt.TRT_FP4QDQ(inputs, block_size=block_size) + + @onnxscript.script(trt) + def _fp4_dynamic_quantize( + inputs: _DYNAMO_T_FLOAT, + scale: _DYNAMO_T_FLOAT, + block_size: int, + ) -> tuple[_DYNAMO_FLOAT4_E2M1, _DYNAMO_FLOAT8_E4M3FN]: + quantized, dynamic_scale = trt.TRT_FP4DynamicQuantize( + inputs, + scale, + axis=-1, + block_size=block_size, + scale_type=17, + ) + return quantized, dynamic_scale + + @onnxscript.script(trt) + def _fp4_dq( + inputs: _DYNAMO_FLOAT4_E2M1, + scale: _DYNAMO_T_FLOAT, + block_size: int, + ) -> _DYNAMO_T_FLOAT: + return trt.DequantizeLinear(inputs, scale, axis=-1, block_size=block_size) + + @onnxscript.script(trt) + def _mxfp8_dynamic_qdq( + inputs: _DYNAMO_T_FLOAT, block_size: int, output_dtype: int + ) -> _DYNAMO_T_FLOAT: + quantized, scale = trt.TRT_MXFP8DynamicQuantize( + inputs, + axis=-1, + block_size=block_size, + output_dtype=17, + ) + return trt.TRT_MXFP8DequantizeLinear( + quantized, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + @onnxscript.script(trt) + def _mxfp8_static_dq( + inputs: _DYNAMO_T_FLOAT, + scale: _DYNAMO_T_FLOAT, + block_size: int, + output_dtype: int, + ) -> _DYNAMO_T_FLOAT: + return trt.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 op.Cast(inputs, to=dtype) + + def _resolve_dtype(inputs, trt_high_precision_dtype: str | None) -> int: + if trt_high_precision_dtype is None: + return int(inputs.dtype) + return onnx_dtype_map[trt_high_precision_dtype] + + def _quantize_op_translation( + inputs, + amax, + num_bits: int, + exponent_bits: int, + unsigned: bool, + narrow_range: bool, + trt_high_precision_dtype: str = None, # noqa: RUF013 + block_size: int = None, # noqa: RUF013 + axis: int = None, # noqa: RUF013 + ): + if num_bits == 8 and exponent_bits == 4: + scale = op.CastLike(op.Div(amax, 448.0), inputs) + return _fp8_qdq(inputs, scale) + + output_dtype = _resolve_dtype(inputs, trt_high_precision_dtype) + if num_bits == 8 and exponent_bits == 0: + if not unsigned: + assert not narrow_range, "ONNX does not support unsigned narrow range INT8." + assert output_dtype in ( + int(inputs.dtype), + onnx.TensorProto.FLOAT, + onnx.TensorProto.BFLOAT16, + ), ( + "TRT StronglyType requires both weights and amax to be in the BF16/FP16, " + "or the QDQ in Float." + ) + inputs = _cast(inputs, output_dtype) + amax = op.Squeeze(op.Cast(amax, to=output_dtype)) + scale = op.Div(amax, float((1 << (7 + int(unsigned))) - 1)) + scale = op.Where(op.Equal(scale, 0.0), op.CastLike(1.0, scale), scale) + zero_point_dtype = onnx.TensorProto.UINT8 if unsigned else onnx.TensorProto.INT8 + zero_point = op.Cast(op.Mul(amax, 0.0), to=zero_point_dtype) + if axis is None: + quantized = op.QuantizeLinear(inputs, scale, zero_point) + return op.DequantizeLinear(quantized, scale, zero_point) + quantized = op.QuantizeLinear(inputs, scale, zero_point, axis=axis) + return op.DequantizeLinear(quantized, scale, zero_point, axis=axis) + + if num_bits == 4 and exponent_bits == 0: + assert block_size is not None and axis is not None, ( + "INT4 ONNX export requires block_size and axis." + ) + input_dtype = int(inputs.dtype) + scale = op.Div(op.Cast(amax, to=output_dtype), 7.0) + output = _int4_dq(inputs, scale, axis, block_size) + return output if input_dtype == output_dtype else op.Cast(output, to=output_dtype) + + raise NotImplementedError( + f"Unsupported num_bits: {num_bits} and exponent_bits: {exponent_bits} for ONNX export." + ) + + def _dynamic_block_quantize_op_translation( + inputs, + block_size: int, + amax, + num_bits: int, + exponent_bits: int, + scale_num_bits: int, + scale_exponent_bits: int, + trt_high_precision_dtype: str = None, # noqa: RUF013 + onnx_quantizer_type: str = None, # noqa: RUF013 + ): + if (num_bits, exponent_bits, scale_num_bits, scale_exponent_bits) == ( + 4, + 2, + 8, + 4, + ): + if onnx_quantizer_type != "dynamic": + return op.Identity(_fp4_qdq(inputs, block_size)) + output_dtype = _resolve_dtype(inputs, trt_high_precision_dtype) + inputs = _cast(inputs, output_dtype) + if amax is None: + scale = op.Constant(value_float=1.0) + else: + scale = op.Div(op.Cast(amax, to=onnx.TensorProto.FLOAT), 2688.0) + scale = op.Where(op.Equal(scale, 0.0), op.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 + dequantized_scale = op.DequantizeLinear(dynamic_scale, scale) + output = _fp4_dq(quantized, dequantized_scale, block_size) + return ( + output + if output_dtype == onnx.TensorProto.FLOAT + else op.Cast(output, to=output_dtype) + ) + + if (num_bits, exponent_bits, scale_num_bits, scale_exponent_bits) == ( + 8, + 4, + 9, + 8, + ): + output_dtype = int(inputs.dtype) + if onnx_quantizer_type == "dynamic": + output = _mxfp8_dynamic_qdq(inputs, block_size, output_dtype) + else: + scale = op.CastLike(1.0, inputs) + output = _mxfp8_static_dq(inputs, scale, block_size, output_dtype) + return op.Identity(output) + + raise NotImplementedError( + f"Unsupported num_bits: ({exponent_bits}, {num_bits - exponent_bits - 1}) " + "and scale_bits: " + f"({scale_exponent_bits}, {scale_num_bits - scale_exponent_bits - 1}) " + "for ONNX export." + ) + + return { + torch.ops.tensorrt.quantize_op.default: _quantize_op_translation, + torch.ops.tensorrt.dynamic_block_quantize_op.default: ( + _dynamic_block_quantize_op_translation + ), + torch.ops.tensorrt.dynamic_block_quantize_op.overload: ( + _dynamic_block_quantize_op_translation + ), + } def export_int8( @@ -206,7 +698,11 @@ def export_int4( if trt_high_precision_dtype is None: trt_high_precision_dtype = otype return g.op( - "trt::DequantizeLinear", inputs, scale_inv_op, axis_i=axis, block_size_i=block_size + "trt::DequantizeLinear", + inputs, + scale_inv_op, + axis_i=axis, + block_size_i=block_size, ).setType( inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) ) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 18b97ac2774..f8e22e7dbb7 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -747,7 +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 + amax = ( + amax.detach() if is_torch_export_mode() or torch.compiler.is_exporting() else amax.data + ) return amax def validate_attr( diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..c12feaf78eb 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) @@ -361,6 +377,28 @@ def forward( axis=None, ): """Forward method.""" + if torch.onnx.is_in_onnx_export() and torch.compiler.is_exporting(): + from .export_onnx import export_int4_dynamo, export_int8_dynamo + + if bias is not None: + inputs = inputs - bias + if num_bits == 4: + outputs = export_int4_dynamo( + inputs, amax, num_bits, trt_high_precision_dtype, block_size, axis + ) + else: + outputs = export_int8_dynamo( + inputs, + amax, + num_bits, + unsigned, + narrow_range, + trt_high_precision_dtype, + ) + if bias is not None: + outputs = outputs + bias + return outputs + if bias is not None: inputs = inputs - bias @@ -371,7 +409,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 +420,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 @@ -435,6 +476,16 @@ def forward( if E != 4 or M != 3: raise NotImplementedError("Only support E=4 & M=3 for now.") + if torch.onnx.is_in_onnx_export() and torch.compiler.is_exporting(): + from .export_onnx import export_fp8_dynamo + + if bias is not None: + inputs = inputs - bias + outputs = export_fp8_dynamo(inputs, amax, trt_high_precision_dtype) + if bias is not None: + outputs = outputs + bias + return outputs + if bias is not None: inputs = inputs - bias @@ -447,6 +498,9 @@ def forward( exponent_bits=4, unsigned=False, narrow_range=False, + trt_high_precision_dtype=trt_high_precision_dtype, + block_size=None, + axis=None, ) if bias is not None: @@ -490,6 +544,8 @@ def _dynamic_block_quantize_forward( exponent_bits, scale_num_bits, scale_exponent_bits, + trt_high_precision_dtype, + onnx_quantizer_type, ) return outputs @@ -549,6 +605,24 @@ def forward( pass_through_bwd=True, ): """Forward method.""" + if torch.onnx.is_in_onnx_export() and torch.compiler.is_exporting(): + from .export_onnx import export_fp4_dynamo, export_mxfp8_dynamo + + if num_bits == (2, 1) and scale_bits == (4, 3): + return export_fp4_dynamo( + inputs, + block_size, + amax, + num_bits, + trt_high_precision_dtype, + onnx_quantizer_type, + ) + if num_bits == (4, 3) and scale_bits == (8, 0): + return export_mxfp8_dynamo(inputs, onnx_quantizer_type, block_size) + raise NotImplementedError( + f"Unsupported num_bits: {num_bits} and scale_bits: {scale_bits} for ONNX export." + ) + _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax) return _dynamic_block_quantize_forward( ctx, diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index fe99cae9bc3..ea9a69fcc53 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -14,6 +14,9 @@ # limitations under the License. +import os +import tempfile + import onnx import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command @@ -75,3 +78,42 @@ def test_torch_onnx_recipe_flag(tmp_path): "TRT_FP8QuantizeLinear", } assert not quantize_ops & {node.op_type for node in onnx.load(onnx_save_path).graph.node} + + +def test_torch_onnx_dynamo_export(tmp_path): + timm_model_name, _ = _MODELS["vit_tiny"] + model_kwargs = ( + '{"depth": 1, "img_size": 32, "embed_dim": 32, "num_heads": 1, "mlp_ratio": 1, ' + '"pretrained_cfg_overlay": {"input_size": [3, 32, 32]}}' + ) + onnx_save_path = tmp_path / "vit_tiny.fp8.dynamo.onnx" + + cmd_parts = extend_cmd_parts( + ["python", "torch_quant_to_onnx.py"], + timm_model_name=timm_model_name, + model_kwargs=model_kwargs, + quantize_mode="fp8", + onnx_save_path=str(onnx_save_path), + calibration_data_size="1", + onnx_opset="24", + ) + cmd_parts.extend(["--no_pretrained", "--dynamo_export", "--trt_build"]) + env = os.environ.copy() + with tempfile.TemporaryDirectory(prefix="modelopt_torch_extensions_") as extension_dir: + env["TORCH_EXTENSIONS_DIR"] = extension_dir + run_example_command(cmd_parts, "torch_onnx", env=env) + + model = onnx.load(onnx_save_path) + onnx.checker.check_model(model, full_check=True) + assert {opset.domain: opset.version for opset in model.opset_import}[""] == 24 + assert not model.functions + node_types = {node.op_type for node in model.graph.node} + assert {"QuantizeLinear", "DequantizeLinear"} <= node_types + assert all( + node.domain not in {"trt", "tensorrt"} and node.op_type != "quantize_op" + for node in model.graph.node + ) + assert any( + initializer.data_type == onnx.TensorProto.FLOAT8E4M3FN + for initializer in model.graph.initializer + ) diff --git a/tests/gpu/torch/quantization/test_onnx_export_dynamo.py b/tests/gpu/torch/quantization/test_onnx_export_dynamo.py new file mode 100644 index 00000000000..b4ed7125bc4 --- /dev/null +++ b/tests/gpu/torch/quantization/test_onnx_export_dynamo.py @@ -0,0 +1,106 @@ +# 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. + +"""GPU integration tests for Dynamo ONNX export of quantized models.""" + +import copy + +import onnx +import pytest +import torch +from _test_utils.torch.misc import minimum_sm + +import modelopt.torch.quantization as mtq +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata + + +class _AlignedLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(128, 128, bias=False) + + def forward(self, inputs): + return self.linear(inputs) + + +_CASES = [ + pytest.param( + "fp8", + mtq.FP8_DEFAULT_CFG, + {"QuantizeLinear", "DequantizeLinear"}, + id="fp8", + ), + pytest.param( + "int8", + mtq.INT8_DEFAULT_CFG, + {"QuantizeLinear", "DequantizeLinear"}, + id="int8", + ), + pytest.param( + "mxfp8", + mtq.MXFP8_DEFAULT_CFG, + {"TRT_MXFP8DynamicQuantize", "TRT_MXFP8DequantizeLinear"}, + id="mxfp8", + ), + pytest.param( + "nvfp4", + mtq.NVFP4_DEFAULT_CFG, + {"TRT_FP4DynamicQuantize", "DequantizeLinear"}, + marks=minimum_sm(100), + id="nvfp4", + ), +] + + +@pytest.mark.timeout(600) +@pytest.mark.parametrize(("quant_format", "config", "expected_ops"), _CASES) +def test_quantized_model_dynamo_export(tmp_path, quant_format, config, expected_ops): + model = _AlignedLinear().eval().cuda() + sample_input = torch.randn(2, 128, device="cuda") + model = mtq.quantize( + model, + copy.deepcopy(config), + forward_loop=lambda candidate: candidate(sample_input), + ) + + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + model_name=f"aligned_linear_{quant_format}_dynamo", + dynamo_export=True, + onnx_opset=24, + weights_dtype="fp16", + ) + onnx_package = OnnxBytes.from_bytes(onnx_bytes) + export_dir = tmp_path / quant_format + onnx_package.write_to_disk(str(export_dir)) + exported = onnx.load( + export_dir / f"{onnx_package.model_name}.onnx", + load_external_data=True, + ) + + onnx.checker.check_model(exported) + ops = {node.op_type for node in exported.graph.node} + assert expected_ops <= ops + assert "quantize_op" not in ops + assert "dynamic_block_quantize_op" not in ops + initializer_dtypes = {initializer.data_type for initializer in exported.graph.initializer} + if quant_format == "mxfp8": + assert onnx.TensorProto.FLOAT8E4M3FN in initializer_dtypes + assert onnx.TensorProto.UINT8 in initializer_dtypes + elif quant_format == "nvfp4": + assert "TRT_FP4QDQ" not in ops + assert onnx.TensorProto.FLOAT4E2M1 in initializer_dtypes + assert onnx.TensorProto.FLOAT8E4M3FN in initializer_dtypes diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 4b1e69ec538..7faeef2d8b6 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -19,7 +19,7 @@ import onnx_graphsurgeon as gs import onnxruntime as ort import pytest -from onnx import TensorProto, helper, numpy_helper +from onnx import TensorProto, checker, helper, numpy_helper from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter from modelopt.onnx.export.nvfp4_exporter import _cast_fp4 @@ -65,7 +65,13 @@ def create_test_model_with_int4_dq_reshape_transpose_matmul(constant_scale: bool # Create nodes dq_inputs = ["weight", "Constant_output_0"] if constant_scale else ["weight", "scale"] dq_node = helper.make_node( - "DequantizeLinear", inputs=dq_inputs, outputs=["dq_output"], name="weight_dq" + "DequantizeLinear", + inputs=dq_inputs, + outputs=["dq_output"], + name="weight_dq", + domain="trt", + axis=0, + block_size=8, ) reshape_constant = helper.make_node( @@ -125,6 +131,171 @@ def create_test_model_with_int4_dq_reshape_transpose_matmul(constant_scale: bool return model +def create_test_model_with_shared_int4_weight(): + weight = numpy_helper.from_array(np.ones((4, 4), dtype=np.float32), "weight") + scale = numpy_helper.from_array(np.full((4, 1), 2.0, dtype=np.float32), "scale") + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [2, 4]) + nodes = [] + outputs = [] + value_info = [] + for index in range(2): + dq_output = f"dq_output_{index}" + output = f"output_{index}" + nodes.extend( + [ + helper.make_node( + "DequantizeLinear", + ["weight", "scale"], + [dq_output], + name=f"weight_dq_{index}", + domain="trt", + axis=1, + block_size=4, + ), + helper.make_node("MatMul", ["input", dq_output], [output], name=f"matmul_{index}"), + ] + ) + outputs.append(helper.make_tensor_value_info(output, TensorProto.FLOAT, [2, 4])) + value_info.append(helper.make_tensor_value_info(dq_output, TensorProto.FLOAT, [4, 4])) + + graph = helper.make_graph( + nodes, + "shared_int4_weight", + [input_tensor], + outputs, + [weight, scale], + value_info=value_info, + ) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("trt", 1)], + ) + + +def create_test_model_with_pre_reshape_only_int4_weight(): + weight = numpy_helper.from_array(np.ones((2, 8), dtype=np.float32), "weight") + blocked_shape = numpy_helper.from_array(np.array([-1, 4], dtype=np.int64), "val_3") + scale = numpy_helper.from_array(np.full((4, 1), 2.0, dtype=np.float32), "scale") + nodes = [ + helper.make_node("Reshape", ["weight", "val_3"], ["view"], name="weight_view"), + helper.make_node( + "DequantizeLinear", + ["view", "scale"], + ["dq_output"], + name="weight_dq", + domain="trt", + axis=0, + block_size=4, + ), + helper.make_node("MatMul", ["input", "dq_output"], ["output"], name="matmul"), + ] + graph = helper.make_graph( + nodes, + "pre_reshape_only_int4_weight", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [2, 4])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [2, 4])], + [weight, blocked_shape, scale], + value_info=[ + helper.make_tensor_value_info("view", TensorProto.FLOAT, [4, 4]), + helper.make_tensor_value_info("dq_output", TensorProto.FLOAT, [4, 4]), + ], + ) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("trt", 1)], + ) + + +def create_test_model_with_pre_and_post_reshape_int4_weight(): + weight = numpy_helper.from_array(np.ones((4, 8), dtype=np.float32), "weight") + blocked_shape = numpy_helper.from_array(np.array([8, 4], dtype=np.int64), "blocked_shape") + target_shape = numpy_helper.from_array(np.array([4, 8], dtype=np.int64), "target_shape") + scale = numpy_helper.from_array(np.full((8, 1), 2.0, dtype=np.float32), "scale") + nodes = [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["view"], name="weight_view"), + helper.make_node( + "DequantizeLinear", + ["view", "scale"], + ["blocked_dq_output"], + name="weight_dq", + domain="trt", + axis=0, + block_size=4, + ), + helper.make_node( + "Reshape", + ["blocked_dq_output", "target_shape"], + ["dq_output"], + name="weight_target_view", + ), + helper.make_node("MatMul", ["input", "dq_output"], ["output"], name="matmul"), + ] + graph = helper.make_graph( + nodes, + "pre_and_post_reshape_int4_weight", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [2, 4])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [2, 8])], + [weight, blocked_shape, target_shape, scale], + value_info=[ + helper.make_tensor_value_info("view", TensorProto.FLOAT, [8, 4]), + helper.make_tensor_value_info("scale", TensorProto.FLOAT, [8, 1]), + helper.make_tensor_value_info("blocked_dq_output", TensorProto.FLOAT, [8, 4]), + helper.make_tensor_value_info("dq_output", TensorProto.FLOAT, [4, 8]), + ], + ) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("trt", 1)], + ) + + +def create_test_model_with_int4_and_dynamic_nvfp4(): + model = create_test_model_with_pre_reshape_only_int4_weight() + graph = model.graph + graph.input.append(helper.make_tensor_value_info("nvfp4_input", TensorProto.FLOAT, [2, 16])) + graph.output.append(helper.make_tensor_value_info("nvfp4_output", TensorProto.FLOAT, [2, 16])) + graph.initializer.append( + numpy_helper.from_array(np.array(1.0, dtype=np.float32), "global_scale") + ) + graph.value_info.extend( + [ + helper.make_tensor_value_info("fp4_quantized", TensorProto.FLOAT4E2M1, [2, 16]), + helper.make_tensor_value_info("dynamic_scale", TensorProto.FLOAT8E4M3FN, [2, 1]), + helper.make_tensor_value_info("dequantized_scale", TensorProto.FLOAT, [2, 1]), + ] + ) + graph.node.extend( + [ + helper.make_node( + "TRT_FP4DynamicQuantize", + ["nvfp4_input", "global_scale"], + ["fp4_quantized", "dynamic_scale"], + name="dynamic_nvfp4_quantize", + domain="trt", + axis=-1, + block_size=16, + scale_type=TensorProto.FLOAT8E4M3FN, + ), + helper.make_node( + "DequantizeLinear", + ["dynamic_scale", "global_scale"], + ["dequantized_scale"], + name="dynamic_nvfp4_scale_dq", + ), + helper.make_node( + "DequantizeLinear", + ["fp4_quantized", "dequantized_scale"], + ["nvfp4_output"], + name="dynamic_nvfp4_dq", + domain="trt", + axis=-1, + block_size=16, + ), + ] + ) + return model + + def create_test_model_with_cast_nodes(): """Create a test model with various Cast nodes to test float32->float16 conversion.""" # Create a simple model with Cast nodes @@ -256,6 +427,55 @@ def create_test_model_with_mxfp8_dq(): return model +def create_test_model_with_shared_mxfp8_weight(num_paths: int): + weight_data = np.linspace(-1.0, 1.0, num=32 * 64, dtype=np.float32).reshape(32, 64) + weight = numpy_helper.from_array(weight_data, "weight") + scale_data = np.array(1.0, dtype=np.float32) + scale = helper.make_node( + "Constant", + [], + ["mx_scale"], + name="scale_constant", + value=numpy_helper.from_array(scale_data), + ) + nodes = [scale] + outputs = [] + value_info = [] + for index in range(num_paths): + dq_output = f"dq_output_{index}" + output = f"output_{index}" + nodes.extend( + [ + helper.make_node( + "TRT_MXFP8DequantizeLinear", + ["weight", "mx_scale"], + [dq_output], + name=f"weight_dq_{index}", + domain="trt", + axis=-1, + block_size=32, + output_dtype=TensorProto.FLOAT, + ), + helper.make_node("MatMul", ["input", dq_output], [output], name=f"matmul_{index}"), + ] + ) + outputs.append(helper.make_tensor_value_info(output, TensorProto.FLOAT, [4, 64])) + value_info.append(helper.make_tensor_value_info(dq_output, TensorProto.FLOAT, [32, 64])) + + graph = helper.make_graph( + nodes, + "shared_mxfp8_weight", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, 32])], + outputs, + [weight], + value_info=value_info, + ) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("trt", 1)], + ) + + def create_test_model_with_nvfp4_qdq(with_transpose: bool = False): """Create a test ONNX model with TRT_FP4QDQ nodes for testing NVFP4. @@ -409,6 +629,73 @@ def test_projection_bias_and_scale_casting(self): ) assert scale_tensor.data_type == TensorProto.FLOAT16 + def test_shared_weight_and_scale_are_processed_independently(self): + model = INT4QuantExporter.pre_process(create_test_model_with_shared_int4_weight()) + dq_nodes = [node for node in model.graph.node if node.domain == "trt"] + + assert len({node.input[0] for node in dq_nodes}) == 2 + assert len({node.input[1] for node in dq_nodes}) == 2 + + model = INT4QuantExporter.compute_scales(model) + initializers = {initializer.name: initializer for initializer in model.graph.initializer} + for node in dq_nodes: + np.testing.assert_array_equal( + numpy_helper.to_array(initializers[node.input[0]]), + np.full((4, 4), 0.5, dtype=np.float32), + ) + + model = INT4QuantExporter.compress_weights(model) + model = INT4QuantExporter.post_process(model) + checker.check_model(model) + + def test_pre_reshape_only_uses_blocked_weight_shape(self): + model = INT4QuantExporter.process_model( + create_test_model_with_pre_reshape_only_int4_weight() + ) + dq_node = next(node for node in model.graph.node if node.domain == "trt") + initializers = {initializer.name: initializer for initializer in model.graph.initializer} + + assert list(initializers[dq_node.input[0]].dims) == [4, 4] + assert list(initializers[dq_node.input[1]].dims) == [4, 1] + assert initializers[dq_node.input[0]].data_type == TensorProto.INT4 + axis = next(attribute.i for attribute in dq_node.attribute if attribute.name == "axis") + assert axis == 1 + assert not any(node.op_type == "Reshape" for node in model.graph.node) + matmul = next(node for node in model.graph.node if node.op_type == "MatMul") + assert matmul.input[1] == dq_node.output[0] + checker.check_model(model) + + def test_pre_and_post_reshape_syncs_scale_metadata(self): + model = INT4QuantExporter.process_model( + create_test_model_with_pre_and_post_reshape_int4_weight() + ) + dq_node = next(node for node in model.graph.node if node.domain == "trt") + scale = next( + initializer + for initializer in model.graph.initializer + if initializer.name == dq_node.input[1] + ) + scale_info = next( + value_info for value_info in model.graph.value_info if value_info.name == scale.name + ) + + assert list(scale.dims) == [4, 2] + assert scale_info.type.tensor_type.elem_type == scale.data_type == TensorProto.FLOAT + assert [dim.dim_value for dim in scale_info.type.tensor_type.shape.dim] == [4, 2] + checker.check_model(model, full_check=True) + + def test_dynamic_nvfp4_dq_is_not_processed_as_int4(self): + model = INT4QuantExporter.process_model(create_test_model_with_int4_and_dynamic_nvfp4()) + initializers = {initializer.name: initializer for initializer in model.graph.initializer} + int4_dq = next(node for node in model.graph.node if node.name == "weight_dq") + nvfp4_dq = next(node for node in model.graph.node if node.name == "dynamic_nvfp4_dq") + + assert initializers[int4_dq.input[0]].data_type == TensorProto.INT4 + assert nvfp4_dq.input[0] == "fp4_quantized" + assert nvfp4_dq.domain == "trt" + assert not any(attribute.name.startswith("_") for attribute in nvfp4_dq.attribute) + checker.check_model(model, full_check=True) + class TestCastFunctions: """Test suite for _cast_fp8 and _cast_fp4 functions.""" @@ -487,6 +774,68 @@ def test_cast_fp4(self, input_array, expected_array): class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" + def test_weight_detection_uses_initializer_membership(self): + model = create_test_model_with_mxfp8_dq() + weight = next(init for init in model.graph.initializer if init.name == "linear.weight") + weight_name = "p_linear_weight" + weight.name = weight_name + dq_node = next( + node for node in model.graph.node if node.op_type == "TRT_MXFP8DequantizeLinear" + ) + dq_node.input[0] = weight_name + + quantized_model = MXFP8QuantExporter.process_model(model) + + weight = next( + init for init in quantized_model.graph.initializer if init.name == weight_name + ) + assert weight.data_type == TensorProto.FLOAT8E4M3FN + + def test_shared_weight_and_scale_match_single_path(self): + baseline = create_test_model_with_shared_mxfp8_weight(1) + baseline.graph.value_info.append( + helper.make_tensor_value_info("mx_scale", TensorProto.FLOAT, []) + ) + baseline = MXFP8QuantExporter.process_model(baseline) + shared = MXFP8QuantExporter.process_model(create_test_model_with_shared_mxfp8_weight(2)) + baseline_dq = next( + node for node in baseline.graph.node if node.op_type == "TRT_MXFP8DequantizeLinear" + ) + shared_dq_nodes = [ + node for node in shared.graph.node if node.op_type == "TRT_MXFP8DequantizeLinear" + ] + baseline_initializers = { + initializer.name: initializer for initializer in baseline.graph.initializer + } + shared_initializers = { + initializer.name: initializer for initializer in shared.graph.initializer + } + baseline_weight = baseline_initializers[baseline_dq.input[0]] + baseline_scale = baseline_initializers[baseline_dq.input[1]] + baseline_scale_info = next( + value_info + for value_info in baseline.graph.value_info + if value_info.name == baseline_scale.name + ) + + assert len({node.input[0] for node in shared_dq_nodes}) == 2 + assert len({node.input[1] for node in shared_dq_nodes}) == 2 + assert baseline_scale_info.type.tensor_type.elem_type == TensorProto.UINT8 + assert [dim.dim_value for dim in baseline_scale_info.type.tensor_type.shape.dim] == [32, 2] + for node in shared_dq_nodes: + weight = shared_initializers[node.input[0]] + scale = shared_initializers[node.input[1]] + assert weight.data_type == baseline_weight.data_type == TensorProto.FLOAT8E4M3FN + assert weight.dims == baseline_weight.dims + assert weight.raw_data == baseline_weight.raw_data + np.testing.assert_array_equal( + numpy_helper.to_array(scale), numpy_helper.to_array(baseline_scale) + ) + + assert not any(node.op_type == "Constant" for node in shared.graph.node) + checker.check_model(baseline, full_check=True) + checker.check_model(shared, full_check=True) + def test_basic_mxfp8_quantization(self): """Test basic MXFP8 quantization with TRT_MXFP8DequantizeLinear nodes.""" model = create_test_model_with_mxfp8_dq() @@ -603,6 +952,7 @@ class TestFP4QDQTo2DQ: def test_fp4qdq_conversion(self, with_transpose): """Test FP4QDQ to 2DQ conversion with and without Transpose node.""" model = create_test_model_with_nvfp4_qdq(with_transpose=with_transpose) + model.opset_import[0].version = 21 # Run FP4QDQ to 2DQ conversion converted_model = NVFP4QuantExporter.process_model(model) @@ -616,6 +966,19 @@ def test_fp4qdq_conversion(self, with_transpose): node for node in converted_model.graph.node if node.op_type == "DequantizeLinear" ] assert len(dq_nodes) == 2 + scale_dq = next(node for node in dq_nodes if len(node.attribute) == 0) + weight_dq = next( + node + for node in dq_nodes + if any(attribute.name == "block_size" for attribute in node.attribute) + ) + assert scale_dq.domain == "" + assert weight_dq.domain == "trt" + assert {opset.domain: opset.version for opset in converted_model.opset_import} == { + "": 21, + "trt": 1, + } + checker.check_model(converted_model, full_check=True) # Verify new initializers are created initializer_names = {init.name for init in converted_model.graph.initializer} 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 399085ad699..e2102f01f22 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -57,6 +57,28 @@ } +class _DefaultScalarModel(nn.Module): + def forward(self, x, flag=False): + if flag: + return -x + return x + 1 + + +def test_onnx_dynamo_export_preserves_default_scalar_as_constant(): + sample_input = torch.ones(2, 3) + + onnx_bytes, metadata = get_onnx_bytes_and_metadata( + _DefaultScalarModel(), (sample_input,), dynamo_export=True, onnx_opset=21 + ) + onnx_model = onnx.load_model_from_string( + OnnxBytes.from_bytes(onnx_bytes).get_onnx_model_file_bytes() + ) + + onnx.checker.check_model(onnx_model, full_check=True) + assert [value.name for value in onnx_model.graph.input] == ["x"] + assert generate_onnx_input(metadata, (sample_input,)).keys() == {"x"} + + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ce2ef626d63..1c7390b4f3a 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -29,12 +29,14 @@ from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester from onnx import TensorProto, helper, numpy_helper +import modelopt.torch._deploy.utils.torch_onnx as torch_onnx import modelopt.torch.quantization as mtq import modelopt.torch.quantization.tensor_quant as tensor_quant from modelopt.onnx import utils from modelopt.onnx.export import NVFP4QuantExporter from modelopt.onnx.export.nvfp4_exporter import _encode_nvfp4_block_scale from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata from modelopt.torch.quantization.qtensor import NVFP4QTensor from modelopt.torch.quantization.utils import is_quantized_linear @@ -105,6 +107,48 @@ def cpu_dynamic_block_quantize(inputs, *args): onnx.checker.check_model(converted_model) +def test_nvfp4_dynamo_helper_uses_quantized_fp16_conversion(monkeypatch, tmp_path): + sample_input = SimpleLinear.get_input() + + def forward_loop(model): + model(sample_input) + + def cpu_dynamic_block_quantize(inputs, *args): + return inputs + + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize) + model = mtq.quantize(SimpleLinear().eval(), mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop) + + def fail_generic_conversion(*args, **kwargs): + pytest.fail("Dynamo NVFP4 export must use the quantization-aware FP16 conversion path") + + monkeypatch.setattr(torch_onnx, "convert_to_f16", fail_generic_conversion) + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + model_name="nvfp4_dynamo_fp16", + dynamo_export=True, + onnx_opset=24, + weights_dtype="fp16", + ) + onnx_package = OnnxBytes.from_bytes(onnx_bytes) + export_dir = tmp_path / "nvfp4_dynamo_fp16" + onnx_package.write_to_disk(str(export_dir)) + exported = onnx.load(export_dir / f"{onnx_package.model_name}.onnx", load_external_data=True) + + onnx.checker.check_model(exported, full_check=True) + node_types = {node.op_type for node in exported.graph.node} + assert "TRT_FP4DynamicQuantize" in node_types + assert "TRT_FP4QDQ" not in node_types + bias_initializers = [ + initializer + for initializer in exported.graph.initializer + if initializer.name.endswith("bias") + ] + assert bias_initializers + assert all(initializer.data_type == TensorProto.FLOAT16 for initializer in bias_initializers) + + @pytest.mark.parametrize( ("convert", "deprecated"), [ diff --git a/tests/unit/torch/quantization/test_onnx_export_dynamo.py b/tests/unit/torch/quantization/test_onnx_export_dynamo.py new file mode 100644 index 00000000000..edcbdbc42ee --- /dev/null +++ b/tests/unit/torch/quantization/test_onnx_export_dynamo.py @@ -0,0 +1,928 @@ +# 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 tests for Dynamo ONNX export of ModelOpt quantization operators.""" + +import copy + +import onnx +import pytest +import torch +from torch import nn + +pytest.importorskip("onnxscript") + +import modelopt.torch.quantization as mtq +from modelopt.onnx.export import MXFP8QuantExporter +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata +from modelopt.torch.quantization.export_onnx import get_dynamo_onnx_translation_table +from modelopt.torch.quantization.tensor_quant import ( + dynamic_block_quant, + fake_tensor_quant, + scaled_e4m3, +) + + +class _TwoLinear(nn.Module): + def __init__(self, features=16): + super().__init__() + self.linear0 = nn.Linear(features, features, bias=False) + self.linear1 = nn.Linear(features, features, bias=False) + + def forward(self, x): + return self.linear1(torch.nn.functional.silu(self.linear0(x))) + + +class _Linear(nn.Module): + def __init__(self, in_features): + super().__init__() + self.linear = nn.Linear(in_features, 128, bias=False) + + def forward(self, x): + return self.linear(x) + + +class _BiasedQuant(nn.Module): + def __init__(self, quant_format): + super().__init__() + self.quant_format = quant_format + + def forward(self, x, amax, bias): + if self.quant_format == "fp8": + return scaled_e4m3(x, amax, bias, 4, 3, "Float", False) + return fake_tensor_quant(x, amax, bias, 8, False, False, "Float", False, None, None) + + +def _all_nodes(model): + yield from model.graph.node + for function in model.functions: + yield from function.node + + +def _node_attributes(node): + return { + attribute.name: onnx.helper.get_attribute_value(attribute) for attribute in node.attribute + } + + +def _tensor_metadata(model, name): + for value in (*model.graph.input, *model.graph.value_info, *model.graph.output): + if value.name == name: + tensor_type = value.type.tensor_type + shape = [dimension.dim_value for dimension in tensor_type.shape.dim] + return tensor_type.elem_type, shape + for initializer in model.graph.initializer: + if initializer.name == name: + return initializer.data_type, list(initializer.dims) + raise AssertionError(f"Missing type information for {name}") + + +def _scalar_initializer_values(model): + values = set() + for initializer in model.graph.initializer: + value = onnx.numpy_helper.to_array(initializer) + if value.size == 1: + values.add(float(value.item())) + return values + + +def _auto_quantize_fp8(model, sample_input): + return mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[copy.deepcopy(mtq.FP8_DEFAULT_CFG)], + 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] + + +@pytest.mark.timeout(120) +def test_autoquant_fp8_dynamo_export_without_translation_table(tmp_path): + sample_input = torch.randn(2, 16) + model = _auto_quantize_fp8(_TwoLinear().eval(), sample_input) + onnx_path = tmp_path / "two_linear_fp8.onnx" + + torch.onnx.export( + model, + (sample_input,), + onnx_path, + dynamo=True, + opset_version=24, + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported) + assert not exported.functions + opset_imports = {opset.domain: opset.version for opset in exported.opset_import} + assert opset_imports[""] == 24 + assert opset_imports["trt"] == 1 + nodes = list(_all_nodes(exported)) + assert sum(node.op_type == "TRT_FP8QuantizeLinear" for node in nodes) == 4 + assert sum(node.op_type == "TRT_FP8DequantizeLinear" for node in nodes) == 4 + for node in nodes: + if node.op_type == "TRT_FP8QuantizeLinear": + assert _tensor_metadata(exported, node.output[0])[0] == onnx.TensorProto.UINT8 + assert not any(node.op_type == "quantize_op" for node in nodes) + + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + model_name="two_linear_fp8", + dynamo_export=True, + onnx_opset=24, + ) + processed = onnx.load_model_from_string( + OnnxBytes.from_bytes(onnx_bytes).get_onnx_model_file_bytes() + ) + onnx.checker.check_model(processed) + processed_nodes = list(_all_nodes(processed)) + assert {"QuantizeLinear", "DequantizeLinear"} <= {node.op_type for node in processed_nodes} + assert not any( + node.op_type.startswith("TRT_FP8") or node.op_type == "quantize_op" + for node in processed_nodes + ) + + +@pytest.mark.timeout(120) +def test_mixed_autoquant_dynamo_export_without_translation_table(tmp_path): + sample_input = torch.randn(2, 128) + model = mtq.auto_quantize( + _TwoLinear(128).eval(), + constraints={"effective_bits": 6.0}, + quantization_formats=[ + copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG), + copy.deepcopy(mtq.FP8_DEFAULT_CFG), + ], + 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] + onnx_path = tmp_path / "two_linear_mixed.onnx" + + torch.onnx.export( + model, + (sample_input,), + onnx_path, + dynamo=True, + opset_version=24, + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported) + assert not exported.functions + nodes = list(_all_nodes(exported)) + assert sum(node.op_type == "TRT_FP8QuantizeLinear" for node in nodes) == 2 + assert sum(node.op_type == "TRT_FP8DequantizeLinear" for node in nodes) == 2 + assert sum(node.domain == "trt" and node.op_type == "DequantizeLinear" for node in nodes) == 1 + assert not any(node.domain == "tensorrt" for node in nodes) + + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + model_name="two_linear_mixed", + dynamo_export=True, + onnx_opset=24, + weights_dtype="fp32", + ) + onnx_package = OnnxBytes.from_bytes(onnx_bytes) + export_dir = tmp_path / "mixed_helper" + onnx_package.write_to_disk(str(export_dir)) + processed = onnx.load( + export_dir / f"{onnx_package.model_name}.onnx", + load_external_data=True, + ) + + onnx.checker.check_model(processed, full_check=True) + assert not processed.functions + processed_nodes = list(_all_nodes(processed)) + assert ( + sum(node.domain == "" and node.op_type == "QuantizeLinear" for node in processed_nodes) == 1 + ) + assert ( + sum(node.domain == "" and node.op_type == "DequantizeLinear" for node in processed_nodes) + == 2 + ) + int4_dq_nodes = [ + node + for node in processed_nodes + if node.domain == "trt" and node.op_type == "DequantizeLinear" + ] + assert len(int4_dq_nodes) == 1 + initializers = {initializer.name: initializer for initializer in processed.graph.initializer} + assert initializers[int4_dq_nodes[0].input[0]].data_type == onnx.TensorProto.INT4 + assert any( + node.domain == "" + and node.op_type == "DequantizeLinear" + and node.input[0] in initializers + and initializers[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN + for node in processed_nodes + ) + assert not any(node.domain == "tensorrt" for node in processed_nodes) + assert not any(node.op_type.startswith("TRT_") for node in processed_nodes) + assert not any(node.op_type == "quantize_op" for node in processed_nodes) + + +@pytest.mark.parametrize( + ("quant_format", "quantize_op", "dequantize_op"), + [ + ("fp8", "TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"), + ("int8", "QuantizeLinear", "DequantizeLinear"), + ], +) +def test_affine_bias_surrounds_dynamo_qdq(tmp_path, quant_format, quantize_op, dequantize_op): + args = (torch.randn(4, 32), torch.tensor(2.0), torch.full((32,), 0.25)) + + for capture in ("direct", "strict"): + model = _BiasedQuant(quant_format).eval() + export_args = args + export_kwargs = {} + if capture == "strict": + model = torch.export.export(model, args, strict=True) + export_args = () + export_kwargs["custom_translation_table"] = get_dynamo_onnx_translation_table() + onnx_path = tmp_path / f"{quant_format}_{capture}_bias.onnx" + + torch.onnx.export( + model, + export_args, + onnx_path, + dynamo=True, + opset_version=24, + **export_kwargs, + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + nodes = list(exported.graph.node) + subtract = next(node for node in nodes if node.op_type == "Sub") + quantize = next(node for node in nodes if node.op_type == quantize_op) + dequantize = next(node for node in nodes if node.op_type == dequantize_op) + add = next(node for node in nodes if node.op_type == "Add") + assert quantize.input[0] == subtract.output[0] + assert dequantize.input[0] == quantize.output[0] + assert dequantize.output[0] in add.input + if quant_format == "fp8": + assert 448.0 in _scalar_initializer_values(exported) + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("in_features", [128, 256]) +def test_int4_awq_helper_dynamo_export(tmp_path, in_features): + sample_input = torch.randn(2, in_features) + model = _Linear(in_features).eval() + model = mtq.quantize( + model, + copy.deepcopy(mtq.INT4_AWQ_CFG), + forward_loop=lambda candidate: candidate(sample_input), + ) + + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + model_name=f"linear_int4_{in_features}", + dynamo_export=True, + onnx_opset=24, + weights_dtype="fp32", + ) + onnx_package = OnnxBytes.from_bytes(onnx_bytes) + export_dir = tmp_path / str(in_features) + onnx_package.write_to_disk(str(export_dir)) + exported = onnx.load( + export_dir / f"{onnx_package.model_name}.onnx", + load_external_data=True, + ) + + onnx.checker.check_model(exported) + int4_dq_nodes = [ + node + for node in exported.graph.node + if node.domain == "trt" and node.op_type == "DequantizeLinear" + ] + assert len(int4_dq_nodes) == 1 + int4_dq = int4_dq_nodes[0] + initializers = {initializer.name: initializer for initializer in exported.graph.initializer} + weight = initializers[int4_dq.input[0]] + scale = initializers[int4_dq.input[1]] + assert weight.data_type == onnx.TensorProto.INT4 + assert list(weight.dims) == [128, in_features] + assert list(scale.dims) == [128, in_features // 128] + + attributes = { + attribute.name: onnx.helper.get_attribute_value(attribute) + for attribute in int4_dq.attribute + } + assert attributes["axis"] == 1 + assert attributes["block_size"] == 128 + dq_output = next( + value_info + for value_info in exported.graph.value_info + if value_info.name == int4_dq.output[0] + ) + assert [dimension.dim_value for dimension in dq_output.type.tensor_type.shape.dim] == [ + 128, + in_features, + ] + assert not any(node.op_type == "Reshape" for node in exported.graph.node) + gemm = next(node for node in exported.graph.node if node.op_type == "Gemm") + assert gemm.input[1] == int4_dq.output[0] + assert not any(node.op_type == "quantize_op" for node in _all_nodes(exported)) + + +class _DirectNVFP4(nn.Module): + def __init__(self, trt_high_precision_dtype="Float"): + super().__init__() + self.trt_high_precision_dtype = trt_high_precision_dtype + + def forward(self, x, amax): + return dynamic_block_quant( + x, + 16, + amax, + None, + (2, 1), + (4, 3), + self.trt_high_precision_dtype, + "dynamic", + True, + ) + + +class _DirectMXFP8(nn.Module): + def forward(self, x): + return dynamic_block_quant( + x, + 32, + None, + None, + (4, 3), + (8, 0), + None, + "dynamic", + True, + ) + + +def test_nvfp4_direct_dynamo_export_has_logical_float4_shape(tmp_path): + sample_input = torch.randn(4, 32) + onnx_path = tmp_path / "nvfp4_dynamic.onnx" + + torch.onnx.export( + _DirectNVFP4().eval(), + (sample_input, torch.tensor(1.0)), + onnx_path, + dynamo=True, + opset_version=24, + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + fp4_node = next( + node for node in exported.graph.node if node.op_type == "TRT_FP4DynamicQuantize" + ) + assert _node_attributes(fp4_node) == { + "axis": -1, + "block_size": 16, + "scale_type": onnx.TensorProto.FLOAT8E4M3FN, + } + assert _tensor_metadata(exported, fp4_node.output[0]) == ( + onnx.TensorProto.FLOAT4E2M1, + list(sample_input.shape), + ) + assert _tensor_metadata(exported, fp4_node.output[1]) == ( + onnx.TensorProto.FLOAT8E4M3FN, + [sample_input.shape[0], sample_input.shape[1] // 16], + ) + fp4_dq = next( + node + for node in exported.graph.node + if node.op_type == "DequantizeLinear" and node.input[0] == fp4_node.output[0] + ) + assert fp4_dq.domain == "trt" + assert _node_attributes(fp4_dq)["axis"] == -1 + assert _node_attributes(fp4_dq)["block_size"] == 16 + assert _tensor_metadata(exported, fp4_dq.output[0]) == ( + onnx.TensorProto.FLOAT, + list(sample_input.shape), + ) + assert 2688.0 in _scalar_initializer_values(exported) + assert any(node.op_type == "Where" for node in exported.graph.node) + + +@pytest.mark.parametrize("capture", ["direct", "strict"]) +def test_nvfp4_dynamic_opset21_full_check(tmp_path, capture): + sample_input = torch.randn(4, 32) + model = _DirectNVFP4().eval() + args = (sample_input, torch.tensor(1.0)) + export_kwargs = {} + if capture == "strict": + model = torch.export.export(model, args, strict=True) + args = () + export_kwargs["custom_translation_table"] = get_dynamo_onnx_translation_table() + onnx_path = tmp_path / f"nvfp4_opset21_{capture}.onnx" + + torch.onnx.export( + model, + args, + onnx_path, + dynamo=True, + opset_version=21, + **export_kwargs, + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + assert not exported.functions + quantize = next( + node for node in exported.graph.node if node.op_type == "TRT_FP4DynamicQuantize" + ) + scale_dq = next( + node + for node in exported.graph.node + if node.op_type == "DequantizeLinear" and node.input[0] == quantize.output[1] + ) + fp4_dq = next( + node + for node in exported.graph.node + if node.op_type == "DequantizeLinear" and node.input[0] == quantize.output[0] + ) + assert scale_dq.domain == "" + assert fp4_dq.domain == "trt" + assert _tensor_metadata(exported, quantize.output[0])[0] == onnx.TensorProto.FLOAT4E2M1 + + +class _StrictNVFP4(nn.Module): + def __init__(self, trt_high_precision_dtype, onnx_quantizer_type): + super().__init__() + self.trt_high_precision_dtype = trt_high_precision_dtype + self.onnx_quantizer_type = onnx_quantizer_type + + def forward(self, x, amax): + return torch.ops.tensorrt.dynamic_block_quantize_op.default( + x, + 16, + amax, + 4, + 2, + 8, + 4, + self.trt_high_precision_dtype, + self.onnx_quantizer_type, + ) + + +_NVFP4_DTYPES = [ + pytest.param("Float", torch.float32, onnx.TensorProto.FLOAT, id="float"), + pytest.param("Half", torch.float16, onnx.TensorProto.FLOAT16, id="half"), + pytest.param("BFloat16", torch.bfloat16, onnx.TensorProto.BFLOAT16, id="bfloat16"), +] + + +@pytest.mark.parametrize("capture", ["direct", "strict"]) +@pytest.mark.parametrize(("trt_high_precision_dtype", "torch_dtype", "onnx_dtype"), _NVFP4_DTYPES) +def test_nvfp4_dynamic_dynamo_output_dtype( + tmp_path, capture, trt_high_precision_dtype, torch_dtype, onnx_dtype +): + sample_input = torch.randn(4, 32, dtype=torch_dtype) + model = _DirectNVFP4(trt_high_precision_dtype).eval() + args = (sample_input, torch.tensor(1.0)) + if capture == "strict": + model = torch.export.export( + _StrictNVFP4(trt_high_precision_dtype, "dynamic"), args, strict=True + ) + args = () + onnx_path = tmp_path / f"nvfp4_{capture}_{trt_high_precision_dtype}.onnx" + + torch.onnx.export( + model, + args, + onnx_path, + dynamo=True, + opset_version=24, + custom_translation_table=get_dynamo_onnx_translation_table(), + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + assert not exported.functions + fp4_node = next( + node for node in exported.graph.node if node.op_type == "TRT_FP4DynamicQuantize" + ) + assert _node_attributes(fp4_node) == { + "axis": -1, + "block_size": 16, + "scale_type": onnx.TensorProto.FLOAT8E4M3FN, + } + assert _tensor_metadata(exported, fp4_node.output[0])[0] == onnx.TensorProto.FLOAT4E2M1 + assert _tensor_metadata(exported, fp4_node.output[1])[0] == onnx.TensorProto.FLOAT8E4M3FN + fp4_dq = next( + node + for node in exported.graph.node + if node.op_type == "DequantizeLinear" and node.input[0] == fp4_node.output[0] + ) + assert fp4_dq.domain == "trt" + assert _node_attributes(fp4_dq)["axis"] == -1 + assert _node_attributes(fp4_dq)["block_size"] == 16 + assert _tensor_metadata(exported, exported.graph.output[0].name) == ( + onnx_dtype, + list(sample_input.shape), + ) + assert 2688.0 in _scalar_initializer_values(exported) + assert any(node.op_type == "Where" for node in exported.graph.node) + + +@pytest.mark.parametrize("capture", ["direct", "strict"]) +def test_mxfp8_dynamic_dynamo_contract(tmp_path, capture): + sample_input = torch.randn(4, 64) + model = _DirectMXFP8().eval() + args = (sample_input,) + export_kwargs = {} + if capture == "strict": + model = torch.export.export(model, args, strict=True) + args = () + export_kwargs["custom_translation_table"] = get_dynamo_onnx_translation_table() + onnx_path = tmp_path / f"mxfp8_{capture}.onnx" + + torch.onnx.export( + model, + args, + onnx_path, + dynamo=True, + opset_version=24, + **export_kwargs, + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + assert not exported.functions + opset_imports = {opset.domain: opset.version for opset in exported.opset_import} + assert opset_imports == {"": 24, "trt": 1} + quantize = next( + node for node in exported.graph.node if node.op_type == "TRT_MXFP8DynamicQuantize" + ) + dequantize = next( + node for node in exported.graph.node if node.op_type == "TRT_MXFP8DequantizeLinear" + ) + assert _node_attributes(quantize) == { + "axis": -1, + "block_size": 32, + "output_dtype": onnx.TensorProto.FLOAT8E4M3FN, + } + assert _node_attributes(dequantize) == { + "axis": -1, + "block_size": 32, + "output_dtype": onnx.TensorProto.FLOAT, + } + assert list(dequantize.input[:2]) == list(quantize.output) + assert _tensor_metadata(exported, exported.graph.output[0].name) == ( + onnx.TensorProto.FLOAT, + list(sample_input.shape), + ) + if capture == "direct": + assert _tensor_metadata(exported, quantize.output[0]) == ( + onnx.TensorProto.FLOAT8E4M3FN, + list(sample_input.shape), + ) + assert _tensor_metadata(exported, quantize.output[1]) == ( + onnx.TensorProto.UINT8, + [sample_input.shape[0], sample_input.shape[1] // 32], + ) + + +def test_nvfp4_strict_none_quantizer_type_uses_static_marker(tmp_path): + sample_input = torch.randn(4, 32) + exported_program = torch.export.export( + _StrictNVFP4("Float", None), + (sample_input, torch.tensor(1.0)), + strict=True, + ) + onnx_path = tmp_path / "nvfp4_static_default.onnx" + + torch.onnx.export( + exported_program, + (), + onnx_path, + dynamo=True, + opset_version=24, + custom_translation_table=get_dynamo_onnx_translation_table(), + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + assert not exported.functions + opset_imports = {opset.domain: opset.version for opset in exported.opset_import} + assert opset_imports == {"": 24, "trt": 1} + node_types = {node.op_type for node in exported.graph.node} + assert "TRT_FP4QDQ" in node_types + assert "TRT_FP4DynamicQuantize" not in node_types + + +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 == "fp8": + return torch.ops.tensorrt.quantize_op.default( + x, amax, 8, 4, False, False, "Float", None, None + ) + 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 == "int4": + return torch.ops.tensorrt.quantize_op.default( + x, amax, 4, 0, False, True, "Float", 32, 0 + ) + if self.quant_format.startswith("nvfp4"): + quantizer_type = self.quant_format.removeprefix("nvfp4_") + return torch.ops.tensorrt.dynamic_block_quantize_op.default( + x, 16, amax, 4, 2, 8, 4, "Float", quantizer_type + ) + if self.quant_format.startswith("mxfp8"): + quantizer_type = self.quant_format.removeprefix("mxfp8_") + return torch.ops.tensorrt.dynamic_block_quantize_op.overload( + x, 32, None, 8, 4, 9, 8, "Float", quantizer_type + ) + raise AssertionError(f"Unknown format: {self.quant_format}") + + +class _StrictINT8(nn.Module): + def __init__(self, unsigned): + super().__init__() + self.unsigned = unsigned + + def forward(self, x, amax): + return torch.ops.tensorrt.quantize_op.default( + x, amax, 8, 0, self.unsigned, False, "Float", None, 0 + ) + + +class _StrictINT4Weight(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(4, 32, dtype=torch.float16), requires_grad=False) + self.register_buffer("amax", torch.ones(4, 1)) + + def forward(self): + return torch.ops.tensorrt.quantize_op.default( + self.weight, + self.amax, + 4, + 0, + False, + True, + "Float", + 32, + 0, + ) + + +class _StrictMXFP8Weight(nn.Module): + def __init__(self): + super().__init__() + weight = torch.linspace(-1.0, 1.0, steps=32 * 64).reshape(32, 64) + self.weight = nn.Parameter(weight, requires_grad=False) + + def forward(self, x): + weight = torch.ops.tensorrt.dynamic_block_quantize_op.overload( + self.weight, + 32, + None, + 8, + 4, + 9, + 8, + "Float", + "static", + ) + return torch.matmul(x, weight) + + +_STRICT_CASES = [ + pytest.param( + "fp8", + 21, + {("trt", "TRT_FP8QuantizeLinear"), ("trt", "TRT_FP8DequantizeLinear")}, + id="fp8-opset21", + ), + pytest.param( + "fp8", + 24, + {("trt", "TRT_FP8QuantizeLinear"), ("trt", "TRT_FP8DequantizeLinear")}, + id="fp8-opset24", + ), + pytest.param( + "int8", + 21, + {("", "QuantizeLinear"), ("", "DequantizeLinear")}, + id="int8", + ), + pytest.param("int4", 21, {("trt", "DequantizeLinear")}, id="int4-awq"), + pytest.param( + "nvfp4_dynamic", + 21, + {("trt", "TRT_FP4DynamicQuantize"), ("", "DequantizeLinear")}, + id="nvfp4-dynamic", + ), + pytest.param( + "nvfp4_static", + 21, + {("trt", "TRT_FP4QDQ")}, + id="nvfp4-static", + ), + pytest.param( + "mxfp8_dynamic", + 21, + { + ("trt", "TRT_MXFP8DynamicQuantize"), + ("trt", "TRT_MXFP8DequantizeLinear"), + }, + id="mxfp8-dynamic", + ), + pytest.param( + "mxfp8_static", + 21, + {("trt", "TRT_MXFP8DequantizeLinear")}, + id="mxfp8-static", + ), +] + + +@pytest.mark.parametrize(("quant_format", "opset", "expected_nodes"), _STRICT_CASES) +def test_strict_exported_program_uses_translation_table( + tmp_path, quant_format, opset, expected_nodes +): + sample_input = torch.randn(4, 32) + amax = torch.ones(4, 1) if quant_format in {"int8", "int4"} else torch.tensor(1.0) + exported_program = torch.export.export( + _StrictQuantOp(quant_format), + (sample_input, amax), + strict=True, + ) + if quant_format.startswith("mxfp8"): + assert any( + node.target == torch.ops.tensorrt.dynamic_block_quantize_op.overload + for node in exported_program.graph.nodes + ) + onnx_path = tmp_path / f"{quant_format}.onnx" + + torch.onnx.export( + exported_program, + (), + onnx_path, + dynamo=True, + opset_version=opset, + custom_translation_table=get_dynamo_onnx_translation_table(), + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported) + assert not exported.functions + actual_nodes = {(node.domain, node.op_type) for node in _all_nodes(exported)} + assert expected_nodes <= actual_nodes + assert ("tensorrt", "quantize_op") not in actual_nodes + assert ("tensorrt", "dynamic_block_quantize_op") not in actual_nodes + + +@pytest.mark.parametrize( + ("unsigned", "zero_point_dtype", "scale_denominator"), + [ + pytest.param(False, onnx.TensorProto.INT8, 127.0, id="signed"), + pytest.param(True, onnx.TensorProto.UINT8, 255.0, id="unsigned"), + ], +) +def test_strict_int8_translation_contract(tmp_path, unsigned, zero_point_dtype, scale_denominator): + sample_input = torch.randn(4, 32) + amax = torch.ones(4, 1) + exported_program = torch.export.export( + _StrictINT8(unsigned), + (sample_input, amax), + strict=True, + ) + onnx_path = tmp_path / f"int8_{'unsigned' if unsigned else 'signed'}.onnx" + + torch.onnx.export( + exported_program, + (), + onnx_path, + dynamo=True, + opset_version=24, + custom_translation_table=get_dynamo_onnx_translation_table(), + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + assert not exported.functions + quantize = next(node for node in exported.graph.node if node.op_type == "QuantizeLinear") + dequantize = next(node for node in exported.graph.node if node.op_type == "DequantizeLinear") + assert _node_attributes(quantize)["axis"] == 0 + assert _node_attributes(dequantize)["axis"] == 0 + assert quantize.input[2] == dequantize.input[2] + assert _tensor_metadata(exported, quantize.input[2]) == (zero_point_dtype, [4]) + assert _tensor_metadata(exported, quantize.output[0]) == ( + zero_point_dtype, + list(sample_input.shape), + ) + assert scale_denominator in _scalar_initializer_values(exported) + assert any(node.op_type == "Where" for node in exported.graph.node) + + +def test_strict_int4_keeps_initializer_as_marker_input(tmp_path): + exported_program = torch.export.export(_StrictINT4Weight(), (), strict=True) + onnx_path = tmp_path / "int4_weight.onnx" + + torch.onnx.export( + exported_program, + (), + onnx_path, + dynamo=True, + opset_version=24, + custom_translation_table=get_dynamo_onnx_translation_table(), + ) + + exported = onnx.load(onnx_path) + onnx.checker.check_model(exported, full_check=True) + assert not exported.functions + int4_dq = next( + node + for node in exported.graph.node + if node.domain == "trt" and node.op_type == "DequantizeLinear" + ) + assert int4_dq.input[0] in {initializer.name for initializer in exported.graph.initializer} + assert exported.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT + + +def test_strict_mxfp8_static_weight_postprocess_syncs_initializer_metadata(tmp_path): + sample_input = torch.randn(4, 32) + exported_program = torch.export.export(_StrictMXFP8Weight(), (sample_input,), strict=True) + onnx_path = tmp_path / "mxfp8_weight.onnx" + + torch.onnx.export( + exported_program, + (), + onnx_path, + dynamo=True, + opset_version=21, + custom_translation_table=get_dynamo_onnx_translation_table(), + ) + + exported = onnx.load(onnx_path) + assert not exported.functions + mxfp8_dq = next( + node for node in exported.graph.node if node.op_type == "TRT_MXFP8DequantizeLinear" + ) + weight_name, scale_name = mxfp8_dq.input + assert _tensor_metadata(exported, weight_name) == (onnx.TensorProto.FLOAT, [32, 64]) + assert _tensor_metadata(exported, scale_name) == (onnx.TensorProto.FLOAT, []) + + exported = MXFP8QuantExporter.process_model(exported) + assert _tensor_metadata(exported, weight_name) == ( + onnx.TensorProto.FLOAT8E4M3FN, + [32, 64], + ) + assert _tensor_metadata(exported, scale_name) == (onnx.TensorProto.UINT8, [32, 2]) + assert exported.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT + onnx.checker.check_model(exported, full_check=True) + + +def test_dynamo_translation_table_covers_custom_op_overloads(): + table = get_dynamo_onnx_translation_table() + + assert { + torch.ops.tensorrt.quantize_op.default, + torch.ops.tensorrt.dynamic_block_quantize_op.default, + torch.ops.tensorrt.dynamic_block_quantize_op.overload, + } <= table.keys() + + +def test_custom_op_schemas_accept_legacy_positional_calls(): + sample_input = torch.randn(4, 32) + amax = torch.tensor(1.0) + + fp8_output = torch.ops.tensorrt.quantize_op(sample_input, amax, 8, 4, False, False) + mxfp8_output = torch.ops.tensorrt.dynamic_block_quantize_op.overload( + sample_input, 32, None, 8, 4, 9, 8 + ) + + assert fp8_output.shape == sample_input.shape + assert fp8_output.dtype == sample_input.dtype + assert mxfp8_output.shape == sample_input.shape + assert mxfp8_output.dtype == sample_input.dtype diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index d72bdca6e1c..3aba0f717a7 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -31,12 +31,72 @@ register_quant_backend, unregister_quant_backend, ) +from modelopt.torch.quantization.tensor_quant import fake_tensor_quant class TestFakeTensorQuantCPU(FakeTensorQuantTester): device = "cpu" +class _ExportedINT8(torch.nn.Module): + def __init__(self, unsigned=False, axis=None): + super().__init__() + self.unsigned = unsigned + self.axis = axis + + def forward(self, inputs, amax): + return fake_tensor_quant( + inputs, + amax, + None, + 8, + self.unsigned, + True, + None, + False, + None, + self.axis, + ) + + +@pytest.mark.parametrize( + ("amax", "axis", "unsigned"), + [ + pytest.param(torch.tensor(2.0), None, False, id="scalar-signed"), + pytest.param(torch.tensor([[2.0], [1.0]]), 0, False, id="per-channel-signed"), + pytest.param(torch.tensor(2.0), None, True, id="scalar-unsigned"), + ], +) +def test_strict_exported_program_int8_cpu_matches_eager(amax, axis, unsigned): + inputs = torch.tensor([[-2.0, -0.4, 0.2, 1.7], [-0.8, -0.1, 0.3, 0.9]]) + if unsigned: + inputs = inputs.abs() + model = _ExportedINT8(unsigned, axis) + expected = model(inputs, amax) + + exported_program = torch.export.export(model, (inputs, amax), strict=True) + + assert any( + node.target == torch.ops.tensorrt.quantize_op.default + for node in exported_program.graph.nodes + ) + torch.testing.assert_close(exported_program.module()(inputs, amax), expected, rtol=0, atol=0) + + +def test_strict_exported_program_int8_cpu_preserves_errors(): + inputs = torch.ones(2, 4) + amax = torch.tensor(2.0) + unsigned_program = torch.export.export( + _ExportedINT8(unsigned=True), (inputs, amax), strict=True + ) + signed_program = torch.export.export(_ExportedINT8(), (inputs, amax), strict=True) + + with pytest.raises(TypeError, match="Negative values encountered in unsigned quantization"): + unsigned_program.module()(-inputs, amax) + with pytest.raises(ValueError, match="Negative values in amax"): + signed_program.module()(inputs, -amax) + + class TestQuantizerAttributeConfig: def test_scaled_mode(self): num_bits = np.random.randint(1, 16)