codingutf8 2018 the google ai language team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch optimization for bert model import math import warnings from functools import partial from typing import callable iterable optional tuple union import torch from torch import nn from torch optim import optimizer from torch optim lrscheduler import lambdalr reducelronplateau from trainerutils import schedulertype from utils import logging from utils versions import requireversion logger logging getloggername def getconstantlambdanone return 1 def getconstantscheduleoptimizer optimizer lastepoch int 1 return lambdalroptimizer getconstantlambda lastepochlastepoch def getreduceonplateauscheduleoptimizer optimizer return reducelronplateauoptimizer def getconstantschedulewithwarmuplrlambdacurrentstep int numwarmupsteps int if currentstep numwarmupsteps return floatcurrentstep floatmax1 0 numwarmupsteps return 1 0 def getconstantschedulewithwarmupoptimizer optimizer numwarmupsteps int lastepoch int 1 lrlambda partialgetconstantschedulewithwarmuplrlambda numwarmupstepsnumwarmupsteps return lambdalroptimizer lrlambda lastepochlastepoch def getlinearschedulewithwarmuplrlambdacurrentstep int numwarmupsteps int numtrainingsteps int if currentstep numwarmupsteps return floatcurrentstep floatmax1 numwarmupsteps return max0 0 floatnumtrainingsteps currentstep floatmax1 numtrainingsteps numwarmupsteps def getlinearschedulewithwarmupoptimizer numwarmupsteps numtrainingsteps lastepoch1 lrlambda partial getlinearschedulewithwarmuplrlambda numwarmupstepsnumwarmupsteps numtrainingstepsnumtrainingsteps return lambdalroptimizer lrlambda lastepoch def getcosineschedulewithwarmuplrlambda currentstep int numwarmupsteps int numtrainingsteps int numcycles float if currentstep numwarmupsteps return floatcurrentstep floatmax1 numwarmupsteps progress floatcurrentstep numwarmupsteps floatmax1 numtrainingsteps numwarmupsteps return max0 0 0 5 1 0 math cosmath pi floatnumcycles 2 0 progress def getcosineschedulewithwarmup optimizer optimizer numwarmupsteps int numtrainingsteps int numcycles float 0 5 lastepoch int 1 lrlambda partial getcosineschedulewithwarmuplrlambda numwarmupstepsnumwarmupsteps numtrainingstepsnumtrainingsteps numcyclesnumcycles return lambdalroptimizer lrlambda lastepoch def getcosinewithhardrestartsschedulewithwarmuplrlambda currentstep int numwarmupsteps int numtrainingsteps int numcycles int if currentstep numwarmupsteps return floatcurrentstep floatmax1 numwarmupsteps progress floatcurrentstep numwarmupsteps floatmax1 numtrainingsteps numwarmupsteps if progress 1 0 return 0 0 return max0 0 0 5 1 0 math cosmath pi floatnumcycles progress 1 0 def getcosinewithhardrestartsschedulewithwarmup optimizer optimizer numwarmupsteps int numtrainingsteps int numcycles int 1 lastepoch int 1 lrlambda partial getcosinewithhardrestartsschedulewithwarmuplrlambda numwarmupstepsnumwarmupsteps numtrainingstepsnumtrainingsteps numcyclesnumcycles return lambdalroptimizer lrlambda lastepoch def getpolynomialdecayschedulewithwarmuplrlambda currentstep int numwarmupsteps int numtrainingsteps int lrend float power float lrinit int if currentstep numwarmupsteps return floatcurrentstep floatmax1 numwarmupsteps elif currentstep numtrainingsteps return lrend lrinit as lambdalr multiplies by lrinit else lrrange lrinit lrend decaysteps numtrainingsteps numwarmupsteps pctremaining 1 currentstep numwarmupsteps decaysteps decay lrrange pctremainingpower lrend return decay lrinit as lambdalr multiplies by lrinit def getpolynomialdecayschedulewithwarmup optimizer numwarmupsteps numtrainingsteps lrend1e7 power1 0 lastepoch1 lrinit optimizer defaultslr if not lrinit lrend raise valueerrorflrend lrend must be be smaller than initial lr lrinit lrlambda partial getpolynomialdecayschedulewithwarmuplrlambda numwarmupstepsnumwarmupsteps numtrainingstepsnumtrainingsteps lrendlrend powerpower lrinitlrinit return lambdalroptimizer lrlambda lastepoch def getinversesqrtschedulelrlambdacurrentstep int numwarmupsteps int timescale int none if currentstep numwarmupsteps return floatcurrentstep floatmax1 numwarmupsteps shift timescale numwarmupsteps decay 1 0 math sqrtcurrentstep shift timescale return decay def getinversesqrtschedule optimizer optimizer numwarmupsteps int timescale int none lastepoch int 1 note this implementation is adapted from https github comgoogleresearchbigvisionblobf071ce68852d56099437004fd70057597a95f6efbigvisionutils pyl930 if timescale is none timescale numwarmupsteps lrlambda partialgetinversesqrtschedulelrlambda numwarmupstepsnumwarmupsteps timescaletimescale return lambdalroptimizer lrlambda lastepochlastepoch typetoschedulerfunction schedulertype linear getlinearschedulewithwarmup schedulertype cosine getcosineschedulewithwarmup schedulertype cosinewithrestarts getcosinewithhardrestartsschedulewithwarmup schedulertype polynomial getpolynomialdecayschedulewithwarmup schedulertype constant getconstantschedule schedulertype constantwithwarmup getconstantschedulewithwarmup schedulertype inversesqrt getinversesqrtschedule schedulertype reduceonplateau getreduceonplateauschedule def getscheduler name unionstr schedulertype optimizer optimizer numwarmupsteps optionalint none numtrainingsteps optionalint none schedulerspecifickwargs optionaldict none name schedulertypename schedulefunc typetoschedulerfunctionname if name schedulertype constant or name schedulertype reduceonplateau return schedulefuncoptimizer all other schedulers require numwarmupsteps if numwarmupsteps is none raise valueerrorfname requires numwarmupsteps please provide that argument if name schedulertype constantwithwarmup return schedulefuncoptimizer numwarmupstepsnumwarmupsteps if name schedulertype inversesqrt return schedulefuncoptimizer numwarmupstepsnumwarmupsteps all other schedulers require numtrainingsteps if numtrainingsteps is none raise valueerrorfname requires numtrainingsteps please provide that argument if schedulerspecifickwargs is none schedulerspecifickwargs return schedulefunc optimizer numwarmupstepsnumwarmupsteps numtrainingstepsnumtrainingsteps schedulerspecifickwargs class adamwoptimizer def init self params iterablenn parameter parameter lr float 1e3 betas tuplefloat float 0 9 0 999 eps float 1e6 weightdecay float 0 0 correctbias bool true nodeprecationwarning bool false if not nodeprecationwarning warnings warn this implementation of adamw is deprecated and will be removed in a future version use the pytorch implementation torch optim adamw instead or set nodeprecationwarningtrue to disable this warning futurewarning requireversiontorch1 5 0 add with alpha if lr 0 0 raise valueerrorfinvalid learning rate lr should be 0 0 if not 0 0 betas0 1 0 raise valueerrorfinvalid beta parameter betas0 should be in 0 0 1 0 if not 0 0 betas1 1 0 raise valueerrorfinvalid beta parameter betas1 should be in 0 0 1 0 if not 0 0 eps raise valueerrorfinvalid epsilon value eps should be 0 0 defaults lr lr betas betas eps eps weightdecay weightdecay correctbias correctbias super initparams defaults torch nograd def stepself closure callable none loss none if closure is not none loss closure for group in self paramgroups for p in groupparams if p grad is none continue grad p grad if grad issparse raise runtimeerroradam does not support sparse gradients please consider sparseadam instead state self statep state initialization if lenstate 0 statestep 0 exponential moving average of gradient values stateexpavg torch zeroslikep exponential moving average of squared gradient values stateexpavgsq torch zeroslikep expavg expavgsq stateexpavg stateexpavgsq beta1 beta2 groupbetas statestep 1 decay the first and second moment running average coefficient inplace operations to update the averages at the same time expavg mulbeta1 addgrad alpha1 0 beta1 expavgsq mulbeta2 addcmulgrad grad value1 0 beta2 denom expavgsq sqrt addgroupeps stepsize grouplr if groupcorrectbias no bias correction for bert biascorrection1 1 0 beta1 statestep biascorrection2 1 0 beta2 statestep stepsize stepsize math sqrtbiascorrection2 biascorrection1 p addcdivexpavg denom valuestepsize just adding the square of the weights to the loss function is not the correct way of using l2 regularizationweight decay with adam since that will interact with the m and v parameters in strange ways instead we want to decay the weights in a manner that doesn t interact with the mv parameters this is equivalent to adding the square of the weights to the loss with plain nonmomentum sgd add weight decay at the end fixed version if groupweightdecay 0 0 p addp alphagrouplr groupweightdecay return loss class adafactoroptimizer replace adamw with adafactor copy from fairseq s adafactor implementation https github comhuggingfacetransformersblob8395f14de6068012787d83989c3627c3df6a252bsrctransformersoptimization pyl505 performs a single optimization step arguments closure callable optional a closure that reevaluates the model and returns the loss state initialization exponential moving average of gradient values approximation of exponential moving average of square of gradient since optimization adafactor performs its own scheduling if the training loop relies on a scheduler e g for logging this class creates a proxy object that retrieves the current lr values from the optimizer it returns initiallr during startup and the actual lr during stepping get a proxy schedule for optimization adafactor args optimizer torch optim optimizer the optimizer for which to schedule the learning rate initiallr float optional defaults to 0 0 initial lr return optimization adafactor proxy schedule object coding utf 8 2018 the google ai language team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch optimization for bert model create a schedule with a constant learning rate using the learning rate set in optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate last_epoch int optional defaults to 1 the index of the last epoch when resuming training return torch optim lr_scheduler lambdalr with the appropriate schedule create a schedule with a constant learning rate that decreases when a metric has stopped improving args optimizer torch optim optimizer the optimizer for which to schedule the learning rate return torch optim lr_scheduler reducelronplateau with the appropriate schedule create a schedule with a constant learning rate preceded by a warmup period during which the learning rate increases linearly between 0 and the initial lr set in the optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate num_warmup_steps int the number of steps for the warmup phase last_epoch int optional defaults to 1 the index of the last epoch when resuming training return torch optim lr_scheduler lambdalr with the appropriate schedule create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0 after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate num_warmup_steps int the number of steps for the warmup phase num_training_steps int the total number of training steps last_epoch int optional defaults to 1 the index of the last epoch when resuming training return torch optim lr_scheduler lambdalr with the appropriate schedule create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to 0 after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate num_warmup_steps int the number of steps for the warmup phase num_training_steps int the total number of training steps num_cycles float optional defaults to 0 5 the number of waves in the cosine schedule the defaults is to just decrease from the max value to 0 following a half cosine last_epoch int optional defaults to 1 the index of the last epoch when resuming training return torch optim lr_scheduler lambdalr with the appropriate schedule create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to 0 with several hard restarts after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate num_warmup_steps int the number of steps for the warmup phase num_training_steps int the total number of training steps num_cycles int optional defaults to 1 the number of hard restarts to use last_epoch int optional defaults to 1 the index of the last epoch when resuming training return torch optim lr_scheduler lambdalr with the appropriate schedule as lambdalr multiplies by lr_init as lambdalr multiplies by lr_init create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the optimizer to end lr defined by lr_end after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate num_warmup_steps int the number of steps for the warmup phase num_training_steps int the total number of training steps lr_end float optional defaults to 1e 7 the end lr power float optional defaults to 1 0 power factor last_epoch int optional defaults to 1 the index of the last epoch when resuming training note power defaults to 1 0 as in the fairseq implementation which in turn is based on the original bert implementation at https github com google research bert blob f39e881b169b9d53bea03d2d341b31707a6c052b optimization py l37 return torch optim lr_scheduler lambdalr with the appropriate schedule create a schedule with an inverse square root learning rate from the initial lr set in the optimizer after a warmup period which increases lr linearly from 0 to the initial lr set in the optimizer args optimizer torch optim optimizer the optimizer for which to schedule the learning rate num_warmup_steps int the number of steps for the warmup phase timescale int optional defaults to num_warmup_steps time scale last_epoch int optional defaults to 1 the index of the last epoch when resuming training return torch optim lr_scheduler lambdalr with the appropriate schedule note this implementation is adapted from https github com google research big_vision blob f071ce68852d56099437004fd70057597a95f6ef big_vision utils py l930 unified api to get any scheduler from its name args name str or schedulertype the name of the scheduler to use optimizer torch optim optimizer the optimizer that will be used during training num_warmup_steps int optional the number of warmup steps to do this is not required by all schedulers hence the argument being optional the function will raise an error if it s unset and the scheduler type requires it num_training_steps int optional the number of training steps to do this is not required by all schedulers hence the argument being optional the function will raise an error if it s unset and the scheduler type requires it scheduler_specific_kwargs dict optional extra parameters for schedulers such as cosine with restarts mismatched scheduler types and scheduler parameters will cause the scheduler function to raise a typeerror all other schedulers require num_warmup_steps all other schedulers require num_training_steps implements adam algorithm with weight decay fix as introduced in decoupled weight decay regularization https arxiv org abs 1711 05101 parameters params iterable nn parameter parameter iterable of parameters to optimize or dictionaries defining parameter groups lr float optional defaults to 0 001 the learning rate to use betas tuple float float optional defaults to 0 9 0 999 adam s betas parameters b1 b2 eps float optional defaults to 1e 06 adam s epsilon for numerical stability weight_decay float optional defaults to 0 0 decoupled weight decay to apply correct_bias bool optional defaults to true whether or not to correct bias in adam for instance in bert tf repository they use false no_deprecation_warning bool optional defaults to false a flag used to disable the deprecation warning set to true to disable the warning add_ with alpha performs a single optimization step arguments closure callable optional a closure that reevaluates the model and returns the loss state initialization exponential moving average of gradient values exponential moving average of squared gradient values decay the first and second moment running average coefficient in place operations to update the averages at the same time no bias correction for bert just adding the square of the weights to the loss function is not the correct way of using l2 regularization weight decay with adam since that will interact with the m and v parameters in strange ways instead we want to decay the weights in a manner that doesn t interact with the m v parameters this is equivalent to adding the square of the weights to the loss with plain non momentum sgd add weight decay at the end fixed version adafactor pytorch implementation can be used as a drop in replacement for adam original fairseq code https github com pytorch fairseq blob master fairseq optim adafactor py paper adafactor adaptive learning rates with sublinear memory cost https arxiv org abs 1804 04235 note that this optimizer internally adjusts the learning rate depending on the scale_parameter relative_step and warmup_init options to use a manual external learning rate schedule you should set scale_parameter false and relative_step false arguments params iterable nn parameter parameter iterable of parameters to optimize or dictionaries defining parameter groups lr float optional the external learning rate eps tuple float float optional defaults to 1e 30 0 001 regularization constants for square gradient and parameter scale respectively clip_threshold float optional defaults to 1 0 threshold of root mean square of final gradient update decay_rate float optional defaults to 0 8 coefficient used to compute running averages of square beta1 float optional coefficient used for computing running averages of gradient weight_decay float optional defaults to 0 0 weight decay l2 penalty scale_parameter bool optional defaults to true if true learning rate is scaled by root mean square relative_step bool optional defaults to true if true time dependent learning rate is computed instead of external learning rate warmup_init bool optional defaults to false time dependent learning rate computation depends on whether warm up initialization is being used this implementation handles low precision fp16 bfloat values but we have not thoroughly tested recommended t5 finetuning settings https discuss huggingface co t t5 finetuning tips 684 3 training without lr warmup or clip_threshold is not recommended use scheduled lr warm up to fixed lr use clip_threshold 1 0 https arxiv org abs 1804 04235 disable relative updates use scale_parameter false additional optimizer operations like gradient clipping should not be used alongside adafactor example python adafactor model parameters scale_parameter false relative_step false warmup_init false lr 1e 3 others reported the following combination to work well python adafactor model parameters scale_parameter true relative_step true warmup_init true lr none when using lr none with trainer you will most likely need to use optimization adafactorschedule scheduler as following python from transformers optimization import adafactor adafactorschedule optimizer adafactor model parameters scale_parameter true relative_step true warmup_init true lr none lr_scheduler adafactorschedule optimizer trainer trainer optimizers optimizer lr_scheduler usage python replace adamw with adafactor optimizer adafactor model parameters lr 1e 3 eps 1e 30 1e 3 clip_threshold 1 0 decay_rate 0 8 beta1 none weight_decay 0 0 relative_step false scale_parameter false warmup_init false add_ with alpha copy from fairseq s adafactor implementation https github com huggingface transformers blob 8395f14de6068012787d83989c3627c3df6a252b src transformers optimization py l505 performs a single optimization step arguments closure callable optional a closure that reevaluates the model and returns the loss state initialization exponential moving average of gradient values approximation of exponential moving average of square of gradient since optimization adafactor performs its own scheduling if the training loop relies on a scheduler e g for logging this class creates a proxy object that retrieves the current lr values from the optimizer it returns initial_lr during startup and the actual lr during stepping if called before stepping get a proxy schedule for optimization adafactor args optimizer torch optim optimizer the optimizer for which to schedule the learning rate initial_lr float optional defaults to 0 0 initial lr return optimization adafactor proxy schedule object
import math
import warnings
from functools import partial
from typing import Callable, Iterable, Optional, Tuple, Union

