kanshan/minimind
1
0
Fork 0
mirror of https://github.com/jingyaogong/minimind.git synced 2026-05-01 19:58:15 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #536 from yuyu5333/fix/attn_forward

Browse Source 
fix: attn_forwad when is_causal=True assert attn_mask is None
This commit is contained in:
jingyaogong 2025-11-18 13:02:46 +08:00 committed by GitHub
commit ce9394670b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
model/model_minimind.py
View File

@ -194,13 +194,7 @@ 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 = ( output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
None
if attention_mask is None
else attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1).bool()
)
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:
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( scores = scores + torch.triu(