github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-01 03:45:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
wangbluo 2024-10-10 18:38:47 +08:00
parent 5ecc27e150
commit efe3042bb2
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

@ -482,7 +482,8 @@ class RingAttention(torch.autograd.Function):
for i in range(num_rings): for i in range(num_rings):
for j in range(num_inner_group): for j in range(num_inner_group):
# find inner ring group in one sp groups # find inner ring group in one sp groups
ranks = list(range(j + i * num_ring_size, j + (i + 1) * num_ring_size, tp_size)) start = j + i * num_ring_size
ranks = list(range(start, start + tp_size * inner_ring_size, tp_size))
group = dist.new_group(ranks) group = dist.new_group(ranks)
if rank in ranks: if rank in ranks:
inner_ring_group = group inner_ring_group = group