PR review

Author: jackzhxng

examples/models/llama/export_llama_lib.py

@@ -23,9 +23,7 @@
 import torch
 
 from executorch.devtools.etrecord import generate_etrecord
-from executorch.exir.passes.cache_pos_init_mutable_pass import (
-    CachePosToInitializedMutableBufferPass,
-)
+from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass
 
 from executorch.extension.llm.export.builder import DType, LLMEdgeManager
 
@@ -765,7 +763,7 @@ def _export_llama(args) -> LLMEdgeManager:  # noqa: C901
 
     additional_passes = []
     if args.model in TORCHTUNE_DEFINED_MODELS:
-        additional_passes = [CachePosToInitializedMutableBufferPass()]
+        additional_passes = [InitializedMutableBufferPass(["cache_pos"])]
     if args.generate_etrecord:
         if not builder_exported_to_edge.edge_manager:
             raise ValueError("Unable to generate etrecord due to missing edge manager.")

exir/passes/cache_pos_init_mutable_pass.py renamed to exir/passes/init_mutable_pass.py

@@ -5,10 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 
 
+from typing import List
+
 from executorch.exir.pass_base import ExportPass
 
 
-class CachePosToInitializedMutableBufferPass(ExportPass):
+class InitializedMutableBufferPass(ExportPass):
     """
     If the buffer has the name "cache_pos", such as in an kv_cache
     module with `self.register_buffer("cache_pos", torch.arange(10))`,
@@ -17,11 +19,13 @@ class CachePosToInitializedMutableBufferPass(ExportPass):
     an initialized state.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, patterns: List[str]) -> None:
         super().__init__()
+        self.patterns = patterns
 
     def placeholder(self, name: str, arg, meta):
-        if "cache_pos" in name:
-            meta["et_init_buffer"] = True
+        for pattern in self.patterns:
+            if pattern in name:
+                meta["et_init_buffer"] = True
 
         return super().placeholder(name, arg, meta)

extension/llm/export/builder.py

@@ -396,7 +396,9 @@ def to_backend(self, partitioners: Optional[List[Partitioner]]) -> "LLMEdgeManag
 
         return self
 
-    def to_executorch(self, passes: Optional[List[PassType]]) -> "LLMEdgeManager":
+    def to_executorch(
+        self, passes: Optional[List[PassType]] = None
+    ) -> "LLMEdgeManager":
         """
         Lower the model to executorch and get an ExecutorchProgram.
         """