Generation fixes (#2787)

Author: SalmanMohammadi

tests/torchtune/generation/test_generation.py

@@ -406,6 +406,7 @@ def test_stop_tokens_batched_uneven_stopping(
         model = request.getfixturevalue(model)
         temperature = 0.6
         top_k = 100
+        pad_id = -100
 
         stop_tokens = [3991, 3987, 3969]
 
@@ -418,13 +419,14 @@ def test_stop_tokens_batched_uneven_stopping(
             temperature=temperature,
             top_k=top_k,
             stop_tokens=stop_tokens,
+            pad_id=pad_id,
         )
 
         expected_output = torch.tensor(
             [
                 [2, 3, 4, 5, 6, 7, 8, 9, 3954, 3920, 3991],
-                [2, 3, 4, 5, 6, 7, 8, 9, 3983, 3987, 0],
-                [2, 3, 4, 5, 6, 7, 8, 9, 3969, 0, 0],
+                [2, 3, 4, 5, 6, 7, 8, 9, 3983, 3987, pad_id],
+                [2, 3, 4, 5, 6, 7, 8, 9, 3969, pad_id, pad_id],
             ]
         )
 
@@ -445,7 +447,7 @@ def test_stop_tokens_batched_uneven_stopping_left_padded(
         model = request.getfixturevalue(model)
         temperature = 0.6
         top_k = 100
-
+        pad_id = -100
         # Updated stop tokens to match the new generated tokens
         stop_tokens = [3991, 3987, 3969]
 
@@ -458,13 +460,14 @@ def test_stop_tokens_batched_uneven_stopping_left_padded(
             temperature=temperature,
             top_k=top_k,
             stop_tokens=stop_tokens,
+            pad_id=pad_id,
         )
 
         expected_output = torch.tensor(
             [
                 [0, 0, 2, 3, 4, 5, 6, 7, 8, 9, 3954, 3920, 3991],
-                [0, 0, 2, 3, 4, 5, 6, 7, 8, 9, 3983, 3987, 0],
-                [0, 0, 2, 3, 4, 5, 6, 7, 8, 9, 3969, 0, 0],
+                [0, 0, 2, 3, 4, 5, 6, 7, 8, 9, 3983, 3987, pad_id],
+                [0, 0, 2, 3, 4, 5, 6, 7, 8, 9, 3969, pad_id, pad_id],
             ]
         )
         assert torch.equal(outputs, expected_output)

torchtune/generation/_generation.py

@@ -206,7 +206,7 @@ def generate(
     top_k: Optional[int] = None,
     stop_tokens: Optional[list[int]] = None,
     rng: Optional[torch.Generator] = None,
-    custom_generate_next_token: Optional[Callable] = None,
+    compiled_generate_next_token: Optional[Callable] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Generates tokens from a model conditioned on a prompt, and also returns logits for the generations.
@@ -223,23 +223,48 @@ def generate(
         stop_tokens (Optional[list[int]]): If specified, generation is stopped when any of these tokens are generated,
             default None.
         rng (Optional[torch.Generator]): random number generator, default None.
-        custom_generate_next_token (Optional[Callable]): If specified, we'll use the
-            ``custom_generate_next_token function``. This is generally only useful if
-            you want to specify a ``torch.compile`` version of the generate next token for
-            performance reasons. If None, we use the default :func:`generate_next_token`.
+        compiled_generate_next_token (Optional[Callable]): This argument is typically a reference to a compiled version of
+            the :func:`generate_next_token` function. During autoregressive decoding, this function is called instead of the default
+            :func:`generate_next_token` in order to accelerate generation. :func:`generate_next_token` will still be used for the
+            first token generation - or "pre-fill" pass.
             Default is None.
 
     Note:
         This function has only been tested with decoder-only models.
 
     Examples:
-        >>> model = torchtune.models.llama3.llama3_8b()
-        >>> tokenizer = torchtune.models.llama3.llama3_tokenizer()
-        >>> prompt = tokenizer.encode("Hi my name is")
-        >>> rng.manual_seed(42)
-        >>> output, logits = generate(model, torch.tensor(prompt), max_generated_tokens=100, pad_id=0)
+        >>> import torch
+        >>> from torchtune.models.llama3 import llama3_tokenizer
+        >>> from torchtune.models.llama3 import llama3_8b
+        >>> from torchtune.generation import generate
+        >>> from torchtune.training.checkpointing import FullModelHFCheckpointer
+        >>> from torchtune.data import Message
+
+        >>> model = llama3_8b().cuda()
+
+        >>> checkpointer = FullModelHFCheckpointer(
+        ...     checkpoint_dir="/tmp/Meta-Llama-3-8B-Instruct",
+        ...     checkpoint_files=[
+        ...         "model-00001-of-00004.safetensors",
+        ...         "model-00002-of-00004.safetensors",
+        ...         "model-00003-of-00004.safetensors",
+        ...         "model-00004-of-00004.safetensors",
+        ...     ],
+        ...     model_type="LLAMA3",
+        ...     output_dir="/tmp/torchtune/llama3_8b",
+        ... )
+        >>> checkpoint = checkpointer.load_checkpoint()
+        >>> model.load_state_dict(checkpoint["model"])
+
+        >>> tokenizer = llama3_tokenizer("/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model")
+        >>> messages = [
+        ...     Message(role="assistant", content="Hi my name is"),
+        ... ]
+        >>> prompt = tokenizer({"messages": messages}, inference=True)
+        >>> output, logits = generate(model, torch.tensor(prompt["tokens"], device='cuda'), max_generated_tokens=100, pad_id=0)
         >>> print(tokenizer.decode(output[0].tolist()))
-        Hi my name is Jeremy and I'm a friendly language model assistant!
+
+        >>> Hi my name is Marley. Nice to meet you, Marley! How are you doing today?... [truncated]
 
     Returns:
         tuple[torch.Tensor, torch.Tensor]: tuple of two tensors:
@@ -251,9 +276,6 @@ def generate(
     """
     prompt = prompt.view(1, -1) if prompt.ndim == 1 else prompt
 
-    if custom_generate_next_token is None:
-        custom_generate_next_token = generate_next_token
-
     bsz, prompt_length = prompt.size()
     total_response_length = prompt_length + max_generated_tokens
 
@@ -356,6 +378,12 @@ def generate(
         if stop_token_reached.all().item():
             return generated_tokens, generated_logits
 
+    next_token_fn = (
+        compiled_generate_next_token
+        if compiled_generate_next_token is not None
+        else generate_next_token
+    )
+
     for _ in range(max_generated_tokens - 1):
         # update stop_token_mask if we reached a stop token in a previous step
         # by appending the logical not of stop_token_reached to the end of the mask
@@ -387,7 +415,7 @@ def generate(
             condition = uniform_val >= 1.0 - epsilon
             q = -torch.where(condition, -epsilon, torch.log(uniform_val))
 
-        tokens, logits = custom_generate_next_token(
+        tokens, logits = next_token_fn(
             model,
             input_pos=curr_input_pos,
             x=tokens.clone(),
@@ -409,7 +437,7 @@ def generate(
 
     # mask out generated tokens in seqs that already hit a stop token
     if stop_tokens is not None:
-        generated_tokens *= stop_token_mask
+        generated_tokens.masked_fill_(~stop_token_mask.bool(), pad_id)
         generated_logits *= stop_token_mask[:, -generated_logits.shape[1] :, None]
 
     return generated_tokens, generated_logits