github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-22 13:41:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
wangbluo 2024-10-15 13:26:44 +08:00
parent 6be9862aaf
commit fd92789af2
4 changed files with 9 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -488,7 +488,7 @@ class RingAttention(torch.autograd.Function):
q, # (B, H, Sq, D) q, # (B, H, Sq, D)
k, k,
v, v,
sp_group, sp_axis,
attention_mask_type, attention_mask_type,
cu_seqlens=None, cu_seqlens=None,
max_seqlen=None, max_seqlen=None,
@ -499,7 +499,6 @@ class RingAttention(torch.autograd.Function):
return_softmax=False, return_softmax=False,
inner_ring_size=None, inner_ring_size=None,
pg_mesh=None, pg_mesh=None,
sp_axis=None,
**kwargs, **kwargs,
): ):
""" """
@ -546,7 +545,7 @@ class RingAttention(torch.autograd.Function):
assert pg_mesh is not None, f"Error: The pg mesh is None! please check the process group initialization." assert pg_mesh is not None, f"Error: The pg mesh is None! please check the process group initialization."
clone_pg = lambda pg: dist.new_group(dist.get_process_group_ranks(pg)) clone_pg = lambda pg: dist.new_group(dist.get_process_group_ranks(pg))
sp_group = pg_mesh.get_group_along_axis(sp_axis)
if inner_ring_size != None: if inner_ring_size != None:
RingAttention.SP_GROUP = sp_group RingAttention.SP_GROUP = sp_group
inner_ring_group, inter_ring_group = RingAttention.get_double_ring_groups(sp_axis, pg_mesh, inner_ring_size) inner_ring_group, inter_ring_group = RingAttention.get_double_ring_groups(sp_axis, pg_mesh, inner_ring_size)

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

@ -857,19 +857,18 @@ def get_gpt2_flash_attention_forward(shard_config: Optional[ShardConfig] = None)
dropout_p = self.attn_dropout.p if self.training else 0.0 dropout_p = self.attn_dropout.p if self.training else 0.0
sp_mode = shard_config.sequence_parallelism_mode sp_mode = shard_config.sequence_parallelism_mode
sp_group = shard_config.sequence_parallel_process_group shard_config.sequence_parallel_process_group
if sp_mode == "ring_attn": if sp_mode == "ring_attn":
attn_output = RingAttention.attention( attn_output = RingAttention.attention(
query, query,
key, key,
value, value,
sp_group, sp_axis=shard_config.sp_axis,
**attention_mask, **attention_mask,
dropout_p=dropout_p, dropout_p=dropout_p,
scale=scale, scale=scale,
inner_ring_size=shard_config.inner_ring_size, inner_ring_size=shard_config.inner_ring_size,
pg_mesh=shard_config.pg_mesh, pg_mesh=shard_config.pg_mesh,
sp_axis=shard_config.sp_axis,
) )
else: else:
attn_output = ColoAttention.attention(query, key, value, **attention_mask, dropout_p=dropout_p, scale=scale) attn_output = ColoAttention.attention(query, key, value, **attention_mask, dropout_p=dropout_p, scale=scale)

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

@ -568,11 +568,10 @@ def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, s
query_states, query_states,
key_states, key_states,
value_states, value_states,
sp_group, sp_axis=shard_config.sp_axis,
**attention_mask, **attention_mask,
inner_ring_size=shard_config.inner_ring_size, inner_ring_size=shard_config.inner_ring_size,
pg_mesh=shard_config.pg_mesh, pg_mesh=shard_config.pg_mesh,
sp_axis=shard_config.sp_axis,
) )
elif shard_config.enable_flash_attention: elif shard_config.enable_flash_attention:

8
tests/test_shardformer/test_layer/test_ring_attn.py
View File

@ -44,12 +44,11 @@ def check_ring_attn(seq_len, bs, nheads, d, dtype, inner_ring_size):
q, q,
k, k,
v, v,
sp_group, sp_axis,
AttnMaskType.CAUSAL, AttnMaskType.CAUSAL,
return_softmax=True, return_softmax=True,
inner_ring_size=inner_ring_size, inner_ring_size=inner_ring_size,
pg_mesh=pg_mesh, pg_mesh=pg_mesh,
sp_axis=sp_axis,
) )
ring_out = ring_out.transpose(1, 2) ring_out = ring_out.transpose(1, 2)
out, lse, _ = flash_attn_qkvpacked_func( out, lse, _ = flash_attn_qkvpacked_func(
@ -88,6 +87,7 @@ def check_packed_seq(seqlen, bs, nheads, d, dtype):
device = get_current_device() device = get_current_device()
sp_group = dist.group.WORLD sp_group = dist.group.WORLD
sp_size = dist.get_world_size() sp_size = dist.get_world_size()
sp_axis = 2
atol = rtol = 7e-3 atol = rtol = 7e-3
torch.cuda.manual_seed(2) torch.cuda.manual_seed(2)
# Prepare varlen attention mask # Prepare varlen attention mask
@ -128,11 +128,11 @@ def check_packed_seq(seqlen, bs, nheads, d, dtype):
q_ring, q_ring,
k_ring, k_ring,
v_ring, v_ring,
sp_group, sp_axis,
**mask_info, **mask_info,
pad_output=False, pad_output=False,
return_softmax=True, return_softmax=True,
pg_mesh=dist.group.WORLD, pg_mesh=ProcessGroupMesh(1, 1, sp_size, 1),
# deterministic=True # deterministic=True
) )
ring_out = ring_out.transpose(1, 2).reshape(-1, nheads, d) ring_out = ring_out.transpose(1, 2).reshape(-1, nheads, d)