ranger21.py

# Ranger21 - @lessw2020  and @NestorDemeure
# with contributions from:
# @BrianPugh
# @Kayuksel
# @TheZothen

# core components based on:

# MADGRAD: https://arxiv.org/abs/2101.11075

# warmup:  https://arxiv.org/abs/1910.04209v3

# stable weight decay: https://arxiv.org/abs/2011.11152v3

# Gradient Centralization: https://arxiv.org/abs/2004.01461v2

# positive negative momentum:  https://arxiv.org/abs/2103.17182

# adaptive gradient clipping: https://arxiv.org/abs/2102.06171)
#   big thanks to @kayuksel for suggestion to include agc, and initial code,
#   @lucidrains and @rwightman for additional code impl reference

# softplus transformation to denom:  https://arxiv.org/abs/1908.00700

# lookahead:
# Lookahead Optimizer: https://arxiv.org/abs/1907.08610

# norm loss: https://arxiv.org/abs/2103.06583v1
# big thanks to Theodoros Georgiou for TF code implementation, and he and team for inventing norm loss

# flat lr + cosine decay: original work 2019 (fastai team)

# Chebyshev fractal steps:

# This space for rent - send in your improvements!


import torch
import torch.optim as TO
import torch.nn.functional as F

import math
import collections

# this is to support showing the lr curves after a training run.
import matplotlib.pyplot as plt

import copy
from torch import linalg as LA

import numpy as np


def cheb_steps(m, M, T):
    C, R = (M + m) / 2.0, (M - m) / 2.0
    thetas = (np.arange(T) + 0.5) / T * np.pi
    return 1.0 / (C - R * np.cos(thetas))


def cheb_perm(T):
    perm = np.array([0])
    while len(perm) < T:
        perm = np.vstack([perm, 2 * len(perm) - 1 - perm]).T.flatten()
    return perm


def get_chebs(num_epochs):
    num_epochs = num_epochs - 2
    steps = cheb_steps(0.1, 1, num_epochs)
    perm = cheb_perm(num_epochs)
    cheb_schedule = steps[perm]
    print(f"cheb schedule made with len {len(cheb_schedule)}")
    return cheb_schedule


def normalize_gradient(x, use_channels=False, epsilon=1e-8):
    """  use stdev to normalize gradients """
    size = x.dim()
    # print(f"size = {size}")

    if (size > 1) and use_channels:
        s = x.std(dim=tuple(range(1, size)), keepdim=True) + epsilon
        # print(f"s = {s}")
        x.div_(s)  # , keepdim=True)

    elif torch.numel(x) > 2:
        s = x.std() + epsilon
        x.div_(s)  # , keepdim=True)
    return x


def centralize_gradient(x, gc_conv_only=False):
    """credit - https://github.com/Yonghongwei/Gradient-Centralization """

    size = x.dim()
    # print(f"size = {size}")

    if gc_conv_only:
        if size > 3:
            x.add_(-x.mean(dim=tuple(range(1, size)), keepdim=True))
    else:
        if size > 1:
            x.add_(-x.mean(dim=tuple(range(1, size)), keepdim=True))
    return x


