github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix

pull/6061/head
wangbluo 2024-09-13 10:24:41 +00:00
parent 0ad3129cb9
commit b582319273
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -210,6 +210,7 @@ class ColoAttention:
} }
) )
if is_causal: if is_causal:
attention_mask = kv_padding_mask[:, None, :].expand(b, s_q, s_kv).to(dtype=dtype, device=device)
outputs["attention_mask_type"] = AttnMaskType.PADDED_CAUSAL outputs["attention_mask_type"] = AttnMaskType.PADDED_CAUSAL
if memory_size < MEMORY_BOUND: if memory_size < MEMORY_BOUND:
if s_q != 1: if s_q != 1: