saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 04:39:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add fusions for re-designed Phi-3 vision and Phi-3.5 vision ONNX models (#22026)

Browse source 
### Description
This PR adds the optimizer logic to fuse the newly designed exported
ONNX models for Phi-3 vision and Phi-3.5 vision.

### Motivation and Context
After the re-designed export of Phi-3 vision and Phi-3.5 vision, the
ONNX models for the vision component and embedding component contain
`If` and `Loop` ops to handle multi-image support.
This commit is contained in:
kunal-vaishnavi 2024-09-10 16:18:05 -07:00 committed by GitHub
parent 19954decaf
commit c5418f35d4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 97 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
onnxruntime/python/tools/transformers/fusion_gelu.py
View file

@ -98,10 +98,13 @@ class FusionGelu(Fusion):
return
self.nodes_to_remove.extend(subgraph_nodes)
fused_node = helper.make_node("Gelu", inputs=[subgraph_input], outputs=[subgraph_output])
fused_node = helper.make_node(
"Gelu", inputs=[subgraph_input], outputs=[subgraph_output], name=self.model.create_node_name("Gelu")
)
fused_node.domain = "com.microsoft"
self.nodes_to_add.append(fused_node)
self.node_name_to_graph_name[fused_node.name] = self.this_graph_name
self.increase_counter("Gelu")
return True
def fuse_2(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]:
@ -172,10 +175,13 @@ class FusionGelu(Fusion):
return
self.nodes_to_remove.extend(subgraph_nodes)
fused_node = helper.make_node("Gelu", inputs=[root_node.output[0]], outputs=[mul.output[0]])
fused_node = helper.make_node(
"Gelu", inputs=[root_node.output[0]], outputs=[mul.output[0]], name=self.model.create_node_name("Gelu")
)
fused_node.domain = "com.microsoft"
self.nodes_to_add.append(fused_node)
self.node_name_to_graph_name[fused_node.name] = self.this_graph_name
self.increase_counter("Gelu")
return True
def fuse_3(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]:
@ -243,8 +249,11 @@ class FusionGelu(Fusion):
return
self.nodes_to_remove.extend(subgraph_nodes)
fused_node = helper.make_node("Gelu", inputs=[root_node.output[0]], outputs=[last_mul.output[0]])
fused_node = helper.make_node(
"Gelu", inputs=[root_node.output[0]], outputs=[last_mul.output[0]], name=self.model.create_node_name("Gelu")
)
fused_node.domain = "com.microsoft"
self.nodes_to_add.append(fused_node)
self.node_name_to_graph_name[fused_node.name] = self.this_graph_name
self.increase_counter("Gelu")
return True

36
onnxruntime/python/tools/transformers/fusion_skiplayernorm.py
View file

@ -24,14 +24,15 @@ class FusionSkipLayerNormalization(Fusion):
model: OnnxModel,
fused_op_type: str = "SkipLayerNormalization",
search_op_types: str = "LayerNormalization",
shape_infer: bool = True,
):
super().__init__(model, fused_op_type, search_op_types)
# Update shape inference is needed since other fusions might add new edge which does not have shape info yet.
self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True)
if self.shape_infer_helper is None:
# TODO(tianleiwu): support subgraph in shape inference or add broadcasting in SkipLayerNormalization op.
logger.warning("symbolic shape inference disabled or failed.")
if shape_infer:
# Update shape inference is needed since other fusions might add new edge which does not have shape info yet.
self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True)
if self.shape_infer_helper is None:
# TODO(tianleiwu): support subgraph in shape inference or add broadcasting in SkipLayerNormalization op.
logger.warning("symbolic shape inference disabled or failed.")
def fuse(self, node, input_name_to_nodes, output_name_to_node):
add = self.model.get_parent(node, 0, output_name_to_node)
@ -56,18 +57,19 @@ class FusionSkipLayerNormalization(Fusion):
# Root Mean Square Layer Normalization
simplified = node.op_type == "SimplifiedLayerNormalization"
if self.shape_infer_helper is not None:
# TODO(tianleiwu): support broadcasting Skip shape (1, sequence_length, hidden_size) or (sequence_length, hidden_size)
if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]):
logger.debug(
"skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same",
add.input[0],
add.input[1],
)
if hasattr(self, "shape_infer_helper"):
if self.shape_infer_helper is not None:
# TODO(tianleiwu): support broadcasting Skip shape (1, sequence_length, hidden_size) or (sequence_length, hidden_size)
if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]):
logger.debug(
"skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same",
add.input[0],
add.input[1],
)
return
else:
logger.debug("skip SkipLayerNormalization fusion since symbolic shape inference failed")
return
else:
logger.debug("skip SkipLayerNormalization fusion since symbolic shape inference failed")
return
gather_path = self.model.match_parent_path(add, ["Gather"], [None])
if gather_path is not None and self.model.find_graph_input(gather_path[0].input[1]) is None:

