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

fix

Browse Source
This commit is contained in:
wangbluo 2024-10-14 14:56:05 +08:00
parent e1e86f9f1f
commit d891e50617
1 changed files with 6 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -445,18 +445,11 @@ class RingAttention(torch.autograd.Function):
tp_size = dist.get_world_size(tp_group) if tp_group is not None else 1 tp_size = dist.get_world_size(tp_group) if tp_group is not None else 1
sp_rank = dist.get_rank(sp_group) sp_rank = dist.get_rank(sp_group)
if inner_ring_size is None: assert inner_ring_size is not None
if torch.cuda.device_count() >= dist.get_world_size():
# single node, no need to consider NICs assert (
return sp_group, sp_group inner_ring_size <= sp_size and sp_size % inner_ring_size == 0
if sp_size <= 4: ), f"Error: sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"
inner_ring_size = min(2, sp_size)
else:
inner_ring_size = min(4, sp_size)
else:
assert (
inner_ring_size <= sp_size and sp_size % inner_ring_size == 0
), f"Error: sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"
if inner_ring_size == sp_size: if inner_ring_size == sp_size:
return sp_group, sp_group return sp_group, sp_group
@ -575,7 +568,7 @@ class RingAttention(torch.autograd.Function):
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))
if RingAttention.SP_GROUP is not sp_group: 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( inner_ring_group, inter_ring_group = RingAttention.get_double_ring_groups(
sp_group, tp_group, inner_ring_size sp_group, tp_group, inner_ring_size