github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 02:26:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix zero3 fp16 and add zero3 model context (#62)

Browse Source
This commit is contained in:
ver217
2021-12-10 17:48:50 +08:00
committed by GitHub
parent 9a0466534c
commit 7d3711058f
5 changed files with 114 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/initialize.py
View File

@@ -220,7 +220,9 @@ def initialize(model: Union[nn.Module, List[nn.Module]],
# first sync model across dp ranks
model.to(get_current_device())
sync_model_param_in_dp(model)
use_zero3 = hasattr(gpc.config, 'zero') and gpc.config.zero.level == 3
if not use_zero3:
sync_model_param_in_dp(model)
# check amp and zero
fp16_cfg = gpc.config.get('fp16', None)