github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 03:52:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer]fix flash attention, when mask is casual, just don't unpad it (#5084)

Browse Source 
* fix flash attn

* fix

fix
This commit is contained in:
flybird11111
2023-11-22 16:00:07 +08:00
committed by GitHub
parent 75af66cd81
commit aae496631c
6 changed files with 16 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/shardformer/modeling/chatglm2.py
View File

@@ -51,7 +51,8 @@ def get_flash_core_attention_forward():
attn_mask_type = AttnMaskType.causal
else:
flash_attention_mask = ~(attention_mask[:, :, -1].squeeze(1).to(torch.bool)).contiguous()
attn_mask_type = AttnMaskType.paddedcausal
if not torch.all(flash_attention_mask):
attn_mask_type = AttnMaskType.paddedcausal
attention = ColoAttention(
embed_dim=self.hidden_size_per_partition,