Repository: hpcaitech/ColossalAI
File: colossalai/zero/low_level/low_level_optim.py (Python)

Code before:
            self._grad_store.reset_grads_by_group_id(group_id)

            # update the params in the optimizer
            self.optim.param_groups[group_id]['params'] = real_master_params[group_id]

        # unscale and clip grads
        global_norm = calculate_global_norm_from_list(norm_list=norm_groups)
        self._unscale_and_clip_grads(grad_partition_groups, global_norm)

        # update the parameters
        self.optim.step()

        # release the grad
        grad_partition_groups = []
        for group_id in range(self.num_param_groups):
            release_param_grad(self._master_param_groups_of_current_rank[group_id])

        # update working partition updated by the current rank
        dtype = real_working_params[0][0].dtype
        device = real_working_params[0][0].device
        for group_id in range(self.num_param_groups):
            master_working_param = self.optim.param_groups[group_id]['params']
            for idx, splited_param in enumerate(master_working_param):
                working_param = real_working_params[group_id][idx]
                all_splited_param = [
                    torch.zeros(splited_param.shape, device=device, dtype=dtype) for _ in range(self._world_size)
                ]
                dist.all_gather(all_splited_param, splited_param.to(device).to(dtype), group=self.dp_pg)
                working_param.data.copy_(flatten(all_splited_param)[:working_param.numel()].reshape_as(working_param))

            self.optim.param_groups[group_id]['params'] = self._master_param_groups_of_current_rank[group_id]

    #############################
    # Mixed Precision Utilities #
    #############################

    def _unscale_and_clip_grads(self, grad_groups_flat, total_norm):
        # compute combined scale factor for this group
        div_scale = 1.0
        if self.mixed_precision_mixin is not None:
            div_scale = self.mixed_precision_mixin.get_grad_div_scale()

        if self._clip_grad_norm > 0.:
            # norm is in fact norm*scale
            clip = ((total_norm / div_scale) + 1e-6) / self._clip_grad_norm
            if clip > 1:
                div_scale = clip * div_scale

        for grad in grad_groups_flat:
            grad.data.mul_(1. / div_scale)


Review comment (suggestion):
You'd better keep `all_splited_param` and `splited_param` on cuda, in case user move model to cpu after forward-backward manually.

Code after:
            self._grad_store.reset_grads_by_group_id(group_id)

            # update the params in the optimizer
            self.optim.param_groups[group_id]['params'] = real_master_params[group_id]

        # unscale and clip grads
        global_norm = calculate_global_norm_from_list(norm_list=norm_groups)
        self._unscale_and_clip_grads(grad_partition_groups, global_norm)

        # update the parameters
        self.optim.step()

        # release the grad
        grad_partition_groups = []
        for group_id in range(self.num_param_groups):
            release_param_grad(self._master_param_groups_of_current_rank[group_id])

        # update working partition updated by the current rank
        dtype = real_working_params[0][0].dtype
        for group_id in range(self.num_param_groups):
            master_working_param = self.optim.param_groups[group_id]['params']
            for idx, splited_param in enumerate(master_working_param):
                working_param = real_working_params[group_id][idx]
                all_splited_param = [
                    torch.zeros(splited_param.shape, device="cuda", dtype=dtype) for _ in range(self._world_size)
                ]
                dist.all_gather(all_splited_param, splited_param.cuda().to(dtype), group=self.dp_pg)
                working_param.data.copy_(flatten(all_splited_param)[:working_param.numel()].reshape_as(working_param))

            self.optim.param_groups[group_id]['params'] = self._master_param_groups_of_current_rank[group_id]

    #############################
    # Mixed Precision Utilities #
    #############################

    def _unscale_and_clip_grads(self, grad_groups_flat, total_norm):
        # compute combined scale factor for this group
        div_scale = 1.0
        if self.mixed_precision_mixin is not None:
            div_scale = self.mixed_precision_mixin.get_grad_div_scale()

        if self._clip_grad_norm > 0.:
            # norm is in fact norm*scale
            clip = ((total_norm / div_scale) + 1e-6) / self._clip_grad_norm
            if clip > 1:
                div_scale = clip * div_scale

        for grad in grad_groups_flat:
            grad.data.mul_(1. / div_scale)

    ############################