def METHOD_NAME(self):
    for index, (param, global_param, smoothed_grad, grad) in enumerate(
        zip(
            self.params,
            self.global_params,
            self.smoothed_grads,
            # all gpus would share the same value of smoothed_grad, since it is
            # always computed on synchronized gradients.
            self.grads,
        )
    ):
        # global_param is basically last syncrhornized parameter. though
        # smoothed_grad is local, all processes will have same value of
        # smoothed_grad and hence param is globally synchronized copy.
        # smoothed_grad(t) = BM * smoothed_grad(t-1) + BM_lr * grad(t)
        smoothed_grad = self.block_momentum * smoothed_grad + self.block_lr * grad
        param.data.copy_(global_param - smoothed_grad)
        # A Nesterov momentum here is to do a partial weight update before
        # calculating the gradient
        if self.use_nbm:
            param.data.copy_(param.data - self.block_momentum * smoothed_grad)
        # backup for the next synchronization.
        self.smoothed_grads[index] = smoothed_grad
        global_param.copy_(param.data)