I'm implementing padding support directly on my LLM model. To do so, I add extra rows to the boolean attention mask with all False values.
import torch
scale = 0.125
# 5 tokens, 1 padding
seq_length = 5
max_seq_length = 10
q = torch.randn(1, 2, seq_length + 1, 8)
k = torch.randn(1, 2, max_seq_length, 8)
v = torch.randn(1, 2, max_seq_length, 8)
mask = torch.tensor([[[
[ True, False, False, False, False, False, False, False, False, False], # regular mask
[ True, True, False, False, False, False, False, False, False, False],
[ True, True, True, False, False, False, False, False, False, False],
[ True, True, True, True, False, False, False, False, False, False],
[ True, True, True, True, True, False, False, False, False, False],
[False, False, False, False, False, False, False, False, False, False]]]]) # padding mask
def torch_sdpa(q, k, v, mask):
return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask, dropout_p=0.0, scale=scale)
def naive_sdpa_1(q, k, v, mask):
att = (q @ k.transpose(-2, -1)) * scale
att = torch.masked_fill(att, ~mask, float("-inf"))
att = torch.nn.functional.softmax(att, dim=-1)
return att @ v
def naive_sdpa_2(q, k, v, mask):
att = (q @ k.transpose(-2, -1)) * scale
att = torch.masked_fill(att, ~mask, torch.finfo(att.dtype).min)
att = torch.nn.functional.softmax(att, dim=-1)
return att @ v
y = torch_sdpa(q, k, v, mask)
print(torch.isnan(y).any())
y = naive_sdpa_1(q, k, v, mask)
print(torch.isnan(y).any())
y = naive_sdpa_2(q, k, v, mask)
print(torch.isnan(y).any())This seems to suggest that there's a precision issue. Is this expected or known? Could scaled_dot_product_attention be fixed?
Thank you for your help.
🐛 Describe the bug
I'm implementing padding support directly on my LLM model. To do so, I add extra rows to the boolean attention mask with all
Falsevalues.However, calling
torch.nn.functional.scaled_dot_product_attentionreturns NaNs.An equivalent and naive implementation also does (see
naive_fsdpa_1).But if I use
torch.finfo(att.dtype).minas the masked fill value, the result is correct (seenaive_fsdpa_2)Here's the repro:
This seems to suggest that there's a precision issue. Is this expected or known? Could
scaled_dot_product_attentionbe fixed?Side question: do I lose flash attn if I append the padding mask like this?
Thank you for your help.
Maybe linked: #101967
Versions
Current master
cc @ezyang @gchanan @zou3519 @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg @msaroufim @wconstab @bdhirsh @anijain2305