From 79e36e17e898e35c00a0c128a413b51d4ccbd536 Mon Sep 17 00:00:00 2001 From: wanghao Date: Sat, 18 Oct 2025 21:02:46 +0800 Subject: [PATCH] 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. --- model/model_minimind.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index ecd99b6..e3285b6 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -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)): - attn_mask = ( - None - if attention_mask is None - else attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1).bool() + # Fix: 不能同时使用 attn_mask 和 is_causal=True + output = F.scaled_dot_product_attention( + xq, xk, xv, + 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: + # 传统注意力机制（fallback） 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), - diagonal=1 - ).unsqueeze(0).unsqueeze(0) # scores+mask + + # 添加因果mask + if seq_len > 1: + 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: extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) extended_attention_mask = (1.0 - extended_attention_mask) * -1e9