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

fix

pull/6071/head
wangbluo 2024-09-26 18:03:09 +08:00
parent 6705dad41b
commit 3fab92166e
1 changed files with 31 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -442,8 +442,7 @@ class RingAttention(torch.autograd.Function):
Tuple[dist.ProcessGroup, dist.ProcessGroup]: Inner-ring process group and inter-ring process group.
"""
sp_size = dist.get_world_size(sp_group)
dist.get_rank(sp_group)
dist.get_world_size(tp_group)
tp_size = dist.get_world_size(tp_group)
if inner_ring_size is None:
if torch.cuda.device_count() >= dist.get_world_size():
@ -476,21 +475,38 @@ class RingAttention(torch.autograd.Function):
rank = dist.get_rank()
groups = int(world_size / sp_size)
# Create inner ring groups
for group_id in range(groups):
for i in range(inner_ring_size):
ranks = list(range(i + (group_id * sp_size), (1 + group_id) * sp_size, inner_ring_size))
group = dist.new_group(ranks)
if tp_size > 1:
for group_id in range(groups):
for i in range(inner_ring_size):
ranks = list(range(i + (group_id * sp_size), (1 + group_id) * sp_size, inner_ring_size))
group = dist.new_group(ranks)
if rank in ranks:
inner_ring_group = group
for group_id in range(groups):
for i in range(num_rings):
ranks = list(range(i + group_id * num_rings, world_size, sp_size))
group = dist.new_group(ranks)
if rank in ranks:
inter_ring_group = group
else:
for i in range(sp_size // 2):
ranks = list(range((i) * num_rings, (i + 1) * num_rings, 1))
if rank in ranks:
print(
"rank:",
rank,
"inner ranks:",
ranks,
)
group = dist.new_group(ranks)
inner_ring_group = group
# Create inter ring groups
for group_id in range(groups):
for i in range(num_rings):
ranks = list(range(i + group_id * num_rings, world_size, sp_size))
group = dist.new_group(ranks)
if rank in ranks:
inter_ring_group = group
for group_id in range(num_rings):
for i in range(num_rings):
ranks = list(range(i + group_id * num_rings, world_size, inner_ring_size))
ranks = [0, 1, 4, 5] if rank == 0 or rank == 1 or rank == 4 or rank == 5 else [2, 3, 6, 7]
if rank in ranks:
group = dist.new_group(ranks)
inter_ring_group = group
return inner_ring_group, inter_ring_group