Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zero/stage_1_and_2.py (Python)

Code before:
                    if param.grad_accum is None:
                        param.grad_accum = param.grad.to(
                            self.gradient_accumulation_dtype)
                    else:
                        param.grad_accum.add_(
                            param.grad.to(self.gradient_accumulation_dtype).view(
                                param.grad_accum.shape))
                    # Don't clear if you still need the native grad for reduction (stage 2)
                    if self.use_grad_accum_for_reduction:
                        param.grad = None

    def set_grad_accum_pointer(self):
        for group in self.bit16_groups:
            for param in group:
                param.grad_accum = param.grad

    # Clear the tensor the reduction gradient attribute is pointing to
    def clear_grad_reduc_pointer(self, param):
        param.grad_reduc = None
        if self.use_grad_accum_for_reduction:
            param.grad_accum = None
        else:
            param.grad = None

    # Point to the gradient tensor to use in reduction
    def set_grad_reduc_pointer(self, param):
        if self.use_grad_accum_for_reduction:
            param.grad_reduc = param.grad_accum
        else:
            param.grad_reduc = param.grad

    # Called after a memory defrag moves tensor object via copy
    def reset_pointers_to_reduc_after_copy(self, param):
        if self.use_grad_accum_for_reduction:
            param.grad_accum = param.grad_reduc
        else:
            param.grad = param.grad_reduc
            if self.partition_gradients or not self.use_separate_grad_accum:
                param.grad_accum = param.grad

    def create_reduce_and_remove_grad_hooks(self):
        self.grad_accs = []
        for i, param_group in enumerate(self.bit16_groups):
            for param in param_group:
                if param.requires_grad:

                    def wrapper(param, i):
                        param_tmp = param.expand_as(param)
                        grad_acc = param_tmp.grad_fn.next_functions[0][0]

                        def reduce_partition_and_remove_grads(*notneeded):

Review comment (suggestion):
I think we can avoid this setter by using a getter to access the gradient to reduce. Something like 

```python
  def get_gradient_for_reduction(self, param):
    return param.grad_accum.to(self.dtype) if self.use_grad_accum_for_reduction else param.grad
```

Then all uses of `param.grad_reduc` can be replaced. 

Code after:

                        grad_acc.register_hook(reduce_partition_and_remove_grads)
                        self.grad_accs.append(grad_acc)

                    wrapper(param, i)

    def get_param_id(self, param):
        unique_id = id(param)
        return self.param_id[unique_id]

    def report_ipg_memory_usage(self, tag, param_elems):
        elem_count = self.elements_in_ipg_bucket + param_elems
        percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size
        see_memory_usage(
            f"{tag}: elems in_bucket {self.elements_in_ipg_bucket} param {param_elems} max_percent {percent_of_bucket_size}"
        )

    # create a flat tensor aligned at the alignment boundary
    def flatten_dense_tensors_aligned(self, tensor_list, alignment):
        return self.flatten(align_dense_tensors(tensor_list, alignment))

    ############### Independent Partition Gradient ########################
    def reduce_independent_p_g_buckets_and_remove_grads(self, param, i):

        grad_reduc = self.get_gradient_for_reduction(param)
        if self.elements_in_ipg_bucket + param.numel() > self.reduce_bucket_size:
            self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel())
            self.reduce_ipg_grads()
            if self.contiguous_gradients and self.overlap_comm:
                # Swap ipg_index between 0 and 1
                self.ipg_index = 1 - self.ipg_index
            self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel())

        param_id = self.get_param_id(param)
        assert self.params_already_reduced[param_id] == False, \
            f"The parameter {param_id} has already been reduced. \
            Gradient computed twice for this partition. \
            Multiple gradient reduction is currently not supported"

        if self.contiguous_gradients:
            if param.numel() > self.reduce_bucket_size:
                self.extra_large_param_to_reduce = param
            else:
                # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening
                new_grad_tensor = self.ipg_buffer[self.ipg_index].narrow(0, self.elements_in_ipg_bucket, param.numel())
                new_grad_tensor.copy_(param.grad.view(-1))
                grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc)

        self.elements_in_ipg_bucket += param.numel()

        assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient"