nemo_automodel.components.optim.scheduler

Learning rate decay and weight decay incr functions.

Module Contents

Classes

Name	Description
`OptimizerParamScheduler`	Anneals learning rate and weight decay.

Functions

Name	Description
`_as_float`	Return `value` as a float, or None when it is not numeric (e.g. relative-step Adafactor’s `lr=None`).
`_set_group_mults`	Store each param group’s constructed `lr` / `weight_decay` as `lr_mult` / `wd_mult` relative to group 0.

Data

_T

logger

API

class nemo_automodel.components.optim.scheduler.OptimizerParamScheduler(
    optimizer: torch.optim.optimizer.Optimizer,
    init_lr: float,
    max_lr: float,
    min_lr: float,
    lr_warmup_steps: int,
    lr_decay_steps: int,
    lr_decay_style: str,
    start_wd: float,
    end_wd: float,
    wd_incr_steps: int,
    wd_incr_style: str,
    use_checkpoint_opt_param_scheduler: bool | None = True,
    override_opt_param_scheduler: bool | None = False,
    wsd_decay_steps: int | None = None,
    lr_wsd_decay_style: str | None = None
)

Anneals learning rate and weight decay.

Parameters:

optimizer

Optimizer

the optimizer to be used

init_lr

float

initial learning rate

max_lr

float

maximum learning rate

min_lr

float

minimum learning rate

lr_warmup_steps

int

number of warmup steps

lr_decay_steps

int

number of decay steps

lr_decay_style

str

decay style for learning rate

start_wd

float

initial weight decay

end_wd

float

final weight decay

wd_incr_steps

int

number of weight decay increment steps

wd_incr_style

str

weight decay increment style

use_checkpoint_opt_param_scheduler

boolDefaults to True

whether to use the checkpoint values for the optimizer param scheduler. Defaults to True.

override_opt_param_scheduler

boolDefaults to False

whether to override the optimizer param scheduler values with the class values. Defaults to False.

wsd_decay_steps

intDefaults to None

number of weight decay decay steps. Defaults to None.

lr_wsd_decay_style

strDefaults to None

decay style for learning rate during weight decay decay steps. Defaults to None.

max_lr

= float(max_lr)

num_steps

= 0

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler.__repr__() -> str

Return a string representation of the OptimizerParamScheduler.

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler._check_and_set(
    cls_value: nemo_automodel.components.optim.scheduler._T,
    sd_value: nemo_automodel.components.optim.scheduler._T,
    name: str
) -> nemo_automodel.components.optim.scheduler._T

Auxiliary function for checking the values in the checkpoint and setting them.

Parameters:

cls_value

class value

sd_value

checkpoint value

name

str

name of the parameter

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler.get_lr(
    param_group: dict[str, typing.Any]
) -> float

Learning rate decay functions from: https://openreview.net/pdf?id=BJYwwY9ll pg. 4.

Argsa: param_group (dict): parameter group from the optimizer.

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler.get_wd() -> float

Weight decay incr functions.

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler.load_state_dict(
    state_dict: dict[str, typing.Any]
) -> None

Load the state dict.

Every restored field is applied before the final step call, so the learning rate and weight decay written into optimizer.param_groups both reflect the checkpoint’s schedule rather than the constructor’s. The step count is treated as absolute, so loading is idempotent.

Parameters:

state_dict

dict

state dict to be load

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler.state_dict() -> dict[str, typing.Any]

Return the state dict.

nemo_automodel.components.optim.scheduler.OptimizerParamScheduler.step(
    increment: int
) -> None

Set lr and weight decay for all parameter groups.

Each group gets the scheduled value times its lr_mult / wd_mult (see _set_group_mults).

Parameters:

increment

int

number of steps to increment

nemo_automodel.components.optim.scheduler._as_float(
    value: typing.Any
) -> float | None

Return value as a float, or None when it is not numeric (e.g. relative-step Adafactor’s lr=None).

nemo_automodel.components.optim.scheduler._set_group_mults(
    optimizer: torch.optim.optimizer.Optimizer
) -> None

Store each param group’s constructed lr / weight_decay as lr_mult / wd_mult relative to group 0.

OptimizerParamScheduler.step rewrites every group’s lr and weight_decay from the schedule, so a group built with its own values (a no-decay group, a separate head LR) would otherwise be reset to the base schedule. LRSchedulerConfig.build reads the schedule’s base values from group 0, so each group keeps its ratio to group 0 as a multiplier. The multipliers live in the param groups, so they are saved with the optimizer and restored on resume instead of being recomputed from the resumed run’s config.

A group that already carries lr_mult (or a per-group max_lr / min_lr) or wd_mult keeps it unchanged. Every other group, group 0 included, gets both keys so the saved param-group layout does not depend on which groups happen to differ from group 0. The multiplier is 1.0 when group 0’s value is zero or non-numeric (e.g. relative-step Adafactor’s lr=None).

Parameters:

optimizer

Optimizer

Optimizer whose param groups still hold their constructed values.

nemo_automodel.components.optim.scheduler._T = TypeVar('_T')

nemo_automodel.components.optim.scheduler.logger = logging.getLogger(__name__)