From 88a66a289bfc5cea3177736f11d80b3425e18b72 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 23 Mar 2023 09:45:16 -0700 Subject: [PATCH] Fix prune_graph and gpt attention fusion scripts (#15147) Fix two issues: (1) GPT attention fusion: get_parent could return None when the input is initializer, add a check (2) ONNX node could have optional inputs and outputs. During prune_graph, we shall exclude empty inputs/outputs. Here we exclude "" from output_name_to_node and input_name_to_nodes. Add an option allow_remove_graph_inputs in prune_graph --- .../transformers/fusion_gpt_attention.py | 8 +-- .../python/tools/transformers/onnx_model.py | 50 +++++++++++-------- 2 files changed, 33 insertions(+), 25 deletions(-) diff --git a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py index 7fe3257950..b8a1cbb9f2 100644 --- a/onnxruntime/python/tools/transformers/fusion_gpt_attention.py +++ b/onnxruntime/python/tools/transformers/fusion_gpt_attention.py @@ -43,7 +43,7 @@ class FusionGptAttentionPastBase(Fusion): # | # {present} gather = self.model.get_parent(concat_v, 0, output_name_to_node) - if gather.op_type != "Gather": + if gather is None or gather.op_type != "Gather": logger.debug("match_past_pattern_1: expect Gather for past") return None @@ -53,7 +53,7 @@ class FusionGptAttentionPastBase(Fusion): past = gather.input[0] parent = self.model.get_parent(concat_k, 0, output_name_to_node) - if parent.op_type == "Gather": + if parent and parent.op_type == "Gather": gather_past_k = parent else: past_k_nodes = self.model.match_parent_path(concat_k, ["Transpose", "Gather"], [0, 0]) @@ -95,12 +95,12 @@ class FusionGptAttentionPastBase(Fusion): # {present} # squeeze = self.model.get_parent(concat_v, 0, output_name_to_node) - if squeeze.op_type != "Squeeze": + if squeeze is None or squeeze.op_type != "Squeeze": logger.debug("match_past_pattern_2: expect Squeeze as parent of concat_v") return None split = self.model.get_parent(squeeze, 0, output_name_to_node) - if split.op_type != "Split": + if split is None or split.op_type != "Split": logger.debug("match_past_pattern_2: expect Split for past path") return None diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index cce2cbe5a4..7cfc6d355c 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -49,17 +49,19 @@ class OnnxModel: input_name_to_nodes = {} for node in self.nodes(): for input_name in node.input: - if input_name not in input_name_to_nodes: - input_name_to_nodes[input_name] = [node] - else: - input_name_to_nodes[input_name].append(node) + if input_name: # could be empty when it is optional + if input_name not in input_name_to_nodes: + input_name_to_nodes[input_name] = [node] + else: + input_name_to_nodes[input_name].append(node) return input_name_to_nodes def output_name_to_node(self): output_name_to_node = {} for node in self.nodes(): for output_name in node.output: - output_name_to_node[output_name] = node + if output_name: # could be empty when it is optional + output_name_to_node[output_name] = node return output_name_to_node def nodes(self): @@ -599,7 +601,7 @@ class OnnxModel: Defaults to True. keep_io_types (Union[bool, List[str]], optional): boolean or a list of float32 input/output names. If True, model inputs/outputs should be left as float32. - Defaults to False. + Defaults to True. op_block_list (List[str], optional): List of operator types to leave as float32. Defaults to None, which will use `float16.DEFAULT_OP_BLOCK_LIST`. node_block_list (List[str], optional): List of node names to leave as float32. Defaults to None. @@ -751,14 +753,18 @@ class OnnxModel: if len(unused_nodes) > 0: logger.debug(f"Removed unused constant nodes: {len(unused_nodes)}") - def prune_graph(self, outputs=None): + def prune_graph(self, outputs=None, allow_remove_graph_inputs=True): """ - Prune graph to keep only required outputs. It removes unnecessary inputs and nodes. - Nodes are not linked (directly or indirectly) to any required output will be removed. + Prune graph to keep only required outputs. It removes unnecessary nodes that are not linked + (directly or indirectly) to any required output. + + There is also an option to remove graph inputs that are not used to generate any required output. Args: outputs (list): a list of graph outputs to retain. If it is None, all graph outputs will be kept. + allow_remove_graph_inputs (bool): allow remove graph inputs. """ + if len(self.graphs()) > 1: logger.debug("Skip prune_graph since graph has subgraph") return @@ -793,13 +799,14 @@ class OnnxModel: self.model.graph.output.remove(output) # remove inputs not used by any node. - input_name_to_nodes = self.input_name_to_nodes() input_to_remove = [] - for input in self.model.graph.input: - if input.name not in input_name_to_nodes: - input_to_remove.append(input) - for input in input_to_remove: - self.model.graph.input.remove(input) + if allow_remove_graph_inputs: + input_name_to_nodes = self.input_name_to_nodes() + for input in self.model.graph.input: + if input.name not in input_name_to_nodes: + input_to_remove.append(input) + for input in input_to_remove: + self.model.graph.input.remove(input) if input_to_remove or output_to_remove or nodes_to_remove: removed = [] @@ -813,7 +820,7 @@ class OnnxModel: self.update_graph() - def update_graph(self, verbose=False): + def update_graph(self, verbose=False, allow_remove_graph_inputs=False): graph = self.model.graph remaining_input_names = [] @@ -831,11 +838,12 @@ class OnnxModel: # remove graph input that is not used inputs_to_remove = [] - for input in graph.input: - if input.name not in remaining_input_names: - inputs_to_remove.append(input) - for input in inputs_to_remove: - graph.input.remove(input) + if allow_remove_graph_inputs: + for input in graph.input: + if input.name not in remaining_input_names: + inputs_to_remove.append(input) + for input in inputs_to_remove: + graph.input.remove(input) names_to_remove = [input.name for input in inputs_to_remove] logger.debug(f"remove {len(inputs_to_remove)} unused inputs: {names_to_remove}")