[6421642][ONNX][Quantization] Fix NVFP4 exporter node ordering

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

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)