Rewrite rules implementation for LLaMA-2/ LLaMA-3

onnxscript/ir/_core.py

@@ -753,7 +753,7 @@
 
 class SymbolicDim(_protocols.SymbolicDimProtocol, _display.PrettyPrintable):
     __slots__ = ("_value",)
-
+    
     def __init__(self, value: str | None) -> None:
         """Initialize a symbolic dimension.
 

onnxscript/rewriter/function_rule.py

@@ -86,22 +86,15 @@
 
 class FunctionRewriteRule(pattern.RewriteRule):
     FUNCTION_KEYWORD: str | tuple[str]
-    """The keyword to match the function name. If a tuple, any keyword will match."""
-
     PACKAGE_NAME: str
-    """The package name to match.
-
-    For example, 'transformers' to match for domain name 'pkg.transformers.4.36.2'.
-    """
-
     _opset_imports: dict[str, int]
     onnx_opset: onnxscript.values.Opset
 
-    def __init__(self, opset: onnxscript.values.Opset = onnxscript.opset18) -> None:  # type: ignore[has-type]
+    def __init__(self, opset: onnxscript.values.Opset = onnxscript.opset18) -> None:
         self.onnx_opset = opset
 
     def _match_function(self, function: ir.Function, pkg_name: str) -> bool:
-        # TODO: Consolidate more checks from `compose_new_function` to here.
+        print("----> Checking function:", function.name, "in package:", pkg_name)
         if pkg_name != self.PACKAGE_NAME:
             logger.info(
                 "Rule %s did not match function %s::%s. Package name mismatch '%s' != '%s'.",
@@ -113,11 +106,15 @@
             )
             return False
         if isinstance(self.FUNCTION_KEYWORD, str):
-            return function.name.find(self.FUNCTION_KEYWORD) != -1
+            match = function.name.find(self.FUNCTION_KEYWORD) != -1
+            print(f"----> Function name '{function.name}' match with '{self.FUNCTION_KEYWORD}': {match}")
+            return match
         elif isinstance(self.FUNCTION_KEYWORD, tuple):
-            return any(function.name.find(keyword) != -1 for keyword in self.FUNCTION_KEYWORD)
+            match = any(function.name.find(keyword) != -1 for keyword in self.FUNCTION_KEYWORD)
+            print(f"----> Function name '{function.name}' match with any of '{self.FUNCTION_KEYWORD}': {match}")
+            return match
         else:
-            raise ValueError(  # noqa: TRY004
+            raise ValueError(
                 f"Function keyword must be str or tuple, got {self.FUNCTION_KEYWORD}"
             )
 
@@ -129,10 +126,17 @@
                 return node
         return None
 
+    def _find_function_by_name(
+        self, function: ir.Function, keyword: str
+    ) -> ir.Function | None:
+        for node in function:
+            if node.name.find(keyword) != -1:
+                return node
+        return None
+
     def _find_node_by_type(
         self, function: ir.Function, domain: str, op_type: str
     ) -> ir.Node | None:
-        # Repeat
         for node in function:
             if node.domain == domain and node.op_type == op_type:
                 return node
@@ -141,18 +145,12 @@
     def compose_new_function(
         self, old_function: ir.Function, pkg_version: version.Version | None
     ) -> ir.Function:
-        """Compose a new function from the old function.
-
-        Returns:
-            A tuple of the new function and the opset imports.
-
-        Raises:
-            FunctionRewriteError: If the rewrite fails.
-        """
-        # self._version_controller is created in the subclass
-        func = self._version_controller.dispatch(pkg_version)  # type: ignore[attr-defined]
+        print("----> (2) pkg_version", pkg_version, "old_function", old_function.name)
+        func = self._version_controller.dispatch(pkg_version)
         if func is not None:
+            print("----> (2.5) Dispatch function found, applying...")
             new_function = func(self, old_function)
+            print("----> (2.6) New function created.")
             return new_function
         raise FunctionRewriteError(
             f"No rewrite implementation for package version {pkg_version}."
@@ -163,6 +161,7 @@
     ) -> tuple[ir.OperatorIdentifier, ir.Function] | None:
         try:
             pkg_name, pkg_version = parse_domain(function.domain)
+            print("----> (1) Parsed domain, pkg_name:", pkg_name, "pkg_version:", pkg_version)
         except FunctionRewriteError as e:
             logger.warning("Could not parse domain: %s", e)
             return None
@@ -177,33 +176,39 @@
             )
 
         if not self._match_function(function, pkg_name):
+            print("----> (1.5) Function does not match.")
             return None
         logger.info(
             "Rule %s matched function %s::%s",
             self.__class__.__name__,
             function.domain,
             function.name,
         )
+        print("----> (1.6) Function matched.")
         try:
             new_function = self.compose_new_function(function, pkg_version)
         except FunctionRewriteError as e:
             logger.warning("Could not rewrite function: %s", e)
             return None
 
+        if not hasattr(new_function, 'name'):
+            logger.error("new_function does not have a 'name' attribute. Received: %s", type(new_function))
+            return None
+
         new_function.name = function.name
         new_function.domain = function.domain
 
         return function.identifier(), new_function
 
     def try_rewrite(self, model: ir.Model, value) -> bool:
         raise NotImplementedError(
-            "Use `try_rewrite_function` instead for function based rewrites."
+            "Use try_rewrite_function instead for function based rewrites."
         )
 
     def apply_to_model(
         self, model: ir.Model, *, commute: bool = False
     ) -> tuple[int, ir.Model]:
-        del commute  # unused
+        del commute
 
         old_function_to_new_function: dict[ir.OperatorIdentifier, ir.Function] = {}
         for function in model.functions.values():

onnxscript/rewriter/onnxruntime/transformers/__init__.py

@@ -18,4 +18,12 @@
     layernorm.LNRewriteRule,
     fastgelu.GeluRewriteRule,
     biassplitgelu.GegluRewriteRule,
+    multihead_attention.MHALlama2RewriteRule,
+    multihead_attention.GQALlama3RewriteRule,
+    multihead_attention.AttentionRewriteRule,
+    multihead_attention.MLP3RewriteRule,
+    multihead_attention.GQALlama3RewriteRuleFirstAttention,
+    multihead_attention.MLPRewriteRule,
+    multihead_attention.GQALlamaRewriteRule,
+
 ]

onnxscript/rewriter/onnxruntime/transformers/layernorm.py

@@ -1,43 +1,38 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from __future__ import annotations
-
-import logging
-
 import onnxscript
 from onnxscript import ir
 from onnxscript.rewriter import _ir_utils, function_rule
-
+import logging
 logger = logging.getLogger(__name__)
 
-
 class LNRewriteRule(function_rule.FunctionRewriteRule):
-    FUNCTION_KEYWORD = "layernorm"
+    FUNCTION_KEYWORD = "norm"
     PACKAGE_NAME = "transformers"
     _version_controller = function_rule.VersionController()
 
-    @_version_controller.register_version()
+    @_version_controller.register_version(min_version="4.40", max_version="4.50")
     def _fusion(self, function: ir.Function) -> ir.Function:
-        # TODO(bowbao): Might be more desirable to annotate as attribute in nn.Module
-        aten_add_node = self._find_node_by_type(function, "", "Add")
+       # depending on graph, you may have to find node by name or type instead of function
+        aten_add_node = self._find_function_by_name(function, "aten_add")
         if aten_add_node is None:
-            raise function_rule.FunctionRewriteError("Could not find Add node")
 
+            raise function_rule.FunctionRewriteError("Could not find Add node")
         eps_ir_value = _ir_utils.propagate_const_value(aten_add_node.inputs[1])
         eps_const_value = eps_ir_value.const_value
         if eps_const_value is None:
+            print("could not find")
             raise function_rule.FunctionRewriteError("Could not find eps")
         eps_numpy_value = eps_const_value.numpy()
         eps = eps_numpy_value.item()
         logger.info("eps: %s", eps)
 
-        # TODO(ORT): SimplifiedLayerNormalization in ort is defined under onnx domain.
-        # https://github.com/microsoft/onnxruntime/issues/7573
-        # msft_op = onnxscript.values.Opset("com.microsoft", 1)
         op = self.onnx_opset
+        msft_op = onnxscript.values.Opset("com.microsoft", 1)
 
         def ln(input, weight):
-            return op.SimplifiedLayerNormalization(
+            return msft_op.SimplifiedLayerNormalization(
                 input, weight, axis=-1, epsilon=eps, stash_type=1
             )