github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-01 20:05:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
wangbluo 2024-09-13 06:00:58 +00:00
parent f393867cff
commit dc032172c3
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/shardformer/layer/attn.py
View File

@ -173,7 +173,7 @@ class ColoAttention:
# no padding # no padding
assert is_causal assert is_causal
outputs["attention_mask_type"] = AttnMaskType.CAUSAL outputs["attention_mask_type"] = AttnMaskType.CAUSAL
if memory_size < MEMORY_BOUND and not is_causal: if memory_size < MEMORY_BOUND:
attention_mask = torch.ones(s_q, s_kv, dtype=dtype, device=device) attention_mask = torch.ones(s_q, s_kv, dtype=dtype, device=device)
if s_q != 1: if s_q != 1:
attention_mask.tril_(diagonal=0) attention_mask.tril_(diagonal=0)