[zero] solve hang

2025-09-07 12:01:39 +00:00 · 2024-07-09 08:14:00 +00:00
parent b5bfeb2efd
commit 13b48ac0aa
8 changed files with 218 additions and 335 deletions
--- a/colossalai/zero/low_level/bookkeeping/gradient_store.py
+++ b/colossalai/zero/low_level/bookkeeping/gradient_store.py
@@ -19,7 +19,6 @@ class GradientStore(BaseStore):
        """
        self._grads_of_params = dict()
        # stage 2
-        self._partition_grads = partition_grad
        self._working_index = 0 if partition_grad else self._local_rank
        # for zero2, it's `param_id: [grad_local_rank]`
        self.grad_to_param_mapping = dict()