diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index e8bdfa2db1f..707550312dc 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -34,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. @@ -421,4 +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") + _topologically_sort_graph_nodes(graph) + return onnx_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)