import torch
from torch import nn
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau

from .trainer_utils import SchedulerType
from .utils import logging
from .utils.versions import require_version


logger = logging.get_logger(__name__)


def _get_constant_lambda(_=None):
    return 1


def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
    

    return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch)


def get_reduce_on_plateau_schedule(optimizer: Optimizer):
    

    return ReduceLROnPlateau(optimizer)


def _get_constant_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1.0, num_warmup_steps))
    return 1.0


def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
    

    lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps)
    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)


def _get_linear_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))


def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    

    lr_lambda = partial(
        _get_linear_schedule_with_warmup_lr_lambda,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
    )
    return LambdaLR(optimizer, lr_lambda, last_epoch)


def _get_cosine_schedule_with_warmup_lr_lambda(
    current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float
):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))


def get_cosine_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
):
    

    lr_lambda = partial(
        _get_cosine_schedule_with_warmup_lr_lambda,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
        num_cycles=num_cycles,
    )
    return LambdaLR(optimizer, lr_lambda, last_epoch)


def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda(
    current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int
):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
    if progress >= 1.0:
        return 0.0
    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))


def get_cosine_with_hard_restarts_schedule_with_warmup(
    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
):
    

    lr_lambda = partial(
        _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
        num_cycles=num_cycles,
    )
    return LambdaLR(optimizer, lr_lambda, last_epoch)


