Add migraphx ep for llama2 scripts

ROCm · Aug 28, 2024 · 29f9c36 · 29f9c36
1 parent c604fe9
commit 29f9c36
Show file tree

Hide file tree

Showing 7 changed files with 922 additions and 41 deletions.
diff --git a/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc b/onnxruntime/core/providers/migraphx/migraphx_execution_provider.cc
@@ -834,6 +834,7 @@ GetUnsupportedNodeIndices(const GraphViewer& graph_viewer,
                                                     "GlobalMaxPool",
                                                     "Greater",
                                                     "GreaterOrEqual",
+                                                    "GroupQueryAttention",
                                                     "HardSigmoid",
                                                     "HardSwish",
                                                     "Identity",
@@ -903,8 +904,10 @@ GetUnsupportedNodeIndices(const GraphViewer& graph_viewer,
                                                     "Shape",
                                                     "Sigmoid",
                                                     "Sign",
+                                                    "SimplifiedLayerNormalization",
                                                     "Sin",
                                                     "Sinh",
+                                                    "SkipSimplifiedLayerNormalization",
                                                     "Slice",
                                                     "Softmax",
                                                     "Softplus",

diff --git a/onnxruntime/python/tools/transformers/models/llama/benchmark.py b/onnxruntime/python/tools/transformers/models/llama/benchmark.py
@@ -593,7 +593,7 @@ def get_args(rank=0):
         "--device",
         type=str,
         default="cuda" if torch.cuda.is_available() else "cpu",
-        choices=["cpu", "cuda", "rocm"],
+        choices=["cpu", "cuda", "rocm", "migraphx"],
     )
     parser.add_argument("-id", "--device-id", type=int, default=0)
     parser.add_argument("-w", "--warmup-runs", type=int, default=5)
@@ -622,9 +622,11 @@ def get_args(rank=0):
     # Set runtime properties
     if "ort" in args.benchmark_type:
         setattr(args, "execution_provider", f"{args.device.upper()}ExecutionProvider")  # noqa: B010
+        if args.device == "migraphx":
+            setattr(args, "execution_provider", "MIGraphXExecutionProvider")
         if args.execution_provider == "CUDAExecutionProvider":
             args.execution_provider = (args.execution_provider, {"device_id": rank})
-        elif args.execution_provider == "ROCMExecutionProvider":
+        elif args.execution_provider == "ROCMExecutionProvider" or args.execution_provider == "MIGraphXExecutionProvider":
             args.execution_provider = (args.execution_provider, {"device_id": rank})
             args.device = "cuda"