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

Code before:
                if self.use_grad_accum_for_reduction:
                    param.grad = None

    # Clear the tensor the reduction gradient attribute is pointing to
    def clear_grad_reduc(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 point_grad_reduc(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 point_grad_to_reduc(self, param):
        if self.use_grad_accum_for_reduction:
            param.grad_accum = param.grad_reduc
        else:
            param.grad = param.grad_reduc

    def create_set_grad_accum_hooks(self):
        self.accum_accs = []
        for param_group in self.bit16_groups:
            for param in param_group:
                if param.requires_grad:

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

                        def accumulation_grad_setter(*notneeded):
                            self.set_grad_accum(param)

                        grad_acc.register_hook(accumulation_grad_setter)
                        self.accum_accs.append(grad_acc)

                    wrapper(param)

    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)

Review comment (performance):
Hooks can incur runtime overheads that could be non-trivial for zero stage 1, so I would like to avoid them. Couldn't we call `set_grad_accum` for all params in the optimizer constructor?

Code after:
        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"

        self.grads_in_ipg_bucket.append(grad_reduc)
        self.params_in_ipg_bucket.append((i, param, param_id))

        #make sure the average tensor function knows how to average the gradients
        if is_moe_param(param):
            self.ipg_bucket_has_moe_params = True

        self.report_ipg_memory_usage("End ipg_remove_grads", 0)

    def print_rank_0(self, message):