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

2025-09-01 01:06:00 +00:00 · 2022-08-09 16:08:12 +08:00
parent 504419d261
commit 6df3e19be9
1 changed files with 1 additions and 0 deletions
--- a/colossalai/zero/zero_optimizer.py
+++ b/colossalai/zero/zero_optimizer.py
@@ -139,6 +139,7 @@ class ZeroOptimizer(ColossalaiOptimizer):
        self._update_params_ptr()
        ret = self.optim.step(*args, **kwargs)
        self._register_states()
+        self.zero_grad()
        self._update_fp16_params()
        return ret