clean up accumulated fp32 grads between data batches

facebookresearch · May 2, 2024 · c97bfd9 · c97bfd9
1 parent 4b5abe2
commit c97bfd9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1729,6 +1729,8 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # else:
             #     param.unsharded_main_grad.add_(param.grad.data)
             param.unsharded_main_grad = new_unsharded_main_grad_in_fp32
+            # Clean up accumulated grads between data batches
+            self._fsdp_wrapped_module.fp32_grads = []
             param.grad = None
 
         if not self._require_backward_grad_sync: