github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-10-11 08:58:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

[model checkpoint] add gloo groups for cpu tensor communication (#589)

Browse Source
This commit is contained in:
アマデウス
2022-04-01 10:15:52 +08:00
committed by GitHub
parent 54e688b623
commit 297b8baae2
10 changed files with 98 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/context/process_group_initializer/initializer_tensor.py
View File

@@ -34,17 +34,20 @@ class Initializer_Tensor(ProcessGroupInitializer):
local_rank = None
ranks_in_group = None
process_group = None
cpu_group = None
group_world_size = None
mode = ParallelMode.TENSOR
for i in range(self.num_tensor_parallel_group):
ranks = [i * self.tensor_parallel_size + j for j in range(self.tensor_parallel_size)]
group = dist.new_group(ranks)
group_cpu = dist.new_group(ranks, backend='gloo') if dist.get_backend() != 'gloo' else group
if self.rank in ranks:
local_rank = ranks.index(self.rank)
group_world_size = len(ranks)
process_group = group
cpu_group = group_cpu
ranks_in_group = ranks
return local_rank, group_world_size, process_group, ranks_in_group, mode
return local_rank, group_world_size, process_group, cpu_group, ranks_in_group, mode