Clip logits if torchtune

Author: jackzhxng

examples/models/llama/runner/generation.py

@@ -127,15 +127,18 @@ def generate(  # noqa: C901
                     tokens=torch.tensor([tokens], dtype=torch.long, device=self.device),
                 )
 
-        if self.has_full_logits:
-            current_token = next_token(logits[:, -1, :], temperature, top_p)
-        else:
-            current_token = next_token(logits, temperature, top_p)
+            # If the logits aren't already clipped to only contain the last logit, clip them.
+            if self.has_full_logits:
+                current_token = next_token(logits[:, -1, :], temperature, top_p)
+            else:
+                current_token = next_token(logits, 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
 

extension/llm/export/builder.py

@@ -194,13 +194,13 @@ def export(self) -> "LLMEdgeManager":
                     strict=True,
                 ).module()
             else:
-                # 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}")
 
+                # pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
+                #  `Module`.
                 self.pre_autograd_graph_module = export_for_training(
                     self.model,
                     self.example_inputs,