Add options --disable_layer_norm, --disable_gelu and --enable_gelu_approximation (#3750)

onnxruntime/python/tools/bert/BertOnnxModel.py

@@ -18,11 +18,14 @@ logger = logging.getLogger(__name__)
 class BertOptimizationOptions:
 
     def __init__(self, model_type):
+        self.enable_gelu = True
+        self.enable_layer_norm = True
         self.enable_attention = True
         self.enable_skip_layer_norm = True
         self.enable_embed_layer_norm = True
         self.enable_bias_skip_layer_norm = True
         self.enable_bias_gelu = True
+        self.enable_gelu_approximation = False
 
         if model_type == 'gpt2':
             self.enable_skip_layer_norm = False
@@ -515,6 +518,23 @@ class BertOnnxModel(OnnxModel):
         self.remove_nodes(nodes_to_remove)
         self.add_nodes(nodes_to_add)
 
+    def gelu_approximation(self):
+        nodes_to_add = []
+        nodes_to_remove = self.get_nodes_by_op_type("Gelu") + self.get_nodes_by_op_type("BiasGelu")
+        for node in nodes_to_remove:
+            new_node = onnx.helper.make_node("FastGelu",
+                                             inputs=node.input,
+                                             outputs=node.output,
+                                             name=self.create_node_name("FastGelu", node.op_type + "_Approximation"))
+            new_node.domain = "com.microsoft"
+            nodes_to_add.append(new_node)
+
+        if len(nodes_to_add) > 0:
+            logger.info(f"Gelu approximation count:{len(nodes_to_add)}")
+
+        self.remove_nodes(nodes_to_remove)
+        self.add_nodes(nodes_to_add)
+
     def fuse_add_bias_skip_layer_norm(self):
         logger.debug(f"start Bias and SkipLayerNormalization fusion...")
         input_name_to_nodes = self.input_name_to_nodes()
@@ -1129,9 +1149,11 @@ class BertOnnxModel(OnnxModel):
         self.prune_graph()
 
     def optimize(self, options: BertOptimizationOptions = None):
-        self.fuse_layer_norm()
+        if (options is None) or options.enable_layer_norm:
+            self.fuse_layer_norm()
 
-        self.fuse_gelu()
+        if (options is None) or options.enable_gelu:
+            self.fuse_gelu()
 
         self.preprocess()
 
@@ -1159,6 +1181,9 @@ class BertOnnxModel(OnnxModel):
             # Fuse SkipLayerNormalization and Add Bias before it.
             self.fuse_add_bias_skip_layer_norm()
 
+        if (options is not None and options.enable_gelu_approximation):
+            self.gelu_approximation()
+
         self.remove_unused_constant()
 
         # Use symbolic batch dimension in input and output.

onnxruntime/python/tools/bert/README.md

@@ -55,7 +55,7 @@ See below for description of some options of bert_model_optimization.py:
 - **input**: input model path
 - **output**: output model path
 - **model_type**: (*defaul: bert*)
-    There are 4 model types: *bert* (exported by PyTorch), *bert_tf* (BERT exported by PyTorch), *bert_keras* (BERT exported by keras2onnx) and *gpt2* (exported by PyTorch) respectively.
+    There are 4 model types: *bert* (exported by PyTorch), *bert_tf* (BERT exported by tf2onnx), *bert_keras* (BERT exported by keras2onnx) and *gpt2* (exported by PyTorch) respectively.
 - **num_heads**: (*default: 12*)
     Number of attention heads. BERT-base and BERT-large has 12 and 16 respectively.
 - **hidden_size**: (*default: 768*)

onnxruntime/python/tools/bert/bert_model_optimization.py

@@ -156,6 +156,24 @@ def parse_arguments():
                         help="disable Add Bias and Gelu/FastGelu fusion")
     parser.set_defaults(disable_bias_gelu=False)
 
+    parser.add_argument('--disable_layer_norm',
+                        required=False,
+                        action='store_true',
+                        help="disable LayerNormalization fusion")
+    parser.set_defaults(disable_layer_norm=False)
+
+    parser.add_argument('--disable_gelu',
+                        required=False,
+                        action='store_true',
+                        help="disable Gelu fusion")
+    parser.set_defaults(disable_gelu=False)
+
+    parser.add_argument('--enable_gelu_approximation',
+                        required=False,
+                        action='store_true',
+                        help="enable Gelu/BiasGelu to FastGelu conversion")
+    parser.set_defaults(enable_gelu_approximation=False)
+
     parser.add_argument('--verbose', required=False, action='store_true')
     parser.set_defaults(verbose=False)
 
@@ -173,6 +191,10 @@ def parse_arguments():
 
 def get_optimization_options(args):
     optimization_options = BertOptimizationOptions(args.model_type)
+    if args.disable_gelu:
+        optimization_options.enable_gelu = False
+    if args.disable_layer_norm:
+        optimization_options.enable_layer_norm = False
     if args.disable_attention:
         optimization_options.enable_attention = False
     if args.disable_skip_layer_norm:
@@ -183,6 +205,8 @@ def get_optimization_options(args):
         optimization_options.enable_bias_skip_layer_norm = False
     if args.disable_bias_gelu:
         optimization_options.enable_bias_gelu = False
+    if args.enable_gelu_approximation:
+        optimization_options.enable_gelu_approximation = True
     return optimization_options