class Ranger21(TO.Optimizer):
    def __init__(
        self,
        params,
        lr,
        lookahead_active=True,
        lookahead_mergetime=5,
        lookahead_blending_alpha=0.5,
        lookahead_load_at_validation=False,
        use_madgrad=False,
        use_adabelief=False,
        softplus=True,
        beta_softplus=50,
        using_gc=True,
        using_normgc=True,
        gc_conv_only=False,
        normloss_active=True,
        normloss_factor=1e-4,
        use_adaptive_gradient_clipping=True,
        agc_clipping_value=1e-2,
        agc_eps=1e-3,
        betas=(0.9, 0.999),  # temp for checking tuned warmups
        momentum_type="pnm",
        pnm_momentum_factor=1.0,
        momentum=0.9,
        eps=1e-8,
        num_batches_per_epoch=None,
        num_epochs=None,
        use_cheb=False,
        use_warmup=True,
        num_warmup_iterations=None,
        warmdown_active=True,
        warmdown_start_pct=0.72,
        warmdown_min_lr=3e-5,
        weight_decay=1e-4,
        decay_type="stable",
        warmup_type="linear",
        warmup_pct_default=0.22,
        logging_active=True,
    ):

        # todo - checks on incoming params
        defaults = dict(
            lr=lr, momentum=momentum, betas=betas, eps=eps, weight_decay=weight_decay
        )
        super().__init__(params, defaults)

        # core
        self.logging = logging_active

        # engine
        self.use_madgrad = use_madgrad

        self.num_batches_per_epoch = num_batches_per_epoch
        self.num_epochs = num_epochs

        if not self.use_madgrad:
            self.core_engine = "AdamW"
        else:
            self.core_engine = "madgrad"

        # ada belief:
        self.use_adabelief = use_adabelief

        # eps
        self.eps = eps

        # softplus for denom
        self.softplus = softplus
        self.beta_softplus = beta_softplus

        # norm loss
        self.normloss_active = normloss_active
        self.normloss_factor = normloss_factor

        # lookahead
        self.lookahead_active = lookahead_active
        self.lookahead_mergetime = lookahead_mergetime
        self.lookahead_step = 0
        self.lookahead_alpha = lookahead_blending_alpha

        self.lookahead_validation_load = lookahead_load_at_validation

        # agc
        self.agc_active = use_adaptive_gradient_clipping
        self.agc_clip_val = agc_clipping_value
        self.agc_eps = agc_eps

        # chebs
        self.use_cheb = use_cheb
        self.cheb_schedule = None
        if self.use_cheb:
            if num_epochs is None:
                raise ValueError(
                    "can't produce chebs without num epochs info being passed in"
                )
            self.cheb_schedule = get_chebs(num_epochs)

        self.total_iterations = num_epochs * num_batches_per_epoch
        if not self.total_iterations:
            raise ValueError(
                "missing total iterations, which is calced from num epochs and num iters per epoch param"
            )

        # lr
        self.starting_lr = lr
        self.current_lr = lr

        # warmup - we'll use default recommended in Ma/Yarats unless user specifies num iterations
        # -=-=-=-=-=-=-=-=-=-=-=-=-=--=-=--=-=-
        self.use_warmup = use_warmup
        self.warmup_complete = False
        self.warmup_type = warmup_type
        self.warmup_pct_default = warmup_pct_default

        if num_warmup_iterations is None:
            beta_warmup_iters = math.ceil(
                (2 / (1 - betas[1]))
            )  # default untuned linear warmup

            beta_pct = beta_warmup_iters / self.total_iterations
            # print(f"beta_warmup_pct = {beta_pct}")

            # this can be unreasonable for short runs...so let's compare vs warmup pct % of total epochs
            if beta_pct > 0.45:
                warmup_auto_pct = int(self.warmup_pct_default * self.total_iterations)
                self.num_warmup_iters = warmup_auto_pct
            else:
                self.num_warmup_iters = beta_warmup_iters

        else:  # user passed in specific num
            self.num_warmup_iters = num_warmup_iterations

        # warm down
        self.min_lr = warmdown_min_lr
        self.warmdown_lr_delta = self.starting_lr - self.min_lr
        self.warmdown_active = warmdown_active

        if self.warmdown_active:
            self.warm_down_start_pct = warmdown_start_pct
            self.start_warm_down = int(
                self.warm_down_start_pct * num_epochs * num_batches_per_epoch
            )
            self.warmdown_total_iterations = (
                self.total_iterations - self.start_warm_down
            )
            self.warmdown_displayed = False  # print when warmdown begins...
            self.warmup_curr_pct = 0.01  # used to verify warmup reaches full set point.

            """
            print(f"debug warmdown:\n")
            print(f"warm_down_start_pct = {self.warm_down_start_pct}")
            print(f"num_epochs = {self.num_epochs}, num_batches per epoch = {self.num_batches_per_epoch}")
            print(f" start warmdown at {self.start_warm_down}")
            print(f" total iterations of warmdown = {self.warmdown_total_iterations}")
            print(f" total lr delta = {self.warmdown_lr_delta}")
            """

        self.current_epoch = 0
        self.current_iter = 0

        self.use_gc = using_gc
        self.use_gcnorm = using_normgc
        self.gc_conv_only = gc_conv_only

        # epochs
        self.epoch_count = 0

        # momentum
        self.momentum_pnm = momentum_type == "pnm"

        self.pnm_momentum = pnm_momentum_factor

        # decay
        self.decay = weight_decay
        self.decay_type = decay_type
        self.param_size = 0

        # logging - need to update things before moving these 2 into self.logging toggle
        self.cheb_logging = []
        self.tracking_lr = []

        if self.logging:
            self.tracking_variance_sum = []
            self.tracking_variance_normalized = []

        # display
        engine = "AdamW" if not self.use_madgrad else "MadGrad"

        # print out initial settings to make usage easier

        self.show_settings()

    def __setstate__(self, state):
        super().__setstate__(state)

    # show settings at init or if called
    def show_schedule(self):
        if not self.tracking_lr:
            print(
                "No data from training yet.  Please train and then use this to show the lr curves"
            )
            return
        x = self.tracking_lr
        plt.plot(x)
        maxlr = max(x)
        minlr = min(x)
        startlr = x[0]
        plt.title(
            f"Ranger21 learning rate schedule\nStart={startlr:.2E}\nMax ={maxlr:.2E}\n,Min={minlr:.2E}\n"
        )
        plt.show()

    def show_settings(self):
        print(f"Ranger21 optimizer ready with following settings:\n")
        print(f"Core optimizer = {self.core_engine}")
        print(f"Learning rate of {self.starting_lr}\n")

        print(
            f"Important - num_epochs of training = ** {self.num_epochs} epochs **\nplease confirm this is correct or warmup and warmdown will be off\n"
        )

        if self.use_adabelief:
            print(f"using AdaBelief for variance computation")
        if self.use_warmup:
            print(
                f"Warm-up: {self.warmup_type} warmup, over {self.num_warmup_iters} iterations\n"
            )
        if self.lookahead_active:
            print(
                f"Lookahead active, merging every {self.lookahead_mergetime} steps, with blend factor of {self.lookahead_alpha}"
            )
        if self.normloss_active:
            print(f"Norm Loss active, factor = {self.normloss_factor}")
        if self.decay:
            print(f"Stable weight decay of {self.decay}")

        if self.use_gc:
            print(f"Gradient Centralization = On\n")
        else:
            print("Gradient Centralization = Off\n")

        print(f"Adaptive Gradient Clipping = {self.agc_active}")
        if self.agc_active:
            print(f"\tclipping value of {self.agc_clip_val}")
            print(f"\tsteps for clipping = {self.agc_eps}")

        if self.warmdown_active:
            print(
                f"\nWarm-down: Linear warmdown, starting at {self.warm_down_start_pct*100}%, iteration {self.start_warm_down} of {self.total_iterations}"
            )
            print(f"warm down will decay until {self.min_lr} lr")

    # lookahead functions
    def clear_cache(self):
        """clears the lookahead cached params """

        print(f"clearing lookahead cache...")
        for group in self.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                try:
                    la_params = param_state["lookahead_params"]
                except:
                    print(f"no lookahead cache present.")
                    return

                if len(la_params):
                    param_state["lookahead_params"] = torch.zeros_like(p.data)
        print(f"lookahead cache cleared")

    def clear_and_load_backup(self):
        for group in self.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                p.data.copy_(param_state["backup_params"])
                del param_state["backup_params"]

    def backup_and_load_cache(self):
        for group in self.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["backup_params"] = torch.zeros_like(p.data)
                param_state["backup_params"].copy_(p.data)
                p.data.copy_(param_state["lookahead_params"])

    def unit_norm(self, x):
        """ axis-based Euclidean norm"""
        # verify shape
        keepdim = True
        dim = None

        xlen = len(x.shape)
        # print(f"xlen = {xlen}")

        if xlen <= 1:
            keepdim = False
        elif xlen in (2, 3):  # linear layers
            dim = 1
        elif xlen == 4:  # conv kernels
            dim = (1, 2, 3)
        else:
            dim = tuple(
                [x for x in range(1, xlen)]
            )  # create 1,..., xlen-1 tuple, while avoiding last dim ...

        return x.norm(dim=dim, keepdim=keepdim, p=2.0)

    def agc(self, p):
        """clip gradient values in excess of the unitwise norm.
        the hardcoded 1e-6 is simple stop from div by zero and no relation to standard optimizer eps
        """

        # params = [p for p in parameters if p.grad is not None]
        # if not params:
        #    return

        # for p in params:
        p_norm = self.unit_norm(p).clamp_(self.agc_eps)
        g_norm = self.unit_norm(p.grad)

        max_norm = p_norm * self.agc_clip_val

        clipped_grad = p.grad * (max_norm / g_norm.clamp(min=1e-6))

        new_grads = torch.where(g_norm > max_norm, clipped_grad, p.grad)
        p.grad.detach().copy_(new_grads)

    def warmup_dampening(self, lr, step):

        style = self.warmup_type
        warmup = self.num_warmup_iters

        if style is None:
            return lr

        if step > warmup:
            if not self.warmup_complete:

                if not self.warmup_curr_pct == 1.0:
                    print(
                        f"Error - lr did not achieve full set point from warmup, currently {self.warmup_curr_pct}"
                    )

                self.warmup_complete = True
                print(f"\n** Ranger21 update = Warmup complete - lr set to {lr}\n")
            return lr

        if style == "linear":
            self.warmup_curr_pct = min(1.0, (step / warmup))
            new_lr = lr * self.warmup_curr_pct
            self.current_lr = new_lr
            return new_lr

        # elif style == "exponential":
        # return lr * (1.0 - math.exp(-step / warmup))

        else:
            raise ValueError(f"warmup type {style} not implemented.")

    def get_warm_down(self, lr, iteration):
        """ linear style warmdown """
        if iteration < self.start_warm_down:
            return lr

        if iteration > self.start_warm_down - 1:
            # print when starting
            if not self.warmdown_displayed:
                print(
                    f"\n** Ranger21 update: Warmdown starting now.  Current iteration = {iteration}....\n"
                )
                self.warmdown_displayed = True

            warmdown_iteration = (
                iteration + 1
            ) - self.start_warm_down  # to force the first iteration to be 1 instead of 0

            if warmdown_iteration < 1:
                print(
                    f" warning - iteration started at {iteration} and {self.start_warm_down} with value {warmdown_iteration}"
                )
                warmdown_iteration = 1
            # print(f"warmdown iteration = {warmdown_iteration}")
            # linear start 3672  5650 total iterations 1972 iterations

            warmdown_pct = warmdown_iteration / (
                self.warmdown_total_iterations + 1
            )  # +1 to offset that we have to include first as an iteration to support 1 index instead of 0 based.
            if warmdown_pct > 1.00:
                print(f"error in warmdown pct calc.  new pct = {warmdown_pct}")
                print(f"auto handled but please report issue")
                warmdown_pct = 1.00

            # .5
            lr_range = self.warmdown_lr_delta

            reduction = lr_range * warmdown_pct
            # print(f"lr reduction = {reduction} for {warmdown_pct} with iter {warmdown_iteration} and total iter {iteration}")
            new_lr = self.starting_lr - reduction
            if new_lr < self.min_lr:
                print(f"error in warmdown - lr below min lr. current lr = {new_lr}")
                print(f"auto handling but please report issue!")
                new_lr = self.min_lr

            self.current_lr = new_lr
            return new_lr

            # new_lr = (
            #    self.min_lr
            #    + self.starting_lr
            #    * (1 + math.cos(math.pi * warmdown_iteration / self.warmdown_total_iterations))
            #    / 2
            # )
            # self.current_lr = new_lr
            # return new_lr

    # def new_epoch_handler(self, iteration):

    #    self.epoch_count +=1

    def track_epochs(self, iteration):
        self.current_iter += 1
        if self.current_iter % self.num_batches_per_epoch == 0:
            self.current_iter = 0
            self.epoch_count += 1
            # print(f"New epoch, current epoch = {self.epoch_count}")
            self.tracking_lr.append(self.current_lr)

            # load lookup params for validation
            if self.lookahead_active and self.lookahead_validation_load:
                self.backup_and_load_cache()

    def get_cheb_lr(self, lr, iteration):

        # first confirm we are done with warmup
        if self.use_warmup:
            if iteration < self.num_warmup_iters + 1:
                return lr

        # compute epoch
        current_epoch = (iteration // self.num_batches) + 1
        # print(f"current epoch for cheb = {current_epoch}")
        self.current_epoch = current_epoch
        index = current_epoch - 2
        if index < 0:
            index = 0
        if index > len(self.cheb_schedule) - 1:
            index = len(self.cheb_schedule) - 1

        cheb_value = self.cheb_schedule[index]

        if self.cheb_logging[:-1] != cheb_value:
            self.cheb_logging.append(cheb_value)

        return lr * cheb_value

    def get_variance(self):
        return self.tracking_variance_sum

    def get_state_values(self, group, state):
        beta1, beta2 = group["betas"]
        mean_avg = state["mean_avg"]
        variance_avg = state["variance_avg"]

        return beta1, beta2, mean_avg, variance_avg

    # @staticmethod
    @torch.no_grad()
    def step(self, closure=None):

        loss = None
        if closure is not None and isinstance(closure, collections.abc.Callable):
            with torch.enable_grad():
                loss = closure()

        param_size = 0
        variance_ma_sum = 0.0

        # phase 1 - accumulate all of the variance_ma_sum to use in stable weight decay

        for i, group in enumerate(self.param_groups):
            for j, p in enumerate(group["params"]):
                if p.grad is None:
                    continue

                # if not self.param_size:
                param_size += p.numel()

                # apply agc if enabled
                if self.agc_active:
                    self.agc(p)

                grad = p.grad

                if grad.is_sparse:
                    raise RuntimeError("sparse matrix not supported atm")

                state = self.state[p]
                momentum = group["momentum"]

                # State initialization
                if len(state) == 0:
                    # print("init state")
                    state["step"] = 0
                    # Exponential moving average of gradient values
                    state["grad_ma"] = torch.zeros_like(
                        p, memory_format=torch.preserve_format
                    )
                    # Exponential moving average of squared gradient values
                    state["variance_ma"] = torch.zeros_like(
                        p, memory_format=torch.preserve_format
                    )

                    if self.lookahead_active:
                        state["lookahead_params"] = torch.zeros_like(p.data)
                        state["lookahead_params"].copy_(p.data)

                    if self.use_adabelief:
                        state["variance_ma_belief"] = torch.zeros_like(
                            p, memory_format=torch.preserve_format
                        )
                    if self.momentum_pnm:
                        state["neg_grad_ma"] = torch.zeros_like(
                            p, memory_format=torch.preserve_format
                        )

                        # Maintains max of all exp. moving avg. of sq. grad. values
                        state["max_variance_ma"] = torch.zeros_like(
                            p, memory_format=torch.preserve_format
                        )

                    # Cumulative products of beta1
                    # state["beta1_prod"] = torch.ones_like(
                    #    p.data, memory_format=torch.preserve_format
                    # )

                # centralize gradients
                if self.use_gc:
                    grad = centralize_gradient(
                        grad,
                        gc_conv_only=self.gc_conv_only,
                    )
                if self.use_gcnorm:
                    grad = normalize_gradient(grad)
                # else:
                #    grad = uncentralized_grad

                # phase 1, variance computations

                state["step"] += 1

                step = state["step"]
                lr = group["lr"]

                beta1, beta2 = group["betas"]
                grad_ma = state["grad_ma"]

                bias_correction2 = 1 - beta2 ** state["step"]
                # print(f"bias2 = {bias_correction2}")

                variance_ma = state["variance_ma"]
                if self.use_adabelief:
                    variance_ma_belief = state["variance_ma_belief"]

                # print(f"variance_ma, upper loop = {variance_ma}")

                # update the exp averages
                if self.use_adabelief:
                    grad_ma.mul_(beta1).add_(grad, alpha=1 - beta1)
                    grad_residual = grad - grad_ma
                    variance_ma_belief.mul_(beta2).addcmul(
                        grad_residual, grad_residual, value=1 - beta2
                    )
                # print(f"upper loop grad = {grad.shape}")
                variance_ma.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                # print(f"variance_ma, grad adjusted")
                variance_ma_debiased = variance_ma / bias_correction2

                variance_ma_sum += variance_ma_debiased.sum()
                # print(f"variance_ma_sum = {variance_ma_sum}")
                # else: #madgrad

                # should we dupe variance_ma since stable is assuming adam style] variance?

                # stable wd
                # variance_ma_sum += grad_sum_sq.sum()

            # print(f"variance hat sum = {exp_avg_sq_hat_sum}")
            # Calculate the sqrt of the mean of all elements in exp_avg_sq_hat

            # we will run this first epoch only and then memoize
        if not self.param_size:
            self.param_size = param_size
            print(f"params size saved")
            print(f"total param groups = {i+1}")
            print(f"total params in groups = {j+1}")

        if not self.param_size:
            raise ValueError("failed to set param size")

        # stable weight decay
        if self.use_madgrad:
            variance_normalized = torch.pow(variance_ma_sum / param_size, 1/3)
        else:
            variance_normalized = math.sqrt(variance_ma_sum / param_size)
        # variance_mean = variance_ma_sum / param_size
        if math.isnan(variance_normalized):
            raise RuntimeError("hit nan for variance_normalized")

        # debugging/logging
        if self.logging:
            self.tracking_variance_sum.append(variance_ma_sum.item())
            self.tracking_variance_normalized.append(variance_normalized)

        # print(f"variance_mean = {variance_mean}")
        # print(f"variance_normalized = {variance_normalized}")
        # else:
        #    variance_normalized = math.pow((variance_ma / self.param_size), .3333)

        # print(f"variance mean sqrt = {variance_normalized}")

        # phase 2 - apply weight decay and step
        # ===========================================
        for group in self.param_groups:
            # print(f"In second phase loop")
            step = state["step"]

            # Perform stable weight decay
            decay = group["weight_decay"]
            eps = group["eps"]
            lr = group["lr"]
            momentum = group["momentum"]

            beta1, beta2 = group["betas"]

            # warmup
            # ======================
            if self.use_warmup and not self.warmup_complete:
                lr = self.warmup_dampening(lr, step)
                # print(f"lr = {lr}")

            # chebyshev
            # ===================
            if self.use_cheb and self.warmup_complete:
                lr = self.get_cheb_lr(lr, step)

            # warmdown
            # ==========
            if self.warmdown_active:
                orig_lr = lr
                lr = self.get_warm_down(lr, step)
                assert lr > 0, "lr went negative"

            # madgrad outer
            if self.use_madgrad:
                ck = 1 - momentum
                lamb = lr * math.pow(step, 0.5)

            # stable decay and / or norm loss
            # ==================================
            if decay:
                if not self.use_madgrad:
                    # stable weight decay
                    p.data.mul_(1 - decay * lr / variance_normalized)
                else:
                    p.data.mul_(1 - decay * lamb / variance_normalized)

            if self.normloss_active:
                # apply norm loss
                unorm = self.unit_norm(p.data)
                correction = (
                    2 * self.normloss_factor * (1 - torch.div(1, unorm + self.eps))
                )
                p.mul_(1 - lr * correction)

            # innner loop, params
            for p in group["params"]:
                if p.grad is None:
                    continue

                state = self.state[p]
                inner_grad = p.grad

                if self.use_madgrad:
                    # ================== madgrad ============================
                    if "grad_sum_sq" not in state:
                        state["grad_sum_sq"] = torch.zeros_like(p.data).detach()
                        state["s"] = torch.zeros_like(p.data).detach()
                        if momentum != 0:
                            state["x0"] = torch.clone(p.data).detach()

                    if momentum != 0.0 and grad.is_sparse:
                        raise RuntimeError(
                            "momentum != 0 is not compatible with sparse gradients"
                        )

                    # centralize gradients
                    if self.use_gc:
                        inner_grad = centralize_gradient(
                            inner_grad,
                            gc_conv_only=self.gc_conv_only,
                        )

                    grad_sum_sq = state["grad_sum_sq"]
                    s = state["s"]
                    if momentum == 0:
                        # Compute x_0 from other known quantities
                        rms = grad_sum_sq.pow(1 / 3)
                        if self.softplus:
                            rms = F.softplus(rms, beta=self.beta_softplus)
                        x0 = p.data.addcdiv(s, rms, value=1)
                    else:
                        x0 = state["x0"]

                    # Accumulate second moments

                    # print(f" grad = {grad}")
                    # print(f"lamb = {lamb}")
                    # print(f"gsumsq = {grad_sum_sq}")

                    grad_sum_sq.addcmul_(inner_grad, inner_grad, value=lamb)
                    rms = grad_sum_sq.pow(1 / 3)
                    if self.softplus:
                        rms = F.softplus(rms, beta=self.beta_softplus)

                    # Update s
                    s.data.add_(inner_grad, alpha=lamb)

                    # Step
                    if momentum == 0:
                        p.data.copy_(x0.addcdiv(s, rms, value=-1))
                    else:
                        z = x0.addcdiv(s, rms, value=-1)

                        # p is a moving average of z
                        p.data.mul_(1 - ck).add_(z, alpha=ck)

                else:  # adam with pnm core
                    # ============= adamW with pnm option ========================

                    grad = p.grad

                    beta1, beta2 = group["betas"]

                    grad_ma = state["grad_ma"]
                    variance_ma = state["variance_ma"]
                    if self.use_adabelief:
                        variance_ma_belief = state["variance_ma_belief"]

                    if self.momentum_pnm:

                        max_variance_ma = state["max_variance_ma"]

                        if state["step"] % 2 == 1:
                            grad_ma, neg_grad_ma = (
                                state["grad_ma"],
                                state["neg_grad_ma"],
                            )
                        else:
                            grad_ma, neg_grad_ma = (
                                state["neg_grad_ma"],
                                state["grad_ma"],
                            )

                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2 ** step

                    if self.momentum_pnm:
                        # Maintains the maximum of all 2nd moment running avg. till now
                        torch.max(max_variance_ma, variance_ma, out=variance_ma)
                        # Use the max. for normalizing running avg. of gradient
                        denom = (variance_ma.sqrt() / math.sqrt(bias_correction2)).add_(
                            group["eps"]
                        )

                    # centralize gradients
                    if self.use_gc:
                        grad = centralize_gradient(
                            grad,
                            gc_conv_only=self.gc_conv_only,
                        )
                    if self.use_gcnorm:
                        grad = normalize_gradient(grad)

                    if not self.use_adabelief:
                        grad_ma.mul_(beta1 ** 2).add_(grad, alpha=1 - beta1 ** 2)

                    noise_norm = math.sqrt((1 + beta2) ** 2 + beta2 ** 2)

                    step_size = lr / bias_correction1

                    # softplus the denom
                    if self.softplus:
                        denom = F.softplus(denom, beta=self.beta_softplus)

                    pnmomentum = (
                        grad_ma.mul(1 + self.momentum_pnm)
                        .add(neg_grad_ma, alpha=-self.momentum_pnm)
                        .mul(1 / noise_norm)
                    )

                    p.addcdiv_(pnmomentum, denom, value=-step_size)

                    # denom = variance_biased_ma.sqrt().add(eps)

                    # step_size = lr / bias_correction1

                    # update weights
                    # p.data.add_(weight_mod, alpha=-step_size)
                    # p.addcdiv_(grad_ma, denom, value=-step_size)
        # print(f"\n End optimizer step\n")

        # end of step processes....

        # lookahead
        # ---------------------
        if self.lookahead_active:
            self.lookahead_process_step()

        self.track_epochs(step)
        return loss

    #   Lookahead merge process
    def lookahead_process_step(self):
        """handles blending of params for lookahead step"""

        if not self.lookahead_active:
            return
        self.lookahead_step += 1

        if self.lookahead_step >= self.lookahead_mergetime:
            self.lookahead_step = 0
            # merge lookahead cached params and save current ones
            for group in self.param_groups:
                for p in group["params"]:
                    if p.grad is None:
                        continue

                    param_state = self.state[p]

                    p.data.mul_(self.lookahead_alpha).add_(
                        param_state["lookahead_params"],
                        alpha=1.0 - self.lookahead_alpha,
                    )
                    # save for next merge
                    param_state["lookahead_params"].copy_(p.data)