cerebras.pytorch.optim


`cerebras.pytorch.optim`	Contains all Cerebras compliant Optimizer classes.

class cerebras.pytorch.optim.`Optimizer`**(params, defaults, enable_global_step=False)**

[source]# Bases: cerebras.pytorch.optim.optimizer.torch.optim.Optimizer, abc.ABC The abstract Cerebras base optimizer class. Enforces that the preinitialize method is implemented wherein the optimizer state should be initialized ahead of time Parameters:

params (Union[Iterabletorch.Tensor,Iterable[Dict[str, Any]]]) – Specifies what Tensors should be optimized.
defaults (Dict[str, Any]) – a dict containing default values of optimization options (used when a parameter group doesn’t specify them).
enable_global_step (bool) – If True, the optimizer will keep track of the global step for each parameter.

increment_global_step`(p)`

[source]# Increases the global steps by 1 and returns the current value of global step tensor in torch.float32 format.

state_dict`(_*args_, _**kwargs_)`

[source]#

load_state_dict`(_state_dict_)`

[source]#

register_zero_grad_pre_hook`(_hook_)`

[source]# Register an optimizer zero_grad pre hook which will be called before optimizer zero_grad. It should have the following signature:

hook(optimizer, args, kwargs) -> None or modified args and kwargs

The optimizer argument is the optimizer instance being used. If args and kwargs are modified by the pre-hook, then the transformed values are returned as a tuple containing the new_args and new_kwargs. Parameters: hook (Callable) – The user defined hook to be registered. Returns: a handle that can be used to remove the added hook by calling handle.remove() Return type: torch.utils.hooks.RemovableHandle

`register_zero_grad_post_hook`**(hook)**

[source]# Register an optimizer zero_grad post hook which will be called after optimizer zero_grad. It should have the following signature:

hook(optimizer, args, kwargs)

The optimizer argument is the optimizer instance being used. Parameters: hook (Callable) – The user defined hook to be registered. Returns: a handle that can be used to remove the added hook by calling handle.remove() Return type: torch.utils.hooks.RemovableHandle

`zero_grad`**(*args, _kwargs_)

[source]# Runs the optimizer zero_grad method and calls any pre and post hooks

`apply`(f)

[source]# Calls the function on self

`visit_state`(fn)

[source]# Applies a lambda to each stateful value.

`_abstract_ preinitialize`()

[source]# The optimizer state must be initialized ahead of time in order to capture the full compute graph in the first iteration. This method must be overriden to perform the state preinitialization

abstract `step`**(closure=None)**

[source]# Perform the optimizer step itself. Note, there should be no new state being created in this function. All state must be created ahead of time in preinitialize and only updated in this method.

