Fix Cuda out of memory issue for eager runner

This PR updates the eager runner to disable grad and save memory usage.

It also update the prompt format to not include bos.

Differential Revision: [D65962743](https://our.internmc.facebook.com/intern/diff/D65962743/)

[ghstack-poisoned]

Author: helunwencser

examples/models/llama/runner/eager.py

@@ -84,18 +84,19 @@ def main() -> None:
     parser = build_args_parser()
     args = parser.parse_args()
 
-    runner = EagerLlamaRunner(args)
-    generated_tokens = (
-        runner.chat_completion(temperature=args.temperature)
-        if args.chat
-        else runner.text_completion(
-            prompt=args.prompt,
-            temperature=args.temperature,
-            echo=True,
+    with torch.no_grad():
+        runner = EagerLlamaRunner(args)
+        generated_tokens = (
+            runner.chat_completion(temperature=args.temperature)
+            if args.chat
+            else runner.text_completion(
+                prompt=args.prompt,
+                temperature=args.temperature,
+                echo=True,
+            )
         )
-    )
-    if args.show_tokens:
-        print(f"Generated {len(generated_tokens)} tokens: {generated_tokens}")
+        if args.show_tokens:
+            print(f"Generated {len(generated_tokens)} tokens: {generated_tokens}")
 
 
 if __name__ == "__main__":

examples/models/llama/runner/generation.py

@@ -135,7 +135,7 @@ def text_completion(
             This method generates text completion for the provided prompt, employing nucleus sampling to introduce controlled randomness.
         """
         return self.generate(
-            prompt_tokens=self.tokenizer.encode(prompt, bos=True, eos=False),
+            prompt_tokens=self.tokenizer.encode(prompt, bos=False, eos=False),
             max_seq_len=self.params.max_seq_len,
             temperature=temperature,
             top_p=top_p,
@@ -169,7 +169,7 @@ def chat_completion(
             print("LLM: ", end="", flush=True)
             new_tokens = self.generate(
                 prompt_tokens=self.tokenizer.encode(
-                    self._format_prompt(prompt), bos=True, eos=False
+                    self._format_prompt(prompt), bos=False, eos=False
                 ),
                 max_seq_len=self.params.max_seq_len,
                 temperature=temperature,
@@ -182,8 +182,7 @@ def chat_completion(
         return tokens
 
     def _format_prompt(self, prompt: str) -> str:
-        return f"""
-<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+        return f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
 
 You are a helpful assistant<|eot_id|><|start_header_id|>user<|end_header_id|>