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

[checkpoint] sharded optim save/load grad scaler (#1350)

Browse Source
This commit is contained in:
ver217 2022-07-21 15:21:21 +08:00 committed by GitHub
parent 05fae1fd56
commit ce470ba37e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/zero/sharded_optim/sharded_optim_v2.py
View File

@ -363,7 +363,18 @@ class ShardedOptimizerV2(ColossalaiOptimizer):
self.master_params[p].trans_state(TensorState.HOLD) self.master_params[p].trans_state(TensorState.HOLD)
def state_dict(self):
optim_state_dict = super().state_dict()
scaler_state_dict = self.grad_scaler.state_dict()
optim_state_dict['scaler'] = scaler_state_dict
return optim_state_dict
def load_state_dict(self, *args, **kwargs): def load_state_dict(self, *args, **kwargs):
if 'scaler' not in args[0]:
self._logger.warning('Missing scaler when loading optimizer state dict', ranks=[0])
else:
scaler_state_dict = args[0].pop('scaler')
self.grad_scaler.load_state_dict(scaler_state_dict)
super().load_state_dict(*args, **kwargs) super().load_state_dict(*args, **kwargs)
for group in self.optim.param_groups: for group in self.optim.param_groups:
for p in group['params']: for p in group['params']: