Add verbose and optimization args for parity tests (Gelu, Layernorm, …

onnxruntime/python/tools/transformers/optimizer.py

@@ -65,6 +65,7 @@ def optimize_by_onnxruntime(
     optimized_model_path: Optional[str] = None,
     opt_level: Optional[int] = 99,
     disabled_optimizers=[],
+    verbose=False,
 ) -> str:
     """
     Use onnxruntime to optimize model.
@@ -103,6 +104,10 @@ def optimize_by_onnxruntime(
 
     sess_options.optimized_model_filepath = optimized_model_path
 
+    if verbose:
+        print("Using onnxruntime to optomize model - Debug level Set to verbose")
+        sess_options.log_severity_level = 0
+
     kwargs = {}
     if disabled_optimizers:
         kwargs["disabled_optimizers"] = disabled_optimizers
@@ -119,7 +124,6 @@ def optimize_by_onnxruntime(
         elif torch_version.hip:
             gpu_ep.append("MIGraphXExecutionProvider")
             gpu_ep.append("ROCMExecutionProvider")
-
         session = onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=gpu_ep, **kwargs)
         assert not set(onnxruntime.get_available_providers()).isdisjoint(
             ["CUDAExecutionProvider", "ROCMExecutionProvider", "MIGraphXExecutionProvider"]
@@ -194,6 +198,7 @@ def optimize_model(
     opt_level: Optional[int] = None,
     use_gpu: bool = False,
     only_onnxruntime: bool = False,
+    verbose=False,
 ):
     """Optimize Model by OnnxRuntime and/or python fusion logic.
 
@@ -265,6 +270,7 @@ def optimize_model(
             use_gpu=use_gpu,
             opt_level=opt_level,
             disabled_optimizers=disabled_optimizers,
+            verbose=verbose,
         )
     elif opt_level == 1:
         # basic optimizations (like constant folding and cast elimination) are not specified to execution provider.
@@ -274,6 +280,7 @@ def optimize_model(
             use_gpu=False,
             opt_level=1,
             disabled_optimizers=disabled_optimizers,
+            verbose=verbose,
         )
 
     if only_onnxruntime and not temp_model_path:

onnxruntime/test/python/transformers/parity_utilities.py

@@ -3,14 +3,43 @@
 # Licensed under the MIT License.  See License.txt in the project root for
 # license information.
 # -------------------------------------------------------------------------
-
+import argparse
 import os
 import sys
 
 import numpy
 import torch
 
 
+def parse_arguments(argv=sys.argv[1:], namespace_filter=None):
+
+    parser = argparse.ArgumentParser()
+
+    # useful EPs that don't require the use of optmizer.py
+    parser.add_argument(
+        "-n",
+        "--no_optimize",
+        required=False,
+        action="store_false",
+        default=True,
+        help="Turn off onnxruntime optimizers (Default off optimizers ON)",
+    )
+
+    # useful for debugging and viewing state during test runs
+    parser.add_argument(
+        "-l",
+        "--log_verbose",
+        required=False,
+        action="store_true",
+        default=False,
+        help="Set Onnxruntime log_serverity_level=0 (VERBOSE) ",
+    )
+
+    args, remaining_args = parser.parse_known_args(namespace=namespace_filter)
+
+    return args, sys.argv[:1] + remaining_args
+
+
 def find_transformers_source(sub_dir_paths=[]):
     source_dir = os.path.join(
         os.path.dirname(__file__),
@@ -74,13 +103,16 @@ def optimize_onnx(
     expected_op=None,
     use_gpu=False,
     opt_level=None,
+    verbose=False,
 ):
     if find_transformers_source():
         from optimizer import optimize_model
     else:
         from onnxruntime.transformers.optimizer import optimize_model
 
-    onnx_model = optimize_model(input_onnx_path, model_type="gpt2", use_gpu=use_gpu, opt_level=opt_level)
+    onnx_model = optimize_model(
+        input_onnx_path, model_type="gpt2", use_gpu=use_gpu, opt_level=opt_level, verbose=verbose
+    )
     onnx_model.save_model_to_file(optimized_onnx_path)
 
     if expected_op is not None:
@@ -130,21 +162,28 @@ def compare_outputs(torch_outputs, ort_outputs, atol=1e-06, verbose=True):
     return is_all_close, max(max_abs_diff)
 
 
-def create_ort_session(onnx_model_path, use_gpu=True):
+def create_ort_session(onnx_model_path, use_gpu=True, optimized=True, verbose=False):
     from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions
     from onnxruntime import __version__ as onnxruntime_version
+    from onnxruntime import get_available_providers
 
     sess_options = SessionOptions()
     sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL
     sess_options.intra_op_num_threads = 2
     sess_options.log_severity_level = 2
+
+    if verbose:
+        sess_options.log_severity_level = 0
+
     execution_providers = []
 
     if use_gpu:
         if torch.version.cuda:
             execution_providers.append("CUDAExecutionProvider")
         elif torch.version.hip:
-            execution_providers.append("MIGraphXExecutionProvider")
+            if not optimized:
+                execution_providers.append("MIGraphXExecutionProvider")
+
             execution_providers.append("ROCMExecutionProvider")
 
     execution_providers.append("CPUExecutionProvider")
@@ -174,7 +213,7 @@ def run_parity(
     passed_cases = 0
     max_diffs = []
     printed = False  # print only one sample
-    ort_session = create_ort_session(onnx_model_path, device.type == "cuda")
+    ort_session = create_ort_session(onnx_model_path, device.type == "cuda", optimized=optimized, verbose=verbose)
     for i in range(test_cases):
         input_hidden_states = create_inputs(batch_size, sequence_length, hidden_size, float16, device)
 

onnxruntime/test/python/transformers/test_parity_gelu.py

@@ -85,6 +85,7 @@ def run(
     formula=0,
     sequence_length=2,
     fp32_gelu_op=True,
+    verbose=False,
 ):
     test_name = f"device={device}, float16={float16}, optimized={optimized}, batch_size={batch_size}, sequence_length={sequence_length}, hidden_size={hidden_size}, formula={formula}, fp32_gelu_op={fp32_gelu_op}"
     print(f"\nTesting: {test_name}")
@@ -108,6 +109,7 @@ def run(
             Gelu.get_fused_op(formula),
             use_gpu=use_gpu,
             opt_level=2 if use_gpu else None,
+            verbose=verbose,
         )
         onnx_path = optimized_onnx_path
     else:
@@ -123,7 +125,7 @@ def run(
         device,
         optimized,
         test_cases,
-        verbose=False,
+        verbose,
     )
 
     # clean up onnx file
@@ -135,8 +137,10 @@ def run(
 
 
 class TestGeluParity(unittest.TestCase):
+    verbose = False
+    optmized = True
+
     def setUp(self):
-        self.optimized = True  # Change it to False if you want to test parity of non optimized ONNX
         self.test_cases = 100  # Number of test cases per test run
         self.sequence_length = 2
         self.hidden_size = 768
@@ -159,6 +163,7 @@ def run_test(
         formula,
         enable_assert=True,
         fp32_gelu_op=True,
+        verbose=False,
     ):
         if float16 and device.type == "cpu":  # CPU does not support FP16
             return
@@ -172,11 +177,12 @@ def run_test(
             formula,
             self.sequence_length,
             fp32_gelu_op,
+            verbose,
         )
         if enable_assert:
             self.assertTrue(num_failure == 0, "Failed: " + test_name)
 
-    def run_one(self, optimized, device, hidden_size=768, formula=0):
+    def run_one(self, optimized, device, hidden_size=768, formula=0, verbose=False):
         for batch_size in [4]:
             self.run_test(
                 batch_size,
@@ -186,6 +192,7 @@ def run_one(self, optimized, device, hidden_size=768, formula=0):
                 device=device,
                 formula=formula,
                 enable_assert=formula in self.formula_must_pass,
+                verbose=verbose,
             )
 
             self.run_test(
@@ -197,6 +204,7 @@ def run_one(self, optimized, device, hidden_size=768, formula=0):
                 formula=formula,
                 enable_assert=formula in self.formula_must_pass,
                 fp32_gelu_op=True,
+                verbose=verbose,
             )
 
             self.run_test(
@@ -208,12 +216,13 @@ def run_one(self, optimized, device, hidden_size=768, formula=0):
                 formula=formula,
                 enable_assert=formula in self.formula_must_pass,
                 fp32_gelu_op=False,
+                verbose=verbose,
             )
 
     def test_cpu(self):
         cpu = torch.device("cpu")
         for i in self.formula_to_test:
-            self.run_one(self.optimized, cpu, hidden_size=self.hidden_size, formula=i)
+            self.run_one(self.optimized, cpu, hidden_size=self.hidden_size, formula=i, verbose=self.verbose)
 
     def test_cuda(self):
         if not torch.cuda.is_available():
@@ -223,8 +232,13 @@ def test_cuda(self):
         else:
             gpu = torch.device("cuda")
             for i in self.formula_to_test:
-                self.run_one(self.optimized, gpu, hidden_size=self.hidden_size, formula=i)
+                self.run_one(self.optimized, gpu, hidden_size=self.hidden_size, formula=i, verbose=self.verbose)
 
 
 if __name__ == "__main__":
-    unittest.main()
+    args, remaining_args = parse_arguments(namespace_filter=unittest)
+
+    TestGeluParity.verbose = args.log_verbose
+    TestGeluParity.optimized = args.no_optimize
+
+    unittest.main(argv=remaining_args)