def _get_polynomial_decay_schedule_with_warmup_lr_lambda(
    current_step: int,
    *,
    num_warmup_steps: int,
    num_training_steps: int,
    lr_end: float,
    power: float,
    lr_init: int,
):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    elif current_step > num_training_steps:
        return lr_end / lr_init  
    else:
        lr_range = lr_init - lr_end
        decay_steps = num_training_steps - num_warmup_steps
        pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
        decay = lr_range * pct_remaining**power + lr_end
        return decay / lr_init  


def get_polynomial_decay_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
):
    

    lr_init = optimizer.defaults["lr"]
    if not (lr_init > lr_end):
        raise ValueError(f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})")

    lr_lambda = partial(
        _get_polynomial_decay_schedule_with_warmup_lr_lambda,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
        lr_end=lr_end,
        power=power,
        lr_init=lr_init,
    )
    return LambdaLR(optimizer, lr_lambda, last_epoch)


def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: int = None):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    shift = timescale - num_warmup_steps
    decay = 1.0 / math.sqrt((current_step + shift) / timescale)
    return decay


def get_inverse_sqrt_schedule(
    optimizer: Optimizer, num_warmup_steps: int, timescale: int = None, last_epoch: int = -1
):
    
    
    

    if timescale is None:
        timescale = num_warmup_steps

    lr_lambda = partial(_get_inverse_sqrt_schedule_lr_lambda, num_warmup_steps=num_warmup_steps, timescale=timescale)
    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)


TYPE_TO_SCHEDULER_FUNCTION = {
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.INVERSE_SQRT: get_inverse_sqrt_schedule,
    SchedulerType.REDUCE_ON_PLATEAU: get_reduce_on_plateau_schedule,
}


def get_scheduler(
    name: Union[str, SchedulerType],
    optimizer: Optimizer,
    num_warmup_steps: Optional[int] = None,
    num_training_steps: Optional[int] = None,
    scheduler_specific_kwargs: Optional[dict] = None,
):
    
    name = SchedulerType(name)
    schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
    if name == SchedulerType.CONSTANT or name == SchedulerType.REDUCE_ON_PLATEAU:
        return schedule_func(optimizer)

    
    if num_warmup_steps is None:
        raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")

    if name == SchedulerType.CONSTANT_WITH_WARMUP:
        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)

    if name == SchedulerType.INVERSE_SQRT:
        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)

    
    if num_training_steps is None:
        raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")

    if scheduler_specific_kwargs is None:
        scheduler_specific_kwargs = {}

    return schedule_func(
        optimizer,
        num_warmup_steps=num_warmup_steps,
        num_training_steps=num_training_steps,
        **scheduler_specific_kwargs,
    )


class AdamW(Optimizer):
    

    def __init__(
        self,
        params: Iterable[nn.parameter.Parameter],
        lr: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-6,
        weight_decay: float = 0.0,
        correct_bias: bool = True,
        no_deprecation_warning: bool = False,
    ):
        if not no_deprecation_warning:
            warnings.warn(
                "This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch"
                " implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this"
                " warning",
                FutureWarning,
            )
        require_version("torch>=1.5.0")  
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")
        if not 0.0 <= eps:
            raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0")
        defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "correct_bias": correct_bias}
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure: Callable = None):
        
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead")

                state = self.state[p]

                
                if len(state) == 0:
                    state["step"] = 0
                    
                    state["exp_avg"] = torch.zeros_like(p)
                    
                    state["exp_avg_sq"] = torch.zeros_like(p)

                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                beta1, beta2 = group["betas"]

                state["step"] += 1

                
                
                exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                denom = exp_avg_sq.sqrt().add_(group["eps"])

                step_size = group["lr"]
                if group["correct_bias"]:  
                    bias_correction1 = 1.0 - beta1 ** state["step"]
                    bias_correction2 = 1.0 - beta2 ** state["step"]
                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1

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

                
                
                
                
                
                
                
                
                if group["weight_decay"] > 0.0:
                    p.add_(p, alpha=(-group["lr"] * group["weight_decay"]))

        return loss


class Adafactor(Optimizer):
    

    def __init__(
        self,
        params,
        lr=None,
        eps=(1e-30, 1e-3),
        clip_threshold=1.0,
        decay_rate=-0.8,
        beta1=None,
        weight_decay=0.0,
        scale_parameter=True,
        relative_step=True,
        warmup_init=False,
    ):
        require_version("torch>=1.5.0")  
        if lr is not None and relative_step:
            raise ValueError("Cannot combine manual `lr` and `relative_step=True` options")
        if warmup_init and not relative_step:
            raise ValueError("`warmup_init=True` requires `relative_step=True`")

        defaults = {
            "lr": lr,
            "eps": eps,
            "clip_threshold": clip_threshold,
            "decay_rate": decay_rate,
            "beta1": beta1,
            "weight_decay": weight_decay,
            "scale_parameter": scale_parameter,
            "relative_step": relative_step,
            "warmup_init": warmup_init,
        }
        super().__init__(params, defaults)

    @staticmethod
    def _get_lr(param_group, param_state):
        rel_step_sz = param_group["lr"]
        if param_group["relative_step"]:
            min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2
            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))
        param_scale = 1.0
        if param_group["scale_parameter"]:
            param_scale = max(param_group["eps"][1], param_state["RMS"])
        return param_scale * rel_step_sz

    @staticmethod
    def _get_options(param_group, param_shape):
        factored = len(param_shape) >= 2
        use_first_moment = param_group["beta1"] is not None
        return factored, use_first_moment

    @staticmethod
    def _rms(tensor):
        return tensor.norm(2) / (tensor.numel() ** 0.5)

    @staticmethod
    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
        
        
        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
        return torch.mul(r_factor, c_factor)

    @torch.no_grad()
    def step(self, closure=None):
        
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.dtype in {torch.float16, torch.bfloat16}:
                    grad = grad.float()
                if grad.is_sparse:
                    raise RuntimeError("Adafactor does not support sparse gradients.")

                state = self.state[p]
                grad_shape = grad.shape

                factored, use_first_moment = self._get_options(group, grad_shape)
                
                if len(state) == 0:
                    state["step"] = 0

                    if use_first_moment:
                        
                        state["exp_avg"] = torch.zeros_like(grad)
                    if factored:
                        state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
                        state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
                    else:
                        state["exp_avg_sq"] = torch.zeros_like(grad)

                    state["RMS"] = 0
                else:
                    if use_first_moment:
                        state["exp_avg"] = state["exp_avg"].to(grad)
                    if factored:
                        state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
                        state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
                    else:
                        state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)

                p_data_fp32 = p
                if p.dtype in {torch.float16, torch.bfloat16}:
                    p_data_fp32 = p_data_fp32.float()

                state["step"] += 1
                state["RMS"] = self._rms(p_data_fp32)
                lr = self._get_lr(group, state)

                beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
                update = (grad**2) + group["eps"][0]
                if factored:
                    exp_avg_sq_row = state["exp_avg_sq_row"]
                    exp_avg_sq_col = state["exp_avg_sq_col"]

                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))

                    
                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
                    update.mul_(grad)
                else:
                    exp_avg_sq = state["exp_avg_sq"]

                    exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
                    update = exp_avg_sq.rsqrt().mul_(grad)

                update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
                update.mul_(lr)

                if use_first_moment:
                    exp_avg = state["exp_avg"]
                    exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"]))
                    update = exp_avg

                if group["weight_decay"] != 0:
                    p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr))

                p_data_fp32.add_(-update)

                if p.dtype in {torch.float16, torch.bfloat16}:
                    p.copy_(p_data_fp32)

        return loss


class AdafactorSchedule(LambdaLR):
    

    def __init__(self, optimizer, initial_lr=0.0):
        def lr_lambda(_):
            return initial_lr

        for group in optimizer.param_groups:
            group["initial_lr"] = initial_lr
        super().__init__(optimizer, lr_lambda)
        for group in optimizer.param_groups:
            del group["initial_lr"]

    def get_lr(self):
        opt = self.optimizer
        lrs = [
            opt._get_lr(group, opt.state[group["params"][0]])
            for group in opt.param_groups
            if group["params"][0].grad is not None
        ]
        if len(lrs) == 0:
            lrs = self.base_lrs  
        return lrs


def get_adafactor_schedule(optimizer, initial_lr=0.0):
    
    return AdafactorSchedule(optimizer, initial_lr)