> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.optim.scheduler

Learning rate decay and weight decay incr functions.

## Module Contents

### Classes

| Name                                                                                            | Description                             |
| ----------------------------------------------------------------------------------------------- | --------------------------------------- |
| [`OptimizerParamScheduler`](#nemo_automodel-components-optim-scheduler-OptimizerParamScheduler) | Anneals learning rate and weight decay. |

### Functions

| Name                                                                              | Description                                                                                              |
| --------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------- |
| [`_as_float`](#nemo_automodel-components-optim-scheduler-_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`](#nemo_automodel-components-optim-scheduler-_set_group_mults) | Store each param group's constructed `lr` / `weight_decay` as `lr_mult` / `wd_mult` relative to group 0. |

### Data

[`_T`](#nemo_automodel-components-optim-scheduler-_T)

[`logger`](#nemo_automodel-components-optim-scheduler-logger)

### API

```python
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`** `bool` — default: True

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

---

**`override_opt_param_scheduler`** `bool` — default: False

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

---

**`wsd_decay_steps`** `int` — default: None

number of weight decay decay steps. Defaults to None.

---

**`lr_wsd_decay_style`** `str` — default: None

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

---

**`max_lr`** `= float(max_lr)`

---

**`num_steps`** `= 0`

---

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

Return a string representation of the OptimizerParamScheduler.

```python
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`** `_T`

class value

---

**`sd_value`** `_T`

checkpoint value

---

**`name`** `str`

name of the parameter

---

```python
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](https://openreview.net/pdf?id=BJYwwY9ll) pg. 4.

Argsa:
param\_group (dict): parameter group from the optimizer.

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

Weight decay incr functions.

```python
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

---

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

Return the state dict.

```python
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

---

```python
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`).

```python
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.

---

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

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