diff --git a/modelopt/onnx/export/nvfp4_exporter.py b/modelopt/onnx/export/nvfp4_exporter.py index e8bdfa2db1f..ee5ee39aec2 100644 --- a/modelopt/onnx/export/nvfp4_exporter.py +++ b/modelopt/onnx/export/nvfp4_exporter.py @@ -331,6 +331,7 @@ 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] = {} def _get_precision_dtype() -> str: # Check initializers to determine the precision of the weights @@ -350,18 +351,22 @@ def _cast_input_dtypes(node: onnx.NodeProto, precision_dtype: str): # Create Cast nodes for each input of the target node except bias for i, input_name in enumerate(node.input[:2]): - cast_output_name = input_name + "_f16" # Unique name for the cast output - - # 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], # Cast to FP16/BF16 - ) - - # Insert the Cast node into the graph - graph.node.extend([cast_node]) + 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_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 + 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 + ) + + # Insert the Cast node into the graph + graph.node.extend([cast_node]) # Update the target node input to use the cast node output node.input[i] = cast_output_name @@ -421,4 +426,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") + utils.topologically_sort_graph_nodes(graph) + return onnx_model diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 3b8edf76a84..bc37c4a9333 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -2020,3 +2020,89 @@ def clear_stale_value_info(model: onnx.ModelProto) -> int: fixed_shapes = _reconcile_stale_output_shapes(model) return fixed_outputs + fixed_shapes + n_cleared + + +def topologically_sort_graph_nodes(graph: onnx.GraphProto) -> None: + """Stable-sort graph nodes so tensor producers precede consumers. + + Unlike GraphSurgeon topological sorting, this operates directly on GraphProto + nodes and does not import/export tensor metadata that may include ONNX enum + dtypes such as BF16, which can be misinterpreted as NumPy dtypes during + conversion. + """ + + def get_graph_outer_scope_inputs(subgraph: onnx.GraphProto) -> set[str]: + local_names = { + value.name for value in (*subgraph.input, *subgraph.initializer) if value.name + } + local_names.update(output_name for node in subgraph.node for output_name in node.output) + + outer_scope_inputs = set() + for node in subgraph.node: + outer_scope_inputs.update( + input_name + for input_name in node.input + if input_name and input_name not in local_names + ) + for attr in node.attribute: + graphs = [] + if attr.type == onnx.AttributeProto.GRAPH: + graphs.append(attr.g) + elif attr.type == onnx.AttributeProto.GRAPHS: + graphs.extend(attr.graphs) + + for nested_graph in graphs: + outer_scope_inputs.update( + input_name + for input_name in get_graph_outer_scope_inputs(nested_graph) + if input_name not in local_names + ) + return outer_scope_inputs + + 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): + input_names = set(node.input) + for attr in node.attribute: + if attr.type == onnx.AttributeProto.GRAPH: + input_names.update(get_graph_outer_scope_inputs(attr.g)) + elif attr.type == onnx.AttributeProto.GRAPHS: + for subgraph in attr.graphs: + input_names.update(get_graph_outer_scope_inputs(subgraph)) + + for input_name in input_names: + 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 ONNX graph nodes.") + + graph.ClearField("node") + graph.node.extend(sorted_nodes) diff --git a/tests/unit/torch/quantization/test_onnx_export_cpu.py b/tests/unit/torch/quantization/test_onnx_export_cpu.py index ed062df8c9b..bb5a6dcb44c 100644 --- a/tests/unit/torch/quantization/test_onnx_export_cpu.py +++ b/tests/unit/torch/quantization/test_onnx_export_cpu.py @@ -15,13 +15,25 @@ """Unit tests for ONNX export for CPU quantization.""" +import inspect +import io + +import numpy as np import pytest import torch +onnx = pytest.importorskip("onnx") 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 +from onnx import TensorProto, helper, numpy_helper + +import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.tensor_quant as tensor_quant +from modelopt.onnx import utils +from modelopt.onnx.export import NVFP4QuantExporter +from modelopt.torch.quantization.utils import is_quantized_linear @pytest.mark.parametrize("model_cls", TEST_MODELS) @@ -42,3 +54,157 @@ 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) + + +def test_nvfp4_shared_activation_reuses_cast(): + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, 32]) + outputs = [ + helper.make_tensor_value_info("output0", TensorProto.FLOAT, [4, 64]), + helper.make_tensor_value_info("output1", TensorProto.FLOAT, [4, 64]), + ] + nodes = [] + initializers = [] + value_info = [] + + for index in range(2): + weight_name = f"linear{index}.weight" + fp4qdq_output = f"fp4qdq_output{index}" + initializers.append( + numpy_helper.from_array( + np.linspace(-1.0, 1.0, num=32 * 64, dtype=np.float32).reshape(32, 64), + weight_name, + ) + ) + value_info.append(helper.make_tensor_value_info(fp4qdq_output, TensorProto.FLOAT, [32, 64])) + nodes.extend( + [ + helper.make_node( + "TRT_FP4QDQ", + inputs=[weight_name], + outputs=[fp4qdq_output], + name=f"weight{index}_fp4qdq", + block_size=16, + ), + helper.make_node( + "MatMul", + inputs=["input", fp4qdq_output], + outputs=[f"output{index}"], + name=f"matmul{index}", + ), + ] + ) + + model = helper.make_model( + helper.make_graph( + nodes, + "shared_activation_nvfp4", + [input_tensor], + outputs, + initializers, + value_info=value_info, + ) + ) + + converted_model = NVFP4QuantExporter.process_model(model) + activation_casts = [ + node + for node in converted_model.graph.node + if node.op_type == "Cast" and node.input == ["input"] + ] + assert len(activation_casts) == 1 + assert activation_casts[0].output == ["input_f16"] + assert all( + node.input[0] == "input_f16" + for node in converted_model.graph.node + if node.op_type == "MatMul" + ) + onnx.checker.check_model(converted_model) + + +def test_topologically_sort_graph_nodes_accounts_for_subgraph_captures(): + input_tensor = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1]) + cond_tensor = helper.make_tensor_value_info("cond", TensorProto.BOOL, []) + output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + then_output = helper.make_tensor_value_info("then_output", TensorProto.FLOAT, [1]) + else_output = helper.make_tensor_value_info("else_output", TensorProto.FLOAT, [1]) + + then_graph = helper.make_graph( + [helper.make_node("Identity", ["captured"], ["then_output"], name="then_use_captured")], + "then_branch", + [], + [then_output], + ) + else_graph = helper.make_graph( + [helper.make_node("Identity", ["captured"], ["else_output"], name="else_use_captured")], + "else_branch", + [], + [else_output], + ) + if_node = helper.make_node( + "If", + ["cond"], + ["output"], + name="if_uses_captured", + then_branch=then_graph, + else_branch=else_graph, + ) + producer = helper.make_node("Identity", ["input"], ["captured"], name="producer") + model = helper.make_model( + helper.make_graph( + [if_node, producer], + "outer_scope_capture", + [input_tensor, cond_tensor], + [output_tensor], + ) + ) + + utils.topologically_sort_graph_nodes(model.graph) + + assert [node.name for node in model.graph.node] == ["producer", "if_uses_captured"] + onnx.checker.check_model(model)