Rename model.py to layers.py

Author: Sanggyu Lee

…yerWithKVCacheAndFusedAttention/model.py …erWithKVCacheAndFusedAttention/layers.pytest/modules/model/LlamaDecoderLayerWithKVCacheAndFusedAttention/model.py renamed to test/modules/model/LlamaDecoderLayerWithKVCacheAndFusedAttention/layers.py test/modules/model/LlamaDecoderLayerWithKVCacheAndFusedAttention/model.py renamed to test/modules/model/LlamaDecoderLayerWithKVCacheAndFusedAttention/layers.py

@@ -2,46 +2,6 @@
 prompt = "Lily picked up a flower."
 model_name = "Maykeye/TinyLLama-v0"
 
-captured_input = ()
-
-import copy, inspect, types
-
-from transformers.models.llama.modeling_llama import LlamaDecoderLayer
-
-forward_org = LlamaDecoderLayer.forward
-
-
-def capture_and_forward(self, *args, **kwargs):
-    global captured_input
-
-    # Prepare args tuple for TICO.convert()
-    # Get arg_names in positional args order using inspect
-    sig = inspect.signature(forward_org)
-    args_names = [
-        # signature includes `self`` and `kwargs``.
-        # Just retrieve the ordinary positional inputs only
-        name
-        for name in sig.parameters.keys()
-        if name
-        not in ("self", "kwargs", "use_cache", "position_ids", "output_attentions")
-    ]
-
-    args_dict = dict(zip(args_names, args))
-    args_dict.update(kwargs)
-
-    def populate_args(args_dict, filter):
-        for key in filter:
-            args_dict.pop(key, None)
-        args_tuple = tuple(args_dict.get(name, None) for name in args_names)
-        return copy.deepcopy(args_tuple)
-
-    if args_dict["past_key_value"].get_seq_length() != 0 and captured_input == ():
-        input_to_remove = []
-        captured_input = populate_args(args_dict, input_to_remove)
-
-    return forward_org(self, *args, **kwargs)
-
-
 # Tokenizer
 from transformers import AutoTokenizer
 
@@ -56,30 +16,35 @@ def populate_args(args_dict, filter):
     truncation=True,
 )
 
-
 # Generator
 import torch
 
 from transformers import AutoModelForCausalLM
 
 model = AutoModelForCausalLM.from_pretrained(model_name)
 model.eval()
-model.model.layers[0].forward = types.MethodType(
-    capture_and_forward, model.model.layers[0]
-)
-with torch.no_grad():
+
+from tico.utils.record_input import RecordingInput
+
+target_model = model.model.layers[0]
+condition_fn = lambda args_dict: args_dict["past_key_value"].get_seq_length() != 0
+
+with torch.no_grad(), RecordingInput(target_model, condition_fn) as rec:
     outputs = model.generate(
         **inputs,
         max_new_tokens=32,
         do_sample=False,
         pad_token_id=tokenizer.eos_token_id,
     )
+    captured_input = rec.captured_input
+
 generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 print(generated_text)
 
+
 # ATTENTION FUSER
 
-from typing import List, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 
 @torch.library.impl("circle::attention.llama", "CPU")
@@ -160,7 +125,6 @@ def forward_adapter(
 
 
 # Tico
-
 import tico
 
 from torch import nn
@@ -179,11 +143,13 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
-        past_key_values: Optional[Cache] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[
-            Tuple[torch.Tensor, torch.Tensor]
-        ] = None,  # necessary, but kept here for BC
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
+        **kwargs: Any,
     ) -> Tuple[
         torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
     ]:
@@ -192,7 +158,7 @@ def forward(
             layer_outputs = decoder_layer(
                 hidden_states,
                 attention_mask=attention_mask,
-                past_key_value=past_key_values,
+                past_key_value=past_key_value,
                 cache_position=cache_position,
                 position_embeddings=position_embeddings,
             )
@@ -205,4 +171,4 @@ def forward(
 LlamaAttention.forward = forward_adapter
 layers.eval()
 circle_model = tico.convert(layers, captured_input)
-circle_model.save(f"tinyllama.model.attn.circle")
+circle_model.save(f"tinyllama.layers.attn.circle")