Use non-buggy onnx-graphsurgeon via onnxslim for toposort

Author: xenova

scripts/float16.py

@@ -29,6 +29,7 @@
 import warnings
 from onnx import helper, numpy_helper
 from onnx import onnx_pb as onnx_proto
+import onnxslim.third_party.onnx_graphsurgeon as gs
 
 
 FLOAT32 = 1
@@ -277,9 +278,14 @@ def convert_float_to_float16(
             is_top_level = False  # Going to process sub-graph
         graph_stack = next_level
 
-    sort_topology(model.graph)
     remove_unnecessary_cast_node(model.graph)
 
+    # Topologically sort the graph
+    # NOTE: We do not perform another round of optimization as the model is already optimized
+    graph = gs.import_onnx(model)
+    graph.toposort()
+    model = gs.export_onnx(graph)
+
     return model
 
 
@@ -693,56 +699,6 @@ def convert_float_to_float16_model_path(
     )
 
 
-def sort_graph_node(graph_proto):
-    # find the "first" node in Nodes that its input is not any node's output
-    def find_first_node(output2node_dict):
-        for node in org_nodes:
-            is_not_first_node = any(item in output2node_dict for item in node.input)
-            if not is_not_first_node:
-                return node
-        return None
-
-    # remove the node from output2node_dict using output as key
-    def remove_first_node_from_dict2(first_node):
-        for output in first_node.output:
-            if output in output2node_dict:
-                del output2node_dict[output]
-
-    org_nodes = graph_proto.node
-    # create a dict to store output as key and node as value
-    output2node_dict = {}
-    for node in org_nodes:
-        for output in node.output:
-            output2node_dict[output] = node
-
-    # save the final node after sorted
-    sorted_node = []
-    # traverse the Nodes to find the first node
-    while len(output2node_dict) > 0:
-        first_node = find_first_node(output2node_dict)
-        sorted_node.append(first_node)
-        remove_first_node_from_dict2(first_node)
-        # del node from original nodes list to avoid duplicate traverse
-        org_nodes.remove(first_node)
-
-    for new_node in sorted_node:
-        graph_proto.node.extend([new_node])
-
-
-# The input graph should be mode.graph
-# Recursively sort the topology for each sub-graph
-def sort_topology(graph_proto):
-    assert isinstance(graph_proto, onnx_proto.GraphProto)
-    sort_graph_node(graph_proto)  # sort global graph
-    for node in graph_proto.node:
-        for attr in node.attribute:
-            if isinstance(attr.g, onnx_proto.GraphProto) and len(attr.g.node) > 0:
-                sort_topology(attr.g)  # sort sub-graph
-            for g in attr.graphs:
-                if isinstance(g, onnx_proto.GraphProto):
-                    sort_topology(g)  # sort sub-graph
-
-
 def remove_unnecessary_cast_node(graph_proto: onnx_proto.GraphProto):
     # 1. find all cast nodes in the graph
     cast_node_list = []

scripts/quantize.py

@@ -14,7 +14,6 @@
 from onnxruntime.quantization.registry import IntegerOpsRegistry
 from onnxruntime.quantization.matmul_4bits_quantizer import MatMul4BitsQuantizer
 from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer
-import onnx_graphsurgeon as gs
 
 from . import float16
 from .utils import check_and_save_model
@@ -221,10 +220,6 @@ def quantize_fp16(
         disable_shape_infer=disable_shape_infer,
         op_block_list=blocked_ops,
     )
-
-    graph = gs.import_onnx(model_fp16)
-    graph.toposort()
-    model_fp16 = gs.export_onnx(graph)
     check_and_save_model(model_fp16, save_path)
 
 

scripts/requirements.txt

@@ -1,7 +1,6 @@
-transformers[torch]==4.48.3
+transformers[torch]==4.49.0
 onnxruntime==1.20.1
-optimum@git+https://github.com/huggingface/optimum.git@ce533cf1a9e144d4040581947f301dc3f454b279
+optimum@git+https://github.com/huggingface/optimum.git@b04feaea78cda58d79b8da67dca3fd0c4ab33435
 onnx==1.17.0
 tqdm==4.67.1
 onnxslim==0.1.48
-onnx-graphsurgeon==0.5.5