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-10-11 06:20:34 +00:00
parent 0002ae5956
commit 1507a7528f
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -478,6 +478,7 @@ class RingAttention(torch.autograd.Function):
num_ring_size = world_size // num_rings num_ring_size = world_size // num_rings
if tp_size > 1: if tp_size > 1:
# Create inner ring groups
ranks = [] ranks = []
for i in range(num_rings): for i in range(num_rings):
start = i * num_ring_size start = i * num_ring_size
@ -494,7 +495,7 @@ class RingAttention(torch.autograd.Function):
group = dist.new_group(inner_rank) group = dist.new_group(inner_rank)
if rank in inner_rank: if rank in inner_rank:
inner_ring_group = group inner_ring_group = group
# Create inter ring groups
for i in range(num_ring_size): for i in range(num_ring_size):
inter_rank = [i + j * num_ring_size for j in range(num_rings)] inter_rank = [i + j * num_ring_size for j in range(num_rings)]
group = dist.new_group(inter_rank) group = dist.new_group(inter_rank)