refactor-attention

Author: lucylq

examples/models/llama/attention.py

@@ -162,7 +162,16 @@ def forward(
 
 @register_attention("mha")
 class AttentionMHA(Attention):
-    def __init__(self, args: ModelArgs, layer_id: int, rope: Rope):
+    def __init__(
+        self,
+        args: ModelArgs,
+        layer_id: int,
+        rope: Rope,
+        # wq: nn.Module,
+        # wk: nn.Module,
+        # wv: nn.Module,
+        # wo: nn.Module,
+    ):
         super().__init__()
         self.use_kv_cache = args.use_kv_cache
         self.n_heads = args.n_heads
@@ -196,6 +205,11 @@ def __init__(self, args: ModelArgs, layer_id: int, rope: Rope):
         )
         self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False)
 
+        # self.wq = wq
+        # self.wk = wk
+        # self.wv = wv
+        # self.wo = wo
+
         self.layer_id = layer_id
 
         self.rope = rope

examples/models/llama/llama_transformer.py

@@ -19,6 +19,7 @@
 
 from executorch.examples.models.llama.model_args import ModelArgs
 from executorch.examples.models.llama.norm import RMSNorm
+
 from executorch.examples.models.llama.rope import Rope
 from torch import nn
 
@@ -83,7 +84,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, layer_id: int, args: ModelArgs, rope: Rope):
+    def __init__(self, args: ModelArgs, attention: Attention):
         super().__init__()
         self.use_kv_cache = args.use_kv_cache
         self.n_heads = args.n_heads
@@ -94,8 +95,8 @@ def __init__(self, layer_id: int, args: ModelArgs, rope: Rope):
                 f"Unknown attention type: {args.attention_type}. "
                 f"Available: {list(ATTENTION_REGISTRY.keys())}"
             )
-        cls = ATTENTION_REGISTRY[args.attention_type]
-        self.attention = cls(args, layer_id, rope)
+
+        self.attention = attention
         if args.moe:
             self.block_sparse_moe = MOEFeedForward(args)
         else:
@@ -117,7 +118,7 @@ def forward(self, x, freqs_cos, freqs_sin, attn_options: ForwardOptions):  # x:
 
 
 class Transformer(nn.Module):
-    def __init__(self, params: ModelArgs):
+    def __init__(self, params: ModelArgs, layers: nn.ModuleList, rope: Rope):
         super().__init__()
         self.params = params
         self.vocab_size = params.vocab_size
@@ -130,10 +131,8 @@ def __init__(self, params: ModelArgs):
             if self.apply_embedding
             else None
         )
-        self.rope = Rope(params)
-        self.layers = torch.nn.ModuleList()
-        for layer_id in range(params.n_layers):
-            self.layers.append(TransformerBlock(layer_id, params, self.rope))
+        self.layers = layers
+        self.rope = rope
         self.norm = RMSNorm(params.dim, eps=params.norm_eps)
         self.output = (
             nn.Linear(params.dim, params.vocab_size, bias=False)

examples/models/llama/model.py

@@ -18,6 +18,7 @@
 from executorch.examples.models.llama.llama_transformer import Transformer
 
 from executorch.examples.models.llama.model_args import ModelArgs
+from executorch.examples.models.llama.rope import Rope
 
 try:
     from .fairseq2 import convert_to_llama_checkpoint
@@ -173,7 +174,19 @@ def __init__(self, **kwargs):
         # They possess all other metadata a tensor carries such as size, stride, requires_grad.
         with torch.device("meta"):
             # Model itself is loaded in default dtype, fp32.
-            self.model_ = Transformer(model_args)
+
+            # Construct attention layers.
+            rope = Rope(model_args)
+            layers = nn.ModuleList()
+            cls = ATTENTION_REGISTRY[model_args.attention_type]
+            for layer_id in range(model_args.n_layers):
+                attention = cls(model_args, layer_id, rope)
+                transformer_block = TransformerBlock(model_args, attention)
+                layers.append(transformer_block)
+
+            # Construct transformer model.
+            self.model_ = Transformer(model_args, layers, rope)
+
             # Get checkpoint dtype.
             if checkpoint:
                 self.model_.checkpoint_dtype = get_checkpoint_dtype(checkpoint)