mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-08-06 10:34:23 +00:00
fix
This commit is contained in:
parent
6fb1322db1
commit
91ed32c256
@ -500,8 +500,8 @@ class RingAttention(torch.autograd.Function):
|
||||
k,
|
||||
v,
|
||||
sp_group,
|
||||
tp_group: Optional[dist.ProcessGroup],
|
||||
attention_mask_type,
|
||||
tp_group=None,
|
||||
cu_seqlens=None,
|
||||
max_seqlen=None,
|
||||
valid_indices=None,
|
||||
|
@ -866,7 +866,7 @@ def get_gpt2_flash_attention_forward(shard_config: Optional[ShardConfig] = None)
|
||||
key,
|
||||
value,
|
||||
sp_group,
|
||||
tp_group,
|
||||
tp_group=tp_group,
|
||||
**attention_mask,
|
||||
dropout_p=dropout_p,
|
||||
scale=scale,
|
||||
|
@ -571,7 +571,7 @@ def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, s
|
||||
key_states,
|
||||
value_states,
|
||||
sp_group,
|
||||
tp_group,
|
||||
tp_group=tp_group,
|
||||
**attention_mask,
|
||||
inner_ring_size=shard_config.inner_ring_size,
|
||||
)
|
||||
|
Loading…
Reference in New Issue
Block a user