github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] fix zero optim save/load state dict (#1381)

Browse Source
This commit is contained in:
ver217
2022-07-28 17:19:39 +08:00
committed by GitHub
parent b6fd165f66
commit 828b9e5e0d
3 changed files with 160 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/tensor/process_group.py
View File

@@ -104,8 +104,8 @@ class ProcessGroup:
def set_cpu_groups(self):
if self.has_cpu_groups:
return
self.logger.info(
f'{self._rank} Gloo initialize TP group on {self._tp_rank_list}, DP group on {self._dp_rank_list}')
# self.logger.info(
# f'{self._rank} Gloo initialize TP group on {self._tp_rank_list}, DP group on {self._dp_rank_list}')
PYTORCHPGDICT_.get(self._tp_rank_list, 'gloo')
PYTORCHPGDICT_.get(self._dp_rank_list, 'gloo')
self._has_cpu_groups = True