class cerebras.pytorch.optim.Adadelta(params, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0, maximize=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer Adadelta optimizer implemented to perform the required pre-initialization of the optimizer state.

preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (Optional_[Callable]_) – A closure that reevaluates the model and returns the loss. class cerebras.pytorch.optim.Adafactor(params, lr, eps=(1e-30, 0.001), clip_threshold=1.0, decay_rate=- 0.8, beta1=None, weight_decay=0.0, scale_parameter=True, relative_step=False, warmup_init=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer Adafactor optimizer implemented to conform to execution within the constraints of the Cerebras WSE. preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters:

closure (Callable, optional) – A closure that reevaluates
loss. (the model and returns the) –

class cerebras.pytorch.optim.Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-06, maximize=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer Adagrad optimizer implemented to conform to execution within the constraints of the Cerebras WSE. Parameters:

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float_,_ optional) – learning rate (default: 1e-2)
lr_decay (float_,_ optional) – learning rate decay (default: 0)
weight_decay (float_,_ optional) – weight decay (L2 penalty) (default: 0)
eps (float_,_ optional) – term added to the denominator to improve numerical stability (default: 1e-10)
maximize (bool_,_ optional) – maximize the params based on the objective, instead of minimizing (default: False)

Adaptive Subgradient Methods for Online Learning and Stochastic Optimization: http://jmlr.org/papers/v12/duchi11a.html preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (callable_,_ optional) – A closure that reevaluates the model and returns the loss. class cerebras.pytorch.optim.Adamax(params, lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0.0, maximize=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer Adamax optimizer implemented to perform the required pre-initialization of the optimizer state. preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (Optional_[Callable]_) – A closure that reevaluates the model and returns the loss. class cerebras.pytorch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0.0, amsgrad=False)[source]# Bases: cerebras.pytorch.optim.AdamBase.AdamBase Adam specific overrides to AdamBase handle\_weight\_decay(param_groups)[source]# load\_state\_dict(_state_dict_)[source]# Loads the optimizer state. Parameters: state_dict (dict) – optimizer state. Should be an object returned from a call to state_dict. Adds checkpoint compatibility with the Adam from PyTorch class cerebras.pytorch.optim.AdamW(params, lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0.0, correct_bias=True, amsgrad=False)[source]# Bases: cerebras.pytorch.optim.AdamBase.AdamBase AdamW specific overrides to AdamBase load\_state\_dict(state_dict)[source]# Loads the optimizer state. Parameters: state_dict (dict) – optimizer state. Should be an object returned from a call to state_dict. Adds checkpoint compatibility with the AdamW from HuggingFace class cerebras.pytorch.optim.ASGD(params, lr=0.01, lambd=0.0001, alpha=0.75, t0=1000000.0, weight_decay=0, maximize=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer ASGD optimizer implemented to conform to execution within the constraints of the Cerebras WSE, including pre-initializing optimizer state. For more details, see https://dl.acm.org/citation.cfm?id=131098 preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (Callable_,_ optional) – A closure that reevaluates the model and returns the loss. class cerebras.pytorch.optim.Lamb(params, lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0, adam=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer Implements Lamb algorithm. It has been proposed in Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. Parameters:

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float_,_ optional) – learning rate (default: 1e-3)
betas (Tuple_[float,_ float_]__,_ optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
eps (float_,_ optional) – term added to the denominator to improve numerical stability (default: 1e-8)
weight_decay (float_,_ optional) – weight decay (L2 penalty) (default: 0)
adam (bool_,_ optional) – always use trust ratio = 1, which turns this into Adam. Useful for comparison purposes.

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float_,_ optional) – learning rate (default: 1e-4)
betas (Tuple_[float,_ float_]__,_ optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.99))
weight_decay (float_,_ optional) – weight decay coefficient (default: 0)

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float_,_ optional) – learning rate (default: 2e-3)
betas (Tuple_[float,_ float_]__,_ optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
eps (float_,_ optional) – term added to the denominator to improve numerical stability (default: 1e-8)
weight_decay (float_,_ optional) – weight decay (L2 penalty) (default: 0)
momentum_decay (float_,_ optional) – momentum momentum_decay (default: 4e-3)
foreach (bool_,_ optional) – whether foreach implementation of optimizer is used (default: None)

For further details regarding the algorithm refer to Incorporating Nesterov Momentum into Adam: https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (callable_,_ optional) – A closure that reevaluates the model and returns the loss. class cerebras.pytorch.optim.RAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-06, weight_decay=0.0)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer RAdam optimizer implemented to conform to execution within the constraints of the Cerebras WSE. Parameters:

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float_,_ optional) – learning rate (default: 1e-3)
betas (Tuple_[float,_ float_]__,_ optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
eps (float_,_ optional) – term added to the denominator to improve numerical stability (default: 1e-6)
weight_decay (float_,_ optional) – weight decay (L2 penalty) (default: 0)

preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (callable_,_ optional) – A closure that reevaluates the model and returns the loss. class cerebras.pytorch.optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer RMSprop optimizer implemented to perform the required pre-initialization of the optimizer state. preinitialize()[source]# Allocates tensors for the optimizer state to allow direct compilation of the model before the first step. step(closure=None)# Performs a single optimization step. Parameters: closure (callable_,_ optional) – A closure that reevaluates the model and returns the loss. **class cerebras.pytorch.optim.**Rprop(params, lr=0.001, etas=(0.5, 1.2), step_sizes=(1e-06, 50.0))[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer Rprop optimizer implemented to conform to execution within the constraints of the Cerebras WSE, including pre-initializing optimizer state Parameters:

params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
lr (float_,_ optional) – learning rate (default: 1e-3)
etas (Tuple_[float,_ float_]__,_ optional) – step size multipliers
step_size (Tuple_[float,_ float_]__,_ optional) – Tuple of min, max step size values. Step size is clamped to be between these values.

class cerebras.pytorch.optim.SGD(params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, maximize=False)

[source]# Bases: cerebras.pytorch.optim.optimizer.Optimizer SGD optimizer implemented to conform to execution within the constraints of the Cerebras WSE, including pre-initializing optimizer state Parameters:

params (Iterable_[torch.nn.Parameter]_) – Model parameters
lr (float) – The learning rate to use
momentum (float) – momentum factor
dampening (float) – dampening for momentum
weight_decay (float) – weight decay (L2 penalty)
nesterov (bool) – enables Nesterov momentum

optim helpers#

Contains all Cerebras compliant Optimizer classes.

cerebras.pytorch.optim.`configure_optimizer`**(optimizer_type, params, _kwargs_)

[source]# Configures and requires an Optimizer specified using the provided optimizer type The optimizer class’s signature is inspected and relevant parameters are extracted from the keyword arguments Parameters: optimizer_type (str) – The name of the optimizer to configure

params – The model parameters passed to the optimizer

For example,


optimizer_params = {
    "optimizer_type": "SGD",
    "lr": 0.001,
    "momentum": 0.5,
}
optimizer = cstorch.optim.configure_optimizer(
    optimizer_type=optimizer_params.pop("optimizer_type"),
    params=model.parameters(),
    **optimizer_params
)

Deprecated since version 2.3: Use configure_scheduler instead.

cerebras.pytorch.optim.configure_lr_scheduler(optimizer, learning_rate, adjust_learning_rate=None)

[source]#Configures a learning rate scheduler specified using the provided lr_scheduler typeThe learning rate scheduler’s class’s signature is inspected and relevant parameters are extracted from the keyword argumentsParameters:

optimizer – The optimizer passed to the lr_scheduler
learning_rate – learning rate schedule
adjust_learning_rate (dict) – key: layer types, val: lr scaling factor

The following list describes the possible learning_rate parameter formats:

learning_rate is a Python scalar (int or float)

In this case, configure_lr_scheduler returns an instance of ConstantLR with the provided value as the constant learning rate.

learning_rate is a dictionary

In this case, the dictionary is expected to contain the key scheduler which contains the name of the scheduler you want to configure.The rest of the parameters in the dictionary are passed in a keyword arguments to the specified schedulers init method.

learning_rate is a list of dictionaries

In this case, we assume what is being configured is a SequentialLR unless the any one of the dictionaries contains the key main_scheduler and the corresponding value is ChainedLR.In either case, each element of the list is expected to be a dictionary that follows the format as outlines in case 2.If what is being configured is indeed a SequentialLR, each dictionary entry is also expected to contain the key total_iters specifying the total number of iterations each scheduler should be applied for.

cerebras.pytorch.optim.configure\_optimizer\_params(optimizer_type, kwargs)[source]# Configures and requires an Optimizer specified using the provided optimizer type The optimizer class’s signature is inspected and relevant parameters are extracted from the keyword arguments. Parameters:

optimizer_type (str) – The name of the optimizer to configure
kwargs (dict) – Flattened optimizer params

Returns: Optimizer cls, and args for initialization cerebras.pytorch.optim.configure\_scheduler\_params(learning_rate)[source]# Get the kwargs and LR class from params Parameters: learning_rate (dict) – learning rate config Returns: LR class and args Return type: cls, kw_args cerebras.pytorch.optim.configure_scheduler(optimizer, schedulers_params)[source]# Configures a generic scheduler from scheduler params. The scheduler class’ signature is inspected and relevant parameters are extracted from the keyword arguments. Parameters:

optimizer – The optimizer passed to each scheduler.
schedulers_params (dict) – A dict of scheduler params.

scheduler_params is expected to be a dictionary with a single key corresponding to the name of a Scheduler. The value at this key is a sub-dictionary containing key-value pairs matching the arguments of the scheduler (except optimizer). Example:


LinearLR:
    initial_learning_rate: 0.01
    end_learning_rate: 0.001
    total_iters: 100

Some schedulers take other schedulers as an argument. In that case, nest the sub-scheduler dictionaries inside. For SequentialLR and SequentialWD milestones is calculated by the function and can be ignored.


SequentialLR:
    - LinearLR:
        initial_learning_rate: 0.01
        end_learning_rate: 0.001
        total_iters: 100
    - ExponentialLR:
        initial_learning_rate: 0.001
        decay_rate: 0.8
        total_iters: 100

Generic Scheduler class in `cerebras.pytorch`#

optim.scheduler.Scheduler#

class cerebras.pytorch.optim.scheduler.Scheduler(optimizer, total_iters, last_epoch=- 1, param_group_tags=None)[source]# Generic scheduler class for various optimizer params. Parameters:

optimizer – The optimizer to schedule
total_iters – Number of steps to perform the decay
last_epoch – the initial step to start at
param_group_tags – param group tags to target update for

abstract \_get\_closed_form()[source]# abstract property param\_group\_key# Key of the param group value to modify. For example, ‘lr’ or ‘weight_decay’. get()[source]# state_dict()[source]# load\_state\_dict(state_dict)[source]# increment\_last\_epoch()[source]# Increments the last epoch by 1 step(*args, _kwargs_)**[source]# Steps the scheduler and computes the latest value Only sets the last_epoch if running on CS update\_last\_value()[source]# update_groups(values)[source]# Update the optimizer groups with the latest values get\_last\_value()[source]# Return last computed value by current scheduler.

Learning Rate Schedulers in `cerebras.pytorch`#

Available learning rate schedulers in the cerebras.pytorch package


`ConstantLR`	`PolynomialLR`
`LinearLR`	`ExponentialLR`
`InverseExponentialTimeDecayLR`	`InverseSquareRootDecayLR`
`CosineDecayLR`	`SequentialLR`
`PiecewiseConstantLR`	`MultiStepLR`
`StepLR`	`CosineAnnealingLR`
`LambdaLR`	`CosineAnnealingWarmRestarts`
`MultiplicativeLR`	`ChainedScheduler`

optim.lr_scheduler.LRScheduler

class cerebras.pytorch.optim.lr_scheduler.`LRScheduler`**(*args, _kwargs_)

[source]# property param\_group\_key# get\_last\_lr()[source]# Return last computed learning rate by current scheduler. get_lr()[source]#

optim.lr_scheduler.ConstantLR#

class cerebras.pytorch.optim.lr_scheduler.ConstantLR(*args, _kwargs_)[source]# Maintains a constant learning rate for each parameter group (no decaying). Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
val – The learning_rate value to maintain
total_iters (int) – The number of steps to decay for

*property *val#

optim.lr_scheduler.PolynomialLR#

class cerebras.pytorch.optim.lr_scheduler.PolynomialLR(*args, _kwargs_)**[source]# Decays the learning rate of each parameter group using a polynomial function in the given total_iters. This class is similar to the Pytorch PolynomialLR LRS. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
end_learning_rate (float) – The final learning rate
total_iters (int) – Number of steps to perform the decay
power (float) – Exponent to apply to “x” (as in y=mx+b), which is ratio of step completion (1 for linear) Default: 1.0 (only Linear supported at the moment)
cycle (bool) – Whether to cycle

property initial_val# property end_val#

optim.lr_scheduler.LinearLR#

class cerebras.pytorch.optim.lr_scheduler.LinearLR(*args, **kwargs)[source]# Alias for Polynomial LR scheduler with a power of 1 property initial_val# property end_val#

optim.lr_scheduler.ExponentialLR#

class cerebras.pytorch.optim.lr_scheduler.ExponentialLR(*args, _kwargs_)**[source]# Decays the learning rate of each parameter group by decay_rate every step. This class is similar to the Pytorch ExponentialLR LRS. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
total_iters (int) – Number of steps to perform the decay
decay_rate (float) – The decay rate
staircase (bool) – If True decay the learning rate at discrete intervals

property initial_val#

optim.lr_scheduler.InverseExponentialTimeDecayLR#

class cerebras.pytorch.optim.lr_scheduler.InverseExponentialTimeDecayLR(*args, **kwargs)[source]# Decays the learning rate inverse-exponentially over time, as described in the Keras InverseTimeDecay class. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
step_exponent (int) – Exponential value.
total_iters (int) – Number of steps to perform the decay.
decay_rate (float) – The decay rate.
staircase (bool) – If True decay the learning rate at discrete intervals.

property initial_val#

optim.lr_scheduler.InverseSquareRootDecayLR#

class cerebras.pytorch.optim.lr_scheduler.InverseSquareRootDecayLR(_*args_, _**kwargs_)[source]# Decays the learning rate inverse-squareroot over time, as described in the following equation:

$$
\begin{align*}
l_{r_{t}} = \frac{scale}{\sqrt{\max\{t, warmup\_steps\}}}
\end{align*}
$$

Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
scale (float) – Multiplicative factor to scale the result.
warmup_steps (int) – use initial_learning_rate for the first warmup_steps.

property initial_val #

optim.lr_scheduler.CosineDecayLR#

**class cerebras.pytorch.optim.lr_scheduler.CosineDecayLR(*args, _kwargs_)[source]# Applies the cosine decay schedule as described in the Keras CosineDecay class. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
end_learning_rate (float) – The final learning rate
total_iters (int) – Number of steps to perform the decay

property initial_val# property end_val#

optim.lr_scheduler.SequentialLR#

class cerebras.pytorch.optim.lr_scheduler.SequentialLR(*args, _kwargs_)**[source]# Receives the list of schedulers that is expected to be called sequentially during optimization process and milestone points that provides exact intervals to reflect which scheduler is supposed to be called at a given step. This class is a wrapper around the Pytorch SequentialLR LRS. Parameters:

optimizer (torch.optim.Optimizer) – Wrapped optimizer
schedulers (list) – List of chained schedulers.
milestones (list) – List of integers that reflects milestone points.
last_epoch (int) – The index of last epoch. Default: -1.

optim.lr_scheduler.PiecewiseConstantLR#

class cerebras.pytorch.optim.lr_scheduler.PiecewiseConstantLR(_*args_, _**kwargs_)[source]# Adjusts the learning rate to a predefined constant at each milestone and holds this value until the next milestone. Notice that such adjustment can happen simultaneously with other changes to the learning rate from outside this scheduler. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
learning_rates (List_[float]_) – List of learning rates to maintain before/during each milestone.
milestones (List_[int]_) – List of step indices. Must be increasing.

optim.lr_scheduler.MultiStepLR#

class cerebras.pytorch.optim.lr_scheduler.MultiStepLR(*args, **kwargs)[source]# Decays the learning rate of each parameter group by gamma once the number of steps reaches one of the milestones. Notice that such decay can happen simultaneously with other changes to the learning rate from outside this scheduler. This class is similar to the Pytorch MultiStepLR LRS. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
gamma (float) – Multiplicative factor of learning rate decay.
milestones (List_[int]_) – List of step indices. Must be increasing.

property initial_val#

optim.lr_scheduler.StepLR#

class cerebras.pytorch.optim.lr_scheduler.StepLR(*args, **kwargs)[source]# Decays the learning rate of each parameter group by gamma every step_size. Notice that such decay can happen simultaneously with other changes to the learning rate from outside this scheduler. This class is similar to the Pytorch StepLR LRS. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
step_size (int) – Period of decay.
gamma (float) – Multiplicative factor of decay.

property initial_val#

optim.lr_scheduler.CosineAnnealingLR#

class cerebras.pytorch.optim.lr_scheduler.CosineAnnealingLR(*args, _kwargs_)**[source]# Set the learning rate of each parameter group using a cosine annealing schedule, where $$\eta_{\text{max}}$$ is set to the initial lr and $$\text{For } T_{\text{cur}}$$ is the number of steps since the last restart in SGDR:

$$
\begin{align*}
\text{For } T_{\text{cur}} &\neq (2k+1)T_{\text{max}}: \\
\eta_t &= \eta_{\text{min}} + \frac{1}{2}(\eta_{\text{max}} - \eta_{\text{min}})\left(1 + \cos\left(\frac{T_{\text{cur}}}{T_{\text{max}}} \pi \right)\right) \\
\text{For } T_{\text{cur}} &= (2k+1)T_{\text{max}}: \\
\eta_{t+1} &= \eta_t + \frac{1}{2}(\eta_{\text{max}} - \eta_{\text{min}})\left(1 - \cos\left(\frac{1}{T_{\text{max}}} \pi \right)\right)
\end{align*}

$$`

Notice that because the schedule is defined recursively, the learning rate can be simultaneously modified outside this scheduler by other operators. If the learning rate is set solely by this scheduler, the learning rate at each step becomes:

$$
\begin{align*}
\eta_t &= \eta_{\text{min}} + \frac{1}{2}(\eta_{\text{max}} - \eta_{\text{min}})\left(1 + \cos\left(\frac{T_{\text{cur}}}{T_{\text{max}}} \pi \right)\right) \\
\end{align*}
$$

It has been proposed in SGDR: Stochastic Gradient Descent with Warm Restarts. Note that this only implements the cosine annealing part of SGDR, and not the restarts. This class is similar to the Pytorch CosineAnnealingLR LRS. Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.
T_max (int) – Maximum number of iterations.
eta_min (float) – Minimum learning rate.

property initial_val#

optim.lr_scheduler.LambdaLR#

class cerebras.pytorch.optim.lr_scheduler.LambdaLR(*args, **kwargs)[source]# Sets the learning rate of each parameter group to the initial lr times a given function (which is specified by overriding set_value_lambda). Parameters:

optimizer (torch.optim.Optimizer) – The optimizer to schedule
initial_learning_rate (float) – The initial learning rate.

property initial_val#

optim.lr_scheduler.CosineAnnealingWarmRestarts#

class cerebras.pytorch.optim.lr_scheduler.CosineAnnealingWarmRestarts(*args, _kwargs_)**[source]# Set the learning rate of each parameter group using a cosine annealing schedule, where $$(\eta_{\max} )$$ is set to the initial lr, $$ T_{\text{cur}} $$ is the number of steps since the last restart and $$T_i {\text{ set }} $$ is the number of steps between two warm restarts in SGDR:

$$
\begin{align*}
\eta_t &= \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{T_{\text{cur}}}{T_i} \pi \right)\right) \\
\text{when } T_{\text{cur}} = T_i {\text{ set }} \eta_t &= \eta_{\min}
 \text{when } T_{\text{cur}} = 0 \text{ after restart}, {\text{ set }}\eta_t = \eta_{\max} \quad
\end{align*}
$$
It has been proposed in [SGDR: Stochastic Gradient Descent with Warm Restarts](https://arxiv.org/abs/1608.03983).

This class is similar to the [Pytorch CosineAnnealingWarmRestarts LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingWarmRestarts.html#torch.optim.lr_scheduler.CosineAnnealingWarmRestarts).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial\_learning\_rate** (_float_) – The initial learning rate.

* **T_0** (_int_) – Number of iterations for the first restart.

* **T_mult** (_int_) – A factor increases Ti after a restart. Currently T_mult must be set to 1.0

* **eta_min** (_float_) – Minimum learning rate.


**_property_** `initial_val`[#](#cerebras.pytorch.optim.lr_scheduler.CosineAnnealingWarmRestarts.initial_val "Permalink to this definition")

### optim.lr_scheduler.MultiplicativeLR[#](#optim-lr-scheduler-multiplicativelr "Permalink to this headline")

**_class_ cerebras.pytorch.optim.lr_scheduler.**`MultiplicativeLR`(_*args_, _**kwargs_)[\[source\]](../../_modules/cerebras/pytorch/optim/lr_scheduler.html#MultiplicativeLR)[#](#cerebras.pytorch.optim.lr_scheduler.MultiplicativeLR "Permalink to this definition")

Multiply the learning rate of each parameter group by the supplied coefficient.

**Parameters:** * **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial\_learning\_rate** (_float_) – The initial learning rate.

* **coefficient** (_float_) – Multiplicative factor of learning rate.


**_property_** `initial_val`[#](#cerebras.pytorch.optim.lr_scheduler.MultiplicativeLR.initial_val "Permalink to this definition")

### optim.lr_scheduler.ChainedScheduler[#](#optim-lr-scheduler-chainedscheduler "Permalink to this headline")

**_class_ cerebras.pytorch.optim.lr_scheduler.**`ChainedScheduler`**(_*args_, _**kwargs_)**[\[source\]](../../_modules/cerebras/pytorch/optim/lr_scheduler.html#ChainedScheduler)[#](#cerebras.pytorch.optim.lr_scheduler.ChainedScheduler "Permalink to this definition")

### optim.lr_scheduler.CyclicLR[#](#optim-lr-scheduler-cycliclr "Permalink to this headline")

**_class_ cerebras.pytorch.optim.lr_scheduler.**`CyclicLR`**(_*args_, _**kwargs_)**[\[source\]](../../_modules/cerebras/pytorch/optim/lr_scheduler.html#CyclicLR)[#](#cerebras.pytorch.optim.lr_scheduler.CyclicLR "Permalink to this definition")

Sets the learning rate of each parameter group according to cyclical learning rate policy (CLR). The policy cycles the learning rate between two boundaries with a constant frequency, as detailed in the paper [Cyclical Learning Rates for Training Neural Networks](https://arxiv.org/abs/1506.01186). The distance between the two boundaries can be scaled on a per-iteration or per-cycle basis.

Cyclical learning rate policy changes the learning rate after every batch. step should be called after a batch has been used for training.

This class has three built-in policies, as put forth in the paper:

* “triangular”: A basic triangular cycle without amplitude scaling.

* “triangular2”: A basic triangular cycle that scales initial amplitude by

    half each cycle.

* “exp_range”: A cycle that scales initial amplitude by

    $$ ({\text{gamma}^{\text{cycle iterations}}})$$ at each cycle iteration.


This class is similar to the [Pytorch CyclicLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CyclicLR.html#torch.optim.lr_scheduler.CyclicLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule.

* **base_lr** (_float_) – Initial learning rate which is the lower boundary in the cycle.

* **max_lr** (_float_) – Upper learning rate boundaries in the cycle.

* **step\_size\_up** (_int_) – Number of training iterations in the increasing half of a cycle.

* **step\_size\_down** (_int_) – Number of training iterations in the decreasing half of a cycle.

* **mode** (_str_) – One of `{‘triangular’, ‘triangular2’, ‘exp_range’}`.

* **gamma** (_float_) – Constant in ‘exp_range’ scaling function: gamma**(cycle iterations).

* **scale_mode** (_str_) – `{‘cycle’, ‘iterations’}` Defines whether scale_fn is evaluated on cycle number or cycle iterations.


**_property_** `base_val`[#](#cerebras.pytorch.optim.lr_scheduler.CyclicLR.base_val "Permalink to this definition")

**_property_** `max_val`[#](#cerebras.pytorch.optim.lr_scheduler.CyclicLR.max_val "Permalink to this definition")

### optim.lr_scheduler.OneCycleLR[#](#optim-lr-scheduler-onecyclelr "Permalink to this headline")

**_class_ cerebras.pytorch.optim.**`lr_scheduler.OneCycleLR`*(_*args_, _**kwargs_)*[\[source\]](../../_modules/cerebras/pytorch/optim/lr_scheduler.html#OneCycleLR)[#](#cerebras.pytorch.optim.lr_scheduler.OneCycleLR "Permalink to this definition")

Sets the learning rate of each parameter group according to the 1cycle learning rate policy. The 1cycle policy anneals the learning rate from an initial learning rate to some maximum learning rate and then from that maximum learning rate to some minimum learning rate much lower than the initial learning rate. This policy was initially described in the paper [Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://arxiv.org/abs/1708.07120).

This scheduler is not chainable.

This class is similar to the [Pytorch OneCycleLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.OneCycleLR.html#torch.optim.lr_scheduler.OneCycleLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial\_learning\_rate** (_float_) – Initial learning rate. Compared with PyTorch, this is equivalent to max\_lr / div\_factor.

* **max_lr** (_float_) – Upper learning rate boundaries in the cycle.

* **total_steps** (_int_) – The total number of steps in the cycle.

* **pct_start** (_float_) – The percentage of the cycle (in number of steps) spent increasing the learning rate.

* **final\_div\_factor** (_float_) – Determines the minimum learning rate via min\_lr = initial\_lr/final\_div\_factor.

* **three_phase** (_bool_) – If True, use a third phase of the schedule to annihilate the learning rate

* **anneal_strategy** (_str_) – Specifies the annealing strategy: “cos” for cosine annealing, “linear” for linear annealing.


**_property_** `initial_val`[#](#cerebras.pytorch.optim.lr_scheduler.OneCycleLR.initial_val "Permalink to this definition")

**_property_** `max_val`[#](#cerebras.pytorch.optim.lr_scheduler.OneCycleLR.max_val "Permalink to this definition")

## Weight Decay Schedulers in `cerebras.pytorch`[#](#weight-decay-schedulers-in-cerebras-pytorch "Permalink to this headline")

Available weight decay schedulers in the `cerebras.pytorch` package

  |     |     |
 | --- | --- |
 | [`ConstantWD`](#cerebras.pytorch.optim.weight_decay_scheduler.ConstantWD "cerebras.pytorch.optim.weight_decay_scheduler.ConstantWD") | [`PolynomialWD`](#cerebras.pytorch.optim.weight_decay_scheduler.PolynomialWD "cerebras.pytorch.optim.weight_decay_scheduler.PolynomialWD") |
 | [`LinearWD`](#cerebras.pytorch.optim.weight_decay_scheduler.LinearWD "cerebras.pytorch.optim.weight_decay_scheduler.LinearWD") | [`ExponentialWD`](#cerebras.pytorch.optim.weight_decay_scheduler.ExponentialWD "cerebras.pytorch.optim.weight_decay_scheduler.ExponentialWD") |
 | [`InverseExponentialTimeDecayWD`](#cerebras.pytorch.optim.weight_decay_scheduler.InverseExponentialTimeDecayWD "cerebras.pytorch.optim.weight_decay_scheduler.InverseExponentialTimeDecayWD") | [`InverseSquareRootDecayWD`](#cerebras.pytorch.optim.weight_decay_scheduler.InverseSquareRootDecayWD "cerebras.pytorch.optim.weight_decay_scheduler.InverseSquareRootDecayWD") |
 | [`CosineDecayWD`](#cerebras.pytorch.optim.weight_decay_scheduler.CosineDecayWD "cerebras.pytorch.optim.weight_decay_scheduler.CosineDecayWD") | [`SequentialWD`](#cerebras.pytorch.optim.weight_decay_scheduler.SequentialWD "cerebras.pytorch.optim.weight_decay_scheduler.SequentialWD") |
 | [`PiecewiseConstantWD`](#cerebras.pytorch.optim.weight_decay_scheduler.PiecewiseConstantWD "cerebras.pytorch.optim.weight_decay_scheduler.PiecewiseConstantWD") | [`MultiStepWD`](#cerebras.pytorch.optim.weight_decay_scheduler.MultiStepWD "cerebras.pytorch.optim.weight_decay_scheduler.MultiStepWD") |
 | [`StepWD`](#cerebras.pytorch.optim.weight_decay_scheduler.StepWD "cerebras.pytorch.optim.weight_decay_scheduler.StepWD") | [`CosineAnnealingWD`](#cerebras.pytorch.optim.weight_decay_scheduler.CosineAnnealingWD "cerebras.pytorch.optim.weight_decay_scheduler.CosineAnnealingWD") |
 | [`LambdaWD`](#cerebras.pytorch.optim.weight_decay_scheduler.LambdaWD "cerebras.pytorch.optim.weight_decay_scheduler.LambdaWD") | [`CosineAnnealingWarmRestartsWD`](#cerebras.pytorch.optim.weight_decay_scheduler.CosineAnnealingWarmRestartsWD "cerebras.pytorch.optim.weight_decay_scheduler.CosineAnnealingWarmRestartsWD") |
 | [`MultiplicativeWD`](#cerebras.pytorch.optim.weight_decay_scheduler.MultiplicativeWD "cerebras.pytorch.optim.weight_decay_scheduler.MultiplicativeWD") | [`ChainedWD`](#cerebras.pytorch.optim.weight_decay_scheduler.ChainedWD "cerebras.pytorch.optim.weight_decay_scheduler.ChainedWD") |

### optim.weight\_decay\_scheduler.WeightDecayScheduler[#](#optim-weight-decay-scheduler-weightdecayscheduler "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`WeightDecayScheduler`**(_optimizer_, _total_iters_, _last_epoch=- 1_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#WeightDecayScheduler)[#](#cerebras.pytorch.optim.weight_decay_scheduler.WeightDecayScheduler "Permalink to this definition")

**_property_** `param\_group\_key`[#](#cerebras.pytorch.optim.weight_decay_scheduler.WeightDecayScheduler.param_group_key "Permalink to this definition")

### optim.weight\_decay\_scheduler.ConstantWD[#](#optim-weight-decay-scheduler-constantwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`ConstantWD`**(_optimizer_, _val_, _total_iters=None_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#ConstantWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.ConstantWD "Permalink to this definition")

Maintains a constant weight decay for each parameter group (no decaying).

**Parameters:**

* **optimizer** – The optimizer to schedule

* **val** (_float_) – The weight decay value to maintain

* **total_iters** (_int_) – The number of steps to decay for


### optim.weight\_decay\_scheduler.PolynomialWD[#](#optim-weight-decay-scheduler-polynomialwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`PolynomialWD`**(_optimizer_, _initial_val_, _end_val_, _total_iters_, _power=1.0_, _cycle=False_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#PolynomialWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.PolynomialWD "Permalink to this definition")

Decays the weight decay of each parameter group using a polynomial function in the given total_iters.

This class is similar to the [Pytorch PolynomialLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.PolynomialLR.html#torch.optim.lr_scheduler.PolynomialLR).

**Parameters:**

* **optimizer** – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay

* **end_val** (_float_) – The final weight decay

* **total_iters** (_int_) – Number of steps to perform the decay

* **power** (_float_) – Exponent to apply to “x” (as in y=mx+b), which is ratio of step completion (1 for linear) Default: 1.0 (only Linear supported at the moment)

* **cycle** (_bool_) – Whether to cycle


### optim.weight\_decay\_scheduler.LinearWD[#](#optim-weight-decay-scheduler-linearwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`LinearWD`**(_optimizer_, _initial_val_, _end_val_, _total_iters_, _cycle=False_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#LinearWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.LinearWD "Permalink to this definition")

Alias for Polynomial Scheduler scheduler with a power of 1

### optim.weight\_decay\_scheduler.ExponentialWD[#](#optim-weight-decay-scheduler-exponentialwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`ExponentialWD`**(_optimizer_, _initial_val_, _total_iters_, _decay_rate_, _staircase=False_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#ExponentialWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.ExponentialWD "Permalink to this definition")

Decays the weight decay of each parameter group by decay_rate every step.

This class is similar to the [Pytorch ExponentialLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.ExponentialLR.html#torch.optim.lr_scheduler.ExponentialLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **total_iters** (_int_) – Number of steps to perform the decay

* **decay_rate** (_float_) – The decay rate

* **staircase** (_bool_) – If True decay the weight decay at discrete intervals


### optim.weight\_decay\_scheduler.InverseExponentialTimeDecayWD[#](#optim-weight-decay-scheduler-inverseexponentialtimedecaywd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.**weight\_decay\_scheduler.InverseExponentialTimeDecayWD(_optimizer_, _initial_val_, _step_exponent_, _total_iters_, _decay_rate_, _staircase=False_, _param\_group\_tags=None_)[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#InverseExponentialTimeDecayWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.InverseExponentialTimeDecayWD "Permalink to this definition")

Decays the weight decay inverse-exponentially over time, as described in the [Keras InverseTimeDecay class](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/InverseTimeDecay).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **step_exponent** (_int_) – Exponential weight decay.

* **total_iters** (_int_) – Number of steps to perform the decay.

* **decay_rate** (_float_) – The decay rate.

* **staircase** (_bool_) – If True decay the weight decay at discrete intervals.


### optim.weight\_decay\_scheduler.InverseSquareRootDecayWD[#](#optim-weight-decay-scheduler-inversesquarerootdecaywd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`InverseSquareRootDecayWD`**(_optimizer_, _initial_val=1.0_, _scale=1.0_, _warmup_steps=1.0_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#InverseSquareRootDecayWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.InverseSquareRootDecayWD "Permalink to this definition")

Decays the weight decay inverse-squareroot over time, as described in the following equation:
$$
wd_t = \frac{\text{scale}}{\sqrt{\max\{t, \text{warmup\_steps}\}}}
$$

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **scale** (_float_) – Multiplicative factor to scale the result.

* **warmup_steps** (_int_) – use initial\_val for the first warmup\_steps.


### optim.weight\_decay\_scheduler.CosineDecayWD[#](#optim-weight-decay-scheduler-cosinedecaywd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler**.`CosineDecayWD`**(_optimizer_, _initial_val_, _end_val_, _total_iters_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#CosineDecayWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.CosineDecayWD "Permalink to this definition")

Applies the cosine decay schedule as described in the [Keras CosineDecay class](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/CosineDecay).

**Parameters:**

* **optimizer** – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay

* **end_val** (_float_) – The final weight decay

* **total_iters** (_int_) – Number of steps to perform the decay


### optim.weight\_decay\_scheduler.SequentialWD[#](#optim-weight-decay-scheduler-sequentialwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`SequentialWD`**(_optimizer_, _schedulers_, _milestones_, _last_epoch=- 1_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#SequentialWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.SequentialWD "Permalink to this definition")

Receives the list of schedulers that is expected to be called sequentially during optimization process and milestone points that provides exact intervals to reflect which scheduler is supposed to be called at a given step.

This class is similar to [Pytorch SequentialLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.SequentialLR.html#torch.optim.lr_scheduler.SequentialLR).

**Parameters:**

* **optimizer** – Wrapped optimizer

* **schedulers** (_list_) – List of chained schedulers.

* **milestones** (_list_) – List of integers that reflects milestone points.

* **last_epoch** (_int_) – The index of last epoch. Default: -1.


### optim.weight\_decay\_scheduler.PiecewiseConstantWD[#](#optim-weight-decay-scheduler-piecewiseconstantwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`PiecewiseConstantWD`**(_optimizer_, _vals_, _milestones_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#PiecewiseConstantWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.PiecewiseConstantWD "Permalink to this definition")

Adjusts the weight decay to a predefined constant at each milestone and holds this value until the next milestone. Notice that such adjustment can happen simultaneously with other changes to the weight decays from outside this scheduler.

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **vals** (_List__\[__float__\]_) – List of weight decays to maintain before/during each milestone.

* **milestones** (_List__\[__int__\]_) – List of step indices. Must be increasing.


### optim.weight\_decay\_scheduler.MultiStepWD[#](#optim-weight-decay-scheduler-multistepwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`MultiStepWD`**(_optimizer_, _initial_val_, _gamma_, _milestones_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#MultiStepWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.MultiStepWD "Permalink to this definition")

Decays the weight decay of each parameter group by gamma once the number of steps reaches one of the milestones. Notice that such decay can happen simultaneously with other changes to the weight decay from outside this scheduler.

This class is similar to the [Pytorch MultiStepLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.MultiStepLR.html#torch.optim.lr_scheduler.MultiStepLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **gamma** (_float_) – Multiplicative factor of weight decay decay.

* **milestones** (_List__\[__int__\]_) – List of step indices. Must be increasing.


### optim.weight\_decay\_scheduler.StepWD[#](#optim-weight-decay-scheduler-stepwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler**`.StepWD`**(_optimizer_, _initial_val_, _step_size_, _gamma_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#StepWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.StepWD "Permalink to this definition")

Decays the weight decay of each parameter group by gamma every step_size. Notice that such decay can happen simultaneously with other changes to the weight decay from outside this scheduler.

This class is similar to the [Pytorch StepLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html#torch.optim.lr_scheduler.StepLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial val.

* **step_size** (_int_) – Period of decay.

* **gamma** (_float_) – Multiplicative factor of decay.


### optim.weight\_decay\_scheduler.CosineAnnealingWD[#](#optim-weight-decay-scheduler-cosineannealingwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.**weight\_decay\_scheduler.CosineAnnealingWD(_optimizer_, _initial_val_, _T_max_, _eta_min=0.0_, _param\_group\_tags=None_)[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#CosineAnnealingWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.CosineAnnealingWD "Permalink to this definition")

Set the weight decay of each parameter group using a cosine annealing schedule, where $$(\eta_{\max})$$ is set to the initial wd and $$T_{\text{cur}}$$ is the number of steps since the last restart in SGDR:

$$
\begin{align*}
\eta_t &= \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{T_{\text{cur}}}{T_{\max}} \pi \right)\right), \quad T_{\text{cur}} \neq (2k + 1)T_{\max} \\
\eta_{t+1} &= \eta_t + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 - \cos\left(\frac{1}{T_{\max}} \pi \right)\right), \quad T_{\text{cur}} = (2k + 1)T_{\max}
\end{align*}
$$

Notice that because the schedule is defined recursively, the weight decay can be simultaneously modified outside this scheduler by other operators. If the weight decay is set solely by this scheduler, the weight decay at each step becomes:

$$
\begin{align*}
\eta_t &= \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\frac{T_{\text{cur}}}{T_{\max}} \pi \right)\right)
\end{align*}
$$
It has been proposed in [SGDR: Stochastic Gradient Descent with Warm Restarts](https://arxiv.org/abs/1608.03983). Note that this only implements the cosine annealing part of SGDR, and not the restarts.

This class is similar to the [Pytorch CosineAnnealingLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html#torch.optim.lr_scheduler.CosineAnnealingLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **T_max** (_int_) – Maximum number of iterations.

* **eta_min** (_float_) – Minimum weight decay.


### optim.weight\_decay\_scheduler.LambdaWD[#](#optim-weight-decay-scheduler-lambdawd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`LambdaWD`**(_optimizer_, _initial_val_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#LambdaWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.LambdaWD "Permalink to this definition")

Sets the weight decay of each parameter group to the initial wd times a given function (which is specified by overriding set\_value\_lambda).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.


### optim.weight\_decay\_scheduler.CosineAnnealingWarmRestartsWD[#](#optim-weight-decay-scheduler-cosineannealingwarmrestartswd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`CosineAnnealingWarmRestartsWD`**(_optimizer_, _initial_val_, _T_0_, _T_mult=1_, _eta_min=0.0_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#CosineAnnealingWarmRestartsWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.CosineAnnealingWarmRestartsWD "Permalink to this definition")

Set the weight decay of each parameter group using a cosine annealing schedule, where $$(\eta_{\max} )$$ is set to the initial wd, $$ T_{\text{cur}} $$ is the number of steps since the last restart and $$T_i {\text{ set }} $$ is the number of steps between two warm restarts in SGDR:
$$
\begin{align*}
\eta_t = \eta_{\min} + \frac{1}{2} (\eta_{\max} - \eta_{\min}) \left( 1 + \cos\left( \frac{T_{\text{cur}}}{T_i} \pi \right) \right) \\
\text{when } T_{\text{cur}} = T_i {\text{ set }} \eta_t &= \eta_{\min}.
 \text{when } T_{\text{cur}} = 0 \text{ after restart}, {\text{ set }}\eta_t = \eta_{\max} \quad
 \end{align*}
$$

It has been proposed in [SGDR: Stochastic Gradient Descent with Warm Restarts](https://arxiv.org/abs/1608.03983).

This class is similar to the [Pytorch CosineAnnealingWarmRestarts LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingWarmRestarts.html#torch.optim.lr_scheduler.CosineAnnealingWarmRestarts).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **T_0** (_int_) – Number of iterations for the first restart.

* **T_mult** (_int_) – A factor increases Ti after a restart. Currently T_mult must be set to 1.0

* **eta_min** (_float_) – Minimum weight decay.


### optim.weight\_decay\_scheduler.MultiplicativeWD[#](#optim-weight-decay-scheduler-multiplicativewd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`MultiplicativeWD`**(_optimizer_, _initial_val_, _coefficient_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#MultiplicativeWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.MultiplicativeWD "Permalink to this definition")

Multiply the weight decay of each parameter group by the supplied coefficient.

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – The initial weight decay.

* **coefficient** (_float_) – Multiplicative factor of weight decay.


### optim.weight\_decay\_scheduler.ChainedWD[#](#optim-weight-decay-scheduler-chainedwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.** `ChainedWD`**(_schedulers_, _param\_group\_tags=None_)[\[**source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#ChainedWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.ChainedWD "Permalink to this definition")

Chains list of weight decay schedulers. It takes a list of chainable weight decay schedulers and performs consecutive step() functions belonging to them by just one call.

### optim.weight\_decay\_scheduler.CyclicWD[#](#optim-weight-decay-scheduler-cyclicwd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`CyclicWD`**(_optimizer_, _base_val_, _max_val_, _step\_size\_up=2000_, _step\_size\_down=None_, _mode='triangular'_, _gamma=1.0_, _scale_mode='cycle'_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#CyclicWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.CyclicWD "Permalink to this definition")

Sets the weight decay of each parameter group according to cyclical weight decay policy (CLR). The policy cycles the learning rate between two boundaries with a constant frequency, as detailed in the paper [Cyclical Learning Rates for Training Neural Networks](https://arxiv.org/abs/1506.01186). The distance between the two boundaries can be scaled on a per-iteration or per-cycle basis.

Cyclical weight decay policy changes the weight decay after every batch. step should be called after a batch has been used for training.

This class has three built-in policies, as put forth in the paper:

* “triangular”: A basic triangular cycle without amplitude scaling.

* “triangular2”: A basic triangular cycle that scales initial amplitude by

    half each cycle.

* “exp_range”: A cycle that scales initial amplitude by

    $$(gama ^{cycle iteration})$$ at each cycle iteration.

This class is similar to the [Pytorch CyclicLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CyclicLR.html#torch.optim.lr_scheduler.CyclicLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule.

* **base_val** (_float_) – Initial weight decay which is the lower boundary in the cycle.

* **max_val** (_float_) – Upper weight decay boundaries in the cycle.

* **step\_size\_up** (_int_) – Number of training iterations in the increasing half of a cycle.

* **step\_size\_down** (_int_) – Number of training iterations in the decreasing half of a cycle.

* **mode** (_str_) – One of `{‘triangular’, ‘triangular2’, ‘exp_range’}`.

* **gamma** (_float_) – Constant in ‘exp_range’ scaling function: gamma**(cycle iterations).

* **scale_mode** (_str_) – `{‘cycle’, ‘iterations’}` Defines whether scale_fn is evaluated on cycle number or cycle iterations.


### optim.weight\_decay\_scheduler.OneCycleWD[#](#optim-weight-decay-scheduler-onecyclewd "Permalink to this headline")

**_class_ cerebras.pytorch.optim.weight\_decay\_scheduler.**`OneCycleWD`**(_optimizer_, _initial_val_, _max_val_, _total_steps=1000_, _pct_start=0.3_, _final\_div\_factor=10000.0_, _three_phase=False_, _anneal_strategy='cos'_, _param\_group\_tags=None_)**[\[source\]](../../_modules/cerebras/pytorch/optim/weight_decay_scheduler.html#OneCycleWD)[#](#cerebras.pytorch.optim.weight_decay_scheduler.OneCycleWD "Permalink to this definition")

Sets the weight decay of each parameter group according to the 1cycle weight decay policy. The 1cycle policy anneals the learning rate from an initial weight decay to some maximum weight decay and then from that maximum weight decay to some minimum weight decay much lower than the initial weight decay. This policy was initially described in the paper [Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates](https://arxiv.org/abs/1708.07120).

This scheduler is not chainable.

This class is similar to the [Pytorch OneCycleLR LRS](https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.OneCycleLR.html#torch.optim.lr_scheduler.OneCycleLR).

**Parameters:**

* **optimizer** ([_torch.optim.Optimizer_](https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer "(in PyTorch v2.4)")) – The optimizer to schedule

* **initial_val** (_float_) – Initial weight decay. Compared with PyTorch, this is equivalent to max\_val / div\_factor.

* **max_val** (_float_) – Upper weight decay boundaries in the cycle.

* **total_steps** (_int_) – The total number of steps in the cycle.

* **pct_start** (_float_) – The percentage of the cycle (in number of steps) spent increasing the weight decay.

* **final\_div\_factor** (_float_) – Determines the minimum weight decay via min\_val = initial\_val/final\_div\_factor.

* **three_phase** (_bool_) – If True, use a third phase of the schedule to annihilate the weight decay

* **anneal_strategy** (_str_) – Specifies the annealing strategy: “cos” for cosine annealing, “linear” for linear annealing.

Updates

Legacy Documentation

cerebras.pytorch.optim

class cerebras.pytorch.optim.`Optimizer`**(params, defaults, enable_global_step=False)**

increment_global_step`(p)`

state_dict`(_*args_, _**kwargs_)`

load_state_dict`(_state_dict_)`

register_zero_grad_pre_hook`(_hook_)`

`register_zero_grad_post_hook`**(hook)**

`zero_grad`**(*args, _kwargs_)

`apply`(f)

`visit_state`(fn)

`_abstract_ preinitialize`()

abstract `step`**(closure=None)**

class cerebras.pytorch.optim.SGD(params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, maximize=False)

optim helpers#

cerebras.pytorch.optim.`configure_optimizer`**(optimizer_type, params, _kwargs_)

cerebras.pytorch.optim.configure_lr_scheduler(optimizer, learning_rate, adjust_learning_rate=None)

Generic Scheduler class in `cerebras.pytorch`#

optim.scheduler.Scheduler#

Learning Rate Schedulers in `cerebras.pytorch`#

optim.lr_scheduler.LRScheduler

class cerebras.pytorch.optim.lr_scheduler.`LRScheduler`**(*args, _kwargs_)

optim.lr_scheduler.ConstantLR#

optim.lr_scheduler.PolynomialLR#

optim.lr_scheduler.LinearLR#

optim.lr_scheduler.ExponentialLR#

optim.lr_scheduler.InverseExponentialTimeDecayLR#

optim.lr_scheduler.InverseSquareRootDecayLR#

optim.lr_scheduler.CosineDecayLR#

optim.lr_scheduler.SequentialLR#

optim.lr_scheduler.PiecewiseConstantLR#

optim.lr_scheduler.MultiStepLR#

optim.lr_scheduler.StepLR#

optim.lr_scheduler.CosineAnnealingLR#

optim.lr_scheduler.LambdaLR#

optim.lr_scheduler.CosineAnnealingWarmRestarts#

Updates

Legacy Documentation

​class cerebras.pytorch.optim.Optimizer(params, defaults, enable_global_step=False)

​increment_global_step(p)

​state_dict(_*args_, _**kwargs_)

​load_state_dict(_state_dict_)

​register_zero_grad_pre_hook(_hook_)

​register_zero_grad_post_hook(hook)

​zero_grad(*args, _kwargs_)**

​apply(f)

​visit_state(fn)

​_abstract_ preinitialize()

​abstract step(closure=None)

​

​

​class cerebras.pytorch.optim.SGD(params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, maximize=False)

​optim helpers#

​cerebras.pytorch.optim.configure_optimizer(optimizer_type, params, _kwargs_)**

​cerebras.pytorch.optim.configure_lr_scheduler(optimizer, learning_rate, adjust_learning_rate=None)

​Generic Scheduler class in cerebras.pytorch#

​optim.scheduler.Scheduler#

​Learning Rate Schedulers in cerebras.pytorch#

​optim.lr_scheduler.LRScheduler

​class cerebras.pytorch.optim.lr_scheduler.LRScheduler(*args, _kwargs_)**

​optim.lr_scheduler.ConstantLR#

​optim.lr_scheduler.PolynomialLR#

​optim.lr_scheduler.LinearLR#

​optim.lr_scheduler.ExponentialLR#

​optim.lr_scheduler.InverseExponentialTimeDecayLR#

​optim.lr_scheduler.InverseSquareRootDecayLR#

​optim.lr_scheduler.CosineDecayLR#

​optim.lr_scheduler.SequentialLR#

​optim.lr_scheduler.PiecewiseConstantLR#

​optim.lr_scheduler.MultiStepLR#

​optim.lr_scheduler.StepLR#

​optim.lr_scheduler.CosineAnnealingLR#

​optim.lr_scheduler.LambdaLR#

​optim.lr_scheduler.CosineAnnealingWarmRestarts#

class cerebras.pytorch.optim.`Optimizer`**(params, defaults, enable_global_step=False)**

increment_global_step`(p)`

state_dict`(_*args_, _**kwargs_)`

load_state_dict`(_state_dict_)`

register_zero_grad_pre_hook`(_hook_)`

`register_zero_grad_post_hook`**(hook)**

`zero_grad`**(*args, _kwargs_)

`apply`(f)

`visit_state`(fn)

`_abstract_ preinitialize`()

abstract `step`**(closure=None)**

class cerebras.pytorch.optim.SGD(params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False, maximize=False)

optim helpers#

cerebras.pytorch.optim.`configure_optimizer`**(optimizer_type, params, _kwargs_)

cerebras.pytorch.optim.configure_lr_scheduler(optimizer, learning_rate, adjust_learning_rate=None)

Generic Scheduler class in `cerebras.pytorch`#

optim.scheduler.Scheduler#

Learning Rate Schedulers in `cerebras.pytorch`#

optim.lr_scheduler.LRScheduler

class cerebras.pytorch.optim.lr_scheduler.`LRScheduler`**(*args, _kwargs_)

optim.lr_scheduler.ConstantLR#

optim.lr_scheduler.PolynomialLR#

optim.lr_scheduler.LinearLR#

optim.lr_scheduler.ExponentialLR#

optim.lr_scheduler.InverseExponentialTimeDecayLR#

optim.lr_scheduler.InverseSquareRootDecayLR#

optim.lr_scheduler.CosineDecayLR#

optim.lr_scheduler.SequentialLR#

optim.lr_scheduler.PiecewiseConstantLR#

optim.lr_scheduler.MultiStepLR#

optim.lr_scheduler.StepLR#

optim.lr_scheduler.CosineAnnealingLR#

optim.lr_scheduler.LambdaLR#

optim.lr_scheduler.CosineAnnealingWarmRestarts#