From 303dda25f18b722e327fcc52478be954516e4ea1 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:39:13 +0000 Subject: [PATCH 1/3] Fix NVFP4 exporter node ordering Ensure NVFP4 ONNX post-processing returns graphs with producers before consumers after TRT_FP4QDQ lowering and Cast insertion. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- modelopt/onnx/export/nvfp4_exporter.py | 5 ++ .../test_nvfp4_onnx_export_cpu.py | 77 +++++++++++++++++++ 2 files changed, 82 insertions(+) create mode 100644 tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index e8bdfa2db1f..ceda1d0b48b 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -17,6 +17,7 @@ import numpy as np import onnx +import onnx_graphsurgeon as gs import torch from onnx import numpy_helper @@ -421,4 +422,8 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): graph.initializer.extend(new_initializers) logger.info(f"Removed {len(initializers_to_delete)} initializers") + gs_graph = gs.import_onnx(onnx_model) + gs_graph.toposort() + onnx_model = gs.export_onnx(gs_graph) + return onnx_model diff --git a/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py b/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py new file mode 100644 index 00000000000..68d6c1c2a24 --- /dev/null +++ b/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for NVFP4 Torch quantization ONNX export.""" + +import inspect +import io + +import onnx +import pytest +import torch +from _test_utils.torch.quantization.models import SimpleLinear + +import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.tensor_quant as tensor_quant +from modelopt.onnx.export import NVFP4QuantExporter +from modelopt.torch.quantization.utils import is_quantized_linear + +pytest.importorskip("onnxruntime") + + +def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch): + def forward_loop(model): + model(sample_input) + + def cpu_dynamic_block_quantize(inputs, *args): + return inputs + + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize) + + model = SimpleLinear().eval() + sample_input = model.get_input() + model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop) + + for module in model.modules(): + assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module) + if isinstance(module, torch.nn.Linear): + module.input_quantizer.disable() + module.weight_quantizer._onnx_quantizer_type = "static" + + buffer = io.BytesIO() + if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: + kwargs = {"enable_onnx_checker": False} + else: + kwargs = {} + + torch.onnx.export( + model, + sample_input, + buffer, + input_names=["input"], + output_names=["output"], + export_params=True, + opset_version=21, + dynamo=False, + **kwargs, + ) + + buffer.seek(0) + exported_model = onnx.load_model_from_string(buffer.read()) + assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) + + converted_model = NVFP4QuantExporter.process_model(exported_model) + assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node) + onnx.checker.check_model(converted_model) From d3f71999c4f111ef8fc18cfefd8e2602684d7fd3 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Thu, 16 Jul 2026 21:11:45 +0000 Subject: [PATCH 2/3] Move NVFP4 ONNX regression into CPU export tests Keep the NVFP4 ONNX exporter regression with the existing CPU Torch ONNX export coverage instead of a standalone test file. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- .../test_nvfp4_onnx_export_cpu.py | 77 ------------------- .../quantization/test_onnx_export_cpu.py | 56 ++++++++++++++ 2 files changed, 56 insertions(+), 77 deletions(-) delete mode 100644 tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py diff --git a/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py b/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py deleted file mode 100644 index 68d6c1c2a24..00000000000 --- a/tests/unit/torch/quantization/test_nvfp4_onnx_export_cpu.py +++ /dev/null @@ -1,77 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CPU tests for NVFP4 Torch quantization ONNX export.""" - -import inspect -import io - -import onnx -import pytest -import torch -from _test_utils.torch.quantization.models import SimpleLinear - -import modelopt.torch.quantization as mtq -import modelopt.torch.quantization.tensor_quant as tensor_quant -from modelopt.onnx.export import NVFP4QuantExporter -from modelopt.torch.quantization.utils import is_quantized_linear - -pytest.importorskip("onnxruntime") - - -def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch): - def forward_loop(model): - model(sample_input) - - def cpu_dynamic_block_quantize(inputs, *args): - return inputs - - monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize) - - model = SimpleLinear().eval() - sample_input = model.get_input() - model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop) - - for module in model.modules(): - assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module) - if isinstance(module, torch.nn.Linear): - module.input_quantizer.disable() - module.weight_quantizer._onnx_quantizer_type = "static" - - buffer = io.BytesIO() - if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: - kwargs = {"enable_onnx_checker": False} - else: - kwargs = {} - - torch.onnx.export( - model, - sample_input, - buffer, - input_names=["input"], - output_names=["output"], - export_params=True, - opset_version=21, - dynamo=False, - **kwargs, - ) - - buffer.seek(0) - exported_model = onnx.load_model_from_string(buffer.read()) - assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) - - converted_model = NVFP4QuantExporter.process_model(exported_model) - assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node) - onnx.checker.check_model(converted_model) diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ed062df8c9b..dcfc78b58e4 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -15,14 +15,24 @@ """Unit tests for ONNX export for CPU quantization.""" +import inspect +import io + +import onnx import pytest import torch pytest.importorskip("onnxruntime") from _test_utils.torch.misc import set_seed +from _test_utils.torch.quantization.models import SimpleLinear from _test_utils.torch.quantization.onnx_export import TEST_MODELS, onnx_export_tester +import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.tensor_quant as tensor_quant +from modelopt.onnx.export import NVFP4QuantExporter +from modelopt.torch.quantization.utils import is_quantized_linear + @pytest.mark.parametrize("model_cls", TEST_MODELS) @pytest.mark.parametrize( @@ -42,3 +52,49 @@ def test_onnx_export_cpu(model_cls, num_bits, per_channel_quantization, constant onnx_export_tester( model_cls(), "cpu", num_bits, per_channel_quantization, constant_folding, dtype ) + + +def test_nvfp4_exported_onnx_is_topologically_sorted(monkeypatch): + def forward_loop(model): + model(sample_input) + + def cpu_dynamic_block_quantize(inputs, *args): + return inputs + + monkeypatch.setattr(tensor_quant, "dynamic_block_quantize_op", cpu_dynamic_block_quantize) + + model = SimpleLinear().eval() + sample_input = model.get_input() + model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=forward_loop) + + for module in model.modules(): + assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module) + if isinstance(module, torch.nn.Linear): + module.input_quantizer.disable() + module.weight_quantizer._onnx_quantizer_type = "static" + + buffer = io.BytesIO() + if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters: + kwargs = {"enable_onnx_checker": False} + else: + kwargs = {} + + torch.onnx.export( + model, + sample_input, + buffer, + input_names=["input"], + output_names=["output"], + export_params=True, + opset_version=21, + dynamo=False, + **kwargs, + ) + + buffer.seek(0) + exported_model = onnx.load_model_from_string(buffer.read()) + assert any(node.op_type == "TRT_FP4QDQ" for node in exported_model.graph.node) + + converted_model = NVFP4QuantExporter.process_model(exported_model) + assert not any(node.op_type == "TRT_FP4QDQ" for node in converted_model.graph.node) + onnx.checker.check_model(converted_model) From 7c448e1b65c2fafdbd1d52776e76fafeee140c85 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:16:00 +0000 Subject: [PATCH 3/3] Avoid GraphSurgeon round-trip for NVFP4 BF16 export Use a protobuf-only graph node topological sort so NVFP4 post-processing preserves BF16 ONNX metadata while still returning producer-before-consumer node order. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- modelopt/onnx/export/nvfp4_exporter.py | 48 +++++++++++++++++++++++--- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index ceda1d0b48b..707550312dc 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -17,7 +17,6 @@ import numpy as np import onnx -import onnx_graphsurgeon as gs import torch from onnx import numpy_helper @@ -35,6 +34,49 @@ from .base_exporter import ONNXQuantExporter +def _topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None: + """Stable-sort graph nodes so tensor producers precede consumers.""" + nodes = list(graph.node) + producer_by_tensor: dict[str, int] = {} + for node_index, node in enumerate(nodes): + for output_name in node.output: + if not output_name: + continue + if output_name in producer_by_tensor: + raise ValueError(f"Duplicate producer for tensor {output_name!r}.") + producer_by_tensor[output_name] = node_index + + dependencies: list[set[int]] = [set() for _ in nodes] + dependents: list[set[int]] = [set() for _ in nodes] + + for consumer_index, node in enumerate(nodes): + for input_name in node.input: + producer_index = producer_by_tensor.get(input_name) + if producer_index is None: + continue + if producer_index == consumer_index: + raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.") + dependencies[consumer_index].add(producer_index) + dependents[producer_index].add(consumer_index) + + ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency] + sorted_nodes: list[onnx.NodeProto] = [] + while ready: + node_index = ready.pop(0) + sorted_nodes.append(nodes[node_index]) + for dependent_index in sorted(dependents[node_index]): + dependencies[dependent_index].remove(node_index) + if not dependencies[dependent_index]: + ready.append(dependent_index) + ready.sort() + + if len(sorted_nodes) != len(nodes): + raise ValueError("Cycle detected while sorting NVFP4 ONNX graph nodes.") + + graph.ClearField("node") + graph.node.extend(sorted_nodes) + + def _cast_fp4(array: np.ndarray) -> np.ndarray: """Cast a numpy array to FLOAT4E2M1 using PyTorch. @@ -422,8 +464,6 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): graph.initializer.extend(new_initializers) logger.info(f"Removed {len(initializers_to_delete)} initializers") - gs_graph = gs.import_onnx(onnx_model) - gs_graph.toposort() - onnx_model = gs.export_onnx(gs_graph) + _topologically_sort_graph_nodes(graph) return onnx_model