diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 612be6aebf..4a553325f9 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -139,7 +139,6 @@ class SymbolicShapeInference: "Gather": self._infer_Gather, "GatherElements": self._infer_GatherElements, "GatherND": self._infer_GatherND, - "Gelu": self._pass_on_shape_and_type, "Identity": self._pass_on_shape_and_type, "If": self._infer_If, "Loop": self._infer_Loop, @@ -1536,7 +1535,7 @@ class SymbolicShapeInference: handle_negative_axis(ax, self._get_shape_rank(node, i + num_scan_states)) for i, ax in enumerate(scan_input_axes) ] - # We may have cases where the subgraph has optionial inputs that appear in both subgraph's input and initializer, + # We may have cases where the subgraph has optional inputs that appear in both subgraph's input and initializer, # but not in the node's input. In such cases, the input model might be invalid, but let's skip those optional inputs. assert len(subgraph.input) >= len(node.input) subgraph_inputs = subgraph.input[: len(node.input)] diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 54037dd601..dcc9966ba8 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -24,6 +24,10 @@ class FusionSkipLayerNormalization(Fusion): # 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) @@ -43,15 +47,14 @@ class FusionSkipLayerNormalization(Fusion): if self.shape_infer_helper is not None: if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]): logger.debug( - f"skip skiplayernorm fusion since shape of inputs ({add.input[0]}, {add.input[1]}) are not same" + "skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same", + add.input[0], + add.input[1], ) return else: - # shape_infer_helper can not handle subgraphs. Current work around is to disable skiplayernorm fusion - # longterm todo: support subgraph in symbolic_shape_infer or support add broadcasting in skiplayernorm op - logger.warning( - "symbolic shape infer failed. it's safe to ignore this message if there is no issue with optimized model" - ) + 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: diff --git a/onnxruntime/python/tools/transformers/shape_infer_helper.py b/onnxruntime/python/tools/transformers/shape_infer_helper.py index bbb84a6bf1..e877497ffb 100644 --- a/onnxruntime/python/tools/transformers/shape_infer_helper.py +++ b/onnxruntime/python/tools/transformers/shape_infer_helper.py @@ -3,10 +3,10 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +import logging import os import sys - -import onnx +from typing import Dict # In ORT Package the symbolic_shape_infer.py is in ../tools file_path = os.path.dirname(__file__) @@ -17,85 +17,104 @@ else: from symbolic_shape_infer import SymbolicShapeInference, get_shape_from_type_proto, sympy +logger = logging.getLogger(__name__) + class SymbolicShapeInferenceHelper(SymbolicShapeInference): - def __init__( - self, - model, - verbose=0, - int_max=2**31 - 1, - auto_merge=True, - guess_output_rank=False, - ): + def __init__(self, model, verbose=0, int_max=2**31 - 1, auto_merge=True, guess_output_rank=False): super().__init__(int_max, auto_merge, guess_output_rank, verbose) - self.model_ = onnx.ModelProto() - self.model_.CopyFrom(model) - self.all_shapes_inferred_ = False - self.inferred_ = False + self.model_ = model + self.all_shapes_inferred_: bool = False + self.is_inferred_: bool = False + self.dynamic_axis_mapping_: Dict[str, int] = {} - # The goal is to remove dynamic_axis_mapping - def infer(self, dynamic_axis_mapping): - if self.inferred_: + def infer(self, dynamic_axis_mapping: Dict[str, int], max_runs: int = 128): + """Run shape inference, and try replace dynamic axis from string to integer when mapping is provided. + + Args: + dynamic_axis_mapping (_type_): a dictionary with name of dynamic axis as key, like {"batch_size" : 4} + max_runs (int, optional): limit maximum number of runs to avoid infinite loop. Defaults to 32. + + Returns: + bool: whether all shapes has been inferred or not. + """ + assert dynamic_axis_mapping is not None + + if self.is_inferred_ and self.dynamic_axis_mapping_ == dynamic_axis_mapping: return self.all_shapes_inferred_ - self.dynamic_axis_mapping_ = dynamic_axis_mapping # e.g {"batch_size" : 4, "seq_len" :7} + self.dynamic_axis_mapping_ = dynamic_axis_mapping self._preprocess(self.model_) - while self.run_: - self.all_shapes_inferred_ = self._infer_impl() - self.inferred_ = True + count = 0 + while self.run_: + logger.debug(f"shape infer run {count}") + self.all_shapes_inferred_ = self._infer_impl() + count += 1 + if max_runs > 0 and count >= max_runs: + break + + self.is_inferred_ = True return self.all_shapes_inferred_ - # override _preprocess() to avoid unnecessary model copy since ctor copies the model - def _preprocess(self, in_mp): - self.out_mp_ = in_mp - self.graph_inputs_ = dict([(i.name, i) for i in list(self.out_mp_.graph.input)]) - self.initializers_ = dict([(i.name, i) for i in self.out_mp_.graph.initializer]) - self.known_vi_ = dict([(i.name, i) for i in list(self.out_mp_.graph.input)]) - self.known_vi_.update( - dict( - [ - ( - i.name, - onnx.helper.make_tensor_value_info(i.name, i.data_type, list(i.dims)), - ) - for i in self.out_mp_.graph.initializer - ] - ) - ) - - # Override _get_sympy_shape() in symbolic_shape_infer.py to ensure shape inference by giving the actual value of dynamic axis def _get_sympy_shape(self, node, idx): + """Override it to ensure shape inference by giving the actual value of dynamic axis.""" sympy_shape = [] - for d in self._get_shape(node, idx): - if type(d) == str: - if d in self.dynamic_axis_mapping_.keys(): - sympy_shape.append(self.dynamic_axis_mapping_[d]) - elif d in self.symbolic_dims_: - sympy_shape.append(self.symbolic_dims_[d]) + + shape = self._get_shape(node, idx) + if shape: + for dim in shape: + if isinstance(dim, str): + if dim in self.dynamic_axis_mapping_: + sympy_shape.append(self.dynamic_axis_mapping_[dim]) + elif dim in self.symbolic_dims_: + sympy_shape.append(self.symbolic_dims_[dim]) + else: + sympy_shape.append(sympy.Symbol(dim, integer=True)) else: - sympy_shape.append(sympy.Symbol(d, integer=True)) - else: - assert None != d - sympy_shape.append(d) + assert dim is not None + sympy_shape.append(dim) return sympy_shape def get_edge_shape(self, edge): - assert self.all_shapes_inferred_ == True + """Get shape of an edge. + + Args: + edge (str): name of edge + + Returns: + Optional[List[int]]: the shape, or None if shape is unknown + """ + assert self.all_shapes_inferred_ if edge not in self.known_vi_: - print("Cannot retrive the shape of " + str(edge)) + print("Cannot retrieve the shape of " + str(edge)) return None + type_proto = self.known_vi_[edge].type shape = get_shape_from_type_proto(type_proto) - for i in range(len(shape)): - d = shape[i] - if type(d) == str and d in self.dynamic_axis_mapping_.keys(): - shape[i] = self.dynamic_axis_mapping_[d] + + if shape is not None: + for i, dim in enumerate(shape): + if isinstance(dim, str) and dim in self.dynamic_axis_mapping_: + shape[i] = self.dynamic_axis_mapping_[dim] + return shape def compare_shape(self, edge, edge_other): - assert self.all_shapes_inferred_ == True + """Compare shape of two edges. + + Args: + edge (str): name of edge + edge_other (str): name of another edge + + Raises: + Exception: At least one shape is missed for edges to compare + + Returns: + bool: whether the shape is same or not + """ + assert self.all_shapes_inferred_ shape = self.get_edge_shape(edge) shape_other = self.get_edge_shape(edge_other) if shape is None or shape_other is None: diff --git a/onnxruntime/test/python/transformers/test_shape_infer_helper.py b/onnxruntime/test/python/transformers/test_shape_infer_helper.py index 6429b4cee5..3c5f3b581f 100644 --- a/onnxruntime/test/python/transformers/test_shape_infer_helper.py +++ b/onnxruntime/test/python/transformers/test_shape_infer_helper.py @@ -1,15 +1,17 @@ import unittest +import onnx import pytest +import torch from parity_utilities import find_transformers_source if find_transformers_source(): - from benchmark_helper import OptimizerInfo, Precision + from benchmark_helper import ConfigModifier, OptimizerInfo, Precision from huggingface_models import MODELS from onnx_exporter import export_onnx_model_from_pt from shape_infer_helper import SymbolicShapeInferenceHelper else: - from onnxruntime.transformers.benchmark_helper import OptimizerInfo, Precision + from onnxruntime.transformers.benchmark_helper import ConfigModifier, OptimizerInfo, Precision from onnxruntime.transformers.huggingface_models import MODELS from onnxruntime.transformers.onnx_exporter import export_onnx_model_from_pt from onnxruntime.transformers.shape_infer_helper import SymbolicShapeInferenceHelper @@ -19,15 +21,18 @@ class SymbolicShapeInferenceHelperTest(unittest.TestCase): def _load_onnx(self, model_name): input_names = MODELS[model_name][0] base_path = "../onnx_models/" - import torch + config_modifier = ConfigModifier(None) + fusion_options = None + model_class = "AutoModel" with torch.no_grad(): export_onnx_model_from_pt( model_name, MODELS[model_name][1], MODELS[model_name][2], MODELS[model_name][3], - None, + model_class, + config_modifier, "../cache_models", base_path, input_names[:1], @@ -38,10 +43,9 @@ class SymbolicShapeInferenceHelperTest(unittest.TestCase): True, False, {}, + fusion_options, ) model_path = base_path + model_name.replace("-", "_") + "_1.onnx" - import onnx - return onnx.load_model(model_path) # TODO: use a static lightweight model for test