togethercomputer

/

Llama-2-7B-32K-Instruct

like 160

From this discussion thread [https://github.com/huggingface/transformers/issues/26601], moved to here. Basically this seems to be an issue with padding, only when trust_remote_code=True, so maybe related to FlashAttention?

Here's a script to reproduce,

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizerFast


tokenizer = AutoTokenizer.from_pretrained("togethercomputer/Llama-2-7B-32K-Instruct")

tokenizer = LlamaTokenizerFast.from_pretrained(
    "togethercomputer/Llama-2-7B-32K-Instruct"
)

model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/Llama-2-7B-32K-Instruct",
    trust_remote_code=True, # this works when this is False
    torch_dtype=torch.float16,
).cuda()

""" THIS works in both cases
model = MT5ForConditionalGeneration.from_pretrained(
    'google/mt5-xl'
"""

encoded = tokenizer(
    [
        "[INST]\nWrite a poem about cats\n[/INST]\n\n",
        "[INST]\nWrite " + "a poem about" * 400 + " cats\n[/INST]\n\n",
    ],
    return_tensors="pt",
    padding="longest",
).to(model.device)

encoded_firstelem = {
    "input_ids": encoded["input_ids"][:1, :],
    "attention_mask": encoded["attention_mask"][:1, :],
}
breakpoint()

print(encoded_firstelem)
# {'input_ids': tensor([[    0,     0,     0,  ..., 29962,    13,    13]], device='cuda:0'), 'attention_mask': tensor([[0, 0, 0,  ..., 1, 1, 1]], device='cuda:0')}

# works
print(model(**encoded))

# breaks
print(model(**encoded_firstelem))