github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-04-28 19:55:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
wangbluo 2024-10-09 14:05:26 +08:00
parent 3532f77b90
commit b635dd0669
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -442,7 +442,7 @@ class RingAttention(torch.autograd.Function):
Tuple[dist.ProcessGroup, dist.ProcessGroup]: Inner-ring process group and inter-ring process group. Tuple[dist.ProcessGroup, dist.ProcessGroup]: Inner-ring process group and inter-ring process group.
""" """
sp_size = dist.get_world_size(sp_group) sp_size = dist.get_world_size(sp_group)
tp_size = dist.get_world_size(tp_group) 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: if inner_ring_size is None: