check mask value; add test for xlm-roberta

microsoft · Nov 20, 2024 · 899ed78 · 899ed78
1 parent 9eed369
commit 899ed78
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/onnxruntime/python/tools/transformers/fusion_attention.py b/onnxruntime/python/tools/transformers/fusion_attention.py
@@ -1140,9 +1140,15 @@ def fuse(self, node, input_name_to_nodes, output_name_to_node):
 
         if not is_no_mask_attention and len(mask_nodes) > 1:
             _, mul_val = self.model.get_constant_input(mask_nodes[0])
-            if mul_val != -10000:
-                assert mul_val < 0
-                self.mask_filter_value = mul_val
+            # The mask value shall be a float scalar (usually is the lowest float value).
+            if (
+                (mul_val is None)
+                or not (isinstance(mul_val, np.ndarray) and mul_val.size == 1)
+                or (float(mul_val) >= 0)
+            ):
+                return
+            if float(mul_val) != -10000:
+                self.mask_filter_value = float(mul_val)
 
         if matmul_v.input[0] == root_input and matmul_q.input[0] == root_input and matmul_k.input[0] == root_input:
             mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0]) if not is_no_mask_attention else None

diff --git a/onnxruntime/test/python/transformers/test_optimizer_huggingface_bert.py b/onnxruntime/test/python/transformers/test_optimizer_huggingface_bert.py
@@ -10,6 +10,7 @@
 
 import shutil
 import unittest
+from pathlib import Path
 
 import torch
 from parity_utilities import find_transformers_source
@@ -36,7 +37,10 @@ def run_optimizer_on_model(
         use_external_data_format=False,
         model_type="bert",
     ):
-        shutil.rmtree("./onnx_models", ignore_errors=True)
+        onnx_dir = Path(".") / "onnx_models" / model_name
+        shutil.rmtree(onnx_dir, ignore_errors=True)
+
+        Path(onnx_dir).mkdir(parents=True, exist_ok=True)
 
         model_fusion_statistics = {}
 
@@ -54,7 +58,7 @@ def run_optimizer_on_model(
                 model_class=model_class,
                 config_modifier=config_modifier,
                 cache_dir=default_cache_path,
-                onnx_dir="./onnx_models",
+                onnx_dir=str(onnx_dir),
                 input_names=input_names[:inputs_count],
                 use_gpu=False,
                 precision=Precision.FLOAT32,
@@ -79,8 +83,10 @@ def run_optimizer_on_model(
             "SkipLayerNormalization": expected_fusion_result_list[6],
         }
 
+        node_count = None
         for value in model_fusion_statistics.values():
             node_count = value
+        self.assertIsNotNone(node_count)
 
         actual_node_count = {}
         for op_type in expected_node_count:
@@ -125,6 +131,7 @@ def test_bert(self):
 
     def test_roberta(self):
         model_name = "hf-internal-testing/tiny-random-roberta"
+        # TODO: EmbedLayerNormalization fusion.
         self.run_optimizer_on_model(model_name, [0, 5, 0, 0, 5, 1, 10], inputs_count=1)
         self.run_optimizer_on_model(model_name, [0, 5, 0, 0, 5, 1, 10], inputs_count=2)
 
@@ -133,6 +140,12 @@ def test_distillbert(self):
         self.run_optimizer_on_model(model_name, [1, 5, 0, 0, 5, 0, 10], inputs_count=1)
         self.run_optimizer_on_model(model_name, [1, 5, 0, 0, 5, 0, 10], inputs_count=2)
 
+    def test_xlm_roberta(self):
+        model_name = "hf-internal-testing/tiny-xlm-roberta"
+        # TODO: EmbedLayerNormalization fusion.
+        self.run_optimizer_on_model(model_name, [0, 2, 0, 0, 2, 1, 4], inputs_count=1)
+        self.run_optimizer_on_model(model_name, [0, 2, 0, 0, 2, 1, 4], inputs_count=2)
+
 
 if __name__ == "__main__":
     unittest.main()