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-14 11:45:35 +08:00
parent 703bb5c18d
commit e1e86f9f1f
1 changed files with 1 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -484,12 +484,7 @@ class RingAttention(torch.autograd.Function):
start = i * num_ring_size start = i * num_ring_size
end = (i + 1) * num_ring_size end = (i + 1) * num_ring_size
for idx in range(start, end): for idx in range(start, end):
inner_rank = [] inner_rank = [idx + k * tp_size for k in range(inner_ring_size) if idx + k * tp_size < end]
for k in range(inner_ring_size):
current_num = idx + k * tp_size
if current_num >= end:
break
inner_rank.append(current_num)
if len(inner_rank) == inner_ring_size and inner_rank not in ranks: if len(inner_rank) == inner_ring_size and inner_rank not in ranks:
ranks.append(inner_rank) ranks.append(inner_rank)
group = dist.new_group(inner_rank) group = dist.new_group(inner_rank)