Fix: Remove attn_mask when using is_causal=True in Flash Attention
PyTorch's scaled_dot_product_attention cannot use attn_mask and is_causal=True simultaneously. This commit sets attn_mask=None when is_causal=True to avoid parameter conflicts.pull/497/head
parent
3812a16d20
commit
79e36e17e8
|
|
@ -194,20 +194,25 @@ class Attention(nn.Module):
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)):
|
if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)):
|
||||||
attn_mask = (
|
# Fix: 不能同时使用 attn_mask 和 is_causal=True
|
||||||
None
|
output = F.scaled_dot_product_attention(
|
||||||
if attention_mask is None
|
xq, xk, xv,
|
||||||
else attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1).bool()
|
attn_mask=None, # 使用 is_causal 时必须设为 None
|
||||||
|
dropout_p=self.dropout if self.training else 0.0,
|
||||||
|
is_causal=True
|
||||||
)
|
)
|
||||||
|
|
||||||
output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
|
|
||||||
else:
|
else:
|
||||||
|
# 传统注意力机制(fallback)
|
||||||
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||||||
scores = scores + torch.triu(
|
|
||||||
torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
|
# 添加因果mask
|
||||||
diagonal=1
|
if seq_len > 1:
|
||||||
).unsqueeze(0).unsqueeze(0) # scores+mask
|
scores = scores + torch.triu(
|
||||||
|
torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
|
||||||
|
diagonal=1
|
||||||
|
).unsqueeze(0).unsqueeze(0)
|
||||||
|
|
||||||
|
# 添加 attention_mask
|
||||||
if attention_mask is not None:
|
if attention_mask is not None:
|
||||||
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
|
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
|
||||||
extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
|
extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue