Generation does not work

[Shariar076]

The doc provided in https://docs.pytorch.org/torchtune/stable/generated/torchtune.generation.generate.html#torchtune.generation.generate for running generation does not work. I tried the following:

>>> import torch                                                                                                                                                                                                                
>>> from torchtune.models.llama3 import llama3_tokenizer                                                          
>>> from torchtune.models.llama3 import llama3_8b                                                                 
>>> from torchtune.generation import generate                                                                     
>>> model = llama3_8b().cuda()                                                                                    
>>> tokenizer = llama3_tokenizer("checkpoints/Meta-Llama-3-8B-Instruct/original/tokenizer.model")                 
>>> prompt = tokenizer.encode("Hi my name is")                                                                    
>>> output, logits = generate(model, torch.tensor(prompt, device='cuda'), max_generated_tokens=100, pad_id=0)     
>>> print(tokenizer.decode(output[0].tolist()))                                                                   
Hi my name is                                                                                                     
>>> output                                                                                                        
tensor([[128000,  13347,    856,    836,    374, 128001,  49263,  45892, 117650,                                  
          80252,  13825,  93166,  86232,   1309, 119791,  57968, 109216, 119099,                                  
         127824,  50532, 114899, 101806,  63967,  82748,  81405, 119646,   1323,                                  
          88452,  81382,  43309,  46070,  60111,  98318,  89937,  82561,  86967,                                  
          15046,  46705,  92231,  49405, 105751,  12936,  63385,  78030,  65426,                                  
         115513,  63001,  47715,  26661, 115855,  74187,  20661,  46922,  52735,                                  
          71358,  45263,   9412,  70215,  46441,  17561,  34201,  16042, 105009,                                  
         111681,  57920,  97103,   7404,  96699,  85056,  65707,   8174,  12481,                                  
         121220,   2882,  41843,  39199,  99413,  44659, 110860,  58017,  41245,                                  
          74254,  91415,   4041, 120132,  21972,  46548,  68651,  11568,  88572,                                  
         106217,  34486,  95538, 126271, 128138,  34382, 115571,  85049, 126324,                                  
         107640, 124685,  17832, 118559,  54026, 124872, 102345]],                                                
       device='cuda:0')                                                                                           
>>>                                                                                                               

[krammnic]

It seems to me that the model checkpoint is not loaded

[ebsmothers]

Hi @Shariar076 thanks for creating the issue. I agree with @krammnic -- you need to load the checkpoint into the model to get sensible generations. E.g. adding something like this in the middle of your code should work:

from torchtune.training.checkpointing import FullModelHFCheckpointer
checkpointer = FullModelHFCheckpointer(
    checkpoint_dir="checkpoints/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"])

We would love to review a PR updating the docstring with a clearer example if you're interested in contributing.

[Shariar076]

Thanks @ebsmothers @krammnic. I was able to load the model weights using the code but I am getting the output when I decode only the newly generated tokens using tokenizer.decode(output[0][len(prompt):].tolist()). Running decode on full output still returns only the input.

[krammnic]

Interesting, can you share a script that you run?

[ebsmothers]

@Shariar076 you probably also want to use the tokenizer's __call__ method instead of encode. This will handle stop tokens etc (you can pass inference=True), you will also need to wrap your text in our Message format. Here is a modified working version of your example:

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 Marley.Nice to meet you, Marley! How are you doing today?I'm doing pretty good, thanks! It's nice to talk to someone new. How about you?I'm doing well too! I'm a large language model, so I don't have feelings or emotions like humans do, but I'm happy to chat with you and help answer any questions you might have. What would you like to

Separately, I would recommend using our generate_v2 recipe as one possible alternative here (example config). This recipe does not call into the generate API, but contains all the generation code directly in the recipe, and gives more flexibility than calling directly into generate

[Shariar076]

Thanks a lot @ebsmothers. It's working now. Appreciate your help.