update cpu adam fp16 case

hpcaitech · Oct 17, 2023 · ebc9e52 · ebc9e52
1 parent c7a6004
commit ebc9e52
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 5 deletions.
diff --git a/colossalai/zero/gemini/gemini_ddp.py b/colossalai/zero/gemini/gemini_ddp.py
@@ -373,7 +373,7 @@ def grad_handle(self, p, grad):
                 if chunk.l2_norm_flag:
                     grad_chunk.set_l2_norm()
                 self.chunk_manager.move_chunk(grad_chunk, self.grads_device[p], force_copy=True)
-                if not self.master_weights:
+                if not (self.master_weights) or (self.enable_gradient_accumulation):
                     self.chunk_manager.move_chunk(chunk, self.grads_device[p], force_copy=True)
         return empty_grad
 

diff --git a/tests/test_zero/test_gemini/test_grad_accum.py b/tests/test_zero/test_gemini/test_grad_accum.py
@@ -114,10 +114,6 @@ def exam_gemini_grad_acc(
         check_grad(gemini_model, torch_model)
 
         if (i + 1) % accum_iter == 0:
-            # TODO(Baizhou): Delete following two lines after cpu_adam for fp16 has been merged into main branch (auto policy put gradients in cpu).
-            if placement_config["placement_policy"] == "auto":
-                break
-
             torch_optim.step()
             gemini_optim.step()
             torch_optim.zero_grad()