BF16 optimizer: Clear lp grads after updating hp grads in hook (deepspeedai#5328)

This fix is to solve:

- Previous iteration's lp grads will still alive during the next
iteration's forward. This increases the memory footprint.
- The hook behavior is not aligned to its name
accumulate_hp_grads_and_remove_lp

Co-authored-by: qunyang <quyang@habana.ai>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>

Author: 2 people

deepspeed/runtime/bf16_optimizer.py

@@ -517,7 +517,7 @@ def param_groups(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=False)
+        self._update_hp_grad(lp_param, group_idx, param_idx, clear_lp_grads=True)
 
     def create_grad_acc_hooks(self):
         self.grad_accs = []