diff --git a/colossalai/zero/low_level/bookkeeping/bucket_store.py b/colossalai/zero/low_level/bookkeeping/bucket_store.py index 98f1b78d0..0ab10e25d 100644 --- 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 @@ class BucketStore(BaseStore): 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 @@ class BucketStore(BaseStore): 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 @@ class BucketStore(BaseStore): 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 index 2b3f50ed4..64d6a5395 100644 --- a/colossalai/zero/low_level/low_level_optim.py +++ b/colossalai/zero/low_level/low_level_optim.py @@ -242,10 +242,19 @@ class LowLevelZeroOptimizer(OptimizerWrapper): 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() diff --git a/tests/test_zero/test_low_level/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py index 5a0609bff..9c4474aff 100644 --- 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)