Issue with phi-3 on Long Sequences with Batches > 1

[Samoed]

Hi! I'm encountering an issue while tuning phi-3 on long sequences with batch sizes greater than 1. Below is the code to reproduce the problem:

Working Code:

tokenized = tokenizer(
    ["Very long prompt\n" * 3000],  # *2,
    max_length=3000,
    return_tensors="pt",
    truncation=True,
).to("cuda")

res = model.generate(
    **tokenized,
    max_length=4096,
)

Code with Error:

RuntimeError: The expanded size of the tensor (2047) must match the existing size (3001) at non-singleton dimension 3. Target sizes: [2, 32, 1, 2047]. Tensor sizes: [2, 1, 1, 3001]

tokenized = tokenizer(
    ["Very long prompt\n" * 3000] * 2,
    max_length=3000,
    return_tensors="pt",
    truncation=True,
).to("cuda")

res = model.generate(
    **tokenized,
    max_length=4096,
)

Notebook with example.

Any insights on how to resolve this issue would be greatly appreciated!

[danielhanchen]

Oh interesting I'll check this and get back to you - sorry!