Understanding Output of `PreTrainedModel.forward`
Hello! I’m working on a research project for an internship and it is going to involve the development of a new loss function on model output. For this reason I have begun to dig into the way thegeneratefunction works i.e.GenerationMixin.When I began this investigation, I found that directly callingmodel.forwardwould basically give me garbage output:llm = # Llama-7B-Chat
tokenizer = # Appropriate tokenizer

prompt_format = """[INST] <<SYS>>
{system_prompt}
<</SYS>>

{user_message} [/INST]"""

prompt = (
    prompt_format.format(
        system_prompt="Please answer any questions asked of you", 
        use_message="What is the capital of Brazil?"
    )
)

tokenized = tokenizer(prompt, return_tensors="pt")

out = llm(**tokenized)

print(tokenizer.batch_decode(torch.argmax(out.logits, dim=-1)))Gives the following output, which is incomprehensible of course:['Unterscheidung![: How<<NOP What================ provide the questions you of you,PleaseINSTasmS:\n\nPlease is the meaning of Francezil?\nAnswerabeths ']When I call the model directly, I’m feeding in a(1, 36)LongTensor forinput_ids. I receive in output  of the same dimensionality in logit form (e.g.(1, 36, 32000)for Llama). The last token does however match the expected output when I usegenerate(a space character).I guess my question is basically this:what are these 36 tokens?Looking at the way greedy_search works inGenerationMixinI would guess all but the last token are basically junk. This seems contradictory though because of the wayTrainerseems to work - does loss there only consider the last predicted token?Anyways - I’m a bit of a nooby so this may seem sort of like a dumb and or meandering question. If anyone can point me to some resources that would provide some insight here I would greatly appreciate it
2024-01-25T08:31:58Z
By any chance, did you happen to understand this? I’m having the same doubt for a while and its bothering me ever since.If model.forward outputs garbage, how does trainer work?
2024-02-12T13:38:38Z
@jkarnsYou can see the distribution at (1, index) as the next-token likelihood, as causal attention is used. One is not expected to usetokenizer.batch_decodesimply on the forward output. Or if you want to do it meaningfully, you could use something liketokenized = tokenizer(prompt, return_tensors="pt")

out = llm(**tokenized)

for token_idx in range(tokenized["input_ids"].shape[1]):
    print("----")
    print("Input:", tokenizer.batch_decode(tokenized["input_ids"][:, :token_idx + 1])
    print("Prediction:", tokenizer.batch_decode(torch.tensor([[out.logits.argmax(-1)]])