github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-16 08:27:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix the attn

Browse Source
This commit is contained in:
wangbluo 2024-09-25 18:51:03 +08:00
parent cfd9eda628
commit 65c8297710
2 changed files with 3 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

@ -500,7 +500,7 @@ class RingAttention(torch.autograd.Function):
k,
v,
sp_group,
tp_group,
tp_group : Optional[dist.ProcessGroup],
attention_mask_type,
cu_seqlens=None,
max_seqlen=None,

2
colossalai/shardformer/modeling/gpt2.py
View File

@ -858,12 +858,14 @@ def get_gpt2_flash_attention_forward(shard_config: Optional[ShardConfig] = None)
sp_mode = shard_config.sequence_parallelism_mode
sp_group = shard_config.sequence_parallel_process_group
tp_group = shard_config.tensor_parallel_process_group
if sp_mode == "ring_attn":
attn_output = RingAttention.attention(
query,
key,
value,
sp_group,
tp_group,
**attention_mask,
dropout_p=dropout_p,
scale=scale,