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

2025-09-03 18:19:58 +00:00 · 2023-11-02 17:59:10 +08:00
parent dc003c304c
commit d99b2c961a
4 changed files with 13 additions and 3 deletions
--- a/colossalai/zero/gemini/chunk/chunk.py
+++ b/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