github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[hotfix] zero optim prevents calling inner optim.zero_grad (#1422)

pull/1423/head
ver217 2022-08-09 16:08:12 +08:00 committed by GitHub
parent 504419d261
commit 6df3e19be9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/zero/zero_optimizer.py
View File

@ -139,6 +139,7 @@ class ZeroOptimizer(ColossalaiOptimizer):
self._update_params_ptr() self._update_params_ptr()
ret = self.optim.step(*args, **kwargs) ret = self.optim.step(*args, **kwargs)
self._register_states() self._register_states()
self.zero_grad()
self._update_fp16_params() self._update_fp16_params()
return ret return ret