79
onnxruntime/python/tools/transformers/onnx_model.py
View file

@ -63,9 +63,10 @@ class OnnxModel:
return None
def input_name_to_nodes(self):
def input_name_to_nodes(self, exclude_subgraphs=False):
input_name_to_nodes = {}
for node in self.nodes():
nodes_to_search = self.nodes() if not exclude_subgraphs else self.model.graph.node
for node in nodes_to_search:
for input_name in node.input:
if input_name: # could be empty when it is optional
if input_name not in input_name_to_nodes:
@ -74,9 +75,10 @@ class OnnxModel:
input_name_to_nodes[input_name].append(node)
return input_name_to_nodes
def output_name_to_node(self):
def output_name_to_node(self, exclude_subgraphs=False):
output_name_to_node = {}
for node in self.nodes():
nodes_to_search = self.nodes() if not exclude_subgraphs else self.model.graph.node
for node in nodes_to_search:
for output_name in node.output:
if output_name: # could be empty when it is optional
output_name_to_node[output_name] = node
@ -906,6 +908,31 @@ class OnnxModel:
if len(unused_nodes) > 0:
logger.debug(f"Removed unused constant nodes: {len(unused_nodes)}")
def _get_subgraph_inputs_of_node(self, node):
"""
Get inputs to all nodes in all subgraphs of a node
"""
# Note: This function only handles one-level subgraphs of child nodes.
subgraph_nodes_inputs = set()
for attr in node.attribute:
if attr.type == AttributeProto.GRAPH:
child_nodes = attr.g.node
for child_node in child_nodes:
subgraph_nodes_inputs.update(child_node.input)
return subgraph_nodes_inputs
def _get_subgraph_nodes_and_inputs(self, ops_with_graph_attrs):
"""
Get input names to all nodes in all subgraphs where subgraphs are
graph attributes of a node in the main graph
"""
subgraph_nodes = list(filter(lambda node: node.op_type in ops_with_graph_attrs, self.model.graph.node))
subgraph_nodes_inputs = set()
for parent_node in subgraph_nodes:
subgraph_inputs_of_parent_node = self._get_subgraph_inputs_of_node(parent_node)
subgraph_nodes_inputs.update(subgraph_inputs_of_parent_node)
return subgraph_nodes, subgraph_nodes_inputs
def prune_graph(self, outputs=None, allow_remove_graph_inputs=True):
"""
Prune graph to keep only required outputs. It removes unnecessary nodes that are not linked
@ -918,13 +945,9 @@ class OnnxModel:
allow_remove_graph_inputs (bool): allow remove graph inputs.
"""
if len(self.graphs()) > 1:
# TODO(tianleiwu): handle subgraph
logger.debug("Skip prune_graph since graph has subgraph")
return
keep_outputs = [output.name for output in self.model.graph.output] if outputs is None else outputs
input_name_to_nodes_for_main_graph = self.input_name_to_nodes(exclude_subgraphs=True)
output_name_to_node = self.output_name_to_node()
def get_first_output(node):
@ -932,6 +955,29 @@ class OnnxModel:
return node.output[0]
return next(iter([o for o in node.output if o]), None)
if len(self.graphs()) > 1:
# Get input names for all nodes in all subgraphs
subgraph_nodes, subgraph_nodes_inputs = self._get_subgraph_nodes_and_inputs(
ops_with_graph_attrs={"Loop", "Scan", "If"}
)
if len(subgraph_nodes) == 0:
# TODO: support other ops such as `BeamSearch` that have subgraphs as op attributes
logger.debug("Skip prune_graph since graph has subgraph")
return
# For graphs with subgraphs, add dangling outputs from parent graph nodes to list of outputs to keep
for node in self.model.graph.node:
# TODO: This for-loop logic currently assumes that Loop/Scan/If nodes will not be
# pruned because their subgraphs are needed for computations. This might not be
# true in all cases.
if node in subgraph_nodes:
continue
# Check if node output is an input of a subgraph node and not an input to a node in the main graph
for output in node.output:
if output in subgraph_nodes_inputs and output not in input_name_to_nodes_for_main_graph:
keep_outputs += [output]
# Keep track of nodes to keep. The key is first output of node, and the value is the node.
output_to_node = {}
@ -956,7 +1002,7 @@ class OnnxModel:
first_output = get_first_output(node)
kept_node = output_to_node.get(first_output)
# Need double check the node since fused node might reuse output name of some nodes to be removed.
# Need to double check the node since fused node might reuse output name of some nodes to be removed.
# It is slow to compare whole node, so we compare op_type first to avoid comparing node in most cases.
if kept_node and kept_node.op_type == node.op_type and kept_node == node:
nodes_to_keep.append(node)
@ -997,16 +1043,15 @@ class OnnxModel:
def update_graph(self, verbose=False, allow_remove_graph_inputs=False):
graph = self.model.graph
remaining_input_names = []
remaining_input_names = set()
for node in graph.node:
if node.op_type in ["Loop", "Scan", "If"]:
# TODO: handle inner graph
logger.debug(f"Skip update_graph since graph has operator: {node.op_type}")
return
# Add input names of nodes in subgraphs
subgraph_inputs_of_node = self._get_subgraph_inputs_of_node(node)
remaining_input_names.update(subgraph_inputs_of_node)
if node.op_type != "Constant":
for input_name in node.input:
if input_name not in remaining_input_names:
remaining_input_names.append(input_name)
remaining_input_names.update(node.input)
if verbose:
logger.debug(f"remaining input names: {remaining_input_names}")

6
onnxruntime/python/tools/transformers/onnx_model_bert.py
View file

@ -115,8 +115,8 @@ class BertOnnxModel(OnnxModel):
fusion = FusionSimplifiedLayerNormalization(self)
fusion.apply()
def fuse_skip_layer_norm(self):
fusion = FusionSkipLayerNormalization(self)
def fuse_skip_layer_norm(self, shape_infer=True):
fusion = FusionSkipLayerNormalization(self, shape_infer=shape_infer)
fusion.apply()
def fuse_skip_simplified_layer_norm(self):
@ -344,7 +344,7 @@ class BertOnnxModel(OnnxModel):
self.fuse_reshape()
if (options is None) or options.enable_skip_layer_norm:
self.fuse_skip_layer_norm()
self.fuse_skip_layer_norm(options.enable_shape_inference)
self.fuse_skip_simplified_layer_norm()
if (options is None) or options.enable_rotary_embeddings:

1
onnxruntime/python/tools/transformers/onnx_model_clip.py
View file

@ -24,6 +24,7 @@ class ClipOnnxModel(BertOnnxModel):
op_count = {}
ops = [
"Attention",
"Gelu",
"LayerNormalization",
"QuickGelu",
"SkipLayerNormalization",