From 509f0e41c030e0aae2daf5fba9b681db63aac60e Mon Sep 17 00:00:00 2001 From: Nadav Elyahu <88962733+nelyahu@users.noreply.github.com> Date: Mon, 9 Sep 2024 11:59:27 +0300 Subject: [PATCH] Revert "BF16 optimizer: Clear lp grads after updating hp grads in hook (#5328)" This reverts commit 40009eb1c7a4ef27d00f36fe4f97aeae2e315c0e. --- deepspeed/runtime/bf16_optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/bf16_optimizer.py b/deepspeed/runtime/bf16_optimizer.py index 9898abce801d..6b63efbb23f7 100644 --- a/deepspeed/runtime/bf16_optimizer.py +++ b/deepspeed/runtime/bf16_optimizer.py @@ -534,7 +534,7 @@ def state(self): def accumulate_hp_grads_and_remove_lp(self, lp_param, group_idx, param_idx): assert self.immediate_grad_update - self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=True) + self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=False) def create_grad_acc_hooks(self): self.grad_accs = []