Things work

Author: jackzhxng

examples/models/llama2/runner/eager.py

@@ -10,7 +10,6 @@
 
 import torch
 
-from examples.models.llama2.llama_transformer import ModelArgs
 from executorch.examples.models.model_factory import EagerModelFactory
 
 from .generation import LlamaRunner
@@ -24,13 +23,13 @@ class EagerLlamaRunner(LlamaRunner):
     def __init__(self, args):
         with open(args.params, "r") as f:
             params = json.loads(f.read())
-        model_args: ModelArgs = ModelArgs(
-            max_seq_len=args.max_len,
+        super().__init__(
+            tokenizer_path=args.tokenizer,
+            max_seq_len=args.max_seq_len,
             max_batch_size=1,
             use_kv_cache=True,
-            **params,
+            vocab_size=params["vocab_size"],
         )
-        super().__init__(tokenizer_path=args.tokenizer, model_args=model_args)
         self.model, _, _, _ = EagerModelFactory.create_model(
             "llama2",
             "Llama2Model",

examples/models/llama2/runner/generation.py

@@ -51,10 +51,19 @@ def next_token(logits: torch.Tensor, temperature: float, top_p: float) -> int:
 
 
 class LlamaRunner(ABC):
-    def __init__(self, tokenizer_path: str, model_args: ModelArgs):
-        self.params = model_args
+    def __init__(
+        self,
+        tokenizer_path: str,
+        max_seq_len: int,
+        max_batch_size: int,
+        use_kv_cache: bool,
+        vocab_size: int,
+    ):
+        self.max_seq_len = max_seq_len
+        self.max_batch_size = max_batch_size
+        self.use_kv_cache = use_kv_cache
         self.tokenizer = Tokenizer(tokenizer_path)
-        assert model_args.vocab_size == self.tokenizer.n_words
+        assert vocab_size == self.tokenizer.n_words
 
     @abstractmethod
     def forward(
@@ -75,27 +84,35 @@ def generate(  # noqa: C901
         logits = self.forward(
             tokens=torch.tensor([prompt_tokens], dtype=torch.long),
             input_pos=(
-                torch.tensor([0], dtype=torch.long)
-                if self.params.use_kv_cache
-                else None
+                torch.tensor([0], dtype=torch.long) if self.use_kv_cache else None
             ),
         )
 
-        current_token = next_token(logits, temperature, top_p)
+        # TODO: accomodate TorchTune model, which doesn't
+        # make an optimization of dropping all logits but the last.
+        current_token = next_token(logits[:, -1, :], temperature, top_p)
         tokens = prompt_tokens + [current_token]
 
-        while len(tokens) < self.params.max_seq_len:
-            if self.params.use_kv_cache:
+        i = 0
+        while len(tokens) < self.max_seq_len:
+            print(f"{i} out of {self.max_seq_len} max tokens generated")
+            if self.use_kv_cache:
                 logits = self.forward(
                     tokens=torch.tensor([[current_token]], dtype=torch.long),
                     input_pos=torch.tensor([len(tokens) - 1], dtype=torch.long),
                 )
             else:
-                logits = self.forward(tokens=torch.tensor([tokens], dtype=torch.long))
-            current_token = next_token(logits, temperature, top_p)
-            if current_token in self.tokenizer.stop_tokens:
+                logits = self.forward(
+                    tokens=torch.tensor([tokens], dtype=torch.long, device=self.device),
+                )
+            current_token = next_token(logits[:, -1, :], temperature, top_p)
+            if current_token == self.tokenizer.eos_id or (
+                hasattr(self.tokenizer, "stop_tokens")
+                and current_token in self.tokenizer.stop_tokens
+            ):
                 break
             tokens.append(current_token)
+            i += 1
 
         return tokens if echo else tokens[len(prompt_tokens) :]
 

examples/models/llama2/runner/native.py

@@ -10,18 +10,17 @@
 
 import torch
 
-from executorch.examples.models.llama2.llama_transformer import ModelArgs
 from executorch.extension.pybindings.portable_lib import _load_for_executorch
 
 # Load custom ops and quantized ops.
 from executorch.extension.pybindings import portable_lib  # noqa # usort: skip
 
+from executorch.examples.models.llama2.runner.generation import LlamaRunner
+
 # Note: import this after portable_lib
 # from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
 from executorch.kernels import quantized  # noqa
 
-from executorch.examples.models.llama2.runner.generation import LlamaRunner
-
 
 class NativeLlamaRunner(LlamaRunner):
     """
@@ -31,13 +30,13 @@ class NativeLlamaRunner(LlamaRunner):
     def __init__(self, args):
         with open(args.params, "r") as f:
             params = json.loads(f.read())
-        model_args: ModelArgs = ModelArgs(
+        super().__init__(
+            tokenizer_path=args.tokenizer,
             max_seq_len=args.max_len,
             max_batch_size=1,
             use_kv_cache=args.kv_cache,
             vocab_size=params["vocab_size"],
         )
-        super().__init__(tokenizer_path=args.tokenizer, model_args=model_args)
         self.model = _load_for_executorch(args.pte)
 
     def forward(

extension/llm/export/builder.py

@@ -193,12 +193,19 @@ def capture_pre_autograd_graph(self) -> "LLMEdgeManager":
                     strict=True,
                 ).module()
             else:
-                self.pre_autograd_graph_module = capture_pre_autograd_graph(
+                # pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
+                #  `Module`.
+                print("Exporting with:")
+                print(f"inputs: {self.example_inputs}")
+                print(f"kwargs: {self.example_kwarg_inputs}")
+                print(f"dynamic shapes: {dynamic_shape}")
+
+                self.pre_autograd_graph_module = export_for_training(
                     self.model,
                     self.example_inputs,
                     kwargs=self.example_kwarg_inputs,
                     dynamic_shapes=dynamic_shape,
-                )
+                ).module()
 
         return self