[ET-VK][ez] Apply rotary embedding as Module (#6422)

Pull Request resolved: #6391 ## Context As title. Wrap the `apply_rotary_emb` function call in a `nn.Module` to make it easy to perform a source module replacement for rotary embedding calculation. The Vulkan delegate will use the source module replacement technique to insert a custom op to calculate rotary embeddings. ghstack-source-id: 249175724 @exported-using-ghexport Differential Revision: [D64697589](https://our.internmc.facebook.com/intern/diff/D64697589/) Co-authored-by: Stephen Jia <[email protected]>
pytorch · Oct 21, 2024 · 0309854 · 0309854
1 parent 1247545
commit 0309854
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py
@@ -15,10 +15,10 @@
 import torch.nn.functional as F
 
 from executorch.examples.models.llama.rope import (
-    apply_rotary_emb,
     hf_apply_rotary_emb,
     hf_precompute_freqs_cis,
     precompute_freqs_cis,
+    RotaryEmbedding,
 )
 
 from torch import nn
@@ -311,7 +311,7 @@ def __init__(self, args: ModelArgs, layer_id: int):
         if args.use_hf_rope:
             self.apply_rotary_emb = hf_apply_rotary_emb
         else:
-            self.apply_rotary_emb = apply_rotary_emb
+            self.apply_rotary_emb = RotaryEmbedding()
 
     def forward(
         self,

diff --git a/examples/models/llama/rope.py b/examples/models/llama/rope.py
@@ -92,6 +92,21 @@ def apply_rotary_emb(
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
 
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+        self,
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+    ):
+        xq_out, xk_out = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
+        return xq_out, xk_out
+
+
 # ======================= HuggingFace Implementation ========================