Repository: hpcaitech/ColossalAI

Issue:
[BUG]: low level zero has unstable performance on numeral accuracy
### 🐛 Describe the bug

When running low-level zero on several empty cards, it works fine.
However, if there are other processes on the GPU, the gradients of zero becomes strange and unstable.

And the problem can be reproduced frequently but not ever time.

### Environment

_No response_
[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/zero/low_level/bookkeeping/bucket_store.py b/colossalai/zero/low_level/bookkeeping/bucket_store.py
--- a/colossalai/zero/low_level/bookkeeping/bucket_store.py
+++ b/colossalai/zero/low_level/bookkeeping/bucket_store.py
@@ -13,15 +13,20 @@ class BucketStore(BaseStore):
     def __init__(self, torch_pg: ProcessGroup):
         super().__init__(torch_pg)
 
-        # init and reset
+        # init
         self.current_group_id = 0
+        self._num_elements_in_bucket = 0
         # mapping gardient slices and parameter
         self.grad_to_param_mapping = dict()
 
+        self._grad_in_bucket = dict()
         self._param_list = []
         self._padding_size = []
+        for rank in range(self._world_size):
+            self._grad_in_bucket[rank] = []
 
-        self.reset()
+        # offset_list records number of tensors in the bucket before each reduction
+        self.offset_list = [0]
 
     def num_elements_in_bucket(self) -> int:
         """Return the total number of elements in bucket
@@ -32,6 +37,12 @@ def num_elements_in_bucket(self) -> int:
 
         return self._num_elements_in_bucket
 
+    def reset_num_elements_in_bucket(self):
+        """Set the number of elements in bucket to zero.
+        """
+
+        self._num_elements_in_bucket = 0
+
     def add_param_grad(self, group_id: int, param: Tensor, padding_size: int):
         """Add a param to bucket and record the padding size of a param for gradient padding
 
@@ -46,28 +57,32 @@ def add_param_grad(self, group_id: int, param: Tensor, padding_size: int):
         self._num_elements_in_bucket += (param.numel() + padding_size)
         self.current_group_id = group_id
 
+        # number of tensors in current bucket
+        self.offset_list[-1] += 1
+
     def build_grad_in_bucket(self):
         """Orgnize parameters' gradient(padding and split), follows the paramters' splitting method
 
         Data structure of self._grad_in_bucket:
         {
         rank0: [grad0_rank0, grad1_rank0, ...]
-        rank1: [grad1_rank1, grad1_rank1, ...]
+        rank1: [grad0_rank1, grad1_rank1, ...]
         }
         """
-
         for param, padding_size in zip(self._param_list, self._padding_size):
-            with torch.no_grad():
-                grad = param.grad.detach().flatten()
-                if padding_size > 0:
-                    grad = torch.nn.functional.pad(grad, [0, padding_size])
-                grad_list = grad.split(grad.numel() // self._world_size)
-                for rank in range(self._world_size):
-                    grad_current_rank = grad_list[rank].detach()
-                    self.grad_to_param_mapping[id(grad_current_rank)] = id(param)
-                    self._grad_in_bucket[rank].append(grad_current_rank)
+            grad = param.grad.clone().detach().flatten()
+            if padding_size > 0:
+                with torch.no_grad():
+                    grad = torch.nn.functional.pad(grad.view(-1), [0, padding_size])
+            grad_list = grad.split(grad.numel() // self._world_size)
+            for rank in range(self._world_size):
+                grad_current_rank = grad_list[rank].clone().detach()
+                self.grad_to_param_mapping[id(grad_current_rank)] = id(param)
+                self._grad_in_bucket[rank].append(grad_current_rank)
             param.grad = None
 
+        self.offset_list.append(0)
+
     def get_grad(self) -> Dict:
         """Return the dictionary of gradients slices, of which the keys are ranks
 
@@ -104,10 +119,12 @@ def get_param_id_of_grad(self, grad: Tensor) -> int:
         return self.grad_to_param_mapping[id(grad)]
 
     def reset(self):
-        self.grad_to_param_mapping = dict()
-        self._num_elements_in_bucket = 0
-        self._param_list = []
-        self._padding_size = []
-        self._grad_in_bucket = dict()
+        """Reset the bucket storage after reduction, only release the tensors have been reduced
+        """
+        cur_offset = self.offset_list.pop(0)
+        self._param_list = self._param_list[cur_offset:]
+        self._padding_size = self._padding_size[cur_offset:]
+        for _ in range(cur_offset):
+            del self.grad_to_param_mapping[next(iter(self.grad_to_param_mapping))]
         for rank in range(self._world_size):
-            self._grad_in_bucket[rank] = []
+            self._grad_in_bucket[rank] = self._grad_in_bucket[rank][cur_offset:]
diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py
--- a/colossalai/zero/low_level/low_level_optim.py
+++ b/colossalai/zero/low_level/low_level_optim.py
@@ -242,10 +242,19 @@ def _attach_reduction_hook(self):
     def _run_reduction(self):
         if self._bucket_store.num_elements_in_bucket() > 0:
             self._bucket_store.build_grad_in_bucket()
+
             flat_grads = self._bucket_store.get_flatten_grad()
             flat_grads /= self._world_size
+
+            # ready to add other tensors to bucket
+            self._bucket_store.reset_num_elements_in_bucket()
+
             if self._overlap_communication:
                 stream = self._comm_stream
+                # in case of the memory being reused in the default stream
+                flat_grads.record_stream(stream)
+                # waiting for ops in the default stream finishing
+                stream.wait_stream(torch.cuda.current_stream())
             else:
                 stream = torch.cuda.current_stream()
 


Test patch:
diff --git a/tests/test_zero/test_low_level/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py
--- a/tests/test_zero/test_low_level/test_zero1_2.py
+++ b/tests/test_zero/test_low_level/test_zero1_2.py
@@ -137,7 +137,7 @@ def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype):
     zero_optimizer = LowLevelZeroOptimizer(zero_optimizer,
                                            overlap_communication=True,
                                            initial_scale=1,
-                                           reduce_bucket_size=262144)
+                                           reduce_bucket_size=1024 * 1024)
 
     torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)