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

microsoft#14739) …GPT_Attention) Some EPs require that onnxruntime and optimum optimizations are turned off in order to run correctly. Allowing this option during test runs allows the EP and library to perform their own optimization and be more representative of actual use case conditions. Important for EPs like MIGraphX which require optimizations to be offer for certain operations  Allow flags to turn off optimizations and add verbose output to confirm which EP is being used for the inference run and validate fallbacks  Related to: microsoft#14702 & microsoft#14700 --------- Signed-off-by: Ted Themistokleous <[email protected]> Co-authored-by: Ted Themistokleous <[email protected]>
ROCm · May 24, 2023 · 099f6d5 · 099f6d5
1 parent 81120e9
commit 099f6d5
Show file tree

Hide file tree

Showing 5 changed files with 119 additions and 38 deletions.
diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py
@@ -58,6 +58,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.
@@ -96,6 +97,10 @@ def optimize_by_onnxruntime(
 
     sess_options.optimized_model_filepath = optimized_model_path
 
+    if verbose:
+        print("Using onnxruntime to optimize model - Debug level Set to verbose")
+        sess_options.log_severity_level = 0
+
     kwargs = {}
     if disabled_optimizers:
         kwargs["disabled_optimizers"] = disabled_optimizers
@@ -112,7 +117,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"]
@@ -187,6 +191,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.
 
@@ -258,6 +263,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.
@@ -267,6 +273,7 @@ def optimize_model(
             use_gpu=False,
             opt_level=1,
             disabled_optimizers=disabled_optimizers,
+            verbose=verbose,
         )
 
     if only_onnxruntime and not temp_model_path:

diff --git a/onnxruntime/test/python/transformers/parity_utilities.py b/onnxruntime/test/python/transformers/parity_utilities.py
@@ -3,14 +3,44 @@
 # 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(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,
+        dest="optimize",
+        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 +104,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 +163,26 @@ 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
 
     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 = ["CPUExecutionProvider"]
 
     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")
 
     return InferenceSession(onnx_model_path, sess_options, providers=execution_providers)
@@ -172,7 +210,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)
 

diff --git a/onnxruntime/test/python/transformers/test_parity_gelu.py b/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
+    optimized = 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.optimize
+
+    unittest.main(argv=remaining_args)