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

fix

Browse Source
This commit is contained in:
wangbluo
2024-09-25 19:02:21 +08:00
parent 91ed32c256
commit 6705dad41b
3 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -571,9 +571,9 @@ def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, s
key_states,
value_states,
sp_group,
tp_group=tp_group,
**attention_mask,
inner_ring_size=shard_config.inner_ring_size,
tp_group=tp_group,
)
elif shard_config.enable_flash_attention: