# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/optimizer.eightbit.ipynb.

# %% ../../nbs/optimizer.eightbit.ipynb 1
# Contains code from:
# bitsandbytes - MIT License - Copyright (c) Facebook, Inc. and its affiliates.
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../../nbs/optimizer.eightbit.ipynb 4
from __future__ import annotations

from bitsandbytes.optim.optimizer import Optimizer1State, Optimizer2State, MockArgs
import bitsandbytes.functional as BF

from fastcore.basics import even_mults

from ..imports import *

# %% auto 0
__all__ = ['SGD8bitOptimizer', 'RMSProp8bitOptimizer', 'AdamW8bitOptimizer', 'LARS8bitOptimizer', 'LAMB8bitOptimizer',
           'Lion8bitOptimizer']

# %% ../../nbs/optimizer.eightbit.ipynb 7
def _convert_params(o:list, **defaults) -> list:
    "Convert fastai param_lists to PyTorch param_groups, adding defaults if group doesn't have it"
    splitter = []
    for group in o:
        if isinstance(group, dict):
            splitter.append({**defaults, **group})
        else:
            splitter.append({'params':group, **defaults})
    return splitter

# %% ../../nbs/optimizer.eightbit.ipynb 8
class EightBitFastaiAdapter:
    "Base for adding fastai optimizer functionality to EightBit Optimizers"
    _keep_on_clear = ['force_train', 'do_wd']
    def get_config(self, gindex, pindex, group):
        config = {}
        config["mom"] = group["mom"]
        config["sqr_mom"] = group["sqr_mom"]
        config["wd"] = group["wd"]
        config["eps"] = group["eps"]
        config["lr"] = group["lr"]
        config["optim_bits"] = self.args.optim_bits
        config["min_8bit_size"] = self.args.min_8bit_size
        config["percentile_clipping"] = self.args.percentile_clipping
        config["block_wise"] = self.args.block_wise
        config["max_unorm"] = self.args.max_unorm
        config["skip_zeros"] = self.args.skip_zeros

        if (gindex, pindex) in self.mng.index2config:
            config.update(self.mng.index2config[(gindex, pindex)])
        return config

    def get_params(self,
        n:slice|int=slice(None), # Extended slicing over the optimizer `param_lists`
        with_grad:bool=False # Get all param tuples. If `True` select only those with a gradient
    ):
        "Slice of parameters and parameter states"
        return L((p, self.state[p]) for pg in self.param_groups[n] for p in pg['params']
                    if (hasattr(p, 'grad') and p.grad is not None) or with_grad==False)

    def clear_state(self):
        "Reset the state of the optimizer"
        for p,state in self.get_params():
            self.state[p] = {k: state[k] for k in self._keep_on_clear if k in state}

    def _set_require_grad(self,
        rg:bool, # Requires grad: if `True` sets gradient for parameters, else uses state `state["force_train"]`
        p:Tensor, # Parameter to set gradient
        state:dict, # Parameter's state dict
    ):
        p.requires_grad_(rg or state.get('force_train', False))

    def freeze_to(self, n:int):
        "Freeze parameter groups up to `n`"
        self.frozen_idx = n if n >= 0 else len(self.param_groups) + n
        if self.frozen_idx >= len(self.param_groups):
            warn(f"Freezing {self.frozen_idx} groups; model has {len(self.param_groups)}; whole model is frozen.")
        for o in self.get_params(slice(n, None)):
            self._set_require_grad(True, *o)
        for o in self.get_params(slice(None, n)):
            self._set_require_grad(False, *o)

    def freeze(self):
        "Freeze up to last parameter group"
        assert(len(self.param_groups) > 1)
        self.freeze_to(-1)

    def unfreeze(self):
        "Unfreeze the entire model"
        self.freeze_to(0)

    @property
    def hypers(self):
        return [{k:v for k,v in pg.items() if k != 'params'} for pg in self.param_groups]

    def set_hypers(self, **kwargs):
        "`set_hyper` for all `kwargs`"
        L(kwargs.items()).starmap(self.set_hyper)

    def _set_hyper(self, k, v):
        "Set the value(s) in `v` for hyper-parameter `k`"
        for v_,h in zip(v, self.param_groups):
            h[k] = v_

    def set_hyper(self, k, v):
        "Set the value(s) in `v` for hyper-parameter `k`"
        if isinstance(v, slice):
            if v.start:
                v = even_mults(v.start, v.stop, len(self.param_groups))
            else:
                v = [v.stop/10]*(len(self.param_groups)-1) + [v.stop]
        v = L(v, use_list=None)
        if len(v)==1:
            v = v*len(self.param_groups)
        assert len(v) == len(self.param_groups), f"Trying to set {len(v)} values for {k} but there are {len(self.param_groups)} parameter groups."
        self._set_hyper(k, v)

    @property
    def param_lists(self):
        return [pg['params'] for pg in self.param_groups]

    @param_lists.setter
    def param_lists(self, v):
        for pg,v_ in zip(self.param_groups,v):
            pg['params'] = v_

# %% ../../nbs/optimizer.eightbit.ipynb 10
class EightBitCommon:
    "Common changes to EightBit Optimizers"
    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        if not self.initialized:
            self.check_overrides()
            self.to_gpu()  # needed for fairseq pure fp16 training
            self.initialized = True

        #if self.is_paged: self.page_mng.prefetch_all()
        for gindex, group in enumerate(self.param_groups):
            for pindex, p in enumerate(group["params"]):
                if p.grad is None:
                    continue
                state = self.state[p]
                if 'step' not in state:
                    self.init_state(group, p, gindex, pindex)

                self.prefetch_state(p)
                self.update_step(group, p, gindex, pindex)
                if self.sync_each_step:
                    torch.cuda.synchronize()
        if self.is_paged or not self.sync_each_step:
            # all paged operation are asynchronous, we need
            # to sync to make sure all tensors are in the right state
            torch.cuda.synchronize()

        return loss

# %% ../../nbs/optimizer.eightbit.ipynb 12
class EightBit1StateOptimizer(EightBitCommon, EightBitFastaiAdapter, Optimizer1State):
    "Adds fastai optimizer functionality & compatability to `Optimizer1State`"
    def __init__(
        self,
        optimizer_name,
        params,
        lr=1e-3,
        mom=0.9,
        sqr_mom=0.0,
        eps=1e-8,
        wd=0.0,
        optim_bits=8,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,
        max_unorm=0.0,
        skip_zeros=False,
        is_paged=False,
        sync_each_step=False
    ):
        if not 0.0 <= lr:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= eps:
            raise ValueError(f"Invalid epsilon value: {eps}")
        if not 0.0 <= mom < 1.0:
            raise ValueError(f"Invalid mom value: {mom}")
        if not 0.0 <= sqr_mom < 1.0:
            raise ValueError(f"Invalid sqr_mom value: {sqr_mom}")
        if not 0.0 <= wd:
            raise ValueError(f"Invalid weight_decay value: {wd}")
        defaults = dict(lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd)
        params = L(params)
        params = _convert_params(params, **defaults) if isinstance(params[0], (L,list)) else params
        super(Optimizer1State, self).__init__(params, defaults, optim_bits, is_paged)

        if args is None:
            args = {}
            args["optim_bits"] = optim_bits
            args["percentile_clipping"] = 100
            args["min_8bit_size"] = min_8bit_size
            args["percentile_clipping"] = percentile_clipping
            args["block_wise"] = block_wise
            args["max_unorm"] = max_unorm
            args["skip_zeros"] = skip_zeros

            self.args = MockArgs(args)
        else:
            self.args = args

        self.optimizer_name = optimizer_name
        self.sync_each_step = sync_each_step

    @torch.no_grad()
    def update_step(self, group, p, gindex, pindex):
        state = self.state[p]
        grad = p.grad

        config = self.get_config(gindex, pindex, group)

        state["step"] += 1
        step = state["step"]

        if config["percentile_clipping"] < 100:
            current_gnorm, clip_value, gnorm_scale = BF.percentile_clipping(
                grad, state["gnorm_vec"], step, config["percentile_clipping"]
            )
        else:
            gnorm_scale = 1.0

        if state["state1"].dtype == torch.float:
            BF.optimizer_update_32bit(
                self.optimizer_name,
                grad,
                p,
                state["state1"],
                config["mom"],
                config["eps"],
                step,
                config["lr"],
                None,
                config['sqr_mom'],
                state.get('wd', config['wd']) if state.get('do_wd', True) else 0.0,
                gnorm_scale,
                state["unorm_vec"] if config["max_unorm"] > 0.0 else None,
                max_unorm=config["max_unorm"],
                skip_zeros=config["skip_zeros"],
            )

        elif state["state1"].dtype == torch.uint8 and not config["block_wise"]:
            BF.optimizer_update_8bit(
                self.optimizer_name,
                grad,
                p,
                state["state1"],
                None,
                config["mom"],
                config['sqr_mom'],
                config["eps"],
                step,
                config["lr"],
                state["qmap1"],
                None,
                state["max1"],
                None,
                state["new_max1"],
                None,
                state.get('wd', config['wd']) if state.get('do_wd', True) else 0.0,
                gnorm_scale,
                state["unorm_vec"] if config["max_unorm"] > 0.0 else None,
                max_unorm=config["max_unorm"],
            )

            state["max1"], state["new_max1"] = state["new_max1"], state["max1"]
        elif state["state1"].dtype == torch.uint8 and config["block_wise"]:
            BF.optimizer_update_8bit_blockwise(
                self.optimizer_name,
                grad,
                p,
                state["state1"],
                None,
                config["mom"],
                config['sqr_mom'],
                config["eps"],
                step,
                config["lr"],
                state["qmap1"],
                None,
                state["absmax1"],
                None,
                state.get('wd', config['wd']) if state.get('do_wd', True) else 0.0,
                gnorm_scale=gnorm_scale,
                skip_zeros=config["skip_zeros"],
            )

# %% ../../nbs/optimizer.eightbit.ipynb 14
class EightBit2StateOptimizer(EightBitCommon, EightBitFastaiAdapter, Optimizer2State):
    "Adds fastai optimizer functionality & compatability to `Optimizer2State`"
    def __init__(
        self,
        optimizer_name,
        params,
        lr=1e-3,
        mom=0.9,
        sqr_mom=0.999,
        eps=1e-8,
        wd=0.0,
        optim_bits=8,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,
        max_unorm=0.0,
        skip_zeros=False,
        is_paged=False,
        sync_each_step=False
    ):
        if not 0.0 <= lr:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= eps:
            raise ValueError(f"Invalid epsilon value: {eps}")
        if not 0.0 <= mom < 1.0:
            raise ValueError(f"Invalid mom value: {mom}")
        if not 0.0 <= sqr_mom < 1.0:
            raise ValueError(f"Invalid sqr_mom value: {sqr_mom}")
        if not 0.0 <= wd:
            raise ValueError(f"Invalid weight_decay value: {wd}")
        defaults = dict(lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd)
        params = L(params)
        params = _convert_params(params, **defaults) if isinstance(params[0], (L,list)) else params
        super(Optimizer2State, self).__init__(params, defaults, optim_bits, is_paged)

        if args is None:
            args = {}
            args["optim_bits"] = optim_bits
            args["percentile_clipping"] = 100
            args["min_8bit_size"] = min_8bit_size
            args["percentile_clipping"] = percentile_clipping
            args["block_wise"] = block_wise
            args["max_unorm"] = max_unorm
            args["skip_zeros"] = skip_zeros

            self.args = MockArgs(args)
        else:
            self.args = args

        self.optimizer_name = optimizer_name
        self.sync_each_step = sync_each_step

    @torch.no_grad()
    def update_step(self, group, p, gindex, pindex):
        state = self.state[p]
        grad = p.grad

        config = self.get_config(gindex, pindex, group)

        state["step"] += 1
        step = state["step"]

        if config["percentile_clipping"] < 100:
            current_gnorm, clip_value, gnorm_scale = BF.percentile_clipping(
                grad, state["gnorm_vec"], step, config["percentile_clipping"]
            )
        else:
            gnorm_scale = 1.0

        if state["state1"].dtype == torch.float:
            BF.optimizer_update_32bit(
                self.optimizer_name,
                grad,
                p,
                state["state1"],
                config["mom"],
                config["eps"],
                step,
                config["lr"],
                state["state2"],
                config["sqr_mom"],
                state.get('wd', config['wd']) if state.get('do_wd', True) else 0.0,
                gnorm_scale,
                state["unorm_vec"] if config["max_unorm"] > 0.0 else None,
                max_unorm=config["max_unorm"],
                skip_zeros=config["skip_zeros"],
            )

        elif state["state1"].dtype == torch.uint8 and not config["block_wise"]:
            BF.optimizer_update_8bit(
                self.optimizer_name,
                grad,
                p,
                state["state1"],
                state["state2"],
                config["mom"],
                config['sqr_mom'],
                config["eps"],
                step,
                config["lr"],
                state["qmap1"],
                state["qmap2"],
                state["max1"],
                state["max2"],
                state["new_max1"],
                state["new_max2"],
                state.get('wd', config['wd']) if state.get('do_wd', True) else 0.0,
                gnorm_scale=gnorm_scale,
                unorm_vec=state["unorm_vec"] if config["max_unorm"] > 0.0 else None,
                max_unorm=config["max_unorm"],
            )

            # swap maxes
            state["max1"], state["new_max1"] = state["new_max1"], state["max1"]
            state["max2"], state["new_max2"] = state["new_max2"], state["max2"]
        elif state["state1"].dtype == torch.uint8 and config["block_wise"]:
            BF.optimizer_update_8bit_blockwise(
                self.optimizer_name,
                grad,
                p,
                state["state1"],
                state["state2"],
                config["mom"],
                config['sqr_mom'],
                config["eps"],
                step,
                config["lr"],
                state["qmap1"],
                state["qmap2"],
                state["absmax1"],
                state["absmax2"],
                state.get('wd', config['wd']) if state.get('do_wd', True) else 0.0,
                gnorm_scale=gnorm_scale,
                skip_zeros=config["skip_zeros"],
            )

# %% ../../nbs/optimizer.eightbit.ipynb 17
class SGD8bitOptimizer(EightBit1StateOptimizer):
    "A fastai-compatible bitsandbytes 8-bit SGD optimizer"
    def __init__(
        self,
        params,
        lr,
        mom,
        wd=0,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,
        sync_each_step=False
    ):
        if mom == 0:
            raise NotImplementedError(f"8-bit SGD without momentum {mom=} is not supported")
        super().__init__("momentum", params, lr, mom, 0.0, 0.0, wd, 8, args,
                         min_8bit_size, percentile_clipping, block_wise,
                         sync_each_step=sync_each_step)

# %% ../../nbs/optimizer.eightbit.ipynb 18
class RMSProp8bitOptimizer(EightBit1StateOptimizer):
    "A fastai-compatible bitsandbytes 8-bit RMSProb optimizer"
    def __init__(
        self,
        params,
        lr=1e-2,
        sqr_mom=0.99,
        eps=1e-8,
        wd=0,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,
        sync_each_step=False
    ):
        if sqr_mom == 0:
            raise NotImplementedError(f"8-bit RMSProp with {sqr_mom=} is not supported")
        super().__init__("rmsprop", params, lr, sqr_mom, 0, eps, wd, 8, args,
                         min_8bit_size, percentile_clipping, block_wise,
                         sync_each_step=sync_each_step)

# %% ../../nbs/optimizer.eightbit.ipynb 19
class AdamW8bitOptimizer(EightBit2StateOptimizer):
    "A fastai-compatible bitsandbytes 8-bit AdamW optimizer"
    def __init__(self,
        params,
        lr=1e-3,
        mom=0.9,
        sqr_mom=0.99,
        eps=1e-8,
        wd=1e-2,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,
        is_paged=False,
        sync_each_step=False
    ):
        super().__init__("adam", params, lr, mom, sqr_mom, eps, wd, 8, args,
                         min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged,
                         sync_each_step=sync_each_step)

# %% ../../nbs/optimizer.eightbit.ipynb 20
class LARS8bitOptimizer(EightBit1StateOptimizer):
    "A fastai-compatible bitsandbytes 8-bit LARS optimizer"
    def __init__(
        self,
        params,
        lr,
        mom=0,
        wd=0,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        trust_coeff=0.02,
        sync_each_step=False
    ):
        if mom == 0:
            raise NotImplementedError(f"8-bit LARS without momentum {mom=} is not supported")
        super().__init__("lars", params, lr, mom, 0.0, 0.0, wd, 8, args,
                         min_8bit_size, percentile_clipping, max_unorm=trust_coeff, block_wise=False,
                         sync_each_step=sync_each_step)

# %% ../../nbs/optimizer.eightbit.ipynb 21
class LAMB8bitOptimizer(EightBit2StateOptimizer):
    "A fastai-compatible bitsandbytes 8-bit LAMB optimizer"
    def __init__(
        self,
        params,
        lr=1e-3,
        mom=0.9,
        sqr_mom=0.999,
        eps=1e-8,
        wd=0,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=False,
        sync_each_step=False
    ):
        super().__init__("lamb", params, lr, mom, sqr_mom, eps, wd, 8, args,
                         min_8bit_size, percentile_clipping, block_wise, max_unorm=1.0,
                         sync_each_step=sync_each_step)

# %% ../../nbs/optimizer.eightbit.ipynb 22
class Lion8bitOptimizer(EightBit1StateOptimizer):
    "A fastai-compatible bitsandbytes 8-bit Lion optimizer"
    def __init__(self,
        params,
        lr=1e-4,
        beta1=0.9,
        beta2=0.99,
        wd=0,
        args=None,
        min_8bit_size=4096,
        percentile_clipping=100,
        block_wise=True,
        is_paged=False,
        sync_each_step=False
    ):
        super().__init__("lion", params, lr, beta1, beta2, 0., wd, 8, args,
                         min_8bit_size, percentile_clipping, block_wise, is_paged=is_paged,
                         sync_each_step=sync_each_step)
