From f82d8e85314a36312e8009aa0f1f4afb86e3c74d Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:50:42 +0000 Subject: [PATCH 1/8] [6508436] Fix BF16 FP8 ONNX export Support BF16 initializers during FP8 weight compression and skip redundant precision conversion when the requested dtype already matches the source model. Add focused exporter and end-to-end regression coverage. Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/onnx/export/fp8_exporter.py | 14 +++- modelopt/torch/_deploy/utils/torch_onnx.py | 13 +++- .../unit/onnx/quantization/test_qdq_utils.py | 57 ++++++++++++++- .../deploy/utils/test_torch_onnx_utils.py | 73 +++++++++++++++++++ 5 files changed, 152 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 68927078e83..ed95f99b0c3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -131,6 +131,7 @@ Changelog **Bug Fixes** +- Fix FP8 ONNX export of BF16 models failing during real weight compression. BF16 initializers now retain their raw representation when converted to PyTorch tensors, and post-export precision conversion is skipped when the requested dtype already matches the source model. - Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` / ``mixer.down_proj``, which the previous ``*mlp*`` selector missed, producing checkpoints with a null ``quant_algo``. - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. Stale output shapes are now reconciled via symbolic shape inference, and AutoCast falls back to schema-based type inference so unresolved ops no longer leave tensors untyped. - Fix fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) skipping modules without an ``act_fn`` attribute. Modules applying a custom gated activation between the two ``F.linear`` calls (e.g. ``MiniMaxM3VLExperts``) were silently skipped, leaving routed experts unquantized and failing HF export. Enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 427a7791f3b..6fa09f38263 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -17,6 +17,7 @@ import time +import ml_dtypes import numpy as np import onnx import onnx_graphsurgeon as gs @@ -33,6 +34,13 @@ _FP8_E4M3_SOFTMAX_SCALE = 1.0 / _FP8_E4M3_MAX +def _torch_from_numpy(array: np.ndarray) -> torch.Tensor: + """Convert a NumPy array to a PyTorch tensor while preserving BF16 values.""" + if array.dtype == ml_dtypes.bfloat16: + return torch.from_numpy(array.view(np.int16)).view(torch.bfloat16) + return torch.from_numpy(array) + + class FP8QuantExporter(ONNXQuantExporter): """Exporter for FP8 quantization.""" @@ -78,8 +86,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights = node.inputs[0] scale = node.inputs[1] - torch_weights = torch.from_numpy(weights.values) - torch_scale = torch.from_numpy(scale.values) + torch_weights = _torch_from_numpy(weights.values) + torch_scale = _torch_from_numpy(scale.values) quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] if dq_op.op != "TRT_FP8DequantizeLinear": @@ -194,7 +202,7 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: if any(out.op == "DequantizeLinear" for out in weight_input.outputs): continue - torch_weights = torch.from_numpy(weight_input.values.copy()) + torch_weights = _torch_from_numpy(weight_input.values.copy()) amax = torch_weights.abs().max().float() if amax == 0: continue diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 01fb754bbae..854545c6108 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -527,6 +527,9 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module + first_parameter = next(model.parameters(), None) + source_weights_dtype = first_parameter.dtype if first_parameter is not None else torch.float32 + # 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!) @@ -634,14 +637,20 @@ def get_onnx_bytes_and_metadata( if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - if weights_dtype in ["fp16", "bf16"]: + target_weights_dtype = { + "fp16": torch.float16, + "bf16": torch.bfloat16, + }.get(weights_dtype) + if target_weights_dtype is not None and target_weights_dtype != source_weights_dtype: if ( is_int4_quantized(model) or is_mxfp8_quantized(model) or is_fp8_quantized(model) or is_int8_quantized(model) ): - assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet" + assert weights_dtype == "fp16", ( + "Converting a quantized ONNX graph to BF16 is not supported yet" + ) onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, keep_io_types=False, diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 4b1e69ec538..bc833622f9b 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -15,13 +15,20 @@ import warnings +import ml_dtypes import numpy as np +import onnx import onnx_graphsurgeon as gs import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper -from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter +from modelopt.onnx.export import ( + FP8QuantExporter, + INT4QuantExporter, + MXFP8QuantExporter, + NVFP4QuantExporter, +) from modelopt.onnx.export.nvfp4_exporter import _cast_fp4 from modelopt.onnx.quantization.qdq_utils import ( _cast_fp8, @@ -484,6 +491,54 @@ def test_cast_fp4(self, input_array, expected_array): assert np.all(result == expected_array) +class TestFP8QuantExporter: + """Test suite for FP8QuantExporter.""" + + def test_bf16_weights_and_scale_are_compressed(self): + weight_data = np.array([[0.25, -0.5], [1.0, -2.0]], dtype=np.float32).astype( + ml_dtypes.bfloat16 + ) + scale_data = np.array(0.25, dtype=np.float32).astype(ml_dtypes.bfloat16) + weight = gs.Constant("weight", weight_data) + scale = gs.Constant("linear/weight_quantizer/scale", scale_data) + quantized = gs.Variable("quantized", dtype=np.uint8, shape=weight_data.shape) + dequantized = gs.Variable("dequantized", dtype=ml_dtypes.bfloat16, shape=weight_data.shape) + graph = gs.Graph( + nodes=[ + gs.Node( + op="TRT_FP8QuantizeLinear", + inputs=[weight, scale], + outputs=[quantized], + ), + gs.Node( + op="TRT_FP8DequantizeLinear", + inputs=[quantized, scale], + outputs=[dequantized], + ), + ], + outputs=[dequantized], + opset=23, + ) + + converted_model = FP8QuantExporter.compress_weights(gs.export_onnx(graph)) + + onnx.checker.check_model(converted_model) + assert [node.op_type for node in converted_model.graph.node] == ["DequantizeLinear"] + fp8_weight = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == "linear/weight_quantizer/fp8_weights" + ) + assert fp8_weight.data_type == TensorProto.FLOAT8E4M3FN + assert fp8_weight.raw_data == bytes.fromhex("38 c0 48 d0") + output_scale = next( + initializer + for initializer in converted_model.graph.initializer + if initializer.name == scale.name + ) + assert output_scale.data_type == TensorProto.BFLOAT16 + + class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" 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..2b50c8f7c40 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -24,6 +24,7 @@ import torch.nn as nn from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models +import modelopt.torch.quantization as mtq from modelopt.onnx.utils import get_batch_size_from_bytes, validate_batch_size from modelopt.torch._deploy.utils import ( OnnxBytes, @@ -57,6 +58,26 @@ } +def _export_fp8_linear(source_dtype, weights_dtype): + model = nn.Sequential(nn.Linear(4, 4, bias=False)).eval().to(source_dtype) + sample_input = torch.arange(4, dtype=source_dtype).reshape(1, 4) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model, + (sample_input,), + model_name="fp8_linear", + weights_dtype=weights_dtype, + dq_only=False, + onnx_opset=23, + ) + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) + + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) @@ -157,6 +178,58 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): ) +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "expected_onnx_dtype"), + [ + pytest.param( + torch.bfloat16, + "bf16", + onnx.TensorProto.BFLOAT16, + id="bf16-no-op", + ), + pytest.param( + torch.float32, + "fp16", + onnx.TensorProto.FLOAT16, + id="fp16-conversion", + ), + ], +) +def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, expected_onnx_dtype): + exported_model = _export_fp8_linear(source_dtype, weights_dtype) + + onnx.checker.check_model(exported_model) + assert not any( + node.op_type in {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} + for node in exported_model.graph.node + ) + initializer_by_name = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + fp8_weight_dq_nodes = [ + node + for node in exported_model.graph.node + if node.op_type == "DequantizeLinear" + and node.input[0] in initializer_by_name + and initializer_by_name[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN + ] + assert fp8_weight_dq_nodes + assert all( + initializer_by_name[node.input[1]].data_type == expected_onnx_dtype + for node in fp8_weight_dq_nodes + ) + graph_io = [*exported_model.graph.input, *exported_model.graph.output] + assert all(value.type.tensor_type.elem_type == expected_onnx_dtype for value in graph_io) + + +def test_fp8_export_rejects_bf16_conversion_from_fp32(): + with pytest.raises( + AssertionError, + match="Converting a quantized ONNX graph to BF16 is not supported", + ): + _export_fp8_linear(torch.float32, "bf16") + + class SingleArgModel(nn.Module): def forward(self, x: torch.Tensor): return torch.add(x, x) - x From db92c5bdb850cf21056002ba2b5689d8e6f677ee Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:21:00 +0000 Subject: [PATCH 2/8] [6508436] Honor requested precision in quantized ONNX export Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/onnx/autocast/convert.py | 170 ++++ modelopt/onnx/export/base_exporter.py | 13 +- modelopt/onnx/export/fp8_exporter.py | 66 +- modelopt/onnx/export/int4_exporter.py | 13 +- modelopt/onnx/export/int8_exporter.py | 4 +- modelopt/onnx/export/mxfp8_exporter.py | 28 +- modelopt/onnx/export/nvfp4_exporter.py | 53 +- modelopt/onnx/quantization/gs_patching.py | 12 +- modelopt/onnx/quantization/qdq_utils.py | 87 +- modelopt/torch/_deploy/utils/torch_onnx.py | 246 +++--- modelopt/torch/quantization/export_onnx.py | 132 +-- modelopt/torch/quantization/tensor_quant.py | 1 + .../quantization/test_fp8_mha_exporter.py | 148 +++- .../onnx/quantization/test_gs_patching.py | 45 ++ .../unit/onnx/quantization/test_qdq_utils.py | 361 ++++++++- .../deploy/utils/test_torch_onnx_utils.py | 757 +++++++++++++++++- 17 files changed, 1869 insertions(+), 269 deletions(-) create mode 100644 tests/unit/onnx/quantization/test_gs_patching.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ed95f99b0c3..f8b077d4bfe 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -131,7 +131,7 @@ Changelog **Bug Fixes** -- Fix FP8 ONNX export of BF16 models failing during real weight compression. BF16 initializers now retain their raw representation when converted to PyTorch tensors, and post-export precision conversion is skipped when the requested dtype already matches the source model. +- Make ``native`` the default ONNX export precision, preserving the model's existing precision, while explicit ``fp32``, ``fp16``, and ``bf16`` requests consistently control graph I/O and high-precision quantization boundaries across FP8, INT4, INT8, MXFP8, and NVFP4. This also fixes BF16 FP8 weight compression. - Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` / ``mixer.down_proj``, which the previous ``*mlp*`` selector missed, producing checkpoints with a null ``quant_algo``. - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. Stale output shapes are now reconciled via symbolic shape inference, and AutoCast falls back to schema-based type inference so unresolved ops no longer leave tensors untyped. - Fix fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) skipping modules without an ``act_fn`` attribute. Modules applying a custom gated activation between the two ``F.linear`` calls (e.g. ``MiniMaxM3VLExperts``) were silently skipped, leaving routed experts unquantized and failing HF export. Enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index 65cabe86974..e2c2e3442df 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -48,6 +48,176 @@ DEFAULT_DATA_MAX = 512 DEFAULT_INIT_MAX = np.finfo(np.float16).max LATEST_IR_VERSION_SUPPORTED_BY_ORT = 10 +_FLOAT_TYPES_TO_FP32 = { + onnx.TensorProto.DOUBLE, + onnx.TensorProto.FLOAT16, + onnx.TensorProto.BFLOAT16, +} +_STANDARD_FLOAT_DTYPE_ATTRIBUTES = { + "dtype", + "output_datatype", + "output_dtype", + "precision", + "softmax_precision", + "stash_type", +} +_TRT_FLOAT_DTYPE_ATTRIBUTES = {"output_dtype"} + + +def _convert_tensor_to_fp32(tensor: onnx.TensorProto) -> None: + if tensor.data_type not in _FLOAT_TYPES_TO_FP32: + return + if tensor.data_location == onnx.TensorProto.EXTERNAL and not tensor.raw_data: + raise ValueError("External tensor data must be loaded before FP32 conversion") + if tensor.HasField("segment"): + raise ValueError("Segmented tensors are not supported for FP32 conversion") + + dims = list(tensor.dims) + name = tensor.name + doc_string = tensor.doc_string + metadata_props = [deepcopy(prop) for prop in getattr(tensor, "metadata_props", ())] + if tensor.data_type in (onnx.TensorProto.FLOAT16, onnx.TensorProto.BFLOAT16): + values = ( + onnx_utils.read_f16_tensor_as_fp32(tensor) + if tensor.raw_data + else onnx.numpy_helper.to_array(tensor).astype(np.float32) + ) + else: + values = onnx.numpy_helper.to_array(tensor).astype(np.float32) + + # Release the low-precision payload before allocating the serialized FP32 payload. + tensor.Clear() + tensor.dims.extend(dims) + tensor.data_type = onnx.TensorProto.FLOAT + tensor.name = name + tensor.doc_string = doc_string + if metadata_props: + tensor.metadata_props.extend(metadata_props) + tensor.raw_data = values.tobytes() + + +def _convert_type_to_fp32(type_proto: onnx.TypeProto) -> None: + if type_proto.HasField("tensor_type"): + if type_proto.tensor_type.elem_type in _FLOAT_TYPES_TO_FP32: + type_proto.tensor_type.elem_type = onnx.TensorProto.FLOAT + elif type_proto.HasField("sparse_tensor_type"): + if type_proto.sparse_tensor_type.elem_type in _FLOAT_TYPES_TO_FP32: + type_proto.sparse_tensor_type.elem_type = onnx.TensorProto.FLOAT + elif type_proto.HasField("sequence_type"): + _convert_type_to_fp32(type_proto.sequence_type.elem_type) + elif type_proto.HasField("optional_type"): + _convert_type_to_fp32(type_proto.optional_type.elem_type) + elif type_proto.HasField("map_type"): + _convert_type_to_fp32(type_proto.map_type.value_type) + + +def _dtype_attributes_for_node(node: onnx.NodeProto) -> set[str]: + if node.domain in ("", "ai.onnx"): + attributes = set(_STANDARD_FLOAT_DTYPE_ATTRIBUTES) + if node.op_type == "Cast": + attributes.add("to") + return attributes + if node.domain == "trt": + return set(_TRT_FLOAT_DTYPE_ATTRIBUTES) + return set() + + +def _convert_attribute_to_fp32( + attribute: onnx.AttributeProto, dtype_attributes: set[str] | None = None +) -> None: + if ( + attribute.type == onnx.AttributeProto.INT + and dtype_attributes is not None + and attribute.name in dtype_attributes + and attribute.i in _FLOAT_TYPES_TO_FP32 + ): + attribute.i = onnx.TensorProto.FLOAT + elif attribute.type == onnx.AttributeProto.TENSOR: + _convert_tensor_to_fp32(attribute.t) + elif attribute.type == onnx.AttributeProto.TENSORS: + for tensor in attribute.tensors: + _convert_tensor_to_fp32(tensor) + elif attribute.type == onnx.AttributeProto.SPARSE_TENSOR: + _convert_tensor_to_fp32(attribute.sparse_tensor.values) + elif attribute.type == onnx.AttributeProto.SPARSE_TENSORS: + for sparse_tensor in attribute.sparse_tensors: + _convert_tensor_to_fp32(sparse_tensor.values) + elif attribute.type == onnx.AttributeProto.TYPE_PROTO: + _convert_type_to_fp32(attribute.tp) + elif attribute.type == onnx.AttributeProto.TYPE_PROTOS: + for type_proto in attribute.type_protos: + _convert_type_to_fp32(type_proto) + elif attribute.type == onnx.AttributeProto.GRAPH: + _convert_graph_to_fp32(attribute.g) + elif attribute.type == onnx.AttributeProto.GRAPHS: + for graph in attribute.graphs: + _convert_graph_to_fp32(graph) + + +def _convert_node_to_fp32(node: onnx.NodeProto) -> None: + is_standard_onnx_node = node.domain in ("", "ai.onnx") + dtype_attributes = _dtype_attributes_for_node(node) + for attribute in node.attribute: + if ( + is_standard_onnx_node + and node.op_type == "BitCast" + and attribute.name == "to" + and attribute.type == onnx.AttributeProto.INT + and attribute.i in _FLOAT_TYPES_TO_FP32 + ): + raise ValueError("BitCast targets cannot be converted safely to FP32") + _convert_attribute_to_fp32(attribute, dtype_attributes) + + +def _convert_function_to_fp32(function: onnx.FunctionProto) -> None: + dtype_attribute_refs = set() + bitcast_attribute_refs = set() + for node in function.node: + dtype_attributes = _dtype_attributes_for_node(node) + is_standard_bitcast = node.domain in ("", "ai.onnx") and node.op_type == "BitCast" + for attribute in node.attribute: + if not attribute.ref_attr_name: + continue + if is_standard_bitcast and attribute.name == "to": + bitcast_attribute_refs.add(attribute.ref_attr_name) + elif attribute.name in dtype_attributes: + dtype_attribute_refs.add(attribute.ref_attr_name) + + for attribute in function.attribute_proto: + if ( + attribute.name in bitcast_attribute_refs + and attribute.type == onnx.AttributeProto.INT + and attribute.i in _FLOAT_TYPES_TO_FP32 + ): + raise ValueError("BitCast targets cannot be converted safely to FP32") + dtype_attributes = {attribute.name} if attribute.name in dtype_attribute_refs else None + _convert_attribute_to_fp32(attribute, dtype_attributes) + for value_info in function.value_info: + _convert_type_to_fp32(value_info.type) + for node in function.node: + _convert_node_to_fp32(node) + + +def _convert_graph_to_fp32(graph: onnx.GraphProto) -> None: + for value_info in (*graph.input, *graph.output, *graph.value_info): + _convert_type_to_fp32(value_info.type) + for initializer in graph.initializer: + _convert_tensor_to_fp32(initializer) + for sparse_initializer in graph.sparse_initializer: + _convert_tensor_to_fp32(sparse_initializer.values) + for node in graph.node: + _convert_node_to_fp32(node) + + +def convert_to_fp32(model: onnx.ModelProto) -> onnx.ModelProto: + """Convert FP16, BF16, and FP64 values in an ONNX model to FP32 in place.""" + _convert_graph_to_fp32(model.graph) + for training_info in model.training_info: + _convert_graph_to_fp32(training_info.initialization) + _convert_graph_to_fp32(training_info.algorithm) + for function in model.functions: + _convert_function_to_fp32(function) + return model def _capture_network_io_metadata( diff --git a/modelopt/onnx/export/base_exporter.py b/modelopt/onnx/export/base_exporter.py index 41d80c0e7ec..da596f8a86a 100644 --- a/modelopt/onnx/export/base_exporter.py +++ b/modelopt/onnx/export/base_exporter.py @@ -24,12 +24,17 @@ class ONNXQuantExporter(ABC): """Base class for ONNX quantizer exporters.""" @classmethod - def process_model(cls, onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def process_model( + cls, onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Processes the ONNX model.""" onnx_model = cls.pre_process(onnx_model) onnx_model = cls.compute_scales(onnx_model) onnx_model = cls.compress_weights(onnx_model) - onnx_model = cls.post_process(onnx_model) + if high_precision_dtype is None: + onnx_model = cls.post_process(onnx_model) + else: + onnx_model = cls.post_process(onnx_model, high_precision_dtype) return onnx_model @staticmethod @@ -49,5 +54,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: @staticmethod @abstractmethod - def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def post_process( + onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Post-processes the ONNX model.""" diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 6fa09f38263..438d2e92f68 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -25,6 +25,7 @@ from onnx_graphsurgeon.ir.tensor import LazyValues from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.qdq_utils import np_dtype_map from .base_exporter import ONNXQuantExporter @@ -56,7 +57,7 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: @staticmethod def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Compresses FP32/FP16 weights to FP8 by folding QDQ nodes to DQ only. + """Compresses FP32/FP16/BF16 weights to FP8 by folding QDQ nodes to DQ only. Even though modelopt supports FP8 onnx export, the weights are represented in fp32 + QDQ. The storage is therefore very bad. In this function, @@ -64,7 +65,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights in the output model. TRT custom ops are converted to native ONNX DequantizeLinear. Parameters: - onnx_model: ONNX model with FP32/FP16 weights and TRT_FP8 QDQ nodes. + onnx_model: ONNX model with FP32/FP16/BF16 weights and TRT_FP8 QDQ nodes. Returns: ONNX model with FP8 weights and native ONNX DQ nodes for weights (QDQ preserved for activations). @@ -167,7 +168,9 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return gs.export_onnx(graph) @staticmethod - def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: + def _quantize_conv_weights_to_fp8( + graph: gs.Graph, high_precision_dtype: str | None = None + ) -> int: """Add FP8 weight DequantizeLinear for Conv layers with unquantized weights. Conv weight quantizers are disabled during TorchScript ONNX export because the @@ -182,6 +185,7 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: Args: graph: The onnx-graphsurgeon graph to modify in-place. + high_precision_dtype: Optional ONNX scalar type for the DQ scale and output. Returns: Number of Conv weight DQ nodes inserted. @@ -203,13 +207,34 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: continue torch_weights = _torch_from_numpy(weight_input.values.copy()) + if high_precision_dtype is not None: + scale_dtype = np_dtype_map[high_precision_dtype] + target_torch_dtype = _torch_from_numpy(np.empty((), dtype=scale_dtype)).dtype + torch_weights = torch_weights.to(target_torch_dtype) + else: + scale_dtype = np.float16 + amax = torch_weights.abs().max().float() if amax == 0: continue - scale_val = (amax / _FP8_E4M3_MAX).item() + + scale_value = (amax / _FP8_E4M3_MAX).item() + scale = np.array(scale_value, dtype=scale_dtype) + if high_precision_dtype is not None and scale == 0: + dtype_info = ( + ml_dtypes.finfo(scale_dtype) + if scale_dtype == ml_dtypes.bfloat16 + else np.finfo(scale_dtype) + ) + scale = np.array(dtype_info.smallest_subnormal, dtype=scale_dtype) + scaled_weights = ( + torch_weights / _torch_from_numpy(scale) + if high_precision_dtype is not None + else torch_weights / scale_value + ) # Quantize weights to FP8 (WAR: numpy doesn't support fp8) - fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_data = scaled_weights.to(torch.float8_e4m3fn).view(torch.uint8).numpy() fp8_tensor = onnx.TensorProto() fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN fp8_tensor.dims.extend(fp8_data.shape) @@ -218,13 +243,15 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) ) - # Scale in FP16 — DQ output type matches scale dtype, must match activation type scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", - np.array(scale_val, dtype=np.float16), + scale, ) - dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") + dq_output = gs.Variable( + node.name + "/weight_quantizer/dq_output", + dtype=scale_dtype if high_precision_dtype is not None else None, + ) dq_node = gs.Node( op="DequantizeLinear", name=node.name + "/weight_quantizer/DequantizeLinear", @@ -378,7 +405,7 @@ def _move_transpose_before_qdq(graph: gs.Graph) -> int: return count @staticmethod - def _insert_qdq_after_softmax(graph: gs.Graph) -> int: + def _insert_qdq_after_softmax(graph: gs.Graph, high_precision_dtype: str | None = None) -> int: """Insert FP8 Q→DQ on Softmax outputs feeding MatMul (required by TRT MHA fusion). Softmax output is data-independently bounded to [0, 1], so we use a fixed scale @@ -400,7 +427,13 @@ def _insert_qdq_after_softmax(graph: gs.Graph) -> int: # Match scale dtype to the graph's current float dtype so TRT stronglyTyped # sees consistent Q/DQ types with the surrounding compute. - scale_dtype = softmax_output.dtype if softmax_output.dtype is not None else np.float32 + scale_dtype = ( + np_dtype_map[high_precision_dtype] + if high_precision_dtype is not None + else softmax_output.dtype + if softmax_output.dtype is not None + else np.float32 + ) scale_val = np.array(_FP8_E4M3_SOFTMAX_SCALE, dtype=scale_dtype) scale_constant = gs.Constant(softmax_node.name + "/softmax_q_scale", scale_val) dq_scale_constant = gs.Constant( @@ -416,7 +449,10 @@ def _insert_qdq_after_softmax(graph: gs.Graph) -> int: ) q_output = gs.Variable(softmax_node.name + "/q_output") - dq_output = gs.Variable(softmax_node.name + "/dq_output", dtype=softmax_output.dtype) + dq_output = gs.Variable( + softmax_node.name + "/dq_output", + dtype=scale_dtype if high_precision_dtype is not None else softmax_output.dtype, + ) q_node = gs.Node( op="QuantizeLinear", name=softmax_node.name + "/QuantizeLinear", @@ -444,7 +480,9 @@ def _insert_qdq_after_softmax(graph: gs.Graph) -> int: return count @staticmethod - def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def post_process( + onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Post-processes the ONNX model for FP8 quantization. Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear, @@ -488,14 +526,14 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: ) # Add FP8 weight DQ for Conv layers that had weight quantizers disabled during export - count = FP8QuantExporter._quantize_conv_weights_to_fp8(graph) + count = FP8QuantExporter._quantize_conv_weights_to_fp8(graph, high_precision_dtype) if count > 0: logger.info(f"Inserted FP8 weight DequantizeLinear for {count} Conv nodes") # Attention-aware rewrites so TRT can fuse DQ into the attention MatMuls. n_mul = FP8QuantExporter._move_mul_before_qdq(graph) n_t = FP8QuantExporter._move_transpose_before_qdq(graph) - n_sm = FP8QuantExporter._insert_qdq_after_softmax(graph) + n_sm = FP8QuantExporter._insert_qdq_after_softmax(graph, high_precision_dtype) if n_mul or n_t or n_sm: logger.info( f"Attention QDQ rewrites: moved {n_mul} Mul, {n_t} Transpose; " diff --git a/modelopt/onnx/export/int4_exporter.py b/modelopt/onnx/export/int4_exporter.py index 0da217ae76f..e2edd0a7a93 100644 --- a/modelopt/onnx/export/int4_exporter.py +++ b/modelopt/onnx/export/int4_exporter.py @@ -223,8 +223,11 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def post_process( + onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Post-processes the ONNX model for INT4 quantization.""" + precision_dtype = high_precision_dtype or "Half" def is_pre_quant_scale_node(node: onnx.NodeProto) -> bool: has_pqs_input = any(input for input in node.input if "_pre_quant_scale" in input) @@ -265,12 +268,12 @@ def is_fp32_cast(node: onnx.NodeProto) -> bool: del graph.node[:] graph.node.extend(new_nodes) - # Cast bias to float16 + # Cast bias to the graph's high-precision dtype for node in graph.node: if node.op_type == "Add" and "proj/Add" in node.name: - cast_initializer_to_dtype(node, "Half", initializer_map) + cast_initializer_to_dtype(node, precision_dtype, initializer_map) - # Cast pre quant scales of o_proj and down_proj to float16 + # Cast pre quant scales of o_proj and down_proj to the high-precision dtype for node in graph.node: if node.op_type == "Mul" and ( any( @@ -278,6 +281,6 @@ def is_fp32_cast(node: onnx.NodeProto) -> bool: for x in ("o_proj/input_quantizer/Mul", "down_proj/input_quantizer/Mul") ) ): - cast_initializer_to_dtype(node, "Half", initializer_map) + cast_initializer_to_dtype(node, precision_dtype, initializer_map) return onnx_model diff --git a/modelopt/onnx/export/int8_exporter.py b/modelopt/onnx/export/int8_exporter.py index 4623279b531..03f0a40dd1b 100644 --- a/modelopt/onnx/export/int8_exporter.py +++ b/modelopt/onnx/export/int8_exporter.py @@ -40,6 +40,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def post_process( + onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Post-processes the ONNX model for INT8 quantization.""" return onnx_model diff --git a/modelopt/onnx/export/mxfp8_exporter.py b/modelopt/onnx/export/mxfp8_exporter.py index 8c1e1f4df4f..92b1005925b 100644 --- a/modelopt/onnx/export/mxfp8_exporter.py +++ b/modelopt/onnx/export/mxfp8_exporter.py @@ -143,20 +143,28 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def post_process( + onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Post-processes the ONNX model for MXFP8 quantization. - Sets DQ output type to FP16 and updates GELU nodes to use tanh approximation. + Sets DQ output type and updates GELU nodes to use tanh approximation. """ logger.info("Post-processing MXFP8 quantized model") graph = onnx_model.graph - - # Set output type of DQ to FP16 + precision_dtype = high_precision_dtype or "Half" + precision_suffix = { + "Float": "fp32", + "Half": "fp16", + "BFloat16": "bf16", + }[precision_dtype] + + # Set output type of DQ to the graph's high-precision dtype for node in graph.node: if node.op_type == "TRT_MXFP8DequantizeLinear": for attr in node.attribute: if attr.name == "output_dtype": - attr.i = onnx_dtype_map["Half"] + attr.i = onnx_dtype_map[precision_dtype] # Currently only tanh approximation is supported for Gelu for node in graph.node: @@ -166,20 +174,20 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: attr.s = b"tanh" logger.debug(f"Updated GELU node {node.name} to use tanh approximation") - # Insert cast to fp16 after Sqrt nodes + # Insert cast to the graph's high-precision dtype after Sqrt nodes cast_nodes_to_insert = [] for idx, node in enumerate(graph.node): if node.op_type == "Sqrt": sqrt_output = node.output[0] - cast_output = f"{sqrt_output}_cast_fp16" + cast_output = f"{sqrt_output}_cast_{precision_suffix}" # Create Cast node cast_node = onnx.helper.make_node( "Cast", inputs=[sqrt_output], outputs=[cast_output], - to=onnx_dtype_map["Half"], - name=f"{node.name}_cast_fp16", + to=onnx_dtype_map[precision_dtype], + name=f"{node.name}_cast_{precision_suffix}", ) cast_nodes_to_insert.append((idx + 1, cast_node)) @@ -194,6 +202,6 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: # Insert Cast nodes in reverse order to preserve indices for offset, (pos, cast_node) in enumerate(cast_nodes_to_insert): graph.node.insert(pos + offset, cast_node) - logger.debug(f"Inserted Cast to FP16 after {cast_node.input[0]}") + logger.debug(f"Inserted Cast to {precision_dtype} after {cast_node.input[0]}") return onnx_model diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 338e2725b14..1cd488ea4c5 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -318,7 +318,9 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + def post_process( + onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None + ) -> onnx.ModelProto: """Post-processes the ONNX model for NVFP4 quantization. Replaces TRT_FP4QDQ nodes with two DequantizeLinear nodes and handles @@ -334,37 +336,60 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: value_info_map = {vi.name: vi for vi in graph.value_info} graph_inputs = {inp.name for inp in graph.input} cast_output_cache: dict[tuple[str, str], str] = {} + casted_node_ids: set[int] = set() def _get_precision_dtype() -> str: # Check initializers to determine the precision of the weights precision_dtype = "Half" for initializer in graph.initializer: - if initializer.data_type == 16: + if initializer.data_type == onnx.TensorProto.BFLOAT16: precision_dtype = "BFloat16" break # Assuming all weights are of the same precision return precision_dtype + def _get_linear_consumers(tensor_name: str) -> list[onnx.NodeProto]: + nodes_to_visit = list(tensor_consumers.get(tensor_name, [])) + visited_node_ids = set() + linear_consumers = {} + + while nodes_to_visit: + node = nodes_to_visit.pop() + node_id = id(node) + if node_id in visited_node_ids: + continue + visited_node_ids.add(node_id) + + if node.op_type in {"Gemm", "MatMul"}: + linear_consumers[node_id] = node + elif node.op_type in {"Cast", "Transpose"}: + for output_name in node.output: + nodes_to_visit.extend(tensor_consumers.get(output_name, [])) + + assert linear_consumers, f"No Gemm or MatMul consumes {tensor_name}" + return list(linear_consumers.values()) + def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): # Change the input types to match weight precision (precision_dtype) - if node.op_type == "Transpose": - maybe_matmul = tensor_consumers[node.output[0]][0] - assert maybe_matmul.op_type == "MatMul" - node = maybe_matmul + assert node.op_type in {"Gemm", "MatMul"} # Create Cast nodes for each input of the target node except bias for i, input_name in enumerate(node.input[:2]): cast_output_name = cast_output_cache.get((input_name, precision_dtype)) if cast_output_name is None: - cast_output_suffix = "bf16" if precision_dtype == "BFloat16" else "f16" + cast_output_suffix = { + "Float": "f32", + "Half": "f16", + "BFloat16": "bf16", + }[precision_dtype] cast_output_name = f"{input_name}_{cast_output_suffix}" cast_output_cache[(input_name, precision_dtype)] = cast_output_name - # Create a Cast node to convert the input to FP16/BF16 + # Create a Cast node to convert the input to the selected precision cast_node = onnx.helper.make_node( "Cast", inputs=[input_name], # Original input of the target node outputs=[cast_output_name], - to=onnx_dtype_map[precision_dtype], # Cast to FP16/BF16 + to=onnx_dtype_map[precision_dtype], ) # Insert the Cast node into the graph @@ -373,7 +398,7 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): # Update the target node input to use the cast node output node.input[i] = cast_output_name - precision_dtype = _get_precision_dtype() + precision_dtype = high_precision_dtype or _get_precision_dtype() logger.debug(f"Using precision dtype: {precision_dtype}") fp4_qdq_nodes = [node for node in graph.node if node.op_type == "TRT_FP4QDQ"] @@ -416,9 +441,11 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): block_size, ) - # Cast input dtypes for the next node - next_node = tensor_consumers[node.output[0]][0] - _cast_input_dtypes(next_node, precision_dtype) + # Cast input dtypes for every linear consumer reached through Cast/Transpose wrappers. + for linear_node in _get_linear_consumers(node.output[0]): + if id(linear_node) not in casted_node_ids: + _cast_input_dtypes(linear_node, precision_dtype) + casted_node_ids.add(id(linear_node)) # Remove old initializers new_initializers = [ diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index a0eea84951e..bbd0dadca7c 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -70,9 +70,9 @@ def _export_tensor_proto(tensor: gs.Constant) -> onnx.TensorProto: onnx_tensor = tensor._values.tensor else: # is numpy array. - dtype = getattr( - tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(tensor.values.dtype) - ) + dtype = getattr(tensor, "explicit_dtype", None) + if dtype is None: + dtype = onnx.helper.np_dtype_to_tensor_dtype(tensor.values.dtype) vals = tensor.values if _onnx_supports_int4() and dtype in [onnx.TensorProto.INT4, onnx.TensorProto.UINT4]: @@ -101,9 +101,9 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V ) if tensor.dtype is not None: - dtype = getattr( - tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) - ) + dtype = getattr(tensor, "explicit_dtype", None) + if dtype is None: + dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) else: onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py index 3b48805439a..b172658d7f9 100644 --- a/modelopt/onnx/quantization/qdq_utils.py +++ b/modelopt/onnx/quantization/qdq_utils.py @@ -19,6 +19,7 @@ from collections.abc import Sequence from typing import Any +import ml_dtypes import numpy as np import onnx import onnx_graphsurgeon as gs @@ -61,6 +62,7 @@ onnx_bit_dtype_unsigned_map = {4: "UINT4", 8: "UINT8"} np_dtype_map = { + "BFloat16": ml_dtypes.bfloat16, "Float": np.float32, "Half": np.float16, "INT8": np.int8, @@ -1020,40 +1022,71 @@ def remove_graph_input_q(onnx_model: onnx.ModelProto) -> onnx.ModelProto: def replace_zero_scale_with_smallest_nonzero(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Replace zero scale values with smallest nonzero fp16 value in the ONNX model.""" - graph = onnx_model.graph - fp16_smallest_nonzero = np.float16(6e-08) + """Replace zero scale values with the smallest nonzero value of their dtype.""" qdq_op_types = { "QuantizeLinear", "DequantizeLinear", "TRT_INT4QuantizeLinear", "TRT_INT4DequantizeLinear", } - scale_tensor_names = { - node.input[1] - for node in graph.node - if node.op_type in qdq_op_types and len(node.input) >= 2 - } - # Scales stored as graph initializers (e.g. INT4_AWQ / TRT_INT4DequantizeLinear exports). - for init in graph.initializer: - if init.name in scale_tensor_names: - tensor = numpy_helper.to_array(init) - if tensor.dtype.kind == "f": - new_tensor = np.where(tensor == 0, fp16_smallest_nonzero, tensor).astype( - tensor.dtype - ) - init.CopyFrom(numpy_helper.from_array(new_tensor, init.name)) - # Scales emitted by Constant nodes (legacy QDQ export path). - for node in graph.node: - if node.op_type == "Constant" and node.output[0] in scale_tensor_names: + + def replace_zeros(tensor_proto: onnx.TensorProto) -> None: + dtype = { + onnx.TensorProto.BFLOAT16: ml_dtypes.bfloat16, + onnx.TensorProto.DOUBLE: np.float64, + onnx.TensorProto.FLOAT: np.float32, + onnx.TensorProto.FLOAT16: np.float16, + }.get(tensor_proto.data_type) + if dtype is None: + return + + tensor = numpy_helper.to_array(tensor_proto) + dtype_info = ml_dtypes.finfo(dtype) if dtype == ml_dtypes.bfloat16 else np.finfo(dtype) + smallest_nonzero = np.array(dtype_info.smallest_subnormal, dtype=dtype) + new_tensor = np.where(tensor == 0, smallest_nonzero, tensor).astype(dtype) + tensor_proto.CopyFrom(numpy_helper.from_array(new_tensor, tensor_proto.name)) + + def replace_zero_scales(graph: onnx.GraphProto) -> set[str]: + scale_tensor_names = { + node.input[1] + for node in graph.node + if node.op_type in qdq_op_types and len(node.input) >= 2 and node.input[1] + } + + for node in graph.node: for attr in node.attribute: - if attr.name == "value": - tensor = numpy_helper.to_array(attr.t) - if tensor.dtype.kind == "f": - new_tensor = np.where(tensor == 0, fp16_smallest_nonzero, tensor).astype( - tensor.dtype - ) - attr.t.CopyFrom(numpy_helper.from_array(new_tensor, attr.t.name)) + if attr.type == onnx.AttributeProto.GRAPH: + scale_tensor_names.update(replace_zero_scales(attr.g)) + elif attr.type == onnx.AttributeProto.GRAPHS: + for subgraph in attr.graphs: + scale_tensor_names.update(replace_zero_scales(subgraph)) + + # Scales stored as graph initializers (e.g. INT4_AWQ / TRT_INT4DequantizeLinear exports). + initializer_names = {init.name for init in graph.initializer} + sparse_initializer_names = { + init.values.name for init in graph.sparse_initializer if init.values.name + } + for init in graph.initializer: + if init.name in scale_tensor_names: + replace_zeros(init) + + # Scales emitted by Constant nodes (legacy QDQ export path). + node_output_names = {output for node in graph.node for output in node.output if output} + for node in graph.node: + if node.op_type == "Constant" and node.output[0] in scale_tensor_names: + for attr in node.attribute: + if attr.name == "value": + replace_zeros(attr.t) + + local_definitions = ( + initializer_names + | sparse_initializer_names + | node_output_names + | {value.name for value in graph.input if value.name} + ) + return scale_tensor_names - local_definitions + + replace_zero_scales(onnx_model.graph) return onnx_model diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 854545c6108..0fddd684ffb 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -19,22 +19,20 @@ import contextlib import inspect import json -import logging import os import shutil import tempfile from contextlib import nullcontext +from itertools import chain from typing import Any import onnx -import onnxconverter_common.float16 as _f16_module import torch import torch.nn as nn from onnx import ModelProto -from onnxconverter_common import convert_float_to_float16 from torch.nn.parallel import DataParallel, DistributedDataParallel -from modelopt.onnx.autocast.convert import convert_to_f16 +from modelopt.onnx.autocast.convert import convert_to_f16, convert_to_fp32 from modelopt.onnx.export import ( FP8QuantExporter, INT4QuantExporter, @@ -45,11 +43,9 @@ ) from modelopt.onnx.quantization.qdq_utils import qdq_to_dq, replace_zero_scale_with_smallest_nonzero from modelopt.onnx.utils import ( - change_casts_to_fp16, check_model_uses_external_data, fold_dq_fp32_to_fp16_casts, fold_q_fp16_to_fp32_casts, - fold_qdq_scale_fp16_to_fp32_casts, get_input_names, get_input_shapes, get_node_names, @@ -60,35 +56,12 @@ remove_redundant_casts, ) from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import flatten_tree, standardize_named_model_args from modelopt.torch.utils._pytree import TreeSpec from ..utils.onnx_optimizer import Optimizer -# Monkey-patch for onnxconverter_common bug in remove_unnecessary_cast_node(): -# cast_node_downstream_dict stores either a single node or a list of nodes, but the -# downstream-node handling at lines ~770/787 always does `downstream_node.input`, -# which raises AttributeError("'list' object has no attribute 'input'") when the -# value is a list (i.e. a Cast output feeds multiple consumers). -# TODO: Remove this patch once onnxconverter-common ships a fix. -# Upstream issue: https://github.com/microsoft/onnxconverter-common/issues/261 -_original_remove_unnecessary_cast_node = _f16_module.remove_unnecessary_cast_node - -_logger = logging.getLogger(__name__) - - -def _patched_remove_unnecessary_cast_node(graph): - try: - _original_remove_unnecessary_cast_node(graph) - except AttributeError as e: - if "'list' object has no attribute 'input'" in str(e): - _logger.debug("Skipping remove_unnecessary_cast_node due to known upstream bug: %s", e) - else: - raise - - -_f16_module.remove_unnecessary_cast_node = _patched_remove_unnecessary_cast_node - ModelMetadata = dict[str, Any] ModelType = Any ValueInfoType = Any @@ -97,6 +70,16 @@ def _patched_remove_unnecessary_cast_node(graph): DEFAULT_ONNX_OPSET = 20 ONNX_EXPORT_OUT_PREFIX = "out" TWO_GB = 2 * 1024 * 1024 * 1024 +WEIGHTS_DTYPE_TO_TORCH_DTYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} +WEIGHTS_DTYPE_TO_ONNX_DTYPE = { + "fp32": "Float", + "fp16": "Half", + "bf16": "BFloat16", +} class OnnxBytes: @@ -211,6 +194,61 @@ def _to_expected_onnx_type(val: Any) -> Any: return val +def _cast_floating_tensors(value: Any, dtype: torch.dtype) -> Any: + flat_values, tree_spec = flatten_tree(value) + flat_values = [ + item.to(dtype=dtype) + if isinstance(item, torch.Tensor) and item.is_floating_point() + else item + for item in flat_values + ] + return tree_spec.generate_pytree(flat_values) + + +def _get_autocast_context( + model: nn.Module, flat_input: list[Any], target_dtype: torch.dtype | None +): + if target_dtype not in (torch.float16, torch.bfloat16): + return nullcontext() + + for item in flat_input: + if isinstance(item, torch.Tensor) and item.is_floating_point(): + return torch.autocast(device_type=item.device.type, dtype=target_dtype) + for tensor in chain(model.parameters(), model.buffers()): + if tensor.is_floating_point(): + return torch.autocast(device_type=tensor.device.type, dtype=target_dtype) + for item in flat_input: + if isinstance(item, torch.Tensor): + return torch.autocast(device_type=item.device.type, dtype=target_dtype) + tensor = next(chain(model.parameters(), model.buffers()), None) + if tensor is not None: + return torch.autocast(device_type=tensor.device.type, dtype=target_dtype) + return torch.autocast(device_type="cpu", dtype=target_dtype) + + +@contextlib.contextmanager +def _override_onnx_quantizer_precision(model: nn.Module, high_precision_dtype: str | None): + if high_precision_dtype is None: + yield + return + + sentinel = object() + originals: list[tuple[TensorQuantizer, Any]] = [] + for module in model.modules(): + if isinstance(module, TensorQuantizer): + original = getattr(module, "_trt_high_precision_dtype", sentinel) + originals.append((module, original)) + module.trt_high_precision_dtype = high_precision_dtype + try: + yield + finally: + for quantizer, original in originals: + if original is sentinel: + del quantizer._trt_high_precision_dtype + else: + quantizer.trt_high_precision_dtype = original + + def generate_onnx_input( model_metadata: ModelMetadata, input: Any | tuple, ignore_nesting: bool = False ) -> dict[str, Any]: @@ -429,11 +467,15 @@ def _disable_fp8_conv_weight_quantizers(model: nn.Module): module.weight_quantizer.enable() -def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.ModelProto: +def quantize_weights( + model: nn.Module, + onnx_model: onnx.ModelProto, + high_precision_dtype: str | None = None, +) -> onnx.ModelProto: """Real quantizes the weights in the onnx model. Applies weight quantization to an ONNX model based on the quantization scheme detected - in the PyTorch model. Supports INT4, FP4, and MXFP8 quantization formats. + in the PyTorch model. Supports INT4, NVFP4, MXFP8, FP8, and INT8 quantization formats. The function performs a four-stage process for each detected quantization type: 1. Pre-process - Restructure the graph for quantization @@ -445,6 +487,7 @@ def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.Mode model (nn.Module): The original PyTorch model used to detect quantization schemes. This model should have been quantized using modelopt's quantization APIs. onnx_model (onnx.ModelProto): The ONNX model whose weights will be quantized. + high_precision_dtype: Optional ONNX scalar type used for the surrounding graph. Returns: onnx.ModelProto: The ONNX model with quantized weights applied. The returned model @@ -453,7 +496,7 @@ def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.Mode Notes: - Multiple quantization formats can be applied sequentially if the model contains different quantization schemes for different layers - - The function checks for INT4, FP4, and MXFP8 quantization in the PyTorch model + - The function checks every supported quantization format in the PyTorch model - Each quantization exporter modifies the ONNX graph in-place before returning """ @@ -474,7 +517,7 @@ def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.Mode return onnx_model for onnx_exporter in onnx_exporters: - onnx_model = onnx_exporter.process_model(onnx_model) + onnx_model = onnx_exporter.process_model(onnx_model, high_precision_dtype) return onnx_model @@ -489,7 +532,7 @@ def get_onnx_bytes_and_metadata( dynamo_export: bool = False, onnx_opset: int = DEFAULT_ONNX_OPSET, dq_only: bool = False, - weights_dtype: str = "fp32", + weights_dtype: str = "native", ) -> tuple[bytes, ModelMetadata]: """Get onnx model in bytes from input pytorch model together with the input/output of model. @@ -507,7 +550,11 @@ def get_onnx_bytes_and_metadata( `torch.onnx.export `_. onnx_opset: The onnx opset version to use for exporting the model. dq_only: If True, the exported onnx model is converted to a dq_only model. - weights_dtype: The dtype of the weights in the onnx model. + weights_dtype: Selects the floating-point graph I/O and high-precision Q/DQ boundary + dtype. ``native`` preserves the precision produced by the PyTorch export; + ``fp32``, ``fp16``, and ``bf16`` force that target while leaving format-native + quantized tensors and scales unchanged. Inference inputs supplied to the exported + ONNX model must use the selected explicit floating-point dtype. Returns: bytes: Onnx model in bytes. @@ -519,22 +566,24 @@ def get_onnx_bytes_and_metadata( if not isinstance(model, nn.Module): raise ValueError("Only PyTorch model compilation is supported.") - assert weights_dtype in ["fp32", "fp16", "bf16"], ( - "weights_dtype must be one of fp32, fp16, or bf16" + assert weights_dtype in ["native", "fp32", "fp16", "bf16"], ( + "weights_dtype must be one of native, fp32, fp16, or bf16" ) + if onnx_load_path and weights_dtype != "native": + raise ValueError("weights_dtype must be 'native' when onnx_load_path is provided") # unwrap DDP and DP models if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module - first_parameter = next(model.parameters(), None) - source_weights_dtype = first_parameter.dtype if first_parameter is not None else torch.float32 - # 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()} + target_torch_dtype = WEIGHTS_DTYPE_TO_TORCH_DTYPE.get(weights_dtype) + if target_torch_dtype in (torch.float16, torch.bfloat16): + named_args = _cast_floating_tensors(named_args, target_torch_dtype) # Also standardize dummy_input again so we can use it dummy_input = tuple(named_args.values()) @@ -551,17 +600,8 @@ def get_onnx_bytes_and_metadata( # during inference. input_none_names = list(set(tree_spec_input.names) - set(input_names)) - use_torch_autocast = not ( - is_fp4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - or weights_dtype == "fp32" - ) - autocast = torch.autocast("cuda") if use_torch_autocast else nullcontext() - # Get output once (we export in inference mode - so also using inference mode here!) - with torch.inference_mode(), autocast: + with torch.inference_mode(), _get_autocast_context(model, flat_input, target_torch_dtype): output = model(*named_args.values()) # Get output tree spec @@ -596,7 +636,14 @@ def get_onnx_bytes_and_metadata( conv_wq_context = ( _disable_fp8_conv_weight_quantizers(model) if is_fp8_quantized(model) else nullcontext() ) - with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: + high_precision_dtype = WEIGHTS_DTYPE_TO_ONNX_DTYPE.get(weights_dtype) + with ( + torch.inference_mode(), + _get_autocast_context(model, flat_input, target_torch_dtype), + _override_onnx_quantizer_precision(model, high_precision_dtype), + quantizer_context, + conv_wq_context, + ): additional_kwargs = {} if not dynamo_export: additional_kwargs["dynamic_axes"] = dynamic_axes @@ -632,49 +679,30 @@ def get_onnx_bytes_and_metadata( tree_spec_input, tree_spec_output, input_none_names, onnx_opt_graph, model ) - onnx_opt_graph = quantize_weights(model, onnx_opt_graph) + onnx_opt_graph = quantize_weights(model, onnx_opt_graph, high_precision_dtype) if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - target_weights_dtype = { - "fp16": torch.float16, - "bf16": torch.bfloat16, - }.get(weights_dtype) - if target_weights_dtype is not None and target_weights_dtype != source_weights_dtype: - if ( - is_int4_quantized(model) - or is_mxfp8_quantized(model) - or is_fp8_quantized(model) - or is_int8_quantized(model) - ): - assert weights_dtype == "fp16", ( - "Converting a quantized ONNX graph to BF16 is not supported yet" - ) - onnx_opt_graph = convert_float_to_float16( - onnx_opt_graph, - keep_io_types=False, - disable_shape_infer=True, - check_fp16_ready=False, - op_block_list=["QuantizeLinear", "DequantizeLinear", "Div"], - ) - # Change FP32 cast nodes feeding into Concat/Add to FP16 - op_list = ["Concat", "Add", "Sqrt", "LayerNormalization", "Clip", "Mul", "Exp"] - onnx_opt_graph = change_casts_to_fp16(onnx_opt_graph, op_list) - # Remove Cast(FP32->FP16) nodes after DQ by setting DQ output to FP16 directly - onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) - # Remove Cast(FP16->FP32) feeding Q/DQ scales so DQ stays FP16 for downstream - # MatMul/Add layers under strongly-typed TRT parsing. - onnx_opt_graph = fold_qdq_scale_fp16_to_fp32_casts(onnx_opt_graph) - else: - onnx_opt_graph = convert_to_f16( - onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False - ) + if weights_dtype == "fp32": + onnx_opt_graph = convert_to_fp32(onnx_opt_graph) + elif weights_dtype in ("fp16", "bf16") and not any( + ( + is_int4_quantized(model), + is_fp4_quantized(model), + is_mxfp8_quantized(model), + is_fp8_quantized(model), + is_int8_quantized(model), + ) + ): + onnx_opt_graph = convert_to_f16( + onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False + ) onnx_opt_graph = remove_redundant_casts(onnx_opt_graph) # Remove Cast nodes around Q/DQ for optimal TRT fusion - if is_fp8_quantized(model): + if is_fp8_quantized(model) and weights_dtype == "fp16": onnx_opt_graph = fold_q_fp16_to_fp32_casts(onnx_opt_graph) onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) @@ -685,22 +713,7 @@ def get_onnx_bytes_and_metadata( # Must be set after all gs.export_onnx() calls as graphsurgeon resets ir_version onnx_opt_graph.ir_version = 10 - # If the onnx model contains external data store the external tensors in one file and save the onnx model - if has_external_data(onnx_save_path): - tensor_paths = get_external_tensor_paths(onnx_path) - onnx.save_model( - onnx_opt_graph, - onnx_save_path, - save_as_external_data=True, - all_tensors_to_one_file=True, - location=f"{model_name}.onnx_data", - size_threshold=1024, - convert_attribute=False, - ) - for path in tensor_paths: - os.remove(path) - else: - onnx.save_model(onnx_opt_graph, onnx_save_path) + _save_onnx_model(onnx_opt_graph, onnx_save_path, model_name) onnx_bytes = OnnxBytes(onnx_save_path) @@ -724,6 +737,33 @@ def has_external_data(onnx_model_path: str): return check_model_uses_external_data(onnx_model) +def _save_onnx_model(onnx_model: onnx.ModelProto, onnx_save_path: str, model_name: str) -> None: + model_dir = os.path.dirname(onnx_save_path) + if not (has_external_data(onnx_save_path) or onnx_model.ByteSize() >= TWO_GB): + onnx.save_model(onnx_model, onnx_save_path) + return + + tensor_paths = get_external_tensor_paths(model_dir) + external_data_name = f"{model_name}.onnx_data" + external_data_path = os.path.join(model_dir, external_data_name) + if os.path.exists(external_data_path): + os.remove(external_data_path) + + onnx.save_model( + onnx_model, + onnx_save_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=external_data_name, + size_threshold=1024, + convert_attribute=True, + ) + external_data_path = os.path.abspath(external_data_path) + for path in tensor_paths: + if os.path.abspath(path) != external_data_path and os.path.exists(path): + os.remove(path) + + def create_model_metadata( tree_spec_input: TreeSpec, tree_spec_output: TreeSpec, diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index e5778c3c96b..a3611df5778 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -103,7 +103,7 @@ """Utility to export a quantized torch model to quantized ONNX.""" import contextlib -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import onnx import torch @@ -124,11 +124,22 @@ "INT8": onnx.TensorProto.INT8, "UINT8": onnx.TensorProto.UINT8, } -mha_valid_precisions = {"Half", "BFloat16"} +mha_fusion_precisions = {"Half", "BFloat16"} +mha_supported_precisions = {"Float", "Half", "BFloat16"} torch_dtype_map = {"Float": torch.float32, "Half": torch.float16, "BFloat16": torch.bfloat16} +def _cast_to_dtype(g: "GraphContext", tensor: torch.Value, dtype: str): + """Cast a graph value only when its dtype differs from the target.""" + if tensor.type().scalarType() == dtype: + return tensor + output_shape = sym_help._get_tensor_sizes(tensor) + return g.op("Cast", tensor, to_i=onnx_dtype_map[dtype]).setType( + tensor.type().with_dtype(torch_dtype_map[dtype]).with_sizes(output_shape) + ) + + def export_int8( g: "GraphContext", inputs: torch.Value, @@ -146,6 +157,9 @@ def export_int8( input_type = inputs.type().scalarType() if trt_high_precision_dtype is None: trt_high_precision_dtype = input_type + assert trt_high_precision_dtype in torch_dtype_map, ( + f"Unsupported high precision dtype: {trt_high_precision_dtype}" + ) if amax.numel() == 1: zero_point, axis = torch.tensor(0.0, device=amax.device), None @@ -169,22 +183,12 @@ def export_int8( scale.masked_fill_(scale == 0, 1.0) scale = g.op("Constant", value_t=scale) - assert trt_high_precision_dtype in (input_type, "Float", "BFloat16"), ( - "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." - ) - - # custom ops, so cast the input if needed. - if trt_high_precision_dtype != input_type: - inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) + inputs = _cast_to_dtype(g, inputs, trt_high_precision_dtype) quantized = g.op("QuantizeLinear", inputs, scale, zero_point, axis_i=axis) out = g.op("DequantizeLinear", quantized, scale, zero_point, axis_i=axis).setType( inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) ) - # custom ops, so cast the output if needed. - if trt_high_precision_dtype != input_type: - inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[input_type]) - return out @@ -199,12 +203,17 @@ def export_int4( ): """Export quantized model to INT4 ONNX.""" assert num_bits == 4, "Number of bits must be 4 for INT4 ONNX export." - scale_inv = amax / 7.0 - scale_inv_op = g.op("Constant", value_t=scale_inv) otype = inputs.type().scalarType() output_shape = sym_help._get_tensor_sizes(inputs) if trt_high_precision_dtype is None: trt_high_precision_dtype = otype + scale_inv = amax / 7.0 + else: + assert trt_high_precision_dtype in torch_dtype_map, ( + f"Unsupported high precision dtype: {trt_high_precision_dtype}" + ) + scale_inv = (amax / 7.0).to(torch_dtype_map[trt_high_precision_dtype]) + scale_inv_op = g.op("Constant", value_t=scale_inv) return g.op( "trt::DequantizeLinear", inputs, scale_inv_op, axis_i=axis, block_size_i=block_size ).setType( @@ -216,14 +225,13 @@ def _fp8_quantize( g: "GraphContext", inputs: torch.Value, scale_inv: float, + output_dtype: str, ): """Helper Function for Quantization.""" - # Emit the scale in the native input dtype so no Cast is inserted between the - # graph and Q/DQ (Cast nodes block TRT from fusing DQ into the MatMul kernel). output_shape = sym_help._get_tensor_sizes(inputs) scale = g.op( "Constant", - value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]), + value_t=torch.tensor(scale_inv, dtype=torch_dtype_map[output_dtype]), ) return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) @@ -255,15 +263,19 @@ def export_fp8( ): """Export quantized model to FP8 ONNX. - ``trt_high_precision_dtype`` is accepted for API compatibility but unused: Q/DQ now - emit scales in the native input dtype, so no intermediate Cast is required. + ``None`` preserves the native input dtype. """ - del trt_high_precision_dtype scale = 1.0 if amax is None else 448.0 / float(amax) - otype = inputs.type().scalarType() + input_dtype = inputs.type().scalarType() + if trt_high_precision_dtype is None: + trt_high_precision_dtype = input_dtype + assert trt_high_precision_dtype in torch_dtype_map, ( + f"Unsupported high precision dtype: {trt_high_precision_dtype}" + ) - q_tensor = _fp8_quantize(g, inputs, 1.0 / scale) - return _fp8_dequantize(g, q_tensor, 1.0 / scale, otype) + inputs = _cast_to_dtype(g, inputs, trt_high_precision_dtype) + q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype) + return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype) def scaled_dot_product_attention( @@ -362,7 +374,7 @@ def export_fp8_mha( q_quantized_scale: float = 1.0, k_quantized_scale: float = 1.0, v_quantized_scale: float = 1.0, - high_precision_flag: str = "Half", + high_precision_flag: str | None = "Half", disable_fp8_mha: bool = True, ): r"""Export quantized fMHA to FP8 ONNX. @@ -408,6 +420,22 @@ def export_fp8_mha( "is_causal and attn_mask cannot be set at the same time" ) + if not disable_fp8_mha: + if high_precision_flag is None: + high_precision_flag = query.type().scalarType() + if high_precision_flag not in mha_supported_precisions: + raise ValueError(f"Unsupported FP8 MHA precision: {high_precision_flag}") + if high_precision_flag == "Float": + query = _cast_to_dtype(g, query, high_precision_flag) + key = _cast_to_dtype(g, key, high_precision_flag) + value = _cast_to_dtype(g, value, high_precision_flag) + elif high_precision_flag in mha_fusion_precisions and { + query.type().scalarType(), + key.type().scalarType(), + value.type().scalarType(), + } != {high_precision_flag}: + raise ValueError("The quantized MHA must have 16-bit inputs.") + scale = sym_help._maybe_get_const(scale, "f") if sym_help._is_none(scale): scale = _attention_scale(g, query) @@ -431,24 +459,15 @@ def export_fp8_mha( query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) if not disable_fp8_mha: - if high_precision_flag not in mha_valid_precisions: - raise ValueError( - "The Quantized config setting doesn't match TRT's fusion pattern; the qdqs must be in 16 bits." - ) - q_input_dtype = query.type().scalarType() - k_input_dtype = key.type().scalarType() - v_input_dtype = value.type().scalarType() - if {q_input_dtype, k_input_dtype, v_input_dtype} != {high_precision_flag}: - raise ValueError("The quantized MHA must have 16-bit inputs.") query_scaled = export_fp8(g, query_scaled, q_quantized_scale, high_precision_flag) - query_scaled = g.op("Cast", query_scaled, to_i=onnx_dtype_map["Float"]) + query_scaled = _cast_to_dtype(g, query_scaled, "Float") key_transposed_scaled = export_fp8( g, key_transposed_scaled, k_quantized_scale, high_precision_flag ) - key_transposed_scaled = g.op("Cast", key_transposed_scaled, to_i=onnx_dtype_map["Float"]) + key_transposed_scaled = _cast_to_dtype(g, key_transposed_scaled, "Float") mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) if not disable_fp8_mha: - mul_qk = g.op("Cast", mul_qk, to_i=onnx_dtype_map[high_precision_flag]) + mul_qk = _cast_to_dtype(g, mul_qk, cast("str", high_precision_flag)) if sym_help._is_none(attn_mask): mul_qk_add = mul_qk @@ -472,7 +491,7 @@ def export_fp8_mha( if not disable_fp8_mha: # Softmax's output scale is hard coded to 1.0 attn_weight = export_fp8(g, attn_weight, 1.0, high_precision_flag) - attn_weight = g.op("Cast", attn_weight, to_i=onnx_dtype_map["Float"]) + attn_weight = _cast_to_dtype(g, attn_weight, "Float") if dropout_p != 0: attn_weight = g.op( @@ -482,11 +501,9 @@ def export_fp8_mha( ) if not disable_fp8_mha: value = export_fp8(g, value, v_quantized_scale, high_precision_flag) - value = g.op("Cast", value, to_i=onnx_dtype_map["Float"]) - return g.op( - "Cast", - g.op("MatMul", attn_weight, value), - to_i=onnx_dtype_map[high_precision_flag], + value = _cast_to_dtype(g, value, "Float") + return _cast_to_dtype( + g, g.op("MatMul", attn_weight, value), cast("str", high_precision_flag) ) else: return g.op("MatMul", attn_weight, value) @@ -502,7 +519,7 @@ def _fp4_dynamic_quantize( scale_type: int = onnx_dtype_map["Float8"], ): """Helper Function for Dynamic Quantization.""" - # TRT StronglyType only supports FP16 QDQ ops, so cast the input if needed. + # Match the input to the requested strongly typed QDQ precision. input_type = inputs.type().scalarType() if trt_high_precision_dtype is None: trt_high_precision_dtype = input_type @@ -596,16 +613,37 @@ def export_mxfp8( onnx_quantizer_type: str, block_size: int, axis: int = -1, + trt_high_precision_dtype: str | None = None, ): """Export quantized model to MXFP8 ONNX.""" input_dtype = inputs.type().scalarType() + output_shape = sym_help._get_tensor_sizes(inputs) + if trt_high_precision_dtype is None: + trt_high_precision_dtype = input_dtype + assert trt_high_precision_dtype in torch_dtype_map, ( + f"Unsupported high precision dtype: {trt_high_precision_dtype}" + ) + if onnx_quantizer_type == "dynamic": + inputs = _cast_to_dtype(g, inputs, trt_high_precision_dtype) x_f8, sx_ui8 = _mxfp8_dynamic_quantize(g, inputs, block_size, axis=axis) - return _mxfp8_dequantize(g, x_f8, sx_ui8, block_size, axis=axis, input_dtype=input_dtype) + output = _mxfp8_dequantize( + g, + x_f8, + sx_ui8, + block_size, + axis=axis, + input_dtype=trt_high_precision_dtype, + ) else: - scale = torch.tensor(1.0, dtype=torch_dtype_map[input_dtype]) - return _mxfp8_dequantize(g, inputs, scale, block_size, axis=axis, input_dtype=input_dtype) + scale = torch.tensor(1.0, dtype=torch_dtype_map[trt_high_precision_dtype]) + output = _mxfp8_dequantize( + g, inputs, scale, block_size, axis=axis, input_dtype=trt_high_precision_dtype + ) + return output.setType( + inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) + ) def export_fp4( diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..aab54ef25a7 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -530,6 +530,7 @@ def symbolic( inputs, onnx_quantizer_type, block_size, + trt_high_precision_dtype=trt_high_precision_dtype, ) raise NotImplementedError( f"Unsupported num_bits: {num_bits} and scale_bits: {scale_bits} for ONNX export." diff --git a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py index 1f7251a9ad9..5900ce43dbf 100644 --- a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py +++ b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py @@ -15,15 +15,52 @@ """Tests for the attention-aware FP8 ONNX graph rewrites in ``FP8QuantExporter``.""" +import io + +import ml_dtypes import numpy as np +import onnx import onnx_graphsurgeon as gs import pytest +import torch +from torch.onnx import symbolic_helper from modelopt.onnx.export.fp8_exporter import FP8QuantExporter +from modelopt.torch.quantization.export_onnx import export_fp8_mha + + +class _FP8MHAFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, query, key, value, high_precision_dtype): + return torch.nn.functional.scaled_dot_product_attention(query, key, value) + + @staticmethod + @symbolic_helper.parse_args("v", "v", "v", "s") + def symbolic(g, query, key, value, high_precision_dtype): + return export_fp8_mha( + g, + query, + key, + value, + q_quantized_scale=1.0, + k_quantized_scale=1.0, + v_quantized_scale=1.0, + high_precision_flag=high_precision_dtype, + disable_fp8_mha=False, + ) + + +class _FP8MHAModule(torch.nn.Module): + def __init__(self, high_precision_dtype): + super().__init__() + self.high_precision_dtype = high_precision_dtype + + def forward(self, query, key, value): + return _FP8MHAFunction.apply(query, key, value, self.high_precision_dtype) -def _var(name): - return gs.Variable(name, dtype=np.float32) +def _var(name, dtype=np.float32, shape=None): + return gs.Variable(name, dtype=dtype, shape=shape) def _qdq(src): @@ -70,16 +107,47 @@ def test_move_transpose_before_qdq_rewrites_dq_transpose_matmul_pattern(): assert q.inputs[0].inputs[0].op == "Transpose" -def test_insert_qdq_after_softmax_adds_fixed_scale_q_dq(): +@pytest.mark.parametrize( + ("high_precision_dtype", "numpy_dtype", "onnx_dtype"), + [ + pytest.param(None, np.float32, onnx.TensorProto.FLOAT, id="legacy"), + pytest.param("Float", np.float32, onnx.TensorProto.FLOAT, id="float"), + pytest.param("Half", np.float16, onnx.TensorProto.FLOAT16, id="half"), + pytest.param("BFloat16", ml_dtypes.bfloat16, onnx.TensorProto.BFLOAT16, id="bfloat16"), + ], +) +def test_insert_qdq_after_softmax_adds_target_scale_q_dq( + high_precision_dtype, numpy_dtype, onnx_dtype +): """Softmax → MatMul picks up ``Q → DQ`` with the fixed ``1/448`` scale.""" - scores, v, y, sm_out = _var("scores"), _var("v"), _var("y"), _var("sm_out") + scores, v, y, sm_out = ( + _var("scores", numpy_dtype, [2, 2]), + _var("v", numpy_dtype, [2, 2]), + _var("y", numpy_dtype, [2, 2]), + _var("sm_out", numpy_dtype, [2, 2]), + ) sm = gs.Node(op="Softmax", inputs=[scores], outputs=[sm_out], attrs={"axis": -1}) mm = gs.Node(op="MatMul", inputs=[sm_out, v], outputs=[y]) graph = _graph([sm, mm], [scores, v], [y]) - assert FP8QuantExporter._insert_qdq_after_softmax(graph) == 1 + count = ( + FP8QuantExporter._insert_qdq_after_softmax(graph) + if high_precision_dtype is None + else FP8QuantExporter._insert_qdq_after_softmax(graph, high_precision_dtype) + ) + assert count == 1 q = next(n for n in graph.nodes if n.op == "QuantizeLinear") - assert np.isclose(float(q.inputs[1].values), 1.0 / 448.0) + dq = next(n for n in graph.nodes if n.op == "DequantizeLinear") + expected_scale = np.array(1.0 / 448.0, dtype=numpy_dtype) + for scale in (q.inputs[1], dq.inputs[1]): + assert scale.values.dtype == expected_scale.dtype + np.testing.assert_array_equal(scale.values, expected_scale) + assert np.dtype(dq.outputs[0].dtype) == np.dtype(numpy_dtype) + assert mm.inputs[0] is dq.outputs[0] + + converted_model = gs.export_onnx(graph) + onnx.checker.check_model(converted_model) + onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) @pytest.mark.parametrize( @@ -116,3 +184,71 @@ def test_rewrites_skip_when_non_matmul_consumer_exists(rewrite): [y_mm, y_side], ) assert getattr(FP8QuantExporter, rewrite)(graph) == 0 + + +@pytest.mark.parametrize( + ("torch_dtype", "high_precision_dtype", "onnx_dtype", "expected_accumulation_casts"), + [ + pytest.param(torch.float32, "Float", onnx.TensorProto.FLOAT, 0, id="float"), + pytest.param(torch.float16, "Half", onnx.TensorProto.FLOAT16, 4, id="half"), + pytest.param(torch.bfloat16, "BFloat16", onnx.TensorProto.BFLOAT16, 4, id="bfloat16"), + pytest.param(torch.float32, None, onnx.TensorProto.FLOAT, 0, id="native-float"), + pytest.param(torch.bfloat16, None, onnx.TensorProto.BFLOAT16, 4, id="native-bfloat16"), + ], +) +def test_fp8_mha_symbolic_preserves_accumulation_contract( + torch_dtype, high_precision_dtype, onnx_dtype, expected_accumulation_casts +): + """FP8-MHA supports FP32 while retaining 16-bit fusion casts.""" + shape = (1, 1, 2, 4) + inputs = tuple(torch.ones(shape, dtype=torch_dtype) for _ in range(3)) + buffer = io.BytesIO() + torch.onnx.export( + _FP8MHAModule(high_precision_dtype), + inputs, + buffer, + opset_version=20, + dynamo=False, + ) + + model = onnx.load_model_from_string(buffer.getvalue()) + onnx.checker.check_model(model) + onnx.shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + + qdq_ops = {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} + qdq_nodes = [node for node in model.graph.node if node.op_type in qdq_ops] + assert len(qdq_nodes) == 8 + + tensor_dtype = { + initializer.name: initializer.data_type for initializer in model.graph.initializer + } + for node in model.graph.node: + if node.op_type == "Constant": + tensor = next((attr.t for attr in node.attribute if attr.name == "value"), None) + if tensor is not None: + tensor_dtype[node.output[0]] = tensor.data_type + assert all(tensor_dtype[node.input[1]] == onnx_dtype for node in qdq_nodes) + + producer_by_output = { + output: node for node in model.graph.node for output in node.output if output + } + accumulation_casts = [ + node + for node in model.graph.node + if node.op_type == "Cast" + and any(attr.name == "to" and attr.i == onnx.TensorProto.FLOAT for attr in node.attribute) + and producer_by_output.get(node.input[0], onnx.NodeProto()).op_type in qdq_ops + ] + assert len(accumulation_casts) == expected_accumulation_casts + back_casts = [ + node + for node in model.graph.node + if node.op_type == "Cast" + and any(attr.name == "to" and attr.i == onnx_dtype for attr in node.attribute) + and producer_by_output.get(node.input[0], onnx.NodeProto()).op_type == "MatMul" + ] + assert len(back_casts) == (0 if onnx_dtype == onnx.TensorProto.FLOAT else 2) + assert all( + value.type.tensor_type.elem_type == onnx_dtype + for value in [*model.graph.input, *model.graph.output] + ) diff --git a/tests/unit/onnx/quantization/test_gs_patching.py b/tests/unit/onnx/quantization/test_gs_patching.py new file mode 100644 index 00000000000..17df218dfe9 --- /dev/null +++ b/tests/unit/onnx/quantization/test_gs_patching.py @@ -0,0 +1,45 @@ +# 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. + +import ml_dtypes +import numpy as np +import onnx +import onnx_graphsurgeon as gs + +from modelopt.onnx.quantization.gs_patching import _export_tensor_proto, _export_value_info_proto + + +def test_export_constant_uses_explicit_bf16_dtype_without_fallback(monkeypatch): + tensor = gs.Constant("scale", np.array([1.0], dtype=ml_dtypes.bfloat16)) + tensor.explicit_dtype = onnx.TensorProto.BFLOAT16 + + def fail_on_fallback(_): + raise AssertionError("NumPy dtype fallback should not be evaluated") + + monkeypatch.setattr(onnx.helper, "np_dtype_to_tensor_dtype", fail_on_fallback) + + tensor_proto = _export_tensor_proto(tensor) + + assert tensor_proto.data_type == onnx.TensorProto.BFLOAT16 + np.testing.assert_array_equal(onnx.numpy_helper.to_array(tensor_proto), tensor.values) + + +def test_export_value_info_uses_explicit_onnx_dtype_without_numpy_conversion(): + tensor = gs.Variable("input", dtype=onnx.TensorProto.BFLOAT16, shape=[1]) + tensor.explicit_dtype = onnx.TensorProto.BFLOAT16 + + value_info = _export_value_info_proto(tensor, do_type_check=True) + + assert value_info.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index bc833622f9b..f2413b82d6a 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -172,7 +172,7 @@ def create_test_model_with_cast_nodes(): return model -def create_test_model_with_proj_nodes(): +def create_test_model_with_proj_nodes(graph_dtype=TensorProto.FLOAT): """Create a test model with projection nodes to test bias and scale casting.""" # Create bias tensor bias_data = np.random.uniform(-1.0, 1.0, size=(16,)).astype(np.float32) @@ -182,7 +182,7 @@ def create_test_model_with_proj_nodes(): scale_data = np.random.uniform(0.1, 1.0, size=(1,)).astype(np.float32) scale_tensor = numpy_helper.from_array(scale_data, "quant_scale") - input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, 16]) + input_tensor = helper.make_tensor_value_info("input", graph_dtype, [4, 16]) # Add node (projection bias) add_node = helper.make_node( @@ -201,7 +201,7 @@ def create_test_model_with_proj_nodes(): nodes=[add_node, mul_node], name="test_graph", inputs=[input_tensor], - outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [4, 16])], + outputs=[helper.make_tensor_value_info("output", graph_dtype, [4, 16])], initializer=[bias_tensor, scale_tensor], ) @@ -397,24 +397,38 @@ def test_quantization_with_constant_scale(self): ) assert any("scale" in input_name for input_name in dq_node.input) - def test_projection_bias_and_scale_casting(self): - """Test that projection biases and quantization scales are cast to float16.""" - model = create_test_model_with_proj_nodes() - - # Run quantization - quantized_model = INT4QuantExporter.process_model(model) + @pytest.mark.parametrize( + ("high_precision_dtype", "onnx_dtype"), + [ + pytest.param(None, TensorProto.FLOAT16, id="legacy"), + pytest.param("Float", TensorProto.FLOAT, id="float"), + pytest.param("Half", TensorProto.FLOAT16, id="half"), + pytest.param("BFloat16", TensorProto.BFLOAT16, id="bfloat16"), + ], + ) + def test_projection_bias_and_scale_casting(self, high_precision_dtype, onnx_dtype): + """Test projection bias and pre-quant scale target casting.""" + graph_dtype = TensorProto.FLOAT if high_precision_dtype is None else onnx_dtype + model = create_test_model_with_proj_nodes(graph_dtype) + + quantized_model = ( + INT4QuantExporter.post_process(model) + if high_precision_dtype is None + else INT4QuantExporter.post_process(model, high_precision_dtype) + ) - # Verify bias tensor is cast to float16 bias_tensor = next( init for init in quantized_model.graph.initializer if "proj_bias" in init.name ) - assert bias_tensor.data_type == TensorProto.FLOAT16 + assert bias_tensor.data_type == onnx_dtype - # Verify quantization scale is cast to float16 scale_tensor = next( init for init in quantized_model.graph.initializer if "quant_scale" in init.name ) - assert scale_tensor.data_type == TensorProto.FLOAT16 + assert scale_tensor.data_type == onnx_dtype + onnx.checker.check_model(quantized_model) + if high_precision_dtype is not None: + onnx.shape_inference.infer_shapes(quantized_model, check_type=True, strict_mode=True) class TestCastFunctions: @@ -538,6 +552,45 @@ def test_bf16_weights_and_scale_are_compressed(self): ) assert output_scale.data_type == TensorProto.BFLOAT16 + def test_conv_uses_target_rounded_weights_and_scale(self): + weight_data = np.array([0.58945024, -13.944608], dtype=np.float32).reshape(1, 1, 1, 2) + input_tensor = gs.Variable("input", dtype=np.float32, shape=[1, 1, 1, 2]) + output_tensor = gs.Variable("output", dtype=np.float32, shape=[1, 1, 1, 1]) + conv = gs.Node( + op="Conv", + name="conv", + inputs=[input_tensor, gs.Constant("weight", weight_data)], + outputs=[output_tensor], + ) + graph = gs.Graph(nodes=[conv], inputs=[input_tensor], outputs=[output_tensor], opset=23) + + assert FP8QuantExporter._quantize_conv_weights_to_fp8(graph, "BFloat16") == 1 + + dq_node = next(node for node in graph.nodes if node.op == "DequantizeLinear") + fp8_weights, scale = dq_node.inputs + assert scale.values.dtype == ml_dtypes.bfloat16 + assert dq_node.outputs[0].dtype == ml_dtypes.bfloat16 + assert fp8_weights._values.tensor.raw_data == bytes([90, 254]) + + def test_conv_clamps_target_rounded_zero_scale_before_quantizing(self): + weight_data = np.array([1e-6, -2e-6], dtype=np.float32).reshape(1, 1, 1, 2) + input_tensor = gs.Variable("input", dtype=np.float16, shape=[1, 1, 1, 2]) + output_tensor = gs.Variable("output", dtype=np.float16, shape=[1, 1, 1, 1]) + conv = gs.Node( + op="Conv", + name="conv", + inputs=[input_tensor, gs.Constant("weight", weight_data)], + outputs=[output_tensor], + ) + graph = gs.Graph(nodes=[conv], inputs=[input_tensor], outputs=[output_tensor], opset=23) + + assert FP8QuantExporter._quantize_conv_weights_to_fp8(graph, "Half") == 1 + + dq_node = next(node for node in graph.nodes if node.op == "DequantizeLinear") + fp8_weights, scale = dq_node.inputs + assert scale.values == np.finfo(np.float16).smallest_subnormal + assert fp8_weights._values.tensor.raw_data == bytes([88, 224]) + class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" @@ -589,6 +642,52 @@ def test_mxfp8_output_dtype_update(self): output_dtype_attr = next(attr for attr in dq_node.attribute if attr.name == "output_dtype") assert output_dtype_attr.i == TensorProto.FLOAT16 + @pytest.mark.parametrize( + ("high_precision_dtype", "onnx_dtype", "suffix"), + [ + pytest.param(None, TensorProto.FLOAT16, "fp16", id="legacy"), + pytest.param("Float", TensorProto.FLOAT, "fp32", id="float"), + pytest.param("Half", TensorProto.FLOAT16, "fp16", id="half"), + pytest.param("BFloat16", TensorProto.BFLOAT16, "bf16", id="bfloat16"), + ], + ) + def test_sqrt_output_cast_uses_target_dtype(self, high_precision_dtype, onnx_dtype, suffix): + graph_dtype = TensorProto.FLOAT if high_precision_dtype is None else onnx_dtype + input_info = helper.make_tensor_value_info("input", graph_dtype, [2]) + output_info = helper.make_tensor_value_info("output", graph_dtype, [2]) + sqrt_node = helper.make_node("Sqrt", inputs=["input"], outputs=["sqrt_output"], name="sqrt") + consumer = helper.make_node( + "Identity", inputs=["sqrt_output"], outputs=["output"], name="consumer" + ) + model = helper.make_model( + helper.make_graph( + [sqrt_node, consumer], + "sqrt_graph", + [input_info], + [output_info], + ) + ) + + converted_model = ( + MXFP8QuantExporter.post_process(model) + if high_precision_dtype is None + else MXFP8QuantExporter.post_process(model, high_precision_dtype) + ) + + cast_node = next(node for node in converted_model.graph.node if node.op_type == "Cast") + cast_to = next(attr.i for attr in cast_node.attribute if attr.name == "to") + assert cast_to == onnx_dtype + assert cast_node.name == f"sqrt_cast_{suffix}" + assert cast_node.input == ["sqrt_output"] + assert cast_node.output == [f"sqrt_output_cast_{suffix}"] + converted_consumer = next( + node for node in converted_model.graph.node if node.name == "consumer" + ) + assert converted_consumer.input == cast_node.output + onnx.checker.check_model(converted_model) + if high_precision_dtype is not None: + onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) + def test_mxfp8_gelu_approximation_update(self): """Test that Gelu nodes are updated to use tanh approximation.""" model = create_test_model_with_mxfp8_dq() @@ -701,6 +800,159 @@ def test_fp4qdq_conversion(self, with_transpose): cast_nodes = [node for node in converted_model.graph.node if node.op_type == "Cast"] assert len(cast_nodes) >= 1 # At least one cast node should be added + @pytest.mark.parametrize( + ("precision_dtype", "onnx_dtype"), + [("Half", TensorProto.FLOAT16), ("BFloat16", TensorProto.BFLOAT16)], + ) + def test_existing_weight_cast_does_not_hide_matmul(self, precision_dtype, onnx_dtype): + weight_data = np.linspace(-1.0, 1.0, 8 * 32, dtype=np.float32).reshape(8, 32) + weight = numpy_helper.from_array(weight_data, "linear.weight") + fp4qdq = helper.make_node( + "TRT_FP4QDQ", + inputs=[weight.name], + outputs=["fp4qdq_output"], + name="weight_fp4qdq", + block_size=16, + ) + cast = helper.make_node( + "Cast", + inputs=["fp4qdq_output"], + outputs=["weight_cast"], + name="weight_cast", + to=onnx_dtype, + ) + transpose = helper.make_node( + "Transpose", + inputs=["weight_cast"], + outputs=["weight_transposed"], + name="weight_transpose", + perm=[1, 0], + ) + matmul = helper.make_node( + "MatMul", + inputs=["activation", "weight_transposed"], + outputs=["output"], + name="matmul", + ) + graph = helper.make_graph( + [fp4qdq, cast, transpose, matmul], + "nvfp4_cast_graph", + [helper.make_tensor_value_info("activation", TensorProto.FLOAT, [1, 32])], + [helper.make_tensor_value_info("output", onnx_dtype, [1, 8])], + [weight], + value_info=[ + helper.make_tensor_value_info("fp4qdq_output", TensorProto.FLOAT, [8, 32]), + helper.make_tensor_value_info("weight_cast", onnx_dtype, [8, 32]), + helper.make_tensor_value_info("weight_transposed", onnx_dtype, [32, 8]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 23)]) + + converted_model = NVFP4QuantExporter.process_model(model, precision_dtype) + + onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) + converted_matmul = next( + node for node in converted_model.graph.node if node.op_type == "MatMul" + ) + producer_map = { + output: node for node in converted_model.graph.node for output in node.output + } + assert all( + producer_map[input_name].op_type == "Cast" + and helper.get_attribute_value(producer_map[input_name].attribute[0]) == onnx_dtype + for input_name in converted_matmul.input + ) + + def test_shared_weight_casts_all_direct_and_wrapped_linear_consumers(self): + weight_data = np.linspace(-1.0, 1.0, 32 * 32, dtype=np.float32).reshape(32, 32) + weight = numpy_helper.from_array(weight_data, "linear.weight") + fp4qdq = helper.make_node( + "TRT_FP4QDQ", + inputs=[weight.name], + outputs=["fp4qdq_output"], + name="weight_fp4qdq", + block_size=16, + ) + cast = helper.make_node( + "Cast", + inputs=["fp4qdq_output"], + outputs=["weight_cast"], + name="weight_cast", + to=TensorProto.FLOAT16, + ) + transpose_nodes = [ + helper.make_node( + "Transpose", + inputs=["weight_cast"], + outputs=[f"weight_transposed_{index}"], + name=f"weight_transpose_{index}", + perm=[1, 0], + ) + for index in range(2) + ] + matmul_nodes = [ + *[ + helper.make_node( + "MatMul", + inputs=[f"activation_{index}", "fp4qdq_output"], + outputs=[f"output_{index}"], + name=f"matmul_{index}", + ) + for index in range(2) + ], + *[ + helper.make_node( + "MatMul", + inputs=[f"activation_{index}", f"weight_transposed_{index - 2}"], + outputs=[f"output_{index}"], + name=f"matmul_{index}", + ) + for index in range(2, 4) + ], + ] + graph = helper.make_graph( + [fp4qdq, cast, *transpose_nodes, *matmul_nodes], + "nvfp4_fanout_graph", + [ + helper.make_tensor_value_info(f"activation_{index}", TensorProto.FLOAT, [1, 32]) + for index in range(4) + ], + [ + helper.make_tensor_value_info(f"output_{index}", TensorProto.FLOAT16, [1, 32]) + for index in range(4) + ], + [weight], + value_info=[ + helper.make_tensor_value_info("fp4qdq_output", TensorProto.FLOAT, [32, 32]), + helper.make_tensor_value_info("weight_cast", TensorProto.FLOAT16, [32, 32]), + *[ + helper.make_tensor_value_info( + f"weight_transposed_{index}", TensorProto.FLOAT16, [32, 32] + ) + for index in range(2) + ], + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 23)]) + + converted_model = NVFP4QuantExporter.process_model(model, "Half") + + onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) + producer_map = { + output: node for node in converted_model.graph.node for output in node.output + } + converted_matmuls = [ + node for node in converted_model.graph.node if node.op_type == "MatMul" + ] + assert len(converted_matmuls) == 4 + assert all( + producer_map[input_name].op_type == "Cast" + and helper.get_attribute_value(producer_map[input_name].attribute[0]) + == TensorProto.FLOAT16 + for node in converted_matmuls + for input_name in node.input + ) + def create_test_model_with_int4_dq_matmul(): """Create a simple test model with INT4 DequantizeLinear -> MatMul pattern. @@ -1072,7 +1324,7 @@ def test_column_major_gemm_trans_b_flip(self): print(f"Transpose nodes: {len(transpose_nodes)}") -def _build_model_with_zero_scale_initializer(dq_op_type: str): +def _build_model_with_zero_scale_initializer(dq_op_type: str, scale_dtype=np.float16): """Build an ONNX model whose scale initializer feeds a (Quantize|Dequantize)Linear node. Mirrors the INT4_AWQ failure mode from NVBug 6110209: scales live in graph initializers @@ -1081,10 +1333,11 @@ def _build_model_with_zero_scale_initializer(dq_op_type: str): weight_data = np.random.randint(-8, 8, size=(6, 8), dtype=np.int8) weight_tensor = numpy_helper.from_array(weight_data, "weight") - scale_data = np.array([1e-3, 0.0, 5e-4, 0.0, 0.0, 2e-3], dtype=np.float16).reshape(6, 1) + scale_data = np.array([1e-3, 0.0, 5e-4, 0.0, 0.0, 2e-3], dtype=scale_dtype).reshape(6, 1) scale_tensor = numpy_helper.from_array(scale_data, "scale") - input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT16, [None, 6]) + high_precision_dtype = helper.np_dtype_to_tensor_dtype(np.dtype(scale_dtype)) + input_tensor = helper.make_tensor_value_info("input", high_precision_dtype, [None, 6]) dq_node = helper.make_node( dq_op_type, inputs=["weight", "scale"], outputs=["dq_output"], name="weight_dq" ) @@ -1095,7 +1348,7 @@ def _build_model_with_zero_scale_initializer(dq_op_type: str): nodes=[dq_node, matmul_node], name="test_graph", inputs=[input_tensor], - outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT16, [None, 8])], + outputs=[helper.make_tensor_value_info("output", high_precision_dtype, [None, 8])], initializer=[weight_tensor, scale_tensor], ) return helper.make_model(graph) @@ -1105,8 +1358,17 @@ class TestReplaceZeroScaleWithSmallestNonzero: """Regression tests for ``replace_zero_scale_with_smallest_nonzero`` (NVBug 6110209).""" @pytest.mark.parametrize("dq_op_type", ["DequantizeLinear", "TRT_INT4DequantizeLinear"]) - def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type): - model = _build_model_with_zero_scale_initializer(dq_op_type) + @pytest.mark.parametrize( + ("scale_dtype", "onnx_dtype"), + [ + (np.float16, TensorProto.FLOAT16), + (ml_dtypes.bfloat16, TensorProto.BFLOAT16), + (np.float32, TensorProto.FLOAT), + (np.float64, TensorProto.DOUBLE), + ], + ) + def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type, scale_dtype, onnx_dtype): + model = _build_model_with_zero_scale_initializer(dq_op_type, scale_dtype) scale_before = numpy_helper.to_array( next(init for init in model.graph.initializer if init.name == "scale") ) @@ -1118,7 +1380,14 @@ def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type): scale_after = numpy_helper.to_array(scale_after_init) assert not (scale_after == 0).any() assert (scale_after > 0).all() - assert scale_after_init.data_type == TensorProto.FLOAT16 + assert scale_after_init.data_type == onnx_dtype + + dtype_info = ( + ml_dtypes.finfo(scale_dtype) + if scale_dtype == ml_dtypes.bfloat16 + else np.finfo(scale_dtype) + ) + assert (scale_after[scale_before == 0] == dtype_info.smallest_subnormal).all() def test_constant_node_scale_path_still_patched(self): """Legacy Constant-node QDQ path must continue to be patched.""" @@ -1155,6 +1424,58 @@ def test_constant_node_scale_path_still_patched(self): assert not (scale_arr == 0).any() assert (scale_arr > 0).all() + def test_captured_parent_scale_is_patched_without_crossing_child_scope(self): + captured_scale = numpy_helper.from_array( + np.array(0.0, dtype=ml_dtypes.bfloat16), "captured_scale" + ) + shadowed_scale = numpy_helper.from_array( + np.array(0.0, dtype=ml_dtypes.bfloat16), "shadowed_scale" + ) + subgraph = helper.make_graph( + [ + helper.make_node( + "QuantizeLinear", + ["data", "captured_scale"], + ["captured_output"], + ), + helper.make_node( + "QuantizeLinear", + ["data", "shadowed_scale"], + ["shadowed_output"], + ), + ], + "subgraph", + [ + helper.make_tensor_value_info("data", TensorProto.BFLOAT16, [1]), + helper.make_tensor_value_info("shadowed_scale", TensorProto.BFLOAT16, []), + ], + [ + helper.make_tensor_value_info("captured_output", TensorProto.UINT8, [1]), + helper.make_tensor_value_info("shadowed_output", TensorProto.UINT8, [1]), + ], + ) + scoped_node = helper.make_node( + "ScopedSubgraph", + [], + [], + domain="test", + body=subgraph, + ) + graph = helper.make_graph( + [scoped_node], + "parent_graph", + [], + [], + [captured_scale, shadowed_scale], + ) + model = helper.make_model(graph) + + patched = replace_zero_scale_with_smallest_nonzero(model) + + scales = {init.name: numpy_helper.to_array(init) for init in patched.graph.initializer} + assert scales["captured_scale"] == ml_dtypes.finfo(ml_dtypes.bfloat16).smallest_subnormal + assert scales["shadowed_scale"] == 0 + class TestQdqToDqValidation: """Regression tests for qdq_to_dq input validation.""" 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 2b50c8f7c40..60073c65b95 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -24,7 +24,11 @@ import torch.nn as nn from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models +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.autocast.convert import convert_to_fp32 +from modelopt.onnx.export.base_exporter import ONNXQuantExporter from modelopt.onnx.utils import get_batch_size_from_bytes, validate_batch_size from modelopt.torch._deploy.utils import ( OnnxBytes, @@ -32,7 +36,12 @@ generate_onnx_input, get_onnx_bytes_and_metadata, ) -from modelopt.torch._deploy.utils.torch_onnx import _to_expected_onnx_type +from modelopt.torch._deploy.utils.torch_onnx import ( + _get_autocast_context, + _override_onnx_quantizer_precision, + _to_expected_onnx_type, +) +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import standardize_model_args, unflatten_tree deploy_benchmark_all = get_deploy_models() @@ -57,6 +66,22 @@ k: v for k, v in deploy_benchmark_dynamo.items() if k in _DYNAMO_REPRESENTATIVE_MODELS } +_ONNX_DTYPE_BY_NAME = { + "fp32": onnx.TensorProto.FLOAT, + "fp16": onnx.TensorProto.FLOAT16, + "bf16": onnx.TensorProto.BFLOAT16, +} +_QUANTIZED_LINEAR_CASES = { + "int4": (mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, 128), + "mxfp8": (mtq.MXFP8_DEFAULT_CFG, 32), + "nvfp4": (mtq.NVFP4_DEFAULT_CFG, 32), + "int8": (mtq.INT8_DEFAULT_CFG, 32), +} +_NATIVE_QUANTIZED_LINEAR_CASES = { + "fp8": (mtq.FP8_DEFAULT_CFG, 32), + **_QUANTIZED_LINEAR_CASES, +} + def _export_fp8_linear(source_dtype, weights_dtype): model = nn.Sequential(nn.Linear(4, 4, bias=False)).eval().to(source_dtype) @@ -78,6 +103,79 @@ def _export_fp8_linear(source_dtype, weights_dtype): return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) +def _export_model(model, sample_input, weights_dtype, onnx_opset=20): + weights_dtype_kwargs = {} if weights_dtype is None else {"weights_dtype": weights_dtype} + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model.eval(), + (sample_input,), + onnx_opset=onnx_opset, + **weights_dtype_kwargs, + ) + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) + + +def _tensor_dtype_map(model): + dtype_map = { + value.name: value.type.tensor_type.elem_type + for value in (*model.graph.input, *model.graph.output, *model.graph.value_info) + if value.type.HasField("tensor_type") + } + dtype_map.update( + {initializer.name: initializer.data_type for initializer in model.graph.initializer} + ) + for node in model.graph.node: + if node.op_type != "Constant" or not node.output: + continue + value = next( + (attribute.t for attribute in node.attribute if attribute.name == "value"), None + ) + if value is not None: + dtype_map[node.output[0]] = value.data_type + return dtype_map + + +def _assert_runtime_io_dtype(model, expected_dtype): + initializer_names = {initializer.name for initializer in model.graph.initializer} + runtime_io = [ + *[value for value in model.graph.input if value.name not in initializer_names], + *model.graph.output, + ] + assert runtime_io + assert all(value.type.tensor_type.elem_type == expected_dtype for value in runtime_io) + + +def test_autocast_prefers_floating_input_device(monkeypatch): + autocast_args = {} + + def capture_autocast(*, device_type, dtype): + autocast_args.update(device_type=device_type, dtype=dtype) + return nullcontext() + + monkeypatch.setattr(torch, "autocast", capture_autocast) + flat_input = [ + torch.ones(1, dtype=torch.int64, device="meta"), + torch.ones(1, dtype=torch.float32), + ] + + _get_autocast_context(nn.Identity(), flat_input, torch.bfloat16) + + assert autocast_args == {"device_type": "cpu", "dtype": torch.bfloat16} + + +def test_quant_exporter_preserves_legacy_post_process_signature(): + class LegacyExporter(ONNXQuantExporter): + pre_process = compute_scales = compress_weights = staticmethod(lambda model: model) + + @staticmethod + def post_process(model): + return model + + model = onnx.helper.make_model(onnx.helper.make_graph([], "graph", [], [])) + + assert LegacyExporter.process_model(model) is model + + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) @@ -179,23 +277,48 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): @pytest.mark.parametrize( - ("source_dtype", "weights_dtype", "expected_onnx_dtype"), + ("source_dtype", "weights_dtype", "expected_scale_dtype", "expected_io_dtypes"), [ pytest.param( torch.bfloat16, - "bf16", - onnx.TensorProto.BFLOAT16, - id="bf16-no-op", + "native", + onnx.TensorProto.FLOAT, + (onnx.TensorProto.BFLOAT16, onnx.TensorProto.FLOAT), + id="native-bf16", ), pytest.param( torch.float32, "fp16", onnx.TensorProto.FLOAT16, - id="fp16-conversion", + (onnx.TensorProto.FLOAT16, onnx.TensorProto.FLOAT16), + id="fp32-to-fp16", + ), + pytest.param( + torch.float32, + "bf16", + onnx.TensorProto.BFLOAT16, + (onnx.TensorProto.BFLOAT16, onnx.TensorProto.BFLOAT16), + id="fp32-to-bf16", + ), + pytest.param( + torch.bfloat16, + "bf16", + onnx.TensorProto.BFLOAT16, + (onnx.TensorProto.BFLOAT16, onnx.TensorProto.BFLOAT16), + id="bf16-to-bf16", + ), + pytest.param( + torch.bfloat16, + "fp32", + onnx.TensorProto.FLOAT, + (onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT), + id="bf16-to-fp32", ), ], ) -def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, expected_onnx_dtype): +def test_fp8_export_with_supported_weights_dtype( + source_dtype, weights_dtype, expected_scale_dtype, expected_io_dtypes +): exported_model = _export_fp8_linear(source_dtype, weights_dtype) onnx.checker.check_model(exported_model) @@ -215,19 +338,627 @@ def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, ex ] assert fp8_weight_dq_nodes assert all( - initializer_by_name[node.input[1]].data_type == expected_onnx_dtype + initializer_by_name[node.input[1]].data_type == expected_scale_dtype for node in fp8_weight_dq_nodes ) + graph_io = [*exported_model.graph.input, *exported_model.graph.output] + assert tuple(value.type.tensor_type.elem_type for value in graph_io) == expected_io_dtypes + + +@pytest.mark.parametrize("format_name", _QUANTIZED_LINEAR_CASES) +@pytest.mark.parametrize("weights_dtype", _ONNX_DTYPE_BY_NAME) +def test_quantized_linear_export_uses_requested_weights_dtype( + monkeypatch, format_name, weights_dtype +): + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", lambda inputs, *args: inputs) + quantization_config, features = _QUANTIZED_LINEAR_CASES[format_name] + source_dtype = torch.bfloat16 if weights_dtype == "fp32" else torch.float32 + model = nn.Sequential(nn.Linear(features, 8, bias=False)).eval().to(source_dtype) + sample_input = torch.ones(1, features, dtype=source_dtype) + model = mtq.quantize( + model, + quantization_config, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + exported_model = _export_model(model, sample_input, weights_dtype, onnx_opset=23) + expected_dtype = _ONNX_DTYPE_BY_NAME[weights_dtype] + + onnx.checker.check_model(exported_model) + inferred_model = onnx.shape_inference.infer_shapes( + exported_model, check_type=True, strict_mode=True + ) + _assert_runtime_io_dtype(exported_model, expected_dtype) + dtype_map = _tensor_dtype_map(inferred_model) + matmul = next(node for node in inferred_model.graph.node if node.op_type == "MatMul") + assert all(dtype_map[input_name] == expected_dtype for input_name in matmul.input) + + initializer_map = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + cast_nodes = [node for node in exported_model.graph.node if node.op_type == "Cast"] + + if format_name == "int4": + weight = next( + initializer + for initializer in initializer_map.values() + if initializer.data_type == onnx.TensorProto.INT4 + ) + weight_dq = next( + node + for node in exported_model.graph.node + if node.op_type == "DequantizeLinear" and node.input[0] == weight.name + ) + assert initializer_map[weight_dq.input[1]].data_type == expected_dtype + assert dtype_map[weight_dq.output[0]] == expected_dtype + assert not cast_nodes + elif format_name == "mxfp8": + initializer_dtypes = {initializer.data_type for initializer in initializer_map.values()} + assert onnx.TensorProto.FLOAT8E4M3FN in initializer_dtypes + assert onnx.TensorProto.UINT8 in initializer_dtypes + dq_nodes = [ + node + for node in exported_model.graph.node + if node.op_type == "TRT_MXFP8DequantizeLinear" + ] + assert dq_nodes + assert all( + next(attribute.i for attribute in node.attribute if attribute.name == "output_dtype") + == expected_dtype + for node in dq_nodes + ) + assert not cast_nodes + elif format_name == "nvfp4": + initializer_dtypes = {initializer.data_type for initializer in initializer_map.values()} + assert { + onnx.TensorProto.FLOAT4E2M1, + onnx.TensorProto.FLOAT8E4M3FN, + onnx.TensorProto.FLOAT, + } <= initializer_dtypes + cast_dtypes = { + next(attribute.i for attribute in node.attribute if attribute.name == "to") + for node in cast_nodes + } + assert cast_dtypes == ({expected_dtype} if weights_dtype != "fp32" else set()) + else: + q_nodes = [node for node in inferred_model.graph.node if node.op_type == "QuantizeLinear"] + dq_nodes = [ + node for node in inferred_model.graph.node if node.op_type == "DequantizeLinear" + ] + assert q_nodes and dq_nodes + assert all(dtype_map[node.output[0]] == onnx.TensorProto.INT8 for node in q_nodes) + assert all(dtype_map[node.input[1]] == expected_dtype for node in dq_nodes) + assert all(dtype_map[node.output[0]] == expected_dtype for node in dq_nodes) + assert not cast_nodes + + +@pytest.mark.parametrize("format_name", _NATIVE_QUANTIZED_LINEAR_CASES) +def test_quantized_linear_default_preserves_native_behavior(monkeypatch, format_name): + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", lambda inputs, *args: inputs) + quantization_config, features = _NATIVE_QUANTIZED_LINEAR_CASES[format_name] + model = nn.Sequential(nn.Linear(features, 8, bias=False)).eval().to(torch.bfloat16) + sample_input = torch.ones(1, features, dtype=torch.bfloat16) + model = mtq.quantize( + model, + quantization_config, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + + exported_model = _export_model(model, sample_input, None, onnx_opset=23) + onnx.checker.check_model(exported_model) + inferred_model = onnx.shape_inference.infer_shapes( + exported_model, check_type=True, strict_mode=True + ) + initializer_names = {initializer.name for initializer in exported_model.graph.initializer} + runtime_inputs = [ + value for value in exported_model.graph.input if value.name not in initializer_names + ] + assert runtime_inputs + assert all( + value.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 for value in runtime_inputs + ) + expected_boundary_dtype = ( + onnx.TensorProto.BFLOAT16 if format_name == "nvfp4" else onnx.TensorProto.FLOAT + ) + assert all( + value.type.tensor_type.elem_type == expected_boundary_dtype + for value in exported_model.graph.output + ) + dtype_map = _tensor_dtype_map(inferred_model) + matmul = next(node for node in inferred_model.graph.node if node.op_type == "MatMul") + assert all(dtype_map[input_name] == expected_boundary_dtype for input_name in matmul.input) + + initializer_map = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + initializer_dtypes = {initializer.data_type for initializer in initializer_map.values()} + if format_name == "fp8": + activation_q_nodes = [ + node for node in inferred_model.graph.node if node.op_type == "QuantizeLinear" + ] + cast_by_output = { + node.output[0]: node for node in inferred_model.graph.node if node.op_type == "Cast" + } + assert activation_q_nodes + assert all( + dtype_map[node.input[0]] == onnx.TensorProto.FLOAT + and dtype_map[node.input[1]] == onnx.TensorProto.FLOAT + for node in activation_q_nodes + ) + assert all(node.input[0] in cast_by_output for node in activation_q_nodes) + assert all( + next( + attribute.i + for attribute in cast_by_output[node.input[0]].attribute + if attribute.name == "to" + ) + == onnx.TensorProto.FLOAT + for node in activation_q_nodes + ) + weight_dq_nodes = [ + node + for node in inferred_model.graph.node + if node.op_type == "DequantizeLinear" + and node.input[0] in initializer_map + and initializer_map[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN + ] + assert weight_dq_nodes + dq_nodes = [ + node for node in inferred_model.graph.node if node.op_type == "DequantizeLinear" + ] + assert all( + dtype_map[node.input[1]] == onnx.TensorProto.FLOAT + and dtype_map[node.output[0]] == onnx.TensorProto.FLOAT + for node in dq_nodes + ) + elif format_name == "int4": + assert onnx.TensorProto.INT4 in initializer_dtypes + weight_dq = next( + node + for node in inferred_model.graph.node + if node.op_type == "DequantizeLinear" + and node.input[0] in initializer_map + and initializer_map[node.input[0]].data_type == onnx.TensorProto.INT4 + ) + assert dtype_map[weight_dq.input[1]] == onnx.TensorProto.FLOAT + elif format_name == "mxfp8": + assert {onnx.TensorProto.FLOAT8E4M3FN, onnx.TensorProto.UINT8} <= initializer_dtypes + dq_nodes = [ + node + for node in inferred_model.graph.node + if node.op_type == "TRT_MXFP8DequantizeLinear" + ] + assert dq_nodes + assert all( + next(attribute.i for attribute in node.attribute if attribute.name == "output_dtype") + == onnx.TensorProto.FLOAT16 + for node in dq_nodes + ) + elif format_name == "nvfp4": + assert { + onnx.TensorProto.FLOAT4E2M1, + onnx.TensorProto.FLOAT8E4M3FN, + onnx.TensorProto.FLOAT, + } <= initializer_dtypes + else: + q_nodes = [node for node in inferred_model.graph.node if node.op_type == "QuantizeLinear"] + dq_nodes = [ + node for node in inferred_model.graph.node if node.op_type == "DequantizeLinear" + ] + assert q_nodes and dq_nodes + assert all(dtype_map[node.output[0]] == onnx.TensorProto.INT8 for node in q_nodes) + assert all(dtype_map[node.input[1]] == onnx.TensorProto.FLOAT for node in dq_nodes) + + +@pytest.mark.parametrize("weights_dtype", _ONNX_DTYPE_BY_NAME) +def test_fp8_conv_export_uses_requested_weights_dtype(weights_dtype): + source_dtype = torch.bfloat16 if weights_dtype == "fp32" else torch.float32 + model = nn.Conv2d(3, 4, kernel_size=3, bias=False).eval().to(source_dtype) + sample_input = torch.ones(1, 3, 8, 8, dtype=source_dtype) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + exported_model = _export_model(model, sample_input, weights_dtype, onnx_opset=23) + expected_dtype = _ONNX_DTYPE_BY_NAME[weights_dtype] + + onnx.checker.check_model(exported_model) + inferred_model = onnx.shape_inference.infer_shapes( + exported_model, check_type=True, strict_mode=True + ) + _assert_runtime_io_dtype(exported_model, expected_dtype) + assert not any( + node.op_type in {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} + for node in exported_model.graph.node + ) + + initializer_map = { + initializer.name: initializer for initializer in exported_model.graph.initializer + } + conv = next(node for node in exported_model.graph.node if node.op_type == "Conv") + weight_dq = next( + node + for node in exported_model.graph.node + if node.op_type == "DequantizeLinear" and node.output[0] == conv.input[1] + ) + weight = initializer_map[weight_dq.input[0]] + assert weight.data_type == onnx.TensorProto.FLOAT8E4M3FN + assert initializer_map[weight_dq.input[1]].data_type == expected_dtype + assert _tensor_dtype_map(inferred_model)[weight_dq.output[0]] == expected_dtype + assert not any(node.op_type == "Cast" for node in exported_model.graph.node) + + +class MixedPrecisionLinear(nn.Module): + def __init__(self): + super().__init__() + self.fp32_weight = nn.Parameter(torch.eye(4, dtype=torch.float32)) + self.bf16_weight = nn.Parameter(torch.eye(4, dtype=torch.bfloat16)) + + def forward(self, x): + return torch.matmul(x, self.fp32_weight) + torch.matmul(x, self.bf16_weight) + + +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "expected_onnx_dtype"), + [ + (torch.float32, "bf16", onnx.TensorProto.BFLOAT16), + (torch.bfloat16, "fp32", onnx.TensorProto.FLOAT), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16), + ], +) +def test_parameterless_model_uses_explicit_weights_dtype( + source_dtype, weights_dtype, expected_onnx_dtype +): + exported_model = _export_model( + nn.Identity(), torch.ones(1, 4, dtype=source_dtype), weights_dtype + ) + graph_io = [*exported_model.graph.input, *exported_model.graph.output] assert all(value.type.tensor_type.elem_type == expected_onnx_dtype for value in graph_io) -def test_fp8_export_rejects_bf16_conversion_from_fp32(): - with pytest.raises( - AssertionError, - match="Converting a quantized ONNX graph to BF16 is not supported", +def test_mixed_parameter_model_uses_explicit_weights_dtype(): + model = MixedPrecisionLinear() + exported_model = _export_model(model, torch.ones(1, 4), "bf16") + + assert model.fp32_weight.dtype == torch.float32 + assert model.bf16_weight.dtype == torch.bfloat16 + assert exported_model.graph.initializer + assert all( + initializer.data_type == onnx.TensorProto.BFLOAT16 + for initializer in exported_model.graph.initializer + ) + + +def test_onnx_quantizer_precision_is_restored_after_failure(): + quantizers = nn.ModuleList([TensorQuantizer(), TensorQuantizer()]) + quantizers[0].trt_high_precision_dtype = "Float" + del quantizers[1]._trt_high_precision_dtype + + with _override_onnx_quantizer_precision(quantizers, None): + assert quantizers[0].trt_high_precision_dtype == "Float" + assert not hasattr(quantizers[1], "_trt_high_precision_dtype") + assert quantizers[0].trt_high_precision_dtype == "Float" + assert not hasattr(quantizers[1], "_trt_high_precision_dtype") + + with _override_onnx_quantizer_precision(quantizers, "Half"): + assert all(q.trt_high_precision_dtype == "Half" for q in quantizers) + assert quantizers[0].trt_high_precision_dtype == "Float" + assert not hasattr(quantizers[1], "_trt_high_precision_dtype") + + with ( + pytest.raises(RuntimeError, match="export failed"), + _override_onnx_quantizer_precision(quantizers, "BFloat16"), ): - _export_fp8_linear(torch.float32, "bf16") + assert all(q.trt_high_precision_dtype == "BFloat16" for q in quantizers) + raise RuntimeError("export failed") + + assert quantizers[0].trt_high_precision_dtype == "Float" + assert not hasattr(quantizers[1], "_trt_high_precision_dtype") + + +def _identity_onnx_model(dtype=onnx.TensorProto.FLOAT): + graph = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["input"], ["output"])], + "identity", + [onnx.helper.make_tensor_value_info("input", dtype, [1, 4])], + [onnx.helper.make_tensor_value_info("output", dtype, [1, 4])], + ) + return onnx.helper.make_model(graph) + + +@pytest.mark.parametrize("weights_dtype", ["fp32", "fp16", "bf16"]) +def test_onnx_load_path_rejects_non_native_weights_dtype(tmp_path, weights_dtype): + onnx_path = tmp_path / "identity.onnx" + onnx.save(_identity_onnx_model(), onnx_path) + + with pytest.raises(ValueError, match="weights_dtype must be 'native'"): + get_onnx_bytes_and_metadata( + nn.Identity(), + (torch.ones(1, 4),), + onnx_load_path=str(onnx_path), + weights_dtype=weights_dtype, + ) + + +def test_onnx_load_path_preserves_native_model(tmp_path): + onnx_path = tmp_path / "identity.onnx" + onnx.save(_identity_onnx_model(), onnx_path) + + onnx_bytes, _ = get_onnx_bytes_and_metadata( + nn.Identity(), + (torch.ones(1, 4),), + onnx_load_path=str(onnx_path), + weights_dtype="native", + ) + + loaded = OnnxBytes.from_bytes(onnx_bytes) + assert onnx.load_model_from_string(loaded.get_onnx_model_file_bytes()) + + +def _make_bf16_tensor(name, values): + values = np.asarray(values, dtype=np.float32) + tensor = onnx.TensorProto(name=name, data_type=onnx.TensorProto.BFLOAT16) + tensor.dims.extend(values.shape) + tensor.raw_data = (values.view(np.uint32) >> 16).astype(np.uint16).tobytes() + return tensor + + +def test_convert_to_fp32_recurses_through_graphs_functions_and_attributes(): + branch_value = _make_bf16_tensor("branch_value", [2.0]) + weight = _make_bf16_tensor("weight", [1.0]) + weight.doc_string = "weight metadata" + weight.metadata_props.add(key="source", value="test") + branch = onnx.helper.make_graph( + [onnx.helper.make_node("Constant", [], ["branch_output"], value=branch_value)], + "branch", + [], + [onnx.helper.make_tensor_value_info("branch_output", onnx.TensorProto.BFLOAT16, [1])], + ) + custom_node = onnx.helper.make_node( + "CustomOp", + ["input"], + ["custom_output"], + domain="test", + dtype=onnx.TensorProto.DOUBLE, + output_dtype=onnx.TensorProto.BFLOAT16, + ) + custom_node.attribute.append( + onnx.helper.make_attribute( + "type", onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT16, [1]) + ) + ) + graph = onnx.helper.make_graph( + [ + onnx.helper.make_node( + "If", ["condition"], ["output"], then_branch=branch, else_branch=branch + ), + onnx.helper.make_node("Cast", ["input"], ["cast_output"], to=onnx.TensorProto.FLOAT16), + custom_node, + ], + "recursive", + [onnx.helper.make_tensor_value_info("input", onnx.TensorProto.BFLOAT16, [1])], + [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.DOUBLE, [1])], + initializer=[ + onnx.numpy_helper.from_array(np.array(True), "condition"), + weight, + onnx.numpy_helper.from_array(np.array([3.0], dtype=np.float64), "double_weight"), + ], + value_info=[ + onnx.helper.make_tensor_value_info("custom_output", onnx.TensorProto.FLOAT16, [1]) + ], + ) + function = onnx.helper.make_function( + "test", + "LocalCast", + ["x"], + ["y"], + [onnx.helper.make_node("Cast", ["x"], ["y"], to=onnx.TensorProto.BFLOAT16)], + [onnx.helper.make_opsetid("", 20)], + value_info=[onnx.helper.make_tensor_value_info("y", onnx.TensorProto.BFLOAT16, [1])], + ) + model = onnx.helper.make_model(graph, functions=[function]) + + assert convert_to_fp32(model) is model + + graph = model.graph + function = model.functions[0] + custom_node = next(node for node in graph.node if node.op_type == "CustomOp") + assert all( + value.type.tensor_type.elem_type == onnx.TensorProto.FLOAT + for value in (*graph.input, *graph.output, *graph.value_info) + ) + initializer_map = {initializer.name: initializer for initializer in graph.initializer} + assert initializer_map["weight"].data_type == onnx.TensorProto.FLOAT + assert initializer_map["double_weight"].data_type == onnx.TensorProto.FLOAT + assert initializer_map["weight"].doc_string == "weight metadata" + assert initializer_map["weight"].metadata_props[0].key == "source" + np.testing.assert_array_equal( + np.frombuffer(initializer_map["weight"].raw_data, dtype=np.float32), + np.array([1.0], dtype=np.float32), + ) + attribute_map = {attribute.name: attribute for attribute in custom_node.attribute} + assert attribute_map["dtype"].i == onnx.TensorProto.DOUBLE + assert attribute_map["output_dtype"].i == onnx.TensorProto.BFLOAT16 + assert attribute_map["type"].tp.tensor_type.elem_type == onnx.TensorProto.FLOAT + for attribute in graph.node[0].attribute: + branch_graph = attribute.g + assert branch_graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT + assert branch_graph.node[0].attribute[0].t.data_type == onnx.TensorProto.FLOAT + assert function.value_info[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT + assert function.node[0].attribute[0].i == onnx.TensorProto.FLOAT + + +def test_convert_to_fp32_handles_vetted_dtype_attributes(): + nodes = [ + onnx.helper.make_node("HannWindow", [], [], output_datatype=onnx.TensorProto.FLOAT16), + onnx.helper.make_node("LayerNormalization", [], [], stash_type=onnx.TensorProto.BFLOAT16), + onnx.helper.make_node("Attention", [], [], softmax_precision=onnx.TensorProto.DOUBLE), + onnx.helper.make_node("QuantizeLinear", [], [], precision=onnx.TensorProto.FLOAT16), + onnx.helper.make_node("Cast", [], [], to=onnx.TensorProto.BFLOAT16), + onnx.helper.make_node( + "TRTCustom", [], [], domain="trt", output_dtype=onnx.TensorProto.BFLOAT16 + ), + onnx.helper.make_node( + "CustomOp", + [], + [], + domain="test", + dtype=onnx.TensorProto.BFLOAT16, + precision=onnx.TensorProto.BFLOAT16, + to=onnx.TensorProto.BFLOAT16, + ), + ] + model = onnx.helper.make_model(onnx.helper.make_graph(nodes, "dtype_attributes", [], [])) + + convert_to_fp32(model) + + for node in model.graph.node[:6]: + assert node.attribute[0].i == onnx.TensorProto.FLOAT + assert all(attr.i == onnx.TensorProto.BFLOAT16 for attr in model.graph.node[6].attribute) + + +def _make_function_with_referenced_to(op_type): + to_attribute = onnx.AttributeProto( + name="to", + ref_attr_name="target_dtype", + type=onnx.AttributeProto.INT, + ) + node = onnx.helper.make_node(op_type, ["x"], ["y"]) + node.attribute.append(to_attribute) + return onnx.helper.make_function( + "test", + f"Referenced{op_type}", + ["x"], + ["y"], + [node], + [onnx.helper.make_opsetid("", 26)], + attribute_protos=[onnx.helper.make_attribute("target_dtype", onnx.TensorProto.BFLOAT16)], + ) + + +def test_convert_to_fp32_converts_function_cast_attribute_default(): + function = _make_function_with_referenced_to("Cast") + model = onnx.helper.make_model( + onnx.helper.make_graph([], "function_cast", [], []), functions=[function] + ) + + convert_to_fp32(model) + + assert model.functions[0].attribute_proto[0].i == onnx.TensorProto.FLOAT + + +def test_convert_to_fp32_rejects_function_bitcast_attribute_default(): + function = _make_function_with_referenced_to("BitCast") + model = onnx.helper.make_model( + onnx.helper.make_graph([], "function_bitcast", [], []), functions=[function] + ) + + with pytest.raises(ValueError, match="BitCast targets cannot be converted safely"): + convert_to_fp32(model) + + +def test_convert_to_fp32_rejects_low_precision_bitcast(): + bitcast = onnx.helper.make_node("BitCast", ["input"], ["output"], to=onnx.TensorProto.FLOAT16) + graph = onnx.helper.make_graph( + [bitcast], + "bitcast", + [onnx.helper.make_tensor_value_info("input", onnx.TensorProto.UINT16, [1])], + [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT16, [1])], + ) + + with pytest.raises(ValueError, match="BitCast targets cannot be converted safely"): + convert_to_fp32(onnx.helper.make_model(graph)) + + +def test_convert_to_fp32_rejects_segmented_tensor(): + weight = onnx.numpy_helper.from_array(np.array([1.0], dtype=np.float16), "weight") + weight.segment.begin = 0 + weight.segment.end = 1 + graph = onnx.helper.make_graph( + [], + "segmented", + [], + [], + initializer=[weight], + ) + + with pytest.raises(ValueError, match="Segmented tensors are not supported"): + convert_to_fp32(onnx.helper.make_model(graph)) + + +def test_convert_to_fp32_handles_loaded_external_data(tmp_path): + tensor = onnx.numpy_helper.from_array(np.array([1.0, -2.0], dtype=np.float16), "weight") + graph = onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["weight"], ["output"])], + "external", + [], + [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT16, [2])], + [tensor], + ) + onnx_path = tmp_path / "external.onnx" + onnx.save_model( + onnx.helper.make_model(graph), + onnx_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="weights.data", + size_threshold=0, + ) + + unloaded_model = onnx.load(onnx_path, load_external_data=False) + with pytest.raises(ValueError, match="External tensor data must be loaded"): + convert_to_fp32(unloaded_model) + + loaded_model = onnx.load(onnx_path, load_external_data=True) + convert_to_fp32(loaded_model) + converted_weight = loaded_model.graph.initializer[0] + assert converted_weight.data_type == onnx.TensorProto.FLOAT + assert converted_weight.data_location == onnx.TensorProto.DEFAULT + assert not converted_weight.external_data + np.testing.assert_array_equal( + onnx.numpy_helper.to_array(converted_weight), np.array([1.0, -2.0], dtype=np.float32) + ) + + +def test_save_onnx_model_externalizes_large_attribute_and_replaces_shards(tmp_path, monkeypatch): + values = np.arange(300, dtype=np.float32) + constant = onnx.helper.make_node( + "Constant", + [], + ["output"], + value=onnx.numpy_helper.from_array(values), + ) + graph = onnx.helper.make_graph( + [constant], + "external_attribute", + [], + [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [300])], + ) + model = onnx.helper.make_model(graph) + onnx_path = tmp_path / "model.onnx" + onnx.save_model(model, onnx_path) + external_data_path = tmp_path / "model.onnx_data" + external_data_path.write_bytes(b"stale") + previous_shard = tmp_path / "previous.data" + previous_shard.write_bytes(b"old") + monkeypatch.setattr(torch_onnx, "TWO_GB", 1) + + torch_onnx._save_onnx_model(model, str(onnx_path), "model") + + assert external_data_path.stat().st_size == values.nbytes + assert not previous_shard.exists() + unloaded_model = onnx.load(onnx_path, load_external_data=False) + value = next(attr.t for attr in unloaded_model.graph.node[0].attribute if attr.name == "value") + assert onnx.external_data_helper.uses_external_data(value) + assert next(prop.value for prop in value.external_data if prop.key == "location") == ( + "model.onnx_data" + ) + loaded_model = onnx.load(onnx_path, load_external_data=True) + loaded_value = next( + attr.t for attr in loaded_model.graph.node[0].attribute if attr.name == "value" + ) + np.testing.assert_array_equal(onnx.numpy_helper.to_array(loaded_value), values) class SingleArgModel(nn.Module): From 51c3af4f4ea37b58561e235fea75236852c08007 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:26:00 +0000 Subject: [PATCH 3/8] [6508436] Handle GraphSurgeon integer dtypes Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/quantization/gs_patching.py | 14 ++++++++++-- .../onnx/quantization/test_gs_patching.py | 22 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index bbd0dadca7c..c8525db36bf 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -65,6 +65,14 @@ def _make_variable( return x +def _to_onnx_dtype(dtype) -> int: + if isinstance(dtype, (int, np.integer)): + dtype = int(dtype) + onnx.TensorProto.DataType.Name(dtype) + return dtype + return onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) + + def _export_tensor_proto(tensor: gs.Constant) -> onnx.TensorProto: if isinstance(tensor._values, LazyValues): onnx_tensor = tensor._values.tensor @@ -72,7 +80,8 @@ def _export_tensor_proto(tensor: gs.Constant) -> onnx.TensorProto: # is numpy array. dtype = getattr(tensor, "explicit_dtype", None) if dtype is None: - dtype = onnx.helper.np_dtype_to_tensor_dtype(tensor.values.dtype) + dtype = tensor.values.dtype + dtype = _to_onnx_dtype(dtype) vals = tensor.values if _onnx_supports_int4() and dtype in [onnx.TensorProto.INT4, onnx.TensorProto.UINT4]: @@ -103,7 +112,8 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V if tensor.dtype is not None: dtype = getattr(tensor, "explicit_dtype", None) if dtype is None: - dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) + dtype = tensor.dtype + dtype = _to_onnx_dtype(dtype) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) else: onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) diff --git a/tests/unit/onnx/quantization/test_gs_patching.py b/tests/unit/onnx/quantization/test_gs_patching.py index 17df218dfe9..95adbde79a0 100644 --- a/tests/unit/onnx/quantization/test_gs_patching.py +++ b/tests/unit/onnx/quantization/test_gs_patching.py @@ -17,6 +17,7 @@ import numpy as np import onnx import onnx_graphsurgeon as gs +import pytest from modelopt.onnx.quantization.gs_patching import _export_tensor_proto, _export_value_info_proto @@ -43,3 +44,24 @@ def test_export_value_info_uses_explicit_onnx_dtype_without_numpy_conversion(): value_info = _export_value_info_proto(tensor, do_type_check=True) assert value_info.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 + + +def test_export_value_info_accepts_onnx_dtype_without_explicit_dtype(): + tensor = gs.Variable("input", dtype=onnx.TensorProto.BFLOAT16, shape=[1]) + + value_info = _export_value_info_proto(tensor, do_type_check=True) + + assert value_info.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 + + +def test_export_rejects_unknown_integer_dtype(): + invalid_dtype = max(onnx.TensorProto.DataType.values()) + 1 + constant = gs.Constant("scale", np.array([1.0], dtype=np.float32)) + constant.explicit_dtype = invalid_dtype + + with pytest.raises(ValueError): + _export_tensor_proto(constant) + + variable = gs.Variable("input", dtype=invalid_dtype, shape=[1]) + with pytest.raises(ValueError): + _export_value_info_proto(variable, do_type_check=True) From 2261b292b44fcd7979aa740221d07efcfb264df2 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:35:46 +0000 Subject: [PATCH 4/8] [6508436] Focus ONNX export fix on BF16 FP8 Revert the universal target-precision expansion while retaining the BF16 FP8 real-weight compression fix and its required GraphSurgeon compatibility. The broader multi-format precision contract will be delivered separately. Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/onnx/autocast/convert.py | 170 ---- modelopt/onnx/export/base_exporter.py | 13 +- modelopt/onnx/export/fp8_exporter.py | 66 +- modelopt/onnx/export/int4_exporter.py | 13 +- modelopt/onnx/export/int8_exporter.py | 4 +- modelopt/onnx/export/mxfp8_exporter.py | 28 +- modelopt/onnx/export/nvfp4_exporter.py | 53 +- modelopt/onnx/quantization/gs_patching.py | 21 +- modelopt/onnx/quantization/qdq_utils.py | 87 +- modelopt/torch/_deploy/utils/torch_onnx.py | 247 +++--- modelopt/torch/quantization/export_onnx.py | 132 ++- modelopt/torch/quantization/tensor_quant.py | 1 - .../quantization/test_fp8_mha_exporter.py | 148 +--- .../onnx/quantization/test_gs_patching.py | 67 -- .../unit/onnx/quantization/test_qdq_utils.py | 377 +-------- .../deploy/utils/test_torch_onnx_utils.py | 780 +----------------- 17 files changed, 300 insertions(+), 1909 deletions(-) delete mode 100644 tests/unit/onnx/quantization/test_gs_patching.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f8b077d4bfe..b688a7bf704 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -131,7 +131,7 @@ Changelog **Bug Fixes** -- Make ``native`` the default ONNX export precision, preserving the model's existing precision, while explicit ``fp32``, ``fp16``, and ``bf16`` requests consistently control graph I/O and high-precision quantization boundaries across FP8, INT4, INT8, MXFP8, and NVFP4. This also fixes BF16 FP8 weight compression. +- Fix FP8 ONNX export of BF16 models during real-weight compression. - Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` / ``mixer.down_proj``, which the previous ``*mlp*`` selector missed, producing checkpoints with a null ``quant_algo``. - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. Stale output shapes are now reconciled via symbolic shape inference, and AutoCast falls back to schema-based type inference so unresolved ops no longer leave tensors untyped. - Fix fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) skipping modules without an ``act_fn`` attribute. Modules applying a custom gated activation between the two ``F.linear`` calls (e.g. ``MiniMaxM3VLExperts``) were silently skipped, leaving routed experts unquantized and failing HF export. Enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py index e2c2e3442df..65cabe86974 100644 --- a/modelopt/onnx/autocast/convert.py +++ b/modelopt/onnx/autocast/convert.py @@ -48,176 +48,6 @@ DEFAULT_DATA_MAX = 512 DEFAULT_INIT_MAX = np.finfo(np.float16).max LATEST_IR_VERSION_SUPPORTED_BY_ORT = 10 -_FLOAT_TYPES_TO_FP32 = { - onnx.TensorProto.DOUBLE, - onnx.TensorProto.FLOAT16, - onnx.TensorProto.BFLOAT16, -} -_STANDARD_FLOAT_DTYPE_ATTRIBUTES = { - "dtype", - "output_datatype", - "output_dtype", - "precision", - "softmax_precision", - "stash_type", -} -_TRT_FLOAT_DTYPE_ATTRIBUTES = {"output_dtype"} - - -def _convert_tensor_to_fp32(tensor: onnx.TensorProto) -> None: - if tensor.data_type not in _FLOAT_TYPES_TO_FP32: - return - if tensor.data_location == onnx.TensorProto.EXTERNAL and not tensor.raw_data: - raise ValueError("External tensor data must be loaded before FP32 conversion") - if tensor.HasField("segment"): - raise ValueError("Segmented tensors are not supported for FP32 conversion") - - dims = list(tensor.dims) - name = tensor.name - doc_string = tensor.doc_string - metadata_props = [deepcopy(prop) for prop in getattr(tensor, "metadata_props", ())] - if tensor.data_type in (onnx.TensorProto.FLOAT16, onnx.TensorProto.BFLOAT16): - values = ( - onnx_utils.read_f16_tensor_as_fp32(tensor) - if tensor.raw_data - else onnx.numpy_helper.to_array(tensor).astype(np.float32) - ) - else: - values = onnx.numpy_helper.to_array(tensor).astype(np.float32) - - # Release the low-precision payload before allocating the serialized FP32 payload. - tensor.Clear() - tensor.dims.extend(dims) - tensor.data_type = onnx.TensorProto.FLOAT - tensor.name = name - tensor.doc_string = doc_string - if metadata_props: - tensor.metadata_props.extend(metadata_props) - tensor.raw_data = values.tobytes() - - -def _convert_type_to_fp32(type_proto: onnx.TypeProto) -> None: - if type_proto.HasField("tensor_type"): - if type_proto.tensor_type.elem_type in _FLOAT_TYPES_TO_FP32: - type_proto.tensor_type.elem_type = onnx.TensorProto.FLOAT - elif type_proto.HasField("sparse_tensor_type"): - if type_proto.sparse_tensor_type.elem_type in _FLOAT_TYPES_TO_FP32: - type_proto.sparse_tensor_type.elem_type = onnx.TensorProto.FLOAT - elif type_proto.HasField("sequence_type"): - _convert_type_to_fp32(type_proto.sequence_type.elem_type) - elif type_proto.HasField("optional_type"): - _convert_type_to_fp32(type_proto.optional_type.elem_type) - elif type_proto.HasField("map_type"): - _convert_type_to_fp32(type_proto.map_type.value_type) - - -def _dtype_attributes_for_node(node: onnx.NodeProto) -> set[str]: - if node.domain in ("", "ai.onnx"): - attributes = set(_STANDARD_FLOAT_DTYPE_ATTRIBUTES) - if node.op_type == "Cast": - attributes.add("to") - return attributes - if node.domain == "trt": - return set(_TRT_FLOAT_DTYPE_ATTRIBUTES) - return set() - - -def _convert_attribute_to_fp32( - attribute: onnx.AttributeProto, dtype_attributes: set[str] | None = None -) -> None: - if ( - attribute.type == onnx.AttributeProto.INT - and dtype_attributes is not None - and attribute.name in dtype_attributes - and attribute.i in _FLOAT_TYPES_TO_FP32 - ): - attribute.i = onnx.TensorProto.FLOAT - elif attribute.type == onnx.AttributeProto.TENSOR: - _convert_tensor_to_fp32(attribute.t) - elif attribute.type == onnx.AttributeProto.TENSORS: - for tensor in attribute.tensors: - _convert_tensor_to_fp32(tensor) - elif attribute.type == onnx.AttributeProto.SPARSE_TENSOR: - _convert_tensor_to_fp32(attribute.sparse_tensor.values) - elif attribute.type == onnx.AttributeProto.SPARSE_TENSORS: - for sparse_tensor in attribute.sparse_tensors: - _convert_tensor_to_fp32(sparse_tensor.values) - elif attribute.type == onnx.AttributeProto.TYPE_PROTO: - _convert_type_to_fp32(attribute.tp) - elif attribute.type == onnx.AttributeProto.TYPE_PROTOS: - for type_proto in attribute.type_protos: - _convert_type_to_fp32(type_proto) - elif attribute.type == onnx.AttributeProto.GRAPH: - _convert_graph_to_fp32(attribute.g) - elif attribute.type == onnx.AttributeProto.GRAPHS: - for graph in attribute.graphs: - _convert_graph_to_fp32(graph) - - -def _convert_node_to_fp32(node: onnx.NodeProto) -> None: - is_standard_onnx_node = node.domain in ("", "ai.onnx") - dtype_attributes = _dtype_attributes_for_node(node) - for attribute in node.attribute: - if ( - is_standard_onnx_node - and node.op_type == "BitCast" - and attribute.name == "to" - and attribute.type == onnx.AttributeProto.INT - and attribute.i in _FLOAT_TYPES_TO_FP32 - ): - raise ValueError("BitCast targets cannot be converted safely to FP32") - _convert_attribute_to_fp32(attribute, dtype_attributes) - - -def _convert_function_to_fp32(function: onnx.FunctionProto) -> None: - dtype_attribute_refs = set() - bitcast_attribute_refs = set() - for node in function.node: - dtype_attributes = _dtype_attributes_for_node(node) - is_standard_bitcast = node.domain in ("", "ai.onnx") and node.op_type == "BitCast" - for attribute in node.attribute: - if not attribute.ref_attr_name: - continue - if is_standard_bitcast and attribute.name == "to": - bitcast_attribute_refs.add(attribute.ref_attr_name) - elif attribute.name in dtype_attributes: - dtype_attribute_refs.add(attribute.ref_attr_name) - - for attribute in function.attribute_proto: - if ( - attribute.name in bitcast_attribute_refs - and attribute.type == onnx.AttributeProto.INT - and attribute.i in _FLOAT_TYPES_TO_FP32 - ): - raise ValueError("BitCast targets cannot be converted safely to FP32") - dtype_attributes = {attribute.name} if attribute.name in dtype_attribute_refs else None - _convert_attribute_to_fp32(attribute, dtype_attributes) - for value_info in function.value_info: - _convert_type_to_fp32(value_info.type) - for node in function.node: - _convert_node_to_fp32(node) - - -def _convert_graph_to_fp32(graph: onnx.GraphProto) -> None: - for value_info in (*graph.input, *graph.output, *graph.value_info): - _convert_type_to_fp32(value_info.type) - for initializer in graph.initializer: - _convert_tensor_to_fp32(initializer) - for sparse_initializer in graph.sparse_initializer: - _convert_tensor_to_fp32(sparse_initializer.values) - for node in graph.node: - _convert_node_to_fp32(node) - - -def convert_to_fp32(model: onnx.ModelProto) -> onnx.ModelProto: - """Convert FP16, BF16, and FP64 values in an ONNX model to FP32 in place.""" - _convert_graph_to_fp32(model.graph) - for training_info in model.training_info: - _convert_graph_to_fp32(training_info.initialization) - _convert_graph_to_fp32(training_info.algorithm) - for function in model.functions: - _convert_function_to_fp32(function) - return model def _capture_network_io_metadata( diff --git a/modelopt/onnx/export/base_exporter.py b/modelopt/onnx/export/base_exporter.py index da596f8a86a..41d80c0e7ec 100644 --- a/modelopt/onnx/export/base_exporter.py +++ b/modelopt/onnx/export/base_exporter.py @@ -24,17 +24,12 @@ class ONNXQuantExporter(ABC): """Base class for ONNX quantizer exporters.""" @classmethod - def process_model( - cls, onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def process_model(cls, onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Processes the ONNX model.""" onnx_model = cls.pre_process(onnx_model) onnx_model = cls.compute_scales(onnx_model) onnx_model = cls.compress_weights(onnx_model) - if high_precision_dtype is None: - onnx_model = cls.post_process(onnx_model) - else: - onnx_model = cls.post_process(onnx_model, high_precision_dtype) + onnx_model = cls.post_process(onnx_model) return onnx_model @staticmethod @@ -54,7 +49,5 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: @staticmethod @abstractmethod - def post_process( - onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model.""" diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 438d2e92f68..6fa09f38263 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -25,7 +25,6 @@ from onnx_graphsurgeon.ir.tensor import LazyValues from modelopt.onnx.logging_config import logger -from modelopt.onnx.quantization.qdq_utils import np_dtype_map from .base_exporter import ONNXQuantExporter @@ -57,7 +56,7 @@ def compute_scales(onnx_model: onnx.ModelProto) -> onnx.ModelProto: @staticmethod def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Compresses FP32/FP16/BF16 weights to FP8 by folding QDQ nodes to DQ only. + """Compresses FP32/FP16 weights to FP8 by folding QDQ nodes to DQ only. Even though modelopt supports FP8 onnx export, the weights are represented in fp32 + QDQ. The storage is therefore very bad. In this function, @@ -65,7 +64,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: weights in the output model. TRT custom ops are converted to native ONNX DequantizeLinear. Parameters: - onnx_model: ONNX model with FP32/FP16/BF16 weights and TRT_FP8 QDQ nodes. + onnx_model: ONNX model with FP32/FP16 weights and TRT_FP8 QDQ nodes. Returns: ONNX model with FP8 weights and native ONNX DQ nodes for weights (QDQ preserved for activations). @@ -168,9 +167,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return gs.export_onnx(graph) @staticmethod - def _quantize_conv_weights_to_fp8( - graph: gs.Graph, high_precision_dtype: str | None = None - ) -> int: + def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: """Add FP8 weight DequantizeLinear for Conv layers with unquantized weights. Conv weight quantizers are disabled during TorchScript ONNX export because the @@ -185,7 +182,6 @@ def _quantize_conv_weights_to_fp8( Args: graph: The onnx-graphsurgeon graph to modify in-place. - high_precision_dtype: Optional ONNX scalar type for the DQ scale and output. Returns: Number of Conv weight DQ nodes inserted. @@ -207,34 +203,13 @@ def _quantize_conv_weights_to_fp8( continue torch_weights = _torch_from_numpy(weight_input.values.copy()) - if high_precision_dtype is not None: - scale_dtype = np_dtype_map[high_precision_dtype] - target_torch_dtype = _torch_from_numpy(np.empty((), dtype=scale_dtype)).dtype - torch_weights = torch_weights.to(target_torch_dtype) - else: - scale_dtype = np.float16 - amax = torch_weights.abs().max().float() if amax == 0: continue - - scale_value = (amax / _FP8_E4M3_MAX).item() - scale = np.array(scale_value, dtype=scale_dtype) - if high_precision_dtype is not None and scale == 0: - dtype_info = ( - ml_dtypes.finfo(scale_dtype) - if scale_dtype == ml_dtypes.bfloat16 - else np.finfo(scale_dtype) - ) - scale = np.array(dtype_info.smallest_subnormal, dtype=scale_dtype) - scaled_weights = ( - torch_weights / _torch_from_numpy(scale) - if high_precision_dtype is not None - else torch_weights / scale_value - ) + scale_val = (amax / _FP8_E4M3_MAX).item() # Quantize weights to FP8 (WAR: numpy doesn't support fp8) - fp8_data = scaled_weights.to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() fp8_tensor = onnx.TensorProto() fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN fp8_tensor.dims.extend(fp8_data.shape) @@ -243,15 +218,13 @@ def _quantize_conv_weights_to_fp8( node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) ) + # Scale in FP16 — DQ output type matches scale dtype, must match activation type scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", - scale, + np.array(scale_val, dtype=np.float16), ) - dq_output = gs.Variable( - node.name + "/weight_quantizer/dq_output", - dtype=scale_dtype if high_precision_dtype is not None else None, - ) + dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") dq_node = gs.Node( op="DequantizeLinear", name=node.name + "/weight_quantizer/DequantizeLinear", @@ -405,7 +378,7 @@ def _move_transpose_before_qdq(graph: gs.Graph) -> int: return count @staticmethod - def _insert_qdq_after_softmax(graph: gs.Graph, high_precision_dtype: str | None = None) -> int: + def _insert_qdq_after_softmax(graph: gs.Graph) -> int: """Insert FP8 Q→DQ on Softmax outputs feeding MatMul (required by TRT MHA fusion). Softmax output is data-independently bounded to [0, 1], so we use a fixed scale @@ -427,13 +400,7 @@ def _insert_qdq_after_softmax(graph: gs.Graph, high_precision_dtype: str | None # Match scale dtype to the graph's current float dtype so TRT stronglyTyped # sees consistent Q/DQ types with the surrounding compute. - scale_dtype = ( - np_dtype_map[high_precision_dtype] - if high_precision_dtype is not None - else softmax_output.dtype - if softmax_output.dtype is not None - else np.float32 - ) + scale_dtype = softmax_output.dtype if softmax_output.dtype is not None else np.float32 scale_val = np.array(_FP8_E4M3_SOFTMAX_SCALE, dtype=scale_dtype) scale_constant = gs.Constant(softmax_node.name + "/softmax_q_scale", scale_val) dq_scale_constant = gs.Constant( @@ -449,10 +416,7 @@ def _insert_qdq_after_softmax(graph: gs.Graph, high_precision_dtype: str | None ) q_output = gs.Variable(softmax_node.name + "/q_output") - dq_output = gs.Variable( - softmax_node.name + "/dq_output", - dtype=scale_dtype if high_precision_dtype is not None else softmax_output.dtype, - ) + dq_output = gs.Variable(softmax_node.name + "/dq_output", dtype=softmax_output.dtype) q_node = gs.Node( op="QuantizeLinear", name=softmax_node.name + "/QuantizeLinear", @@ -480,9 +444,7 @@ def _insert_qdq_after_softmax(graph: gs.Graph, high_precision_dtype: str | None return count @staticmethod - def post_process( - onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for FP8 quantization. Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear, @@ -526,14 +488,14 @@ def post_process( ) # Add FP8 weight DQ for Conv layers that had weight quantizers disabled during export - count = FP8QuantExporter._quantize_conv_weights_to_fp8(graph, high_precision_dtype) + count = FP8QuantExporter._quantize_conv_weights_to_fp8(graph) if count > 0: logger.info(f"Inserted FP8 weight DequantizeLinear for {count} Conv nodes") # Attention-aware rewrites so TRT can fuse DQ into the attention MatMuls. n_mul = FP8QuantExporter._move_mul_before_qdq(graph) n_t = FP8QuantExporter._move_transpose_before_qdq(graph) - n_sm = FP8QuantExporter._insert_qdq_after_softmax(graph, high_precision_dtype) + n_sm = FP8QuantExporter._insert_qdq_after_softmax(graph) if n_mul or n_t or n_sm: logger.info( f"Attention QDQ rewrites: moved {n_mul} Mul, {n_t} Transpose; " diff --git a/modelopt/onnx/export/int4_exporter.py b/modelopt/onnx/export/int4_exporter.py index e2edd0a7a93..0da217ae76f 100644 --- a/modelopt/onnx/export/int4_exporter.py +++ b/modelopt/onnx/export/int4_exporter.py @@ -223,11 +223,8 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process( - onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for INT4 quantization.""" - precision_dtype = high_precision_dtype or "Half" def is_pre_quant_scale_node(node: onnx.NodeProto) -> bool: has_pqs_input = any(input for input in node.input if "_pre_quant_scale" in input) @@ -268,12 +265,12 @@ def is_fp32_cast(node: onnx.NodeProto) -> bool: del graph.node[:] graph.node.extend(new_nodes) - # Cast bias to the graph's high-precision dtype + # Cast bias to float16 for node in graph.node: if node.op_type == "Add" and "proj/Add" in node.name: - cast_initializer_to_dtype(node, precision_dtype, initializer_map) + cast_initializer_to_dtype(node, "Half", initializer_map) - # Cast pre quant scales of o_proj and down_proj to the high-precision dtype + # Cast pre quant scales of o_proj and down_proj to float16 for node in graph.node: if node.op_type == "Mul" and ( any( @@ -281,6 +278,6 @@ def is_fp32_cast(node: onnx.NodeProto) -> bool: for x in ("o_proj/input_quantizer/Mul", "down_proj/input_quantizer/Mul") ) ): - cast_initializer_to_dtype(node, precision_dtype, initializer_map) + cast_initializer_to_dtype(node, "Half", initializer_map) return onnx_model diff --git a/modelopt/onnx/export/int8_exporter.py b/modelopt/onnx/export/int8_exporter.py index 03f0a40dd1b..4623279b531 100644 --- a/modelopt/onnx/export/int8_exporter.py +++ b/modelopt/onnx/export/int8_exporter.py @@ -40,8 +40,6 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process( - onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for INT8 quantization.""" return onnx_model diff --git a/modelopt/onnx/export/mxfp8_exporter.py b/modelopt/onnx/export/mxfp8_exporter.py index 92b1005925b..8c1e1f4df4f 100644 --- a/modelopt/onnx/export/mxfp8_exporter.py +++ b/modelopt/onnx/export/mxfp8_exporter.py @@ -143,28 +143,20 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process( - onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for MXFP8 quantization. - Sets DQ output type and updates GELU nodes to use tanh approximation. + Sets DQ output type to FP16 and updates GELU nodes to use tanh approximation. """ logger.info("Post-processing MXFP8 quantized model") graph = onnx_model.graph - precision_dtype = high_precision_dtype or "Half" - precision_suffix = { - "Float": "fp32", - "Half": "fp16", - "BFloat16": "bf16", - }[precision_dtype] - - # Set output type of DQ to the graph's high-precision dtype + + # Set output type of DQ to FP16 for node in graph.node: if node.op_type == "TRT_MXFP8DequantizeLinear": for attr in node.attribute: if attr.name == "output_dtype": - attr.i = onnx_dtype_map[precision_dtype] + attr.i = onnx_dtype_map["Half"] # Currently only tanh approximation is supported for Gelu for node in graph.node: @@ -174,20 +166,20 @@ def post_process( attr.s = b"tanh" logger.debug(f"Updated GELU node {node.name} to use tanh approximation") - # Insert cast to the graph's high-precision dtype after Sqrt nodes + # Insert cast to fp16 after Sqrt nodes cast_nodes_to_insert = [] for idx, node in enumerate(graph.node): if node.op_type == "Sqrt": sqrt_output = node.output[0] - cast_output = f"{sqrt_output}_cast_{precision_suffix}" + cast_output = f"{sqrt_output}_cast_fp16" # Create Cast node cast_node = onnx.helper.make_node( "Cast", inputs=[sqrt_output], outputs=[cast_output], - to=onnx_dtype_map[precision_dtype], - name=f"{node.name}_cast_{precision_suffix}", + to=onnx_dtype_map["Half"], + name=f"{node.name}_cast_fp16", ) cast_nodes_to_insert.append((idx + 1, cast_node)) @@ -202,6 +194,6 @@ def post_process( # Insert Cast nodes in reverse order to preserve indices for offset, (pos, cast_node) in enumerate(cast_nodes_to_insert): graph.node.insert(pos + offset, cast_node) - logger.debug(f"Inserted Cast to {precision_dtype} after {cast_node.input[0]}") + logger.debug(f"Inserted Cast to FP16 after {cast_node.input[0]}") return onnx_model diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index 1cd488ea4c5..338e2725b14 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -318,9 +318,7 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model @staticmethod - def post_process( - onnx_model: onnx.ModelProto, high_precision_dtype: str | None = None - ) -> onnx.ModelProto: + def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for NVFP4 quantization. Replaces TRT_FP4QDQ nodes with two DequantizeLinear nodes and handles @@ -336,60 +334,37 @@ def post_process( value_info_map = {vi.name: vi for vi in graph.value_info} graph_inputs = {inp.name for inp in graph.input} cast_output_cache: dict[tuple[str, str], str] = {} - casted_node_ids: set[int] = set() def _get_precision_dtype() -> str: # Check initializers to determine the precision of the weights precision_dtype = "Half" for initializer in graph.initializer: - if initializer.data_type == onnx.TensorProto.BFLOAT16: + if initializer.data_type == 16: precision_dtype = "BFloat16" break # Assuming all weights are of the same precision return precision_dtype - def _get_linear_consumers(tensor_name: str) -> list[onnx.NodeProto]: - nodes_to_visit = list(tensor_consumers.get(tensor_name, [])) - visited_node_ids = set() - linear_consumers = {} - - while nodes_to_visit: - node = nodes_to_visit.pop() - node_id = id(node) - if node_id in visited_node_ids: - continue - visited_node_ids.add(node_id) - - if node.op_type in {"Gemm", "MatMul"}: - linear_consumers[node_id] = node - elif node.op_type in {"Cast", "Transpose"}: - for output_name in node.output: - nodes_to_visit.extend(tensor_consumers.get(output_name, [])) - - assert linear_consumers, f"No Gemm or MatMul consumes {tensor_name}" - return list(linear_consumers.values()) - def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): # Change the input types to match weight precision (precision_dtype) - assert node.op_type in {"Gemm", "MatMul"} + if node.op_type == "Transpose": + maybe_matmul = tensor_consumers[node.output[0]][0] + assert maybe_matmul.op_type == "MatMul" + node = maybe_matmul # Create Cast nodes for each input of the target node except bias for i, input_name in enumerate(node.input[:2]): cast_output_name = cast_output_cache.get((input_name, precision_dtype)) if cast_output_name is None: - cast_output_suffix = { - "Float": "f32", - "Half": "f16", - "BFloat16": "bf16", - }[precision_dtype] + cast_output_suffix = "bf16" if precision_dtype == "BFloat16" else "f16" cast_output_name = f"{input_name}_{cast_output_suffix}" cast_output_cache[(input_name, precision_dtype)] = cast_output_name - # Create a Cast node to convert the input to the selected precision + # Create a Cast node to convert the input to FP16/BF16 cast_node = onnx.helper.make_node( "Cast", inputs=[input_name], # Original input of the target node outputs=[cast_output_name], - to=onnx_dtype_map[precision_dtype], + to=onnx_dtype_map[precision_dtype], # Cast to FP16/BF16 ) # Insert the Cast node into the graph @@ -398,7 +373,7 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): # Update the target node input to use the cast node output node.input[i] = cast_output_name - precision_dtype = high_precision_dtype or _get_precision_dtype() + precision_dtype = _get_precision_dtype() logger.debug(f"Using precision dtype: {precision_dtype}") fp4_qdq_nodes = [node for node in graph.node if node.op_type == "TRT_FP4QDQ"] @@ -441,11 +416,9 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): block_size, ) - # Cast input dtypes for every linear consumer reached through Cast/Transpose wrappers. - for linear_node in _get_linear_consumers(node.output[0]): - if id(linear_node) not in casted_node_ids: - _cast_input_dtypes(linear_node, precision_dtype) - casted_node_ids.add(id(linear_node)) + # Cast input dtypes for the next node + next_node = tensor_consumers[node.output[0]][0] + _cast_input_dtypes(next_node, precision_dtype) # Remove old initializers new_initializers = [ diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index c8525db36bf..453f7af46bc 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -65,23 +65,14 @@ def _make_variable( return x -def _to_onnx_dtype(dtype) -> int: - if isinstance(dtype, (int, np.integer)): - dtype = int(dtype) - onnx.TensorProto.DataType.Name(dtype) - return dtype - return onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) - - def _export_tensor_proto(tensor: gs.Constant) -> onnx.TensorProto: if isinstance(tensor._values, LazyValues): onnx_tensor = tensor._values.tensor else: # is numpy array. - dtype = getattr(tensor, "explicit_dtype", None) - if dtype is None: - dtype = tensor.values.dtype - dtype = _to_onnx_dtype(dtype) + dtype = getattr( + tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(tensor.values.dtype) + ) vals = tensor.values if _onnx_supports_int4() and dtype in [onnx.TensorProto.INT4, onnx.TensorProto.UINT4]: @@ -113,7 +104,11 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V dtype = getattr(tensor, "explicit_dtype", None) if dtype is None: dtype = tensor.dtype - dtype = _to_onnx_dtype(dtype) + if isinstance(dtype, (int, np.integer)): + dtype = int(dtype) + onnx.TensorProto.DataType.Name(dtype) + else: + dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) else: onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py index b172658d7f9..3b48805439a 100644 --- a/modelopt/onnx/quantization/qdq_utils.py +++ b/modelopt/onnx/quantization/qdq_utils.py @@ -19,7 +19,6 @@ from collections.abc import Sequence from typing import Any -import ml_dtypes import numpy as np import onnx import onnx_graphsurgeon as gs @@ -62,7 +61,6 @@ onnx_bit_dtype_unsigned_map = {4: "UINT4", 8: "UINT8"} np_dtype_map = { - "BFloat16": ml_dtypes.bfloat16, "Float": np.float32, "Half": np.float16, "INT8": np.int8, @@ -1022,71 +1020,40 @@ def remove_graph_input_q(onnx_model: onnx.ModelProto) -> onnx.ModelProto: def replace_zero_scale_with_smallest_nonzero(onnx_model: onnx.ModelProto) -> onnx.ModelProto: - """Replace zero scale values with the smallest nonzero value of their dtype.""" + """Replace zero scale values with smallest nonzero fp16 value in the ONNX model.""" + graph = onnx_model.graph + fp16_smallest_nonzero = np.float16(6e-08) qdq_op_types = { "QuantizeLinear", "DequantizeLinear", "TRT_INT4QuantizeLinear", "TRT_INT4DequantizeLinear", } - - def replace_zeros(tensor_proto: onnx.TensorProto) -> None: - dtype = { - onnx.TensorProto.BFLOAT16: ml_dtypes.bfloat16, - onnx.TensorProto.DOUBLE: np.float64, - onnx.TensorProto.FLOAT: np.float32, - onnx.TensorProto.FLOAT16: np.float16, - }.get(tensor_proto.data_type) - if dtype is None: - return - - tensor = numpy_helper.to_array(tensor_proto) - dtype_info = ml_dtypes.finfo(dtype) if dtype == ml_dtypes.bfloat16 else np.finfo(dtype) - smallest_nonzero = np.array(dtype_info.smallest_subnormal, dtype=dtype) - new_tensor = np.where(tensor == 0, smallest_nonzero, tensor).astype(dtype) - tensor_proto.CopyFrom(numpy_helper.from_array(new_tensor, tensor_proto.name)) - - def replace_zero_scales(graph: onnx.GraphProto) -> set[str]: - scale_tensor_names = { - node.input[1] - for node in graph.node - if node.op_type in qdq_op_types and len(node.input) >= 2 and node.input[1] - } - - for node in graph.node: + scale_tensor_names = { + node.input[1] + for node in graph.node + if node.op_type in qdq_op_types and len(node.input) >= 2 + } + # Scales stored as graph initializers (e.g. INT4_AWQ / TRT_INT4DequantizeLinear exports). + for init in graph.initializer: + if init.name in scale_tensor_names: + tensor = numpy_helper.to_array(init) + if tensor.dtype.kind == "f": + new_tensor = np.where(tensor == 0, fp16_smallest_nonzero, tensor).astype( + tensor.dtype + ) + init.CopyFrom(numpy_helper.from_array(new_tensor, init.name)) + # Scales emitted by Constant nodes (legacy QDQ export path). + for node in graph.node: + if node.op_type == "Constant" and node.output[0] in scale_tensor_names: for attr in node.attribute: - if attr.type == onnx.AttributeProto.GRAPH: - scale_tensor_names.update(replace_zero_scales(attr.g)) - elif attr.type == onnx.AttributeProto.GRAPHS: - for subgraph in attr.graphs: - scale_tensor_names.update(replace_zero_scales(subgraph)) - - # Scales stored as graph initializers (e.g. INT4_AWQ / TRT_INT4DequantizeLinear exports). - initializer_names = {init.name for init in graph.initializer} - sparse_initializer_names = { - init.values.name for init in graph.sparse_initializer if init.values.name - } - for init in graph.initializer: - if init.name in scale_tensor_names: - replace_zeros(init) - - # Scales emitted by Constant nodes (legacy QDQ export path). - node_output_names = {output for node in graph.node for output in node.output if output} - for node in graph.node: - if node.op_type == "Constant" and node.output[0] in scale_tensor_names: - for attr in node.attribute: - if attr.name == "value": - replace_zeros(attr.t) - - local_definitions = ( - initializer_names - | sparse_initializer_names - | node_output_names - | {value.name for value in graph.input if value.name} - ) - return scale_tensor_names - local_definitions - - replace_zero_scales(onnx_model.graph) + if attr.name == "value": + tensor = numpy_helper.to_array(attr.t) + if tensor.dtype.kind == "f": + new_tensor = np.where(tensor == 0, fp16_smallest_nonzero, tensor).astype( + tensor.dtype + ) + attr.t.CopyFrom(numpy_helper.from_array(new_tensor, attr.t.name)) return onnx_model diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 0fddd684ffb..9134df0f2ed 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -19,20 +19,22 @@ import contextlib import inspect import json +import logging import os import shutil import tempfile from contextlib import nullcontext -from itertools import chain from typing import Any import onnx +import onnxconverter_common.float16 as _f16_module import torch import torch.nn as nn from onnx import ModelProto +from onnxconverter_common import convert_float_to_float16 from torch.nn.parallel import DataParallel, DistributedDataParallel -from modelopt.onnx.autocast.convert import convert_to_f16, convert_to_fp32 +from modelopt.onnx.autocast.convert import convert_to_f16 from modelopt.onnx.export import ( FP8QuantExporter, INT4QuantExporter, @@ -43,9 +45,11 @@ ) from modelopt.onnx.quantization.qdq_utils import qdq_to_dq, replace_zero_scale_with_smallest_nonzero from modelopt.onnx.utils import ( + change_casts_to_fp16, check_model_uses_external_data, fold_dq_fp32_to_fp16_casts, fold_q_fp16_to_fp32_casts, + fold_qdq_scale_fp16_to_fp32_casts, get_input_names, get_input_shapes, get_node_names, @@ -56,12 +60,35 @@ remove_redundant_casts, ) from modelopt.torch.quantization.export_onnx import configure_linear_module_onnx_quantizers -from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import flatten_tree, standardize_named_model_args from modelopt.torch.utils._pytree import TreeSpec from ..utils.onnx_optimizer import Optimizer +# Monkey-patch for onnxconverter_common bug in remove_unnecessary_cast_node(): +# cast_node_downstream_dict stores either a single node or a list of nodes, but the +# downstream-node handling at lines ~770/787 always does `downstream_node.input`, +# which raises AttributeError("'list' object has no attribute 'input'") when the +# value is a list (i.e. a Cast output feeds multiple consumers). +# TODO: Remove this patch once onnxconverter-common ships a fix. +# Upstream issue: https://github.com/microsoft/onnxconverter-common/issues/261 +_original_remove_unnecessary_cast_node = _f16_module.remove_unnecessary_cast_node + +_logger = logging.getLogger(__name__) + + +def _patched_remove_unnecessary_cast_node(graph): + try: + _original_remove_unnecessary_cast_node(graph) + except AttributeError as e: + if "'list' object has no attribute 'input'" in str(e): + _logger.debug("Skipping remove_unnecessary_cast_node due to known upstream bug: %s", e) + else: + raise + + +_f16_module.remove_unnecessary_cast_node = _patched_remove_unnecessary_cast_node + ModelMetadata = dict[str, Any] ModelType = Any ValueInfoType = Any @@ -70,16 +97,6 @@ DEFAULT_ONNX_OPSET = 20 ONNX_EXPORT_OUT_PREFIX = "out" TWO_GB = 2 * 1024 * 1024 * 1024 -WEIGHTS_DTYPE_TO_TORCH_DTYPE = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} -WEIGHTS_DTYPE_TO_ONNX_DTYPE = { - "fp32": "Float", - "fp16": "Half", - "bf16": "BFloat16", -} class OnnxBytes: @@ -194,61 +211,6 @@ def _to_expected_onnx_type(val: Any) -> Any: return val -def _cast_floating_tensors(value: Any, dtype: torch.dtype) -> Any: - flat_values, tree_spec = flatten_tree(value) - flat_values = [ - item.to(dtype=dtype) - if isinstance(item, torch.Tensor) and item.is_floating_point() - else item - for item in flat_values - ] - return tree_spec.generate_pytree(flat_values) - - -def _get_autocast_context( - model: nn.Module, flat_input: list[Any], target_dtype: torch.dtype | None -): - if target_dtype not in (torch.float16, torch.bfloat16): - return nullcontext() - - for item in flat_input: - if isinstance(item, torch.Tensor) and item.is_floating_point(): - return torch.autocast(device_type=item.device.type, dtype=target_dtype) - for tensor in chain(model.parameters(), model.buffers()): - if tensor.is_floating_point(): - return torch.autocast(device_type=tensor.device.type, dtype=target_dtype) - for item in flat_input: - if isinstance(item, torch.Tensor): - return torch.autocast(device_type=item.device.type, dtype=target_dtype) - tensor = next(chain(model.parameters(), model.buffers()), None) - if tensor is not None: - return torch.autocast(device_type=tensor.device.type, dtype=target_dtype) - return torch.autocast(device_type="cpu", dtype=target_dtype) - - -@contextlib.contextmanager -def _override_onnx_quantizer_precision(model: nn.Module, high_precision_dtype: str | None): - if high_precision_dtype is None: - yield - return - - sentinel = object() - originals: list[tuple[TensorQuantizer, Any]] = [] - for module in model.modules(): - if isinstance(module, TensorQuantizer): - original = getattr(module, "_trt_high_precision_dtype", sentinel) - originals.append((module, original)) - module.trt_high_precision_dtype = high_precision_dtype - try: - yield - finally: - for quantizer, original in originals: - if original is sentinel: - del quantizer._trt_high_precision_dtype - else: - quantizer.trt_high_precision_dtype = original - - def generate_onnx_input( model_metadata: ModelMetadata, input: Any | tuple, ignore_nesting: bool = False ) -> dict[str, Any]: @@ -467,15 +429,11 @@ def _disable_fp8_conv_weight_quantizers(model: nn.Module): module.weight_quantizer.enable() -def quantize_weights( - model: nn.Module, - onnx_model: onnx.ModelProto, - high_precision_dtype: str | None = None, -) -> onnx.ModelProto: +def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Real quantizes the weights in the onnx model. Applies weight quantization to an ONNX model based on the quantization scheme detected - in the PyTorch model. Supports INT4, NVFP4, MXFP8, FP8, and INT8 quantization formats. + in the PyTorch model. Supports INT4, FP4, and MXFP8 quantization formats. The function performs a four-stage process for each detected quantization type: 1. Pre-process - Restructure the graph for quantization @@ -487,7 +445,6 @@ def quantize_weights( model (nn.Module): The original PyTorch model used to detect quantization schemes. This model should have been quantized using modelopt's quantization APIs. onnx_model (onnx.ModelProto): The ONNX model whose weights will be quantized. - high_precision_dtype: Optional ONNX scalar type used for the surrounding graph. Returns: onnx.ModelProto: The ONNX model with quantized weights applied. The returned model @@ -496,7 +453,7 @@ def quantize_weights( Notes: - Multiple quantization formats can be applied sequentially if the model contains different quantization schemes for different layers - - The function checks every supported quantization format in the PyTorch model + - The function checks for INT4, FP4, and MXFP8 quantization in the PyTorch model - Each quantization exporter modifies the ONNX graph in-place before returning """ @@ -517,7 +474,7 @@ def quantize_weights( return onnx_model for onnx_exporter in onnx_exporters: - onnx_model = onnx_exporter.process_model(onnx_model, high_precision_dtype) + onnx_model = onnx_exporter.process_model(onnx_model) return onnx_model @@ -532,7 +489,7 @@ def get_onnx_bytes_and_metadata( dynamo_export: bool = False, onnx_opset: int = DEFAULT_ONNX_OPSET, dq_only: bool = False, - weights_dtype: str = "native", + weights_dtype: str = "fp32", ) -> tuple[bytes, ModelMetadata]: """Get onnx model in bytes from input pytorch model together with the input/output of model. @@ -550,11 +507,7 @@ def get_onnx_bytes_and_metadata( `torch.onnx.export `_. onnx_opset: The onnx opset version to use for exporting the model. dq_only: If True, the exported onnx model is converted to a dq_only model. - weights_dtype: Selects the floating-point graph I/O and high-precision Q/DQ boundary - dtype. ``native`` preserves the precision produced by the PyTorch export; - ``fp32``, ``fp16``, and ``bf16`` force that target while leaving format-native - quantized tensors and scales unchanged. Inference inputs supplied to the exported - ONNX model must use the selected explicit floating-point dtype. + weights_dtype: The dtype of the weights in the onnx model. Returns: bytes: Onnx model in bytes. @@ -566,24 +519,22 @@ def get_onnx_bytes_and_metadata( if not isinstance(model, nn.Module): raise ValueError("Only PyTorch model compilation is supported.") - assert weights_dtype in ["native", "fp32", "fp16", "bf16"], ( - "weights_dtype must be one of native, fp32, fp16, or bf16" + assert weights_dtype in ["fp32", "fp16", "bf16"], ( + "weights_dtype must be one of fp32, fp16, or bf16" ) - if onnx_load_path and weights_dtype != "native": - raise ValueError("weights_dtype must be 'native' when onnx_load_path is provided") # unwrap DDP and DP models if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module + first_parameter = next(model.parameters(), None) + source_weights_dtype = first_parameter.dtype if first_parameter is not None else torch.float32 + # 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()} - target_torch_dtype = WEIGHTS_DTYPE_TO_TORCH_DTYPE.get(weights_dtype) - if target_torch_dtype in (torch.float16, torch.bfloat16): - named_args = _cast_floating_tensors(named_args, target_torch_dtype) # Also standardize dummy_input again so we can use it dummy_input = tuple(named_args.values()) @@ -600,8 +551,17 @@ def get_onnx_bytes_and_metadata( # during inference. input_none_names = list(set(tree_spec_input.names) - set(input_names)) + use_torch_autocast = not ( + is_fp4_quantized(model) + or is_mxfp8_quantized(model) + or is_fp8_quantized(model) + or is_int8_quantized(model) + or weights_dtype == "fp32" + ) + autocast = torch.autocast("cuda") if use_torch_autocast else nullcontext() + # Get output once (we export in inference mode - so also using inference mode here!) - with torch.inference_mode(), _get_autocast_context(model, flat_input, target_torch_dtype): + with torch.inference_mode(), autocast: output = model(*named_args.values()) # Get output tree spec @@ -636,14 +596,7 @@ def get_onnx_bytes_and_metadata( conv_wq_context = ( _disable_fp8_conv_weight_quantizers(model) if is_fp8_quantized(model) else nullcontext() ) - high_precision_dtype = WEIGHTS_DTYPE_TO_ONNX_DTYPE.get(weights_dtype) - with ( - torch.inference_mode(), - _get_autocast_context(model, flat_input, target_torch_dtype), - _override_onnx_quantizer_precision(model, high_precision_dtype), - quantizer_context, - conv_wq_context, - ): + with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: additional_kwargs = {} if not dynamo_export: additional_kwargs["dynamic_axes"] = dynamic_axes @@ -679,30 +632,50 @@ def get_onnx_bytes_and_metadata( tree_spec_input, tree_spec_output, input_none_names, onnx_opt_graph, model ) - onnx_opt_graph = quantize_weights(model, onnx_opt_graph, high_precision_dtype) + onnx_opt_graph = quantize_weights(model, onnx_opt_graph) if dq_only: onnx_opt_graph = qdq_to_dq(onnx_opt_graph) - if weights_dtype == "fp32": - onnx_opt_graph = convert_to_fp32(onnx_opt_graph) - elif weights_dtype in ("fp16", "bf16") and not any( - ( - is_int4_quantized(model), - is_fp4_quantized(model), - is_mxfp8_quantized(model), - is_fp8_quantized(model), - is_int8_quantized(model), - ) - ): - onnx_opt_graph = convert_to_f16( - onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False - ) + uses_fp8 = is_fp8_quantized(model) + uses_other_unsupported_quantizer = ( + is_int4_quantized(model) or is_mxfp8_quantized(model) or is_int8_quantized(model) + ) + is_bf16_fp8_noop = ( + weights_dtype == "bf16" + and source_weights_dtype == torch.bfloat16 + and uses_fp8 + and not uses_other_unsupported_quantizer + ) + if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop: + if uses_other_unsupported_quantizer or uses_fp8: + assert weights_dtype == "fp16", ( + "Converting a quantized ONNX graph to BF16 is not supported yet" + ) + onnx_opt_graph = convert_float_to_float16( + onnx_opt_graph, + keep_io_types=False, + disable_shape_infer=True, + check_fp16_ready=False, + op_block_list=["QuantizeLinear", "DequantizeLinear", "Div"], + ) + # Change FP32 cast nodes feeding into Concat/Add to FP16 + op_list = ["Concat", "Add", "Sqrt", "LayerNormalization", "Clip", "Mul", "Exp"] + onnx_opt_graph = change_casts_to_fp16(onnx_opt_graph, op_list) + # Remove Cast(FP32->FP16) nodes after DQ by setting DQ output to FP16 directly + onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) + # Remove Cast(FP16->FP32) feeding Q/DQ scales so DQ stays FP16 for downstream + # MatMul/Add layers under strongly-typed TRT parsing. + onnx_opt_graph = fold_qdq_scale_fp16_to_fp32_casts(onnx_opt_graph) + else: + onnx_opt_graph = convert_to_f16( + onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False + ) onnx_opt_graph = remove_redundant_casts(onnx_opt_graph) # Remove Cast nodes around Q/DQ for optimal TRT fusion - if is_fp8_quantized(model) and weights_dtype == "fp16": + if is_fp8_quantized(model): onnx_opt_graph = fold_q_fp16_to_fp32_casts(onnx_opt_graph) onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) @@ -713,7 +686,22 @@ def get_onnx_bytes_and_metadata( # Must be set after all gs.export_onnx() calls as graphsurgeon resets ir_version onnx_opt_graph.ir_version = 10 - _save_onnx_model(onnx_opt_graph, onnx_save_path, model_name) + # If the onnx model contains external data store the external tensors in one file and save the onnx model + if has_external_data(onnx_save_path): + tensor_paths = get_external_tensor_paths(onnx_path) + onnx.save_model( + onnx_opt_graph, + onnx_save_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=f"{model_name}.onnx_data", + size_threshold=1024, + convert_attribute=False, + ) + for path in tensor_paths: + os.remove(path) + else: + onnx.save_model(onnx_opt_graph, onnx_save_path) onnx_bytes = OnnxBytes(onnx_save_path) @@ -737,33 +725,6 @@ def has_external_data(onnx_model_path: str): return check_model_uses_external_data(onnx_model) -def _save_onnx_model(onnx_model: onnx.ModelProto, onnx_save_path: str, model_name: str) -> None: - model_dir = os.path.dirname(onnx_save_path) - if not (has_external_data(onnx_save_path) or onnx_model.ByteSize() >= TWO_GB): - onnx.save_model(onnx_model, onnx_save_path) - return - - tensor_paths = get_external_tensor_paths(model_dir) - external_data_name = f"{model_name}.onnx_data" - external_data_path = os.path.join(model_dir, external_data_name) - if os.path.exists(external_data_path): - os.remove(external_data_path) - - onnx.save_model( - onnx_model, - onnx_save_path, - save_as_external_data=True, - all_tensors_to_one_file=True, - location=external_data_name, - size_threshold=1024, - convert_attribute=True, - ) - external_data_path = os.path.abspath(external_data_path) - for path in tensor_paths: - if os.path.abspath(path) != external_data_path and os.path.exists(path): - os.remove(path) - - def create_model_metadata( tree_spec_input: TreeSpec, tree_spec_output: TreeSpec, diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index a3611df5778..e5778c3c96b 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -103,7 +103,7 @@ """Utility to export a quantized torch model to quantized ONNX.""" import contextlib -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import onnx import torch @@ -124,22 +124,11 @@ "INT8": onnx.TensorProto.INT8, "UINT8": onnx.TensorProto.UINT8, } -mha_fusion_precisions = {"Half", "BFloat16"} -mha_supported_precisions = {"Float", "Half", "BFloat16"} +mha_valid_precisions = {"Half", "BFloat16"} torch_dtype_map = {"Float": torch.float32, "Half": torch.float16, "BFloat16": torch.bfloat16} -def _cast_to_dtype(g: "GraphContext", tensor: torch.Value, dtype: str): - """Cast a graph value only when its dtype differs from the target.""" - if tensor.type().scalarType() == dtype: - return tensor - output_shape = sym_help._get_tensor_sizes(tensor) - return g.op("Cast", tensor, to_i=onnx_dtype_map[dtype]).setType( - tensor.type().with_dtype(torch_dtype_map[dtype]).with_sizes(output_shape) - ) - - def export_int8( g: "GraphContext", inputs: torch.Value, @@ -157,9 +146,6 @@ def export_int8( input_type = inputs.type().scalarType() if trt_high_precision_dtype is None: trt_high_precision_dtype = input_type - assert trt_high_precision_dtype in torch_dtype_map, ( - f"Unsupported high precision dtype: {trt_high_precision_dtype}" - ) if amax.numel() == 1: zero_point, axis = torch.tensor(0.0, device=amax.device), None @@ -183,12 +169,22 @@ def export_int8( scale.masked_fill_(scale == 0, 1.0) scale = g.op("Constant", value_t=scale) - inputs = _cast_to_dtype(g, inputs, trt_high_precision_dtype) + assert trt_high_precision_dtype in (input_type, "Float", "BFloat16"), ( + "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." + ) + + # custom ops, so cast the input if needed. + if trt_high_precision_dtype != input_type: + inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) quantized = g.op("QuantizeLinear", inputs, scale, zero_point, axis_i=axis) out = g.op("DequantizeLinear", quantized, scale, zero_point, axis_i=axis).setType( inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) ) + # custom ops, so cast the output if needed. + if trt_high_precision_dtype != input_type: + inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[input_type]) + return out @@ -203,17 +199,12 @@ def export_int4( ): """Export quantized model to INT4 ONNX.""" assert num_bits == 4, "Number of bits must be 4 for INT4 ONNX export." + scale_inv = amax / 7.0 + scale_inv_op = g.op("Constant", value_t=scale_inv) otype = inputs.type().scalarType() output_shape = sym_help._get_tensor_sizes(inputs) if trt_high_precision_dtype is None: trt_high_precision_dtype = otype - scale_inv = amax / 7.0 - else: - assert trt_high_precision_dtype in torch_dtype_map, ( - f"Unsupported high precision dtype: {trt_high_precision_dtype}" - ) - scale_inv = (amax / 7.0).to(torch_dtype_map[trt_high_precision_dtype]) - scale_inv_op = g.op("Constant", value_t=scale_inv) return g.op( "trt::DequantizeLinear", inputs, scale_inv_op, axis_i=axis, block_size_i=block_size ).setType( @@ -225,13 +216,14 @@ def _fp8_quantize( g: "GraphContext", inputs: torch.Value, scale_inv: float, - output_dtype: str, ): """Helper Function for Quantization.""" + # Emit the scale in the native input dtype so no Cast is inserted between the + # graph and Q/DQ (Cast nodes block TRT from fusing DQ into the MatMul kernel). output_shape = sym_help._get_tensor_sizes(inputs) scale = g.op( "Constant", - value_t=torch.tensor(scale_inv, dtype=torch_dtype_map[output_dtype]), + value_t=torch.tensor(scale_inv).to(torch_dtype_map[inputs.type().scalarType()]), ) return g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) @@ -263,19 +255,15 @@ def export_fp8( ): """Export quantized model to FP8 ONNX. - ``None`` preserves the native input dtype. + ``trt_high_precision_dtype`` is accepted for API compatibility but unused: Q/DQ now + emit scales in the native input dtype, so no intermediate Cast is required. """ + del trt_high_precision_dtype scale = 1.0 if amax is None else 448.0 / float(amax) - input_dtype = inputs.type().scalarType() - if trt_high_precision_dtype is None: - trt_high_precision_dtype = input_dtype - assert trt_high_precision_dtype in torch_dtype_map, ( - f"Unsupported high precision dtype: {trt_high_precision_dtype}" - ) + otype = inputs.type().scalarType() - inputs = _cast_to_dtype(g, inputs, trt_high_precision_dtype) - q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype) - return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype) + q_tensor = _fp8_quantize(g, inputs, 1.0 / scale) + return _fp8_dequantize(g, q_tensor, 1.0 / scale, otype) def scaled_dot_product_attention( @@ -374,7 +362,7 @@ def export_fp8_mha( q_quantized_scale: float = 1.0, k_quantized_scale: float = 1.0, v_quantized_scale: float = 1.0, - high_precision_flag: str | None = "Half", + high_precision_flag: str = "Half", disable_fp8_mha: bool = True, ): r"""Export quantized fMHA to FP8 ONNX. @@ -420,22 +408,6 @@ def export_fp8_mha( "is_causal and attn_mask cannot be set at the same time" ) - if not disable_fp8_mha: - if high_precision_flag is None: - high_precision_flag = query.type().scalarType() - if high_precision_flag not in mha_supported_precisions: - raise ValueError(f"Unsupported FP8 MHA precision: {high_precision_flag}") - if high_precision_flag == "Float": - query = _cast_to_dtype(g, query, high_precision_flag) - key = _cast_to_dtype(g, key, high_precision_flag) - value = _cast_to_dtype(g, value, high_precision_flag) - elif high_precision_flag in mha_fusion_precisions and { - query.type().scalarType(), - key.type().scalarType(), - value.type().scalarType(), - } != {high_precision_flag}: - raise ValueError("The quantized MHA must have 16-bit inputs.") - scale = sym_help._maybe_get_const(scale, "f") if sym_help._is_none(scale): scale = _attention_scale(g, query) @@ -459,15 +431,24 @@ def export_fp8_mha( query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) if not disable_fp8_mha: + if high_precision_flag not in mha_valid_precisions: + raise ValueError( + "The Quantized config setting doesn't match TRT's fusion pattern; the qdqs must be in 16 bits." + ) + q_input_dtype = query.type().scalarType() + k_input_dtype = key.type().scalarType() + v_input_dtype = value.type().scalarType() + if {q_input_dtype, k_input_dtype, v_input_dtype} != {high_precision_flag}: + raise ValueError("The quantized MHA must have 16-bit inputs.") query_scaled = export_fp8(g, query_scaled, q_quantized_scale, high_precision_flag) - query_scaled = _cast_to_dtype(g, query_scaled, "Float") + query_scaled = g.op("Cast", query_scaled, to_i=onnx_dtype_map["Float"]) key_transposed_scaled = export_fp8( g, key_transposed_scaled, k_quantized_scale, high_precision_flag ) - key_transposed_scaled = _cast_to_dtype(g, key_transposed_scaled, "Float") + key_transposed_scaled = g.op("Cast", key_transposed_scaled, to_i=onnx_dtype_map["Float"]) mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) if not disable_fp8_mha: - mul_qk = _cast_to_dtype(g, mul_qk, cast("str", high_precision_flag)) + mul_qk = g.op("Cast", mul_qk, to_i=onnx_dtype_map[high_precision_flag]) if sym_help._is_none(attn_mask): mul_qk_add = mul_qk @@ -491,7 +472,7 @@ def export_fp8_mha( if not disable_fp8_mha: # Softmax's output scale is hard coded to 1.0 attn_weight = export_fp8(g, attn_weight, 1.0, high_precision_flag) - attn_weight = _cast_to_dtype(g, attn_weight, "Float") + attn_weight = g.op("Cast", attn_weight, to_i=onnx_dtype_map["Float"]) if dropout_p != 0: attn_weight = g.op( @@ -501,9 +482,11 @@ def export_fp8_mha( ) if not disable_fp8_mha: value = export_fp8(g, value, v_quantized_scale, high_precision_flag) - value = _cast_to_dtype(g, value, "Float") - return _cast_to_dtype( - g, g.op("MatMul", attn_weight, value), cast("str", high_precision_flag) + value = g.op("Cast", value, to_i=onnx_dtype_map["Float"]) + return g.op( + "Cast", + g.op("MatMul", attn_weight, value), + to_i=onnx_dtype_map[high_precision_flag], ) else: return g.op("MatMul", attn_weight, value) @@ -519,7 +502,7 @@ def _fp4_dynamic_quantize( scale_type: int = onnx_dtype_map["Float8"], ): """Helper Function for Dynamic Quantization.""" - # Match the input to the requested strongly typed QDQ precision. + # TRT StronglyType only supports FP16 QDQ ops, so cast the input if needed. input_type = inputs.type().scalarType() if trt_high_precision_dtype is None: trt_high_precision_dtype = input_type @@ -613,37 +596,16 @@ def export_mxfp8( onnx_quantizer_type: str, block_size: int, axis: int = -1, - trt_high_precision_dtype: str | None = None, ): """Export quantized model to MXFP8 ONNX.""" input_dtype = inputs.type().scalarType() - output_shape = sym_help._get_tensor_sizes(inputs) - if trt_high_precision_dtype is None: - trt_high_precision_dtype = input_dtype - assert trt_high_precision_dtype in torch_dtype_map, ( - f"Unsupported high precision dtype: {trt_high_precision_dtype}" - ) - if onnx_quantizer_type == "dynamic": - inputs = _cast_to_dtype(g, inputs, trt_high_precision_dtype) x_f8, sx_ui8 = _mxfp8_dynamic_quantize(g, inputs, block_size, axis=axis) - output = _mxfp8_dequantize( - g, - x_f8, - sx_ui8, - block_size, - axis=axis, - input_dtype=trt_high_precision_dtype, - ) + return _mxfp8_dequantize(g, x_f8, sx_ui8, block_size, axis=axis, input_dtype=input_dtype) else: - scale = torch.tensor(1.0, dtype=torch_dtype_map[trt_high_precision_dtype]) - output = _mxfp8_dequantize( - g, inputs, scale, block_size, axis=axis, input_dtype=trt_high_precision_dtype - ) - return output.setType( - inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) - ) + scale = torch.tensor(1.0, dtype=torch_dtype_map[input_dtype]) + return _mxfp8_dequantize(g, inputs, scale, block_size, axis=axis, input_dtype=input_dtype) def export_fp4( diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index aab54ef25a7..20e083491aa 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -530,7 +530,6 @@ def symbolic( inputs, onnx_quantizer_type, block_size, - trt_high_precision_dtype=trt_high_precision_dtype, ) raise NotImplementedError( f"Unsupported num_bits: {num_bits} and scale_bits: {scale_bits} for ONNX export." diff --git a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py index 5900ce43dbf..1f7251a9ad9 100644 --- a/tests/unit/onnx/quantization/test_fp8_mha_exporter.py +++ b/tests/unit/onnx/quantization/test_fp8_mha_exporter.py @@ -15,52 +15,15 @@ """Tests for the attention-aware FP8 ONNX graph rewrites in ``FP8QuantExporter``.""" -import io - -import ml_dtypes import numpy as np -import onnx import onnx_graphsurgeon as gs import pytest -import torch -from torch.onnx import symbolic_helper from modelopt.onnx.export.fp8_exporter import FP8QuantExporter -from modelopt.torch.quantization.export_onnx import export_fp8_mha - - -class _FP8MHAFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, query, key, value, high_precision_dtype): - return torch.nn.functional.scaled_dot_product_attention(query, key, value) - - @staticmethod - @symbolic_helper.parse_args("v", "v", "v", "s") - def symbolic(g, query, key, value, high_precision_dtype): - return export_fp8_mha( - g, - query, - key, - value, - q_quantized_scale=1.0, - k_quantized_scale=1.0, - v_quantized_scale=1.0, - high_precision_flag=high_precision_dtype, - disable_fp8_mha=False, - ) - - -class _FP8MHAModule(torch.nn.Module): - def __init__(self, high_precision_dtype): - super().__init__() - self.high_precision_dtype = high_precision_dtype - - def forward(self, query, key, value): - return _FP8MHAFunction.apply(query, key, value, self.high_precision_dtype) -def _var(name, dtype=np.float32, shape=None): - return gs.Variable(name, dtype=dtype, shape=shape) +def _var(name): + return gs.Variable(name, dtype=np.float32) def _qdq(src): @@ -107,47 +70,16 @@ def test_move_transpose_before_qdq_rewrites_dq_transpose_matmul_pattern(): assert q.inputs[0].inputs[0].op == "Transpose" -@pytest.mark.parametrize( - ("high_precision_dtype", "numpy_dtype", "onnx_dtype"), - [ - pytest.param(None, np.float32, onnx.TensorProto.FLOAT, id="legacy"), - pytest.param("Float", np.float32, onnx.TensorProto.FLOAT, id="float"), - pytest.param("Half", np.float16, onnx.TensorProto.FLOAT16, id="half"), - pytest.param("BFloat16", ml_dtypes.bfloat16, onnx.TensorProto.BFLOAT16, id="bfloat16"), - ], -) -def test_insert_qdq_after_softmax_adds_target_scale_q_dq( - high_precision_dtype, numpy_dtype, onnx_dtype -): +def test_insert_qdq_after_softmax_adds_fixed_scale_q_dq(): """Softmax → MatMul picks up ``Q → DQ`` with the fixed ``1/448`` scale.""" - scores, v, y, sm_out = ( - _var("scores", numpy_dtype, [2, 2]), - _var("v", numpy_dtype, [2, 2]), - _var("y", numpy_dtype, [2, 2]), - _var("sm_out", numpy_dtype, [2, 2]), - ) + scores, v, y, sm_out = _var("scores"), _var("v"), _var("y"), _var("sm_out") sm = gs.Node(op="Softmax", inputs=[scores], outputs=[sm_out], attrs={"axis": -1}) mm = gs.Node(op="MatMul", inputs=[sm_out, v], outputs=[y]) graph = _graph([sm, mm], [scores, v], [y]) - count = ( - FP8QuantExporter._insert_qdq_after_softmax(graph) - if high_precision_dtype is None - else FP8QuantExporter._insert_qdq_after_softmax(graph, high_precision_dtype) - ) - assert count == 1 + assert FP8QuantExporter._insert_qdq_after_softmax(graph) == 1 q = next(n for n in graph.nodes if n.op == "QuantizeLinear") - dq = next(n for n in graph.nodes if n.op == "DequantizeLinear") - expected_scale = np.array(1.0 / 448.0, dtype=numpy_dtype) - for scale in (q.inputs[1], dq.inputs[1]): - assert scale.values.dtype == expected_scale.dtype - np.testing.assert_array_equal(scale.values, expected_scale) - assert np.dtype(dq.outputs[0].dtype) == np.dtype(numpy_dtype) - assert mm.inputs[0] is dq.outputs[0] - - converted_model = gs.export_onnx(graph) - onnx.checker.check_model(converted_model) - onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) + assert np.isclose(float(q.inputs[1].values), 1.0 / 448.0) @pytest.mark.parametrize( @@ -184,71 +116,3 @@ def test_rewrites_skip_when_non_matmul_consumer_exists(rewrite): [y_mm, y_side], ) assert getattr(FP8QuantExporter, rewrite)(graph) == 0 - - -@pytest.mark.parametrize( - ("torch_dtype", "high_precision_dtype", "onnx_dtype", "expected_accumulation_casts"), - [ - pytest.param(torch.float32, "Float", onnx.TensorProto.FLOAT, 0, id="float"), - pytest.param(torch.float16, "Half", onnx.TensorProto.FLOAT16, 4, id="half"), - pytest.param(torch.bfloat16, "BFloat16", onnx.TensorProto.BFLOAT16, 4, id="bfloat16"), - pytest.param(torch.float32, None, onnx.TensorProto.FLOAT, 0, id="native-float"), - pytest.param(torch.bfloat16, None, onnx.TensorProto.BFLOAT16, 4, id="native-bfloat16"), - ], -) -def test_fp8_mha_symbolic_preserves_accumulation_contract( - torch_dtype, high_precision_dtype, onnx_dtype, expected_accumulation_casts -): - """FP8-MHA supports FP32 while retaining 16-bit fusion casts.""" - shape = (1, 1, 2, 4) - inputs = tuple(torch.ones(shape, dtype=torch_dtype) for _ in range(3)) - buffer = io.BytesIO() - torch.onnx.export( - _FP8MHAModule(high_precision_dtype), - inputs, - buffer, - opset_version=20, - dynamo=False, - ) - - model = onnx.load_model_from_string(buffer.getvalue()) - onnx.checker.check_model(model) - onnx.shape_inference.infer_shapes(model, check_type=True, strict_mode=True) - - qdq_ops = {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} - qdq_nodes = [node for node in model.graph.node if node.op_type in qdq_ops] - assert len(qdq_nodes) == 8 - - tensor_dtype = { - initializer.name: initializer.data_type for initializer in model.graph.initializer - } - for node in model.graph.node: - if node.op_type == "Constant": - tensor = next((attr.t for attr in node.attribute if attr.name == "value"), None) - if tensor is not None: - tensor_dtype[node.output[0]] = tensor.data_type - assert all(tensor_dtype[node.input[1]] == onnx_dtype for node in qdq_nodes) - - producer_by_output = { - output: node for node in model.graph.node for output in node.output if output - } - accumulation_casts = [ - node - for node in model.graph.node - if node.op_type == "Cast" - and any(attr.name == "to" and attr.i == onnx.TensorProto.FLOAT for attr in node.attribute) - and producer_by_output.get(node.input[0], onnx.NodeProto()).op_type in qdq_ops - ] - assert len(accumulation_casts) == expected_accumulation_casts - back_casts = [ - node - for node in model.graph.node - if node.op_type == "Cast" - and any(attr.name == "to" and attr.i == onnx_dtype for attr in node.attribute) - and producer_by_output.get(node.input[0], onnx.NodeProto()).op_type == "MatMul" - ] - assert len(back_casts) == (0 if onnx_dtype == onnx.TensorProto.FLOAT else 2) - assert all( - value.type.tensor_type.elem_type == onnx_dtype - for value in [*model.graph.input, *model.graph.output] - ) diff --git a/tests/unit/onnx/quantization/test_gs_patching.py b/tests/unit/onnx/quantization/test_gs_patching.py deleted file mode 100644 index 95adbde79a0..00000000000 --- a/tests/unit/onnx/quantization/test_gs_patching.py +++ /dev/null @@ -1,67 +0,0 @@ -# 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. - -import ml_dtypes -import numpy as np -import onnx -import onnx_graphsurgeon as gs -import pytest - -from modelopt.onnx.quantization.gs_patching import _export_tensor_proto, _export_value_info_proto - - -def test_export_constant_uses_explicit_bf16_dtype_without_fallback(monkeypatch): - tensor = gs.Constant("scale", np.array([1.0], dtype=ml_dtypes.bfloat16)) - tensor.explicit_dtype = onnx.TensorProto.BFLOAT16 - - def fail_on_fallback(_): - raise AssertionError("NumPy dtype fallback should not be evaluated") - - monkeypatch.setattr(onnx.helper, "np_dtype_to_tensor_dtype", fail_on_fallback) - - tensor_proto = _export_tensor_proto(tensor) - - assert tensor_proto.data_type == onnx.TensorProto.BFLOAT16 - np.testing.assert_array_equal(onnx.numpy_helper.to_array(tensor_proto), tensor.values) - - -def test_export_value_info_uses_explicit_onnx_dtype_without_numpy_conversion(): - tensor = gs.Variable("input", dtype=onnx.TensorProto.BFLOAT16, shape=[1]) - tensor.explicit_dtype = onnx.TensorProto.BFLOAT16 - - value_info = _export_value_info_proto(tensor, do_type_check=True) - - assert value_info.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 - - -def test_export_value_info_accepts_onnx_dtype_without_explicit_dtype(): - tensor = gs.Variable("input", dtype=onnx.TensorProto.BFLOAT16, shape=[1]) - - value_info = _export_value_info_proto(tensor, do_type_check=True) - - assert value_info.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 - - -def test_export_rejects_unknown_integer_dtype(): - invalid_dtype = max(onnx.TensorProto.DataType.values()) + 1 - constant = gs.Constant("scale", np.array([1.0], dtype=np.float32)) - constant.explicit_dtype = invalid_dtype - - with pytest.raises(ValueError): - _export_tensor_proto(constant) - - variable = gs.Variable("input", dtype=invalid_dtype, shape=[1]) - with pytest.raises(ValueError): - _export_value_info_proto(variable, do_type_check=True) diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index f2413b82d6a..7411c280306 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -30,6 +30,7 @@ NVFP4QuantExporter, ) from modelopt.onnx.export.nvfp4_exporter import _cast_fp4 +from modelopt.onnx.quantization.gs_patching import _export_value_info_proto from modelopt.onnx.quantization.qdq_utils import ( _cast_fp8, apply_column_major_transformation, @@ -172,7 +173,7 @@ def create_test_model_with_cast_nodes(): return model -def create_test_model_with_proj_nodes(graph_dtype=TensorProto.FLOAT): +def create_test_model_with_proj_nodes(): """Create a test model with projection nodes to test bias and scale casting.""" # Create bias tensor bias_data = np.random.uniform(-1.0, 1.0, size=(16,)).astype(np.float32) @@ -182,7 +183,7 @@ def create_test_model_with_proj_nodes(graph_dtype=TensorProto.FLOAT): scale_data = np.random.uniform(0.1, 1.0, size=(1,)).astype(np.float32) scale_tensor = numpy_helper.from_array(scale_data, "quant_scale") - input_tensor = helper.make_tensor_value_info("input", graph_dtype, [4, 16]) + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, 16]) # Add node (projection bias) add_node = helper.make_node( @@ -201,7 +202,7 @@ def create_test_model_with_proj_nodes(graph_dtype=TensorProto.FLOAT): nodes=[add_node, mul_node], name="test_graph", inputs=[input_tensor], - outputs=[helper.make_tensor_value_info("output", graph_dtype, [4, 16])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [4, 16])], initializer=[bias_tensor, scale_tensor], ) @@ -397,38 +398,24 @@ def test_quantization_with_constant_scale(self): ) assert any("scale" in input_name for input_name in dq_node.input) - @pytest.mark.parametrize( - ("high_precision_dtype", "onnx_dtype"), - [ - pytest.param(None, TensorProto.FLOAT16, id="legacy"), - pytest.param("Float", TensorProto.FLOAT, id="float"), - pytest.param("Half", TensorProto.FLOAT16, id="half"), - pytest.param("BFloat16", TensorProto.BFLOAT16, id="bfloat16"), - ], - ) - def test_projection_bias_and_scale_casting(self, high_precision_dtype, onnx_dtype): - """Test projection bias and pre-quant scale target casting.""" - graph_dtype = TensorProto.FLOAT if high_precision_dtype is None else onnx_dtype - model = create_test_model_with_proj_nodes(graph_dtype) - - quantized_model = ( - INT4QuantExporter.post_process(model) - if high_precision_dtype is None - else INT4QuantExporter.post_process(model, high_precision_dtype) - ) + def test_projection_bias_and_scale_casting(self): + """Test that projection biases and quantization scales are cast to float16.""" + model = create_test_model_with_proj_nodes() + # Run quantization + quantized_model = INT4QuantExporter.process_model(model) + + # Verify bias tensor is cast to float16 bias_tensor = next( init for init in quantized_model.graph.initializer if "proj_bias" in init.name ) - assert bias_tensor.data_type == onnx_dtype + assert bias_tensor.data_type == TensorProto.FLOAT16 + # Verify quantization scale is cast to float16 scale_tensor = next( init for init in quantized_model.graph.initializer if "quant_scale" in init.name ) - assert scale_tensor.data_type == onnx_dtype - onnx.checker.check_model(quantized_model) - if high_precision_dtype is not None: - onnx.shape_inference.infer_shapes(quantized_model, check_type=True, strict_mode=True) + assert scale_tensor.data_type == TensorProto.FLOAT16 class TestCastFunctions: @@ -505,14 +492,19 @@ def test_cast_fp4(self, input_array, expected_array): assert np.all(result == expected_array) +def test_graphsurgeon_value_info_accepts_onnx_dtype(): + tensor = gs.Variable("input", dtype=TensorProto.BFLOAT16, shape=[1]) + value_info = _export_value_info_proto(tensor, do_type_check=True) + + assert value_info.type.tensor_type.elem_type == TensorProto.BFLOAT16 + + class TestFP8QuantExporter: """Test suite for FP8QuantExporter.""" def test_bf16_weights_and_scale_are_compressed(self): - weight_data = np.array([[0.25, -0.5], [1.0, -2.0]], dtype=np.float32).astype( - ml_dtypes.bfloat16 - ) - scale_data = np.array(0.25, dtype=np.float32).astype(ml_dtypes.bfloat16) + weight_data = np.array([0.001312255859375], dtype=ml_dtypes.bfloat16) + scale_data = np.array(0.00099945068359375, dtype=ml_dtypes.bfloat16) weight = gs.Constant("weight", weight_data) scale = gs.Constant("linear/weight_quantizer/scale", scale_data) quantized = gs.Variable("quantized", dtype=np.uint8, shape=weight_data.shape) @@ -544,7 +536,7 @@ def test_bf16_weights_and_scale_are_compressed(self): if initializer.name == "linear/weight_quantizer/fp8_weights" ) assert fp8_weight.data_type == TensorProto.FLOAT8E4M3FN - assert fp8_weight.raw_data == bytes.fromhex("38 c0 48 d0") + assert fp8_weight.raw_data == b"\x3a" output_scale = next( initializer for initializer in converted_model.graph.initializer @@ -552,45 +544,6 @@ def test_bf16_weights_and_scale_are_compressed(self): ) assert output_scale.data_type == TensorProto.BFLOAT16 - def test_conv_uses_target_rounded_weights_and_scale(self): - weight_data = np.array([0.58945024, -13.944608], dtype=np.float32).reshape(1, 1, 1, 2) - input_tensor = gs.Variable("input", dtype=np.float32, shape=[1, 1, 1, 2]) - output_tensor = gs.Variable("output", dtype=np.float32, shape=[1, 1, 1, 1]) - conv = gs.Node( - op="Conv", - name="conv", - inputs=[input_tensor, gs.Constant("weight", weight_data)], - outputs=[output_tensor], - ) - graph = gs.Graph(nodes=[conv], inputs=[input_tensor], outputs=[output_tensor], opset=23) - - assert FP8QuantExporter._quantize_conv_weights_to_fp8(graph, "BFloat16") == 1 - - dq_node = next(node for node in graph.nodes if node.op == "DequantizeLinear") - fp8_weights, scale = dq_node.inputs - assert scale.values.dtype == ml_dtypes.bfloat16 - assert dq_node.outputs[0].dtype == ml_dtypes.bfloat16 - assert fp8_weights._values.tensor.raw_data == bytes([90, 254]) - - def test_conv_clamps_target_rounded_zero_scale_before_quantizing(self): - weight_data = np.array([1e-6, -2e-6], dtype=np.float32).reshape(1, 1, 1, 2) - input_tensor = gs.Variable("input", dtype=np.float16, shape=[1, 1, 1, 2]) - output_tensor = gs.Variable("output", dtype=np.float16, shape=[1, 1, 1, 1]) - conv = gs.Node( - op="Conv", - name="conv", - inputs=[input_tensor, gs.Constant("weight", weight_data)], - outputs=[output_tensor], - ) - graph = gs.Graph(nodes=[conv], inputs=[input_tensor], outputs=[output_tensor], opset=23) - - assert FP8QuantExporter._quantize_conv_weights_to_fp8(graph, "Half") == 1 - - dq_node = next(node for node in graph.nodes if node.op == "DequantizeLinear") - fp8_weights, scale = dq_node.inputs - assert scale.values == np.finfo(np.float16).smallest_subnormal - assert fp8_weights._values.tensor.raw_data == bytes([88, 224]) - class TestMXFP8QuantExporter: """Test suite for MXFP8QuantExporter.""" @@ -642,52 +595,6 @@ def test_mxfp8_output_dtype_update(self): output_dtype_attr = next(attr for attr in dq_node.attribute if attr.name == "output_dtype") assert output_dtype_attr.i == TensorProto.FLOAT16 - @pytest.mark.parametrize( - ("high_precision_dtype", "onnx_dtype", "suffix"), - [ - pytest.param(None, TensorProto.FLOAT16, "fp16", id="legacy"), - pytest.param("Float", TensorProto.FLOAT, "fp32", id="float"), - pytest.param("Half", TensorProto.FLOAT16, "fp16", id="half"), - pytest.param("BFloat16", TensorProto.BFLOAT16, "bf16", id="bfloat16"), - ], - ) - def test_sqrt_output_cast_uses_target_dtype(self, high_precision_dtype, onnx_dtype, suffix): - graph_dtype = TensorProto.FLOAT if high_precision_dtype is None else onnx_dtype - input_info = helper.make_tensor_value_info("input", graph_dtype, [2]) - output_info = helper.make_tensor_value_info("output", graph_dtype, [2]) - sqrt_node = helper.make_node("Sqrt", inputs=["input"], outputs=["sqrt_output"], name="sqrt") - consumer = helper.make_node( - "Identity", inputs=["sqrt_output"], outputs=["output"], name="consumer" - ) - model = helper.make_model( - helper.make_graph( - [sqrt_node, consumer], - "sqrt_graph", - [input_info], - [output_info], - ) - ) - - converted_model = ( - MXFP8QuantExporter.post_process(model) - if high_precision_dtype is None - else MXFP8QuantExporter.post_process(model, high_precision_dtype) - ) - - cast_node = next(node for node in converted_model.graph.node if node.op_type == "Cast") - cast_to = next(attr.i for attr in cast_node.attribute if attr.name == "to") - assert cast_to == onnx_dtype - assert cast_node.name == f"sqrt_cast_{suffix}" - assert cast_node.input == ["sqrt_output"] - assert cast_node.output == [f"sqrt_output_cast_{suffix}"] - converted_consumer = next( - node for node in converted_model.graph.node if node.name == "consumer" - ) - assert converted_consumer.input == cast_node.output - onnx.checker.check_model(converted_model) - if high_precision_dtype is not None: - onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) - def test_mxfp8_gelu_approximation_update(self): """Test that Gelu nodes are updated to use tanh approximation.""" model = create_test_model_with_mxfp8_dq() @@ -800,159 +707,6 @@ def test_fp4qdq_conversion(self, with_transpose): cast_nodes = [node for node in converted_model.graph.node if node.op_type == "Cast"] assert len(cast_nodes) >= 1 # At least one cast node should be added - @pytest.mark.parametrize( - ("precision_dtype", "onnx_dtype"), - [("Half", TensorProto.FLOAT16), ("BFloat16", TensorProto.BFLOAT16)], - ) - def test_existing_weight_cast_does_not_hide_matmul(self, precision_dtype, onnx_dtype): - weight_data = np.linspace(-1.0, 1.0, 8 * 32, dtype=np.float32).reshape(8, 32) - weight = numpy_helper.from_array(weight_data, "linear.weight") - fp4qdq = helper.make_node( - "TRT_FP4QDQ", - inputs=[weight.name], - outputs=["fp4qdq_output"], - name="weight_fp4qdq", - block_size=16, - ) - cast = helper.make_node( - "Cast", - inputs=["fp4qdq_output"], - outputs=["weight_cast"], - name="weight_cast", - to=onnx_dtype, - ) - transpose = helper.make_node( - "Transpose", - inputs=["weight_cast"], - outputs=["weight_transposed"], - name="weight_transpose", - perm=[1, 0], - ) - matmul = helper.make_node( - "MatMul", - inputs=["activation", "weight_transposed"], - outputs=["output"], - name="matmul", - ) - graph = helper.make_graph( - [fp4qdq, cast, transpose, matmul], - "nvfp4_cast_graph", - [helper.make_tensor_value_info("activation", TensorProto.FLOAT, [1, 32])], - [helper.make_tensor_value_info("output", onnx_dtype, [1, 8])], - [weight], - value_info=[ - helper.make_tensor_value_info("fp4qdq_output", TensorProto.FLOAT, [8, 32]), - helper.make_tensor_value_info("weight_cast", onnx_dtype, [8, 32]), - helper.make_tensor_value_info("weight_transposed", onnx_dtype, [32, 8]), - ], - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 23)]) - - converted_model = NVFP4QuantExporter.process_model(model, precision_dtype) - - onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) - converted_matmul = next( - node for node in converted_model.graph.node if node.op_type == "MatMul" - ) - producer_map = { - output: node for node in converted_model.graph.node for output in node.output - } - assert all( - producer_map[input_name].op_type == "Cast" - and helper.get_attribute_value(producer_map[input_name].attribute[0]) == onnx_dtype - for input_name in converted_matmul.input - ) - - def test_shared_weight_casts_all_direct_and_wrapped_linear_consumers(self): - weight_data = np.linspace(-1.0, 1.0, 32 * 32, dtype=np.float32).reshape(32, 32) - weight = numpy_helper.from_array(weight_data, "linear.weight") - fp4qdq = helper.make_node( - "TRT_FP4QDQ", - inputs=[weight.name], - outputs=["fp4qdq_output"], - name="weight_fp4qdq", - block_size=16, - ) - cast = helper.make_node( - "Cast", - inputs=["fp4qdq_output"], - outputs=["weight_cast"], - name="weight_cast", - to=TensorProto.FLOAT16, - ) - transpose_nodes = [ - helper.make_node( - "Transpose", - inputs=["weight_cast"], - outputs=[f"weight_transposed_{index}"], - name=f"weight_transpose_{index}", - perm=[1, 0], - ) - for index in range(2) - ] - matmul_nodes = [ - *[ - helper.make_node( - "MatMul", - inputs=[f"activation_{index}", "fp4qdq_output"], - outputs=[f"output_{index}"], - name=f"matmul_{index}", - ) - for index in range(2) - ], - *[ - helper.make_node( - "MatMul", - inputs=[f"activation_{index}", f"weight_transposed_{index - 2}"], - outputs=[f"output_{index}"], - name=f"matmul_{index}", - ) - for index in range(2, 4) - ], - ] - graph = helper.make_graph( - [fp4qdq, cast, *transpose_nodes, *matmul_nodes], - "nvfp4_fanout_graph", - [ - helper.make_tensor_value_info(f"activation_{index}", TensorProto.FLOAT, [1, 32]) - for index in range(4) - ], - [ - helper.make_tensor_value_info(f"output_{index}", TensorProto.FLOAT16, [1, 32]) - for index in range(4) - ], - [weight], - value_info=[ - helper.make_tensor_value_info("fp4qdq_output", TensorProto.FLOAT, [32, 32]), - helper.make_tensor_value_info("weight_cast", TensorProto.FLOAT16, [32, 32]), - *[ - helper.make_tensor_value_info( - f"weight_transposed_{index}", TensorProto.FLOAT16, [32, 32] - ) - for index in range(2) - ], - ], - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 23)]) - - converted_model = NVFP4QuantExporter.process_model(model, "Half") - - onnx.shape_inference.infer_shapes(converted_model, check_type=True, strict_mode=True) - producer_map = { - output: node for node in converted_model.graph.node for output in node.output - } - converted_matmuls = [ - node for node in converted_model.graph.node if node.op_type == "MatMul" - ] - assert len(converted_matmuls) == 4 - assert all( - producer_map[input_name].op_type == "Cast" - and helper.get_attribute_value(producer_map[input_name].attribute[0]) - == TensorProto.FLOAT16 - for node in converted_matmuls - for input_name in node.input - ) - def create_test_model_with_int4_dq_matmul(): """Create a simple test model with INT4 DequantizeLinear -> MatMul pattern. @@ -1324,7 +1078,7 @@ def test_column_major_gemm_trans_b_flip(self): print(f"Transpose nodes: {len(transpose_nodes)}") -def _build_model_with_zero_scale_initializer(dq_op_type: str, scale_dtype=np.float16): +def _build_model_with_zero_scale_initializer(dq_op_type: str): """Build an ONNX model whose scale initializer feeds a (Quantize|Dequantize)Linear node. Mirrors the INT4_AWQ failure mode from NVBug 6110209: scales live in graph initializers @@ -1333,11 +1087,10 @@ def _build_model_with_zero_scale_initializer(dq_op_type: str, scale_dtype=np.flo weight_data = np.random.randint(-8, 8, size=(6, 8), dtype=np.int8) weight_tensor = numpy_helper.from_array(weight_data, "weight") - scale_data = np.array([1e-3, 0.0, 5e-4, 0.0, 0.0, 2e-3], dtype=scale_dtype).reshape(6, 1) + scale_data = np.array([1e-3, 0.0, 5e-4, 0.0, 0.0, 2e-3], dtype=np.float16).reshape(6, 1) scale_tensor = numpy_helper.from_array(scale_data, "scale") - high_precision_dtype = helper.np_dtype_to_tensor_dtype(np.dtype(scale_dtype)) - input_tensor = helper.make_tensor_value_info("input", high_precision_dtype, [None, 6]) + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT16, [None, 6]) dq_node = helper.make_node( dq_op_type, inputs=["weight", "scale"], outputs=["dq_output"], name="weight_dq" ) @@ -1348,7 +1101,7 @@ def _build_model_with_zero_scale_initializer(dq_op_type: str, scale_dtype=np.flo nodes=[dq_node, matmul_node], name="test_graph", inputs=[input_tensor], - outputs=[helper.make_tensor_value_info("output", high_precision_dtype, [None, 8])], + outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT16, [None, 8])], initializer=[weight_tensor, scale_tensor], ) return helper.make_model(graph) @@ -1358,17 +1111,8 @@ class TestReplaceZeroScaleWithSmallestNonzero: """Regression tests for ``replace_zero_scale_with_smallest_nonzero`` (NVBug 6110209).""" @pytest.mark.parametrize("dq_op_type", ["DequantizeLinear", "TRT_INT4DequantizeLinear"]) - @pytest.mark.parametrize( - ("scale_dtype", "onnx_dtype"), - [ - (np.float16, TensorProto.FLOAT16), - (ml_dtypes.bfloat16, TensorProto.BFLOAT16), - (np.float32, TensorProto.FLOAT), - (np.float64, TensorProto.DOUBLE), - ], - ) - def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type, scale_dtype, onnx_dtype): - model = _build_model_with_zero_scale_initializer(dq_op_type, scale_dtype) + def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type): + model = _build_model_with_zero_scale_initializer(dq_op_type) scale_before = numpy_helper.to_array( next(init for init in model.graph.initializer if init.name == "scale") ) @@ -1380,14 +1124,7 @@ def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type, scale_dty scale_after = numpy_helper.to_array(scale_after_init) assert not (scale_after == 0).any() assert (scale_after > 0).all() - assert scale_after_init.data_type == onnx_dtype - - dtype_info = ( - ml_dtypes.finfo(scale_dtype) - if scale_dtype == ml_dtypes.bfloat16 - else np.finfo(scale_dtype) - ) - assert (scale_after[scale_before == 0] == dtype_info.smallest_subnormal).all() + assert scale_after_init.data_type == TensorProto.FLOAT16 def test_constant_node_scale_path_still_patched(self): """Legacy Constant-node QDQ path must continue to be patched.""" @@ -1424,58 +1161,6 @@ def test_constant_node_scale_path_still_patched(self): assert not (scale_arr == 0).any() assert (scale_arr > 0).all() - def test_captured_parent_scale_is_patched_without_crossing_child_scope(self): - captured_scale = numpy_helper.from_array( - np.array(0.0, dtype=ml_dtypes.bfloat16), "captured_scale" - ) - shadowed_scale = numpy_helper.from_array( - np.array(0.0, dtype=ml_dtypes.bfloat16), "shadowed_scale" - ) - subgraph = helper.make_graph( - [ - helper.make_node( - "QuantizeLinear", - ["data", "captured_scale"], - ["captured_output"], - ), - helper.make_node( - "QuantizeLinear", - ["data", "shadowed_scale"], - ["shadowed_output"], - ), - ], - "subgraph", - [ - helper.make_tensor_value_info("data", TensorProto.BFLOAT16, [1]), - helper.make_tensor_value_info("shadowed_scale", TensorProto.BFLOAT16, []), - ], - [ - helper.make_tensor_value_info("captured_output", TensorProto.UINT8, [1]), - helper.make_tensor_value_info("shadowed_output", TensorProto.UINT8, [1]), - ], - ) - scoped_node = helper.make_node( - "ScopedSubgraph", - [], - [], - domain="test", - body=subgraph, - ) - graph = helper.make_graph( - [scoped_node], - "parent_graph", - [], - [], - [captured_scale, shadowed_scale], - ) - model = helper.make_model(graph) - - patched = replace_zero_scale_with_smallest_nonzero(model) - - scales = {init.name: numpy_helper.to_array(init) for init in patched.graph.initializer} - assert scales["captured_scale"] == ml_dtypes.finfo(ml_dtypes.bfloat16).smallest_subnormal - assert scales["shadowed_scale"] == 0 - class TestQdqToDqValidation: """Regression tests for qdq_to_dq input validation.""" 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 60073c65b95..5ca548c645d 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json from contextlib import nullcontext @@ -24,11 +25,7 @@ import torch.nn as nn from _test_utils.torch.deploy.lib_test_models import BaseDeployModel, get_deploy_models -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.autocast.convert import convert_to_fp32 -from modelopt.onnx.export.base_exporter import ONNXQuantExporter from modelopt.onnx.utils import get_batch_size_from_bytes, validate_batch_size from modelopt.torch._deploy.utils import ( OnnxBytes, @@ -36,12 +33,7 @@ generate_onnx_input, get_onnx_bytes_and_metadata, ) -from modelopt.torch._deploy.utils.torch_onnx import ( - _get_autocast_context, - _override_onnx_quantizer_precision, - _to_expected_onnx_type, -) -from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch._deploy.utils.torch_onnx import _to_expected_onnx_type from modelopt.torch.utils import standardize_model_args, unflatten_tree deploy_benchmark_all = get_deploy_models() @@ -66,22 +58,6 @@ k: v for k, v in deploy_benchmark_dynamo.items() if k in _DYNAMO_REPRESENTATIVE_MODELS } -_ONNX_DTYPE_BY_NAME = { - "fp32": onnx.TensorProto.FLOAT, - "fp16": onnx.TensorProto.FLOAT16, - "bf16": onnx.TensorProto.BFLOAT16, -} -_QUANTIZED_LINEAR_CASES = { - "int4": (mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, 128), - "mxfp8": (mtq.MXFP8_DEFAULT_CFG, 32), - "nvfp4": (mtq.NVFP4_DEFAULT_CFG, 32), - "int8": (mtq.INT8_DEFAULT_CFG, 32), -} -_NATIVE_QUANTIZED_LINEAR_CASES = { - "fp8": (mtq.FP8_DEFAULT_CFG, 32), - **_QUANTIZED_LINEAR_CASES, -} - def _export_fp8_linear(source_dtype, weights_dtype): model = nn.Sequential(nn.Linear(4, 4, bias=False)).eval().to(source_dtype) @@ -94,88 +70,13 @@ def _export_fp8_linear(source_dtype, weights_dtype): onnx_bytes, _ = get_onnx_bytes_and_metadata( model, (sample_input,), - model_name="fp8_linear", weights_dtype=weights_dtype, - dq_only=False, onnx_opset=23, ) onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) -def _export_model(model, sample_input, weights_dtype, onnx_opset=20): - weights_dtype_kwargs = {} if weights_dtype is None else {"weights_dtype": weights_dtype} - onnx_bytes, _ = get_onnx_bytes_and_metadata( - model.eval(), - (sample_input,), - onnx_opset=onnx_opset, - **weights_dtype_kwargs, - ) - onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) - return onnx.load_model_from_string(onnx_bytes_obj.get_onnx_model_file_bytes()) - - -def _tensor_dtype_map(model): - dtype_map = { - value.name: value.type.tensor_type.elem_type - for value in (*model.graph.input, *model.graph.output, *model.graph.value_info) - if value.type.HasField("tensor_type") - } - dtype_map.update( - {initializer.name: initializer.data_type for initializer in model.graph.initializer} - ) - for node in model.graph.node: - if node.op_type != "Constant" or not node.output: - continue - value = next( - (attribute.t for attribute in node.attribute if attribute.name == "value"), None - ) - if value is not None: - dtype_map[node.output[0]] = value.data_type - return dtype_map - - -def _assert_runtime_io_dtype(model, expected_dtype): - initializer_names = {initializer.name for initializer in model.graph.initializer} - runtime_io = [ - *[value for value in model.graph.input if value.name not in initializer_names], - *model.graph.output, - ] - assert runtime_io - assert all(value.type.tensor_type.elem_type == expected_dtype for value in runtime_io) - - -def test_autocast_prefers_floating_input_device(monkeypatch): - autocast_args = {} - - def capture_autocast(*, device_type, dtype): - autocast_args.update(device_type=device_type, dtype=dtype) - return nullcontext() - - monkeypatch.setattr(torch, "autocast", capture_autocast) - flat_input = [ - torch.ones(1, dtype=torch.int64, device="meta"), - torch.ones(1, dtype=torch.float32), - ] - - _get_autocast_context(nn.Identity(), flat_input, torch.bfloat16) - - assert autocast_args == {"device_type": "cpu", "dtype": torch.bfloat16} - - -def test_quant_exporter_preserves_legacy_post_process_signature(): - class LegacyExporter(ONNXQuantExporter): - pre_process = compute_scales = compress_weights = staticmethod(lambda model: model) - - @staticmethod - def post_process(model): - return model - - model = onnx.helper.make_model(onnx.helper.make_graph([], "graph", [], [])) - - assert LegacyExporter.process_model(model) is model - - @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) @@ -277,48 +178,13 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): @pytest.mark.parametrize( - ("source_dtype", "weights_dtype", "expected_scale_dtype", "expected_io_dtypes"), + ("source_dtype", "weights_dtype", "expected_onnx_dtype"), [ - pytest.param( - torch.bfloat16, - "native", - onnx.TensorProto.FLOAT, - (onnx.TensorProto.BFLOAT16, onnx.TensorProto.FLOAT), - id="native-bf16", - ), - pytest.param( - torch.float32, - "fp16", - onnx.TensorProto.FLOAT16, - (onnx.TensorProto.FLOAT16, onnx.TensorProto.FLOAT16), - id="fp32-to-fp16", - ), - pytest.param( - torch.float32, - "bf16", - onnx.TensorProto.BFLOAT16, - (onnx.TensorProto.BFLOAT16, onnx.TensorProto.BFLOAT16), - id="fp32-to-bf16", - ), - pytest.param( - torch.bfloat16, - "bf16", - onnx.TensorProto.BFLOAT16, - (onnx.TensorProto.BFLOAT16, onnx.TensorProto.BFLOAT16), - id="bf16-to-bf16", - ), - pytest.param( - torch.bfloat16, - "fp32", - onnx.TensorProto.FLOAT, - (onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT), - id="bf16-to-fp32", - ), + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16), ], ) -def test_fp8_export_with_supported_weights_dtype( - source_dtype, weights_dtype, expected_scale_dtype, expected_io_dtypes -): +def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, expected_onnx_dtype): exported_model = _export_fp8_linear(source_dtype, weights_dtype) onnx.checker.check_model(exported_model) @@ -338,627 +204,41 @@ def test_fp8_export_with_supported_weights_dtype( ] assert fp8_weight_dq_nodes assert all( - initializer_by_name[node.input[1]].data_type == expected_scale_dtype + initializer_by_name[node.input[1]].data_type == expected_onnx_dtype for node in fp8_weight_dq_nodes ) - graph_io = [*exported_model.graph.input, *exported_model.graph.output] - assert tuple(value.type.tensor_type.elem_type for value in graph_io) == expected_io_dtypes - - -@pytest.mark.parametrize("format_name", _QUANTIZED_LINEAR_CASES) -@pytest.mark.parametrize("weights_dtype", _ONNX_DTYPE_BY_NAME) -def test_quantized_linear_export_uses_requested_weights_dtype( - monkeypatch, format_name, weights_dtype -): - monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", lambda inputs, *args: inputs) - quantization_config, features = _QUANTIZED_LINEAR_CASES[format_name] - source_dtype = torch.bfloat16 if weights_dtype == "fp32" else torch.float32 - model = nn.Sequential(nn.Linear(features, 8, bias=False)).eval().to(source_dtype) - sample_input = torch.ones(1, features, dtype=source_dtype) - model = mtq.quantize( - model, - quantization_config, - forward_loop=lambda quantized_model: quantized_model(sample_input), - ) - exported_model = _export_model(model, sample_input, weights_dtype, onnx_opset=23) - expected_dtype = _ONNX_DTYPE_BY_NAME[weights_dtype] - - onnx.checker.check_model(exported_model) - inferred_model = onnx.shape_inference.infer_shapes( - exported_model, check_type=True, strict_mode=True - ) - _assert_runtime_io_dtype(exported_model, expected_dtype) - dtype_map = _tensor_dtype_map(inferred_model) - matmul = next(node for node in inferred_model.graph.node if node.op_type == "MatMul") - assert all(dtype_map[input_name] == expected_dtype for input_name in matmul.input) - - initializer_map = { - initializer.name: initializer for initializer in exported_model.graph.initializer - } - cast_nodes = [node for node in exported_model.graph.node if node.op_type == "Cast"] - - if format_name == "int4": - weight = next( - initializer - for initializer in initializer_map.values() - if initializer.data_type == onnx.TensorProto.INT4 - ) - weight_dq = next( - node - for node in exported_model.graph.node - if node.op_type == "DequantizeLinear" and node.input[0] == weight.name - ) - assert initializer_map[weight_dq.input[1]].data_type == expected_dtype - assert dtype_map[weight_dq.output[0]] == expected_dtype - assert not cast_nodes - elif format_name == "mxfp8": - initializer_dtypes = {initializer.data_type for initializer in initializer_map.values()} - assert onnx.TensorProto.FLOAT8E4M3FN in initializer_dtypes - assert onnx.TensorProto.UINT8 in initializer_dtypes - dq_nodes = [ - node - for node in exported_model.graph.node - if node.op_type == "TRT_MXFP8DequantizeLinear" - ] - assert dq_nodes - assert all( - next(attribute.i for attribute in node.attribute if attribute.name == "output_dtype") - == expected_dtype - for node in dq_nodes - ) - assert not cast_nodes - elif format_name == "nvfp4": - initializer_dtypes = {initializer.data_type for initializer in initializer_map.values()} - assert { - onnx.TensorProto.FLOAT4E2M1, - onnx.TensorProto.FLOAT8E4M3FN, - onnx.TensorProto.FLOAT, - } <= initializer_dtypes - cast_dtypes = { - next(attribute.i for attribute in node.attribute if attribute.name == "to") - for node in cast_nodes - } - assert cast_dtypes == ({expected_dtype} if weights_dtype != "fp32" else set()) - else: - q_nodes = [node for node in inferred_model.graph.node if node.op_type == "QuantizeLinear"] - dq_nodes = [ - node for node in inferred_model.graph.node if node.op_type == "DequantizeLinear" - ] - assert q_nodes and dq_nodes - assert all(dtype_map[node.output[0]] == onnx.TensorProto.INT8 for node in q_nodes) - assert all(dtype_map[node.input[1]] == expected_dtype for node in dq_nodes) - assert all(dtype_map[node.output[0]] == expected_dtype for node in dq_nodes) - assert not cast_nodes - - -@pytest.mark.parametrize("format_name", _NATIVE_QUANTIZED_LINEAR_CASES) -def test_quantized_linear_default_preserves_native_behavior(monkeypatch, format_name): - monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", lambda inputs, *args: inputs) - quantization_config, features = _NATIVE_QUANTIZED_LINEAR_CASES[format_name] - model = nn.Sequential(nn.Linear(features, 8, bias=False)).eval().to(torch.bfloat16) - sample_input = torch.ones(1, features, dtype=torch.bfloat16) - model = mtq.quantize( - model, - quantization_config, - forward_loop=lambda quantized_model: quantized_model(sample_input), - ) - - exported_model = _export_model(model, sample_input, None, onnx_opset=23) - onnx.checker.check_model(exported_model) - inferred_model = onnx.shape_inference.infer_shapes( - exported_model, check_type=True, strict_mode=True - ) - initializer_names = {initializer.name for initializer in exported_model.graph.initializer} - runtime_inputs = [ - value for value in exported_model.graph.input if value.name not in initializer_names - ] - assert runtime_inputs - assert all( - value.type.tensor_type.elem_type == onnx.TensorProto.BFLOAT16 for value in runtime_inputs - ) - expected_boundary_dtype = ( - onnx.TensorProto.BFLOAT16 if format_name == "nvfp4" else onnx.TensorProto.FLOAT - ) - assert all( - value.type.tensor_type.elem_type == expected_boundary_dtype - for value in exported_model.graph.output - ) - dtype_map = _tensor_dtype_map(inferred_model) - matmul = next(node for node in inferred_model.graph.node if node.op_type == "MatMul") - assert all(dtype_map[input_name] == expected_boundary_dtype for input_name in matmul.input) - - initializer_map = { - initializer.name: initializer for initializer in exported_model.graph.initializer - } - initializer_dtypes = {initializer.data_type for initializer in initializer_map.values()} - if format_name == "fp8": - activation_q_nodes = [ - node for node in inferred_model.graph.node if node.op_type == "QuantizeLinear" - ] - cast_by_output = { - node.output[0]: node for node in inferred_model.graph.node if node.op_type == "Cast" - } - assert activation_q_nodes - assert all( - dtype_map[node.input[0]] == onnx.TensorProto.FLOAT - and dtype_map[node.input[1]] == onnx.TensorProto.FLOAT - for node in activation_q_nodes - ) - assert all(node.input[0] in cast_by_output for node in activation_q_nodes) - assert all( - next( - attribute.i - for attribute in cast_by_output[node.input[0]].attribute - if attribute.name == "to" - ) - == onnx.TensorProto.FLOAT - for node in activation_q_nodes - ) - weight_dq_nodes = [ - node - for node in inferred_model.graph.node - if node.op_type == "DequantizeLinear" - and node.input[0] in initializer_map - and initializer_map[node.input[0]].data_type == onnx.TensorProto.FLOAT8E4M3FN - ] - assert weight_dq_nodes - dq_nodes = [ - node for node in inferred_model.graph.node if node.op_type == "DequantizeLinear" - ] - assert all( - dtype_map[node.input[1]] == onnx.TensorProto.FLOAT - and dtype_map[node.output[0]] == onnx.TensorProto.FLOAT - for node in dq_nodes - ) - elif format_name == "int4": - assert onnx.TensorProto.INT4 in initializer_dtypes - weight_dq = next( - node - for node in inferred_model.graph.node - if node.op_type == "DequantizeLinear" - and node.input[0] in initializer_map - and initializer_map[node.input[0]].data_type == onnx.TensorProto.INT4 - ) - assert dtype_map[weight_dq.input[1]] == onnx.TensorProto.FLOAT - elif format_name == "mxfp8": - assert {onnx.TensorProto.FLOAT8E4M3FN, onnx.TensorProto.UINT8} <= initializer_dtypes - dq_nodes = [ - node - for node in inferred_model.graph.node - if node.op_type == "TRT_MXFP8DequantizeLinear" - ] - assert dq_nodes - assert all( - next(attribute.i for attribute in node.attribute if attribute.name == "output_dtype") - == onnx.TensorProto.FLOAT16 - for node in dq_nodes - ) - elif format_name == "nvfp4": - assert { - onnx.TensorProto.FLOAT4E2M1, - onnx.TensorProto.FLOAT8E4M3FN, - onnx.TensorProto.FLOAT, - } <= initializer_dtypes - else: - q_nodes = [node for node in inferred_model.graph.node if node.op_type == "QuantizeLinear"] - dq_nodes = [ - node for node in inferred_model.graph.node if node.op_type == "DequantizeLinear" - ] - assert q_nodes and dq_nodes - assert all(dtype_map[node.output[0]] == onnx.TensorProto.INT8 for node in q_nodes) - assert all(dtype_map[node.input[1]] == onnx.TensorProto.FLOAT for node in dq_nodes) - - -@pytest.mark.parametrize("weights_dtype", _ONNX_DTYPE_BY_NAME) -def test_fp8_conv_export_uses_requested_weights_dtype(weights_dtype): - source_dtype = torch.bfloat16 if weights_dtype == "fp32" else torch.float32 - model = nn.Conv2d(3, 4, kernel_size=3, bias=False).eval().to(source_dtype) - sample_input = torch.ones(1, 3, 8, 8, dtype=source_dtype) - model = mtq.quantize( - model, - mtq.FP8_DEFAULT_CFG, - forward_loop=lambda quantized_model: quantized_model(sample_input), - ) - exported_model = _export_model(model, sample_input, weights_dtype, onnx_opset=23) - expected_dtype = _ONNX_DTYPE_BY_NAME[weights_dtype] - - onnx.checker.check_model(exported_model) - inferred_model = onnx.shape_inference.infer_shapes( - exported_model, check_type=True, strict_mode=True - ) - _assert_runtime_io_dtype(exported_model, expected_dtype) - assert not any( - node.op_type in {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} - for node in exported_model.graph.node - ) - - initializer_map = { - initializer.name: initializer for initializer in exported_model.graph.initializer - } - conv = next(node for node in exported_model.graph.node if node.op_type == "Conv") - weight_dq = next( - node - for node in exported_model.graph.node - if node.op_type == "DequantizeLinear" and node.output[0] == conv.input[1] - ) - weight = initializer_map[weight_dq.input[0]] - assert weight.data_type == onnx.TensorProto.FLOAT8E4M3FN - assert initializer_map[weight_dq.input[1]].data_type == expected_dtype - assert _tensor_dtype_map(inferred_model)[weight_dq.output[0]] == expected_dtype - assert not any(node.op_type == "Cast" for node in exported_model.graph.node) - - -class MixedPrecisionLinear(nn.Module): - def __init__(self): - super().__init__() - self.fp32_weight = nn.Parameter(torch.eye(4, dtype=torch.float32)) - self.bf16_weight = nn.Parameter(torch.eye(4, dtype=torch.bfloat16)) - - def forward(self, x): - return torch.matmul(x, self.fp32_weight) + torch.matmul(x, self.bf16_weight) - - -@pytest.mark.parametrize( - ("source_dtype", "weights_dtype", "expected_onnx_dtype"), - [ - (torch.float32, "bf16", onnx.TensorProto.BFLOAT16), - (torch.bfloat16, "fp32", onnx.TensorProto.FLOAT), - (torch.float32, "fp16", onnx.TensorProto.FLOAT16), - ], -) -def test_parameterless_model_uses_explicit_weights_dtype( - source_dtype, weights_dtype, expected_onnx_dtype -): - exported_model = _export_model( - nn.Identity(), torch.ones(1, 4, dtype=source_dtype), weights_dtype - ) - graph_io = [*exported_model.graph.input, *exported_model.graph.output] assert all(value.type.tensor_type.elem_type == expected_onnx_dtype for value in graph_io) -def test_mixed_parameter_model_uses_explicit_weights_dtype(): - model = MixedPrecisionLinear() - exported_model = _export_model(model, torch.ones(1, 4), "bf16") - - assert model.fp32_weight.dtype == torch.float32 - assert model.bf16_weight.dtype == torch.bfloat16 - assert exported_model.graph.initializer - assert all( - initializer.data_type == onnx.TensorProto.BFLOAT16 - for initializer in exported_model.graph.initializer - ) - - -def test_onnx_quantizer_precision_is_restored_after_failure(): - quantizers = nn.ModuleList([TensorQuantizer(), TensorQuantizer()]) - quantizers[0].trt_high_precision_dtype = "Float" - del quantizers[1]._trt_high_precision_dtype - - with _override_onnx_quantizer_precision(quantizers, None): - assert quantizers[0].trt_high_precision_dtype == "Float" - assert not hasattr(quantizers[1], "_trt_high_precision_dtype") - assert quantizers[0].trt_high_precision_dtype == "Float" - assert not hasattr(quantizers[1], "_trt_high_precision_dtype") - - with _override_onnx_quantizer_precision(quantizers, "Half"): - assert all(q.trt_high_precision_dtype == "Half" for q in quantizers) - assert quantizers[0].trt_high_precision_dtype == "Float" - assert not hasattr(quantizers[1], "_trt_high_precision_dtype") - - with ( - pytest.raises(RuntimeError, match="export failed"), - _override_onnx_quantizer_precision(quantizers, "BFloat16"), +def test_fp8_export_rejects_bf16_conversion_from_fp32(): + with pytest.raises( + AssertionError, + match="Converting a quantized ONNX graph to BF16 is not supported", ): - assert all(q.trt_high_precision_dtype == "BFloat16" for q in quantizers) - raise RuntimeError("export failed") + _export_fp8_linear(torch.float32, "bf16") - assert quantizers[0].trt_high_precision_dtype == "Float" - assert not hasattr(quantizers[1], "_trt_high_precision_dtype") - -def _identity_onnx_model(dtype=onnx.TensorProto.FLOAT): - graph = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["input"], ["output"])], - "identity", - [onnx.helper.make_tensor_value_info("input", dtype, [1, 4])], - [onnx.helper.make_tensor_value_info("output", dtype, [1, 4])], - ) - return onnx.helper.make_model(graph) - - -@pytest.mark.parametrize("weights_dtype", ["fp32", "fp16", "bf16"]) -def test_onnx_load_path_rejects_non_native_weights_dtype(tmp_path, weights_dtype): - onnx_path = tmp_path / "identity.onnx" - onnx.save(_identity_onnx_model(), onnx_path) - - with pytest.raises(ValueError, match="weights_dtype must be 'native'"): - get_onnx_bytes_and_metadata( - nn.Identity(), - (torch.ones(1, 4),), - onnx_load_path=str(onnx_path), - weights_dtype=weights_dtype, - ) - - -def test_onnx_load_path_preserves_native_model(tmp_path): - onnx_path = tmp_path / "identity.onnx" - onnx.save(_identity_onnx_model(), onnx_path) - - onnx_bytes, _ = get_onnx_bytes_and_metadata( - nn.Identity(), - (torch.ones(1, 4),), - onnx_load_path=str(onnx_path), - weights_dtype="native", - ) - - loaded = OnnxBytes.from_bytes(onnx_bytes) - assert onnx.load_model_from_string(loaded.get_onnx_model_file_bytes()) - - -def _make_bf16_tensor(name, values): - values = np.asarray(values, dtype=np.float32) - tensor = onnx.TensorProto(name=name, data_type=onnx.TensorProto.BFLOAT16) - tensor.dims.extend(values.shape) - tensor.raw_data = (values.view(np.uint32) >> 16).astype(np.uint16).tobytes() - return tensor - - -def test_convert_to_fp32_recurses_through_graphs_functions_and_attributes(): - branch_value = _make_bf16_tensor("branch_value", [2.0]) - weight = _make_bf16_tensor("weight", [1.0]) - weight.doc_string = "weight metadata" - weight.metadata_props.add(key="source", value="test") - branch = onnx.helper.make_graph( - [onnx.helper.make_node("Constant", [], ["branch_output"], value=branch_value)], - "branch", - [], - [onnx.helper.make_tensor_value_info("branch_output", onnx.TensorProto.BFLOAT16, [1])], - ) - custom_node = onnx.helper.make_node( - "CustomOp", - ["input"], - ["custom_output"], - domain="test", - dtype=onnx.TensorProto.DOUBLE, - output_dtype=onnx.TensorProto.BFLOAT16, - ) - custom_node.attribute.append( - onnx.helper.make_attribute( - "type", onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT16, [1]) - ) - ) - graph = onnx.helper.make_graph( +def test_fp8_bf16_noop_rejects_incompatible_mixed_format(): + model = nn.Sequential(*(nn.Linear(128, 128, bias=False) for _ in range(2))) + model = model.eval().to(torch.bfloat16) + sample_input = torch.ones(1, 128, dtype=torch.bfloat16) + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"].extend( [ - onnx.helper.make_node( - "If", ["condition"], ["output"], then_branch=branch, else_branch=branch - ), - onnx.helper.make_node("Cast", ["input"], ["cast_output"], to=onnx.TensorProto.FLOAT16), - custom_node, - ], - "recursive", - [onnx.helper.make_tensor_value_info("input", onnx.TensorProto.BFLOAT16, [1])], - [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.DOUBLE, [1])], - initializer=[ - onnx.numpy_helper.from_array(np.array(True), "condition"), - weight, - onnx.numpy_helper.from_array(np.array([3.0], dtype=np.float64), "double_weight"), - ], - value_info=[ - onnx.helper.make_tensor_value_info("custom_output", onnx.TensorProto.FLOAT16, [1]) - ], - ) - function = onnx.helper.make_function( - "test", - "LocalCast", - ["x"], - ["y"], - [onnx.helper.make_node("Cast", ["x"], ["y"], to=onnx.TensorProto.BFLOAT16)], - [onnx.helper.make_opsetid("", 20)], - value_info=[onnx.helper.make_tensor_value_info("y", onnx.TensorProto.BFLOAT16, [1])], - ) - model = onnx.helper.make_model(graph, functions=[function]) - - assert convert_to_fp32(model) is model - - graph = model.graph - function = model.functions[0] - custom_node = next(node for node in graph.node if node.op_type == "CustomOp") - assert all( - value.type.tensor_type.elem_type == onnx.TensorProto.FLOAT - for value in (*graph.input, *graph.output, *graph.value_info) - ) - initializer_map = {initializer.name: initializer for initializer in graph.initializer} - assert initializer_map["weight"].data_type == onnx.TensorProto.FLOAT - assert initializer_map["double_weight"].data_type == onnx.TensorProto.FLOAT - assert initializer_map["weight"].doc_string == "weight metadata" - assert initializer_map["weight"].metadata_props[0].key == "source" - np.testing.assert_array_equal( - np.frombuffer(initializer_map["weight"].raw_data, dtype=np.float32), - np.array([1.0], dtype=np.float32), - ) - attribute_map = {attribute.name: attribute for attribute in custom_node.attribute} - assert attribute_map["dtype"].i == onnx.TensorProto.DOUBLE - assert attribute_map["output_dtype"].i == onnx.TensorProto.BFLOAT16 - assert attribute_map["type"].tp.tensor_type.elem_type == onnx.TensorProto.FLOAT - for attribute in graph.node[0].attribute: - branch_graph = attribute.g - assert branch_graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT - assert branch_graph.node[0].attribute[0].t.data_type == onnx.TensorProto.FLOAT - assert function.value_info[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT - assert function.node[0].attribute[0].i == onnx.TensorProto.FLOAT - - -def test_convert_to_fp32_handles_vetted_dtype_attributes(): - nodes = [ - onnx.helper.make_node("HannWindow", [], [], output_datatype=onnx.TensorProto.FLOAT16), - onnx.helper.make_node("LayerNormalization", [], [], stash_type=onnx.TensorProto.BFLOAT16), - onnx.helper.make_node("Attention", [], [], softmax_precision=onnx.TensorProto.DOUBLE), - onnx.helper.make_node("QuantizeLinear", [], [], precision=onnx.TensorProto.FLOAT16), - onnx.helper.make_node("Cast", [], [], to=onnx.TensorProto.BFLOAT16), - onnx.helper.make_node( - "TRTCustom", [], [], domain="trt", output_dtype=onnx.TensorProto.BFLOAT16 - ), - onnx.helper.make_node( - "CustomOp", - [], - [], - domain="test", - dtype=onnx.TensorProto.BFLOAT16, - precision=onnx.TensorProto.BFLOAT16, - to=onnx.TensorProto.BFLOAT16, - ), - ] - model = onnx.helper.make_model(onnx.helper.make_graph(nodes, "dtype_attributes", [], [])) - - convert_to_fp32(model) - - for node in model.graph.node[:6]: - assert node.attribute[0].i == onnx.TensorProto.FLOAT - assert all(attr.i == onnx.TensorProto.BFLOAT16 for attr in model.graph.node[6].attribute) - - -def _make_function_with_referenced_to(op_type): - to_attribute = onnx.AttributeProto( - name="to", - ref_attr_name="target_dtype", - type=onnx.AttributeProto.INT, - ) - node = onnx.helper.make_node(op_type, ["x"], ["y"]) - node.attribute.append(to_attribute) - return onnx.helper.make_function( - "test", - f"Referenced{op_type}", - ["x"], - ["y"], - [node], - [onnx.helper.make_opsetid("", 26)], - attribute_protos=[onnx.helper.make_attribute("target_dtype", onnx.TensorProto.BFLOAT16)], - ) - - -def test_convert_to_fp32_converts_function_cast_attribute_default(): - function = _make_function_with_referenced_to("Cast") - model = onnx.helper.make_model( - onnx.helper.make_graph([], "function_cast", [], []), functions=[function] - ) - - convert_to_fp32(model) - - assert model.functions[0].attribute_proto[0].i == onnx.TensorProto.FLOAT - - -def test_convert_to_fp32_rejects_function_bitcast_attribute_default(): - function = _make_function_with_referenced_to("BitCast") - model = onnx.helper.make_model( - onnx.helper.make_graph([], "function_bitcast", [], []), functions=[function] - ) - - with pytest.raises(ValueError, match="BitCast targets cannot be converted safely"): - convert_to_fp32(model) - - -def test_convert_to_fp32_rejects_low_precision_bitcast(): - bitcast = onnx.helper.make_node("BitCast", ["input"], ["output"], to=onnx.TensorProto.FLOAT16) - graph = onnx.helper.make_graph( - [bitcast], - "bitcast", - [onnx.helper.make_tensor_value_info("input", onnx.TensorProto.UINT16, [1])], - [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT16, [1])], - ) - - with pytest.raises(ValueError, match="BitCast targets cannot be converted safely"): - convert_to_fp32(onnx.helper.make_model(graph)) - - -def test_convert_to_fp32_rejects_segmented_tensor(): - weight = onnx.numpy_helper.from_array(np.array([1.0], dtype=np.float16), "weight") - weight.segment.begin = 0 - weight.segment.end = 1 - graph = onnx.helper.make_graph( - [], - "segmented", - [], - [], - initializer=[weight], - ) - - with pytest.raises(ValueError, match="Segmented tensors are not supported"): - convert_to_fp32(onnx.helper.make_model(graph)) - - -def test_convert_to_fp32_handles_loaded_external_data(tmp_path): - tensor = onnx.numpy_helper.from_array(np.array([1.0, -2.0], dtype=np.float16), "weight") - graph = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["weight"], ["output"])], - "external", - [], - [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT16, [2])], - [tensor], - ) - onnx_path = tmp_path / "external.onnx" - onnx.save_model( - onnx.helper.make_model(graph), - onnx_path, - save_as_external_data=True, - all_tensors_to_one_file=True, - location="weights.data", - size_threshold=0, - ) - - unloaded_model = onnx.load(onnx_path, load_external_data=False) - with pytest.raises(ValueError, match="External tensor data must be loaded"): - convert_to_fp32(unloaded_model) - - loaded_model = onnx.load(onnx_path, load_external_data=True) - convert_to_fp32(loaded_model) - converted_weight = loaded_model.graph.initializer[0] - assert converted_weight.data_type == onnx.TensorProto.FLOAT - assert converted_weight.data_location == onnx.TensorProto.DEFAULT - assert not converted_weight.external_data - np.testing.assert_array_equal( - onnx.numpy_helper.to_array(converted_weight), np.array([1.0, -2.0], dtype=np.float32) - ) - - -def test_save_onnx_model_externalizes_large_attribute_and_replaces_shards(tmp_path, monkeypatch): - values = np.arange(300, dtype=np.float32) - constant = onnx.helper.make_node( - "Constant", - [], - ["output"], - value=onnx.numpy_helper.from_array(values), - ) - graph = onnx.helper.make_graph( - [constant], - "external_attribute", - [], - [onnx.helper.make_tensor_value_info("output", onnx.TensorProto.FLOAT, [300])], - ) - model = onnx.helper.make_model(graph) - onnx_path = tmp_path / "model.onnx" - onnx.save_model(model, onnx_path) - external_data_path = tmp_path / "model.onnx_data" - external_data_path.write_bytes(b"stale") - previous_shard = tmp_path / "previous.data" - previous_shard.write_bytes(b"old") - monkeypatch.setattr(torch_onnx, "TWO_GB", 1) - - torch_onnx._save_onnx_model(model, str(onnx_path), "model") - - assert external_data_path.stat().st_size == values.nbytes - assert not previous_shard.exists() - unloaded_model = onnx.load(onnx_path, load_external_data=False) - value = next(attr.t for attr in unloaded_model.graph.node[0].attribute if attr.name == "value") - assert onnx.external_data_helper.uses_external_data(value) - assert next(prop.value for prop in value.external_data if prop.key == "location") == ( - "model.onnx_data" - ) - loaded_model = onnx.load(onnx_path, load_external_data=True) - loaded_value = next( - attr.t for attr in loaded_model.graph.node[0].attribute if attr.name == "value" + { + "quantizer_name": "1.weight_quantizer", + "cfg": {"num_bits": 4, "block_sizes": {-1: 128, "type": "static"}}, + }, + {"quantizer_name": "1.input_quantizer", "enable": False}, + ] ) - np.testing.assert_array_equal(onnx.numpy_helper.to_array(loaded_value), values) + model = mtq.quantize(model, config, forward_loop=lambda model: model(sample_input)) + with pytest.raises( + AssertionError, + match="Converting a quantized ONNX graph to BF16 is not supported", + ): + get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) class SingleArgModel(nn.Module): From a2afc683ff442b5c37ae9b481794aa9fb2f58279 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:29:17 +0000 Subject: [PATCH 5/8] [6508436] Address BF16 FP8 export review feedback Keep Conv dequantization type-consistent, normalize BF16 values in FP32, validate ONNX dtype enums explicitly, and reject mixed-source precision no-ops. Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/export/fp8_exporter.py | 27 +++++++-- modelopt/onnx/quantization/gs_patching.py | 3 +- modelopt/torch/_deploy/utils/torch_onnx.py | 7 ++- .../unit/onnx/quantization/test_qdq_utils.py | 2 +- .../deploy/utils/test_torch_onnx_utils.py | 60 +++++++++++++++---- 5 files changed, 80 insertions(+), 19 deletions(-) diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 6fa09f38263..5edc18fd5cc 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -88,6 +88,9 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: scale = node.inputs[1] torch_weights = _torch_from_numpy(weights.values) torch_scale = _torch_from_numpy(scale.values) + if torch.bfloat16 in (torch_weights.dtype, torch_scale.dtype): + torch_weights = torch_weights.float() + torch_scale = torch_scale.float() quantizer_name = scale.name.rsplit("/", 1)[0] dq_op = node.outputs[0].outputs[0] if dq_op.op != "TRT_FP8DequantizeLinear": @@ -207,9 +210,22 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: if amax == 0: continue scale_val = (amax / _FP8_E4M3_MAX).item() + scale_data = np.array(scale_val, dtype=weight_input.values.dtype) + if scale_data < scale_val: + np.nextafter( + scale_data, + np.array(np.inf, dtype=scale_data.dtype), + out=scale_data, + ) + torch_scale = _torch_from_numpy(scale_data) + if torch.bfloat16 in (torch_weights.dtype, torch_scale.dtype): + torch_weights = torch_weights.float() + torch_scale = torch_scale.float() # Quantize weights to FP8 (WAR: numpy doesn't support fp8) - fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_data = ( + (torch_weights / torch_scale).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + ) fp8_tensor = onnx.TensorProto() fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN fp8_tensor.dims.extend(fp8_data.shape) @@ -218,13 +234,16 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) ) - # Scale in FP16 — DQ output type matches scale dtype, must match activation type scale_constant = gs.Constant( node.name + "/weight_quantizer/scale", - np.array(scale_val, dtype=np.float16), + scale_data, ) - dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") + dq_output = gs.Variable( + node.name + "/weight_quantizer/dq_output", + scale_data.dtype, + weight_input.values.shape, + ) dq_node = gs.Node( op="DequantizeLinear", name=node.name + "/weight_quantizer/DequantizeLinear", diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py index 453f7af46bc..503a91f9721 100644 --- a/modelopt/onnx/quantization/gs_patching.py +++ b/modelopt/onnx/quantization/gs_patching.py @@ -106,7 +106,8 @@ def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.V dtype = tensor.dtype if isinstance(dtype, (int, np.integer)): dtype = int(dtype) - onnx.TensorProto.DataType.Name(dtype) + if dtype not in onnx.TensorProto.DataType.values(): + raise ValueError(f"Unknown ONNX tensor dtype for {tensor.name}: {dtype}") else: dtype = onnx.helper.np_dtype_to_tensor_dtype(np.dtype(dtype)) onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 9134df0f2ed..e6614f69e5f 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -527,8 +527,9 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module - first_parameter = next(model.parameters(), None) - source_weights_dtype = first_parameter.dtype if first_parameter is not None else torch.float32 + source_parameter_dtypes = { + parameter.dtype for parameter in model.parameters() if parameter.is_floating_point() + } # 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 @@ -643,7 +644,7 @@ def get_onnx_bytes_and_metadata( ) is_bf16_fp8_noop = ( weights_dtype == "bf16" - and source_weights_dtype == torch.bfloat16 + and source_parameter_dtypes == {torch.bfloat16} and uses_fp8 and not uses_other_unsupported_quantizer ) diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 7411c280306..f370b15211e 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -536,7 +536,7 @@ def test_bf16_weights_and_scale_are_compressed(self): if initializer.name == "linear/weight_quantizer/fp8_weights" ) assert fp8_weight.data_type == TensorProto.FLOAT8E4M3FN - assert fp8_weight.raw_data == b"\x3a" + assert fp8_weight.raw_data == b"\x3b" output_scale = next( initializer for initializer 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 5ca548c645d..9256746279a 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -59,9 +59,17 @@ } -def _export_fp8_linear(source_dtype, weights_dtype): - model = nn.Sequential(nn.Linear(4, 4, bias=False)).eval().to(source_dtype) - sample_input = torch.arange(4, dtype=source_dtype).reshape(1, 4) +def _export_fp8_model(source_dtype, weights_dtype, conv=False): + if conv: + model = nn.Sequential(nn.Conv2d(1, 1, 1, bias=False)) + sample_input = torch.ones(1, 1, 2, 2, dtype=source_dtype) + else: + model = nn.Sequential(nn.Linear(4, 4, bias=False)) + sample_input = torch.arange(4, dtype=source_dtype).reshape(1, 4) + model = model.eval().to(source_dtype) + if conv: + with torch.no_grad(): + model[0].weight.fill_(1e-38) model = mtq.quantize( model, mtq.FP8_DEFAULT_CFG, @@ -178,16 +186,19 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): @pytest.mark.parametrize( - ("source_dtype", "weights_dtype", "expected_onnx_dtype"), + ("source_dtype", "weights_dtype", "expected_onnx_dtype", "conv"), [ - (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16), - (torch.float32, "fp16", onnx.TensorProto.FLOAT16), + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16, False), + (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16, True), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16, False), ], ) -def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, expected_onnx_dtype): - exported_model = _export_fp8_linear(source_dtype, weights_dtype) +def test_fp8_export_with_supported_weights_dtype( + source_dtype, weights_dtype, expected_onnx_dtype, conv +): + exported_model = _export_fp8_model(source_dtype, weights_dtype, conv) - onnx.checker.check_model(exported_model) + onnx.checker.check_model(exported_model, full_check=True) assert not any( node.op_type in {"TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"} for node in exported_model.graph.node @@ -207,6 +218,11 @@ def test_fp8_export_with_supported_weights_dtype(source_dtype, weights_dtype, ex initializer_by_name[node.input[1]].data_type == expected_onnx_dtype for node in fp8_weight_dq_nodes ) + if conv: + assert all( + set(initializer_by_name[node.input[0]].raw_data).isdisjoint({0x7F, 0xFF}) + for node in fp8_weight_dq_nodes + ) graph_io = [*exported_model.graph.input, *exported_model.graph.output] assert all(value.type.tensor_type.elem_type == expected_onnx_dtype for value in graph_io) @@ -216,7 +232,7 @@ def test_fp8_export_rejects_bf16_conversion_from_fp32(): AssertionError, match="Converting a quantized ONNX graph to BF16 is not supported", ): - _export_fp8_linear(torch.float32, "bf16") + _export_fp8_model(torch.float32, "bf16") def test_fp8_bf16_noop_rejects_incompatible_mixed_format(): @@ -241,6 +257,30 @@ def test_fp8_bf16_noop_rejects_incompatible_mixed_format(): get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) +def test_fp8_bf16_noop_rejects_mixed_source_dtypes(): + class MixedDtypeModel(nn.Module): + def __init__(self): + super().__init__() + self.bf16_layer = nn.Linear(4, 4, bias=False).to(torch.bfloat16) + self.fp32_layer = nn.Linear(4, 4, bias=False) + + def forward(self, inputs): + hidden = self.bf16_layer(inputs) + return self.fp32_layer(hidden.float()) + + model = MixedDtypeModel().eval() + sample_input = torch.ones(1, 4, dtype=torch.bfloat16) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + with pytest.raises( + AssertionError, match="Converting a quantized ONNX graph to BF16 is not supported" + ): + get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) + + class SingleArgModel(nn.Module): def forward(self, x: torch.Tensor): return torch.add(x, x) - x From 97302606f09a49cf70002fee52c926a30aacc05d Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:39:35 +0000 Subject: [PATCH 6/8] [6508436] Include floating buffers in BF16 validation Require both floating parameters and registered buffers to be BF16 before treating FP8 export as a no-op conversion. Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/torch/_deploy/utils/torch_onnx.py | 9 +++++--- .../deploy/utils/test_torch_onnx_utils.py | 23 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index e6614f69e5f..8d80f823328 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -24,6 +24,7 @@ import shutil import tempfile from contextlib import nullcontext +from itertools import chain from typing import Any import onnx @@ -527,8 +528,10 @@ def get_onnx_bytes_and_metadata( if isinstance(model, (DataParallel, DistributedDataParallel)): model = model.module - source_parameter_dtypes = { - parameter.dtype for parameter in model.parameters() if parameter.is_floating_point() + source_floating_dtypes = { + tensor.dtype + for tensor in chain(model.parameters(), model.buffers()) + if tensor.is_floating_point() } # Standardize model args and also tensorize them so they also appear in the onnx graph! @@ -644,7 +647,7 @@ def get_onnx_bytes_and_metadata( ) is_bf16_fp8_noop = ( weights_dtype == "bf16" - and source_parameter_dtypes == {torch.bfloat16} + and source_floating_dtypes == {torch.bfloat16} and uses_fp8 and not uses_other_unsupported_quantizer ) 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 9256746279a..e7f8cbb603b 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -281,6 +281,29 @@ def forward(self, inputs): get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) +def test_fp8_bf16_noop_rejects_mixed_buffer_dtype(): + class MixedBufferModel(nn.Module): + def __init__(self): + super().__init__() + self.bf16_layer = nn.Linear(4, 4, bias=False).to(torch.bfloat16) + self.register_buffer("fp32_offset", torch.ones(4)) + + def forward(self, inputs): + return self.bf16_layer(inputs).float() + self.fp32_offset + + model = MixedBufferModel().eval() + sample_input = torch.ones(1, 4, dtype=torch.bfloat16) + model = mtq.quantize( + model, + mtq.FP8_DEFAULT_CFG, + forward_loop=lambda quantized_model: quantized_model(sample_input), + ) + with pytest.raises( + AssertionError, match="Converting a quantized ONNX graph to BF16 is not supported" + ): + get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23) + + class SingleArgModel(nn.Module): def forward(self, x: torch.Tensor): return torch.add(x, x) - x From b53ef035a3885676bd9af291cf09d3f59013a278 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:32:28 +0000 Subject: [PATCH 7/8] [6508436] Move changelog entry to 0.47 Place the BF16 FP8 ONNX export fix under the correct release section. Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b688a7bf704..f3666f1bee4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -54,6 +54,7 @@ Changelog - Avoid querying CUDA/Blackwell capability when ``NVFP4QTensor.quantize`` uses its CPU path or has the optional TensorRT-LLM fast path disabled. - Fix NVFP4 ONNX export to quantize FP4 weights with the published FP8 block scales, matching eager ModelOpt packed weights. Block scales below ``2**-9`` are now clamped to that minimum, and non-finite or negative scales raise an error. +- Fix FP8 ONNX export of BF16 models during real-weight compression. - Fix Megatron-Bridge Quantization Aware Distillation of a vision-language model silently discarding the ModelOpt state, so the distilled checkpoint restored no quantizers and exported as an unquantized model. Re-run QAD to regenerate any affected checkpoint. - Fix Megatron-Core HuggingFace export silently omitting fused (grouped GEMM) MoE experts for architectures without an ``experts.linear_fc1`` rule (e.g. ``Qwen3MoeForCausalLM``), which produced a valid-looking checkpoint containing no expert weights. The exporter now raises instead of writing that checkpoint; the scripts also avoid the situation by selecting ``SequentialMLP`` for those architectures. - Fix GatedDeltaNet (Qwen3.5) quantizer exclusions on Megatron-Core: the recipe patterns name the HuggingFace ``linear_attn`` module, so the ``conv1d`` was calibrated and the alpha / beta gate projections were exported in FP8. ``conv1d`` now has a ``self_attention`` alias in the default disabled-quantizer units, and the alpha / beta projections are exported in BF16 (they share Megatron's fused ``in_proj`` quantizer and cannot be disabled by name). @@ -131,7 +132,6 @@ Changelog **Bug Fixes** -- Fix FP8 ONNX export of BF16 models during real-weight compression. - Fix NemotronH dense MLP quantization with the ``nvfp4_mlp_only`` and ``nvfp4_omlp_only`` recipe families. NemotronH registers these projections as ``mixer.up_proj`` / ``mixer.down_proj``, which the previous ``*mlp*`` selector missed, producing checkpoints with a null ``quant_algo``. - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. Stale output shapes are now reconciled via symbolic shape inference, and AutoCast falls back to schema-based type inference so unresolved ops no longer leave tensors untyped. - Fix fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) skipping modules without an ``act_fn`` attribute. Modules applying a custom gated activation between the two ``F.linear`` calls (e.g. ``MiniMaxM3VLExperts``) were silently skipped, leaving routed experts unquantized and failing HF export. Enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. From 4c6a904a7c841669012be32ca95b4ec7af62a7a2 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:11:23 +0000 Subject: [PATCH 8/8] [6508436] Address BF16 FP8 review feedback Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/onnx/export/fp8_exporter.py | 1 + modelopt/torch/_deploy/utils/torch_onnx.py | 10 ++++++- .../deploy/utils/test_torch_onnx_utils.py | 30 ++++++++++++++----- 3 files changed, 32 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index 5edc18fd5cc..b23019a2c03 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -211,6 +211,7 @@ def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: continue scale_val = (amax / _FP8_E4M3_MAX).item() scale_data = np.array(scale_val, dtype=weight_input.values.dtype) + # Round up so normalizing by the serialized scale stays within the FP8 range. if scale_data < scale_val: np.nextafter( scale_data, diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 8d80f823328..6d828225986 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -533,6 +533,7 @@ def get_onnx_bytes_and_metadata( for tensor in chain(model.parameters(), model.buffers()) if tensor.is_floating_point() } + source_floating_dtype_names = ", ".join(sorted(map(str, source_floating_dtypes))) or "none" # 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 @@ -645,6 +646,12 @@ def get_onnx_bytes_and_metadata( uses_other_unsupported_quantizer = ( is_int4_quantized(model) or is_mxfp8_quantized(model) or is_int8_quantized(model) ) + if weights_dtype == "fp16" and uses_fp8 and torch.bfloat16 in source_floating_dtypes: + raise AssertionError( + "Converting a BF16 FP8 ONNX graph to FP16 is not supported yet " + f"(source floating dtypes: {source_floating_dtype_names})" + ) + is_bf16_fp8_noop = ( weights_dtype == "bf16" and source_floating_dtypes == {torch.bfloat16} @@ -654,7 +661,8 @@ def get_onnx_bytes_and_metadata( if weights_dtype in ["fp16", "bf16"] and not is_bf16_fp8_noop: if uses_other_unsupported_quantizer or uses_fp8: assert weights_dtype == "fp16", ( - "Converting a quantized ONNX graph to BF16 is not supported yet" + "Converting a quantized ONNX graph to BF16 is not supported yet " + f"(source floating dtypes: {source_floating_dtype_names})" ) onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, 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 e7f8cbb603b..24135c8f7f8 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -69,7 +69,7 @@ def _export_fp8_model(source_dtype, weights_dtype, conv=False): model = model.eval().to(source_dtype) if conv: with torch.no_grad(): - model[0].weight.fill_(1e-38) + model[0].weight.fill_(1e-38 if source_dtype == torch.bfloat16 else 1.0) model = mtq.quantize( model, mtq.FP8_DEFAULT_CFG, @@ -191,6 +191,7 @@ def test_onnx_export_and_inputs(model: BaseDeployModel): (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16, False), (torch.bfloat16, "bf16", onnx.TensorProto.BFLOAT16, True), (torch.float32, "fp16", onnx.TensorProto.FLOAT16, False), + (torch.float32, "fp16", onnx.TensorProto.FLOAT16, True), ], ) def test_fp8_export_with_supported_weights_dtype( @@ -227,12 +228,24 @@ def test_fp8_export_with_supported_weights_dtype( assert all(value.type.tensor_type.elem_type == expected_onnx_dtype for value in graph_io) -def test_fp8_export_rejects_bf16_conversion_from_fp32(): - with pytest.raises( - AssertionError, - match="Converting a quantized ONNX graph to BF16 is not supported", - ): - _export_fp8_model(torch.float32, "bf16") +@pytest.mark.parametrize( + ("source_dtype", "weights_dtype", "error"), + [ + ( + torch.float32, + "bf16", + r"Converting a quantized ONNX graph to BF16.*source floating dtypes: torch.float32", + ), + ( + torch.bfloat16, + "fp16", + r"Converting a BF16 FP8 ONNX graph to FP16.*source floating dtypes: torch.bfloat16", + ), + ], +) +def test_fp8_export_rejects_unsupported_dtype_conversion(source_dtype, weights_dtype, error): + with pytest.raises(AssertionError, match=error): + _export_fp8_model(source_dtype, weights_dtype) def test_fp8_bf16_noop_rejects_incompatible_mixed_format(): @@ -299,7 +312,8 @@ def forward(self, inputs): forward_loop=lambda quantized_model: quantized_model(sample_input), ) with pytest.raises( - AssertionError, match="Converting a quantized ONNX graph to BF16 is not supported" + AssertionError, + match=r"source floating dtypes: torch.bfloat16, torch.float32", ): get_onnx_bytes_and_metadata(model, (sample_input,), weights_dtype="bf16", onnx_opset=23)