[hotfix] fix grad accumulation plus clipping for gemini (#5002)

colossalai/zero/gemini/chunk/chunk.py

@@ -637,6 +637,7 @@ class Chunk:
             # grad chunk is initialized, just reallocate cuda global chunk
             self.grad_chunk.cuda_shard = None
             self.grad_chunk.is_gathered = True
+            self.grad_chunk.l2_norm = None
             alloc_storage(self.grad_chunk.cuda_global_chunk)
 
         return self.grad_chunk

colossalai/zero/gemini/gemini_ddp.py

@@ -343,6 +343,7 @@ class GeminiDDP(ModelWrapper):
                         grad_chunk = self.chunk_manager.rearrange_accumulated_grad_chunk(chunk)
                     else:
                         grad_chunk = chunk.grad_chunk
+                        chunk.grad_chunk.l2_norm = None
 
                 # hold -> compute -> hold after bwd
                 grad_chunk.tensor_trans_state(p, TensorState.COMPUTE)