> 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

## Submodules

* **[`nemo_automodel.components.optim.dion`](/nemo-automodel/nemo_automodel/components/optim/dion)**
* **[`nemo_automodel.components.optim.optimizer`](/nemo-automodel/nemo_automodel/components/optim/optimizer)**
* **[`nemo_automodel.components.optim.precision_warnings`](/nemo-automodel/nemo_automodel/components/optim/precision_warnings)**
* **[`nemo_automodel.components.optim.scheduler`](/nemo-automodel/nemo_automodel/components/optim/scheduler)**

## Package Contents

### Classes

| Name                                                                                                  | Description                                                               |
| ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`AdamConfig`](#nemo_automodel-components-optim-optimizer-AdamConfig)                                 | `torch.optim.Adam`.                                                       |
| [`AdamWConfig`](#nemo_automodel-components-optim-optimizer-AdamWConfig)                               | `torch.optim.AdamW`.                                                      |
| [`Dion2Config`](#nemo_automodel-components-optim-optimizer-Dion2Config)                               | `dion.Dion2` — recommended successor to the legacy Dion optimizer.        |
| [`DionConfig`](#nemo_automodel-components-optim-optimizer-DionConfig)                                 | `dion.Dion` — legacy low-rank optimizer (prefer `Dion2Config`).           |
| [`FlashAdamWConfig`](#nemo_automodel-components-optim-optimizer-FlashAdamWConfig)                     | `flashoptim.FlashAdamW`.                                                  |
| [`FusedAdamConfig`](#nemo_automodel-components-optim-optimizer-FusedAdamConfig)                       | `transformer_engine.pytorch.optimizers.FusedAdam`.                        |
| [`LRSchedulerConfig`](#nemo_automodel-components-optim-optimizer-LRSchedulerConfig)                   | LR scheduler configuration.  `None` fields are computed by                |
| [`MuonConfig`](#nemo_automodel-components-optim-optimizer-MuonConfig)                                 | `dion.Muon` — matrix-aware update for 2D+ params, scalar fallback for 1D. |
| [`NorMuonConfig`](#nemo_automodel-components-optim-optimizer-NorMuonConfig)                           | `dion.NorMuon` — Muon variant with neuron-wise normalization.             |
| [`OptimizerConfig`](#nemo_automodel-components-optim-optimizer-OptimizerConfig)                       | Base optimizer config.                                                    |
| [`OptimizerFromFactoryConfig`](#nemo_automodel-components-optim-optimizer-OptimizerFromFactoryConfig) | Build an optimizer from an arbitrary factory callable plus kwargs.        |
| [`OptimizerParamScheduler`](#nemo_automodel-components-optim-scheduler-OptimizerParamScheduler)       | Anneals learning rate and weight decay.                                   |
| [`ParamGroupOverride`](#nemo_automodel-components-optim-optimizer-ParamGroupOverride)                 | Per-parameter-group learning-rate / weight-decay override.                |

### Functions

| Name                                                                                          | Description                                                              |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`build_dion_optimizer`](#nemo_automodel-components-optim-dion-build_dion_optimizer)          | Build the parameter groups and resolve the device mesh for a Dion-family |
| [`build_optimizer`](#nemo_automodel-components-optim-optimizer-build_optimizer)               | Build one optimizer per `model.parts` (or `[model]`).                    |
| [`build_optimizer_config`](#nemo_automodel-components-optim-optimizer-build_optimizer_config) | Normalize an optimizer `target` plus `kwargs` into an `OptimizerConfig`. |
| [`is_dion_optimizer`](#nemo_automodel-components-optim-dion-is_dion_optimizer)                | Return whether an optimizer factory targets a Dion-family optimizer.     |

### Data

[`OPTIMIZER_CONFIG_REGISTRY`](#nemo_automodel-components-optim-optimizer-OPTIMIZER_CONFIG_REGISTRY)

### API

```python
class nemo_automodel.components.optim.optimizer.AdamConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0001,
    weight_decay: float = 0.01,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-08,
    amsgrad: bool = False
)
```

Dataclass

**Bases:** [OptimizerConfig](#nemo_automodel-components-optim-optimizer-OptimizerConfig)

`torch.optim.Adam`.

**`amsgrad`** `bool = False`

---

**`betas`** `tuple[float, float] = (0.9, 0.999)`

---

**`eps`** `float = 1e-08`

---

**`lr`** `float = 0.0001`

---

**`weight_decay`** `float = 0.01`

---

```python
nemo_automodel.components.optim.optimizer.AdamConfig._build_optimizer(
    params,
    foreach: bool | None = None
) -> torch.optim.Optimizer
```

```python
class nemo_automodel.components.optim.optimizer.AdamWConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0001,
    weight_decay: float = 0.01,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-08,
    amsgrad: bool = False,
    fused: bool = False
)
```

Dataclass

**Bases:** [OptimizerConfig](#nemo_automodel-components-optim-optimizer-OptimizerConfig)

`torch.optim.AdamW`.

**`amsgrad`** `bool = False`

---

**`betas`** `tuple[float, float] = (0.9, 0.999)`

---

**`eps`** `float = 1e-08`

---

**`fused`** `bool = False`

---

**`lr`** `float = 0.0001`

---

**`weight_decay`** `float = 0.01`

---

```python
nemo_automodel.components.optim.optimizer.AdamWConfig._build_optimizer(
    params,
    foreach: bool | None = None
) -> torch.optim.Optimizer
```

```python
class nemo_automodel.components.optim.optimizer.Dion2Config(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0005,
    weight_decay: float = 0.0,
    scalar_opt: str = 'adamw',
    scalar_betas: tuple[float, float] = (0.9, 0.999),
    scalar_eps: float = 1e-08,
    scalar_lr: float | None = None,
    embed_lr: float | None = None,
    lm_head_lr: float | None = None,
    fraction: float = 0.25,
    ef_decay: float = 0.95,
    betas: tuple[float, float] = (0.9, 0.95),
    epsilon: float = 1e-08,
    adjust_lr: str = 'spectral_norm'
)
```

Dataclass

**Bases:** [\_DionConfigBase](#nemo_automodel-components-optim-optimizer-_DionConfigBase)

`dion.Dion2` — recommended successor to the legacy Dion optimizer.

**`adjust_lr`** `str = 'spectral_norm'`

---

**`betas`** `tuple[float, float] = (0.9, 0.95)`

---

**`ef_decay`** `float = 0.95`

---

**`epsilon`** `float = 1e-08`

---

**`fraction`** `float = 0.25`

---

```python
nemo_automodel.components.optim.optimizer.Dion2Config._make_optimizer(
    param_groups,
    ctor_kwargs
)
```

```python
class nemo_automodel.components.optim.optimizer.DionConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0005,
    weight_decay: float = 0.0,
    scalar_opt: str = 'adamw',
    scalar_betas: tuple[float, float] = (0.9, 0.999),
    scalar_eps: float = 1e-08,
    scalar_lr: float | None = None,
    embed_lr: float | None = None,
    lm_head_lr: float | None = None,
    mu: float = 0.95,
    betas: tuple[float, float] = (0.9, 0.95),
    epsilon: float = 1e-08,
    rank_fraction: float = 1.0,
    rank_multiple_of: int = 1,
    power_iters: int = 1,
    qr_method: str = 'rcqr'
)
```

Dataclass

**Bases:** [\_DionConfigBase](#nemo_automodel-components-optim-optimizer-_DionConfigBase)

`dion.Dion` — legacy low-rank optimizer (prefer `Dion2Config`).

Legacy Dion takes separate replicate/outer/inner shard meshes; for FSDP2 the
resolved 1-D shard submesh maps to `outer_shard_mesh`.

**`_mesh_kwarg`** `str = 'outer_shard_mesh'`

---

**`betas`** `tuple[float, float] = (0.9, 0.95)`

---

**`epsilon`** `float = 1e-08`

---

**`mu`** `float = 0.95`

---

**`power_iters`** `int = 1`

---

**`qr_method`** `str = 'rcqr'`

---

**`rank_fraction`** `float = 1.0`

---

**`rank_multiple_of`** `int = 1`

---

```python
nemo_automodel.components.optim.optimizer.DionConfig._make_optimizer(
    param_groups,
    ctor_kwargs
)
```

```python
class nemo_automodel.components.optim.optimizer.FlashAdamWConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0001,
    weight_decay: float = 0.01,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-08,
    master_weight_bits: int = 24
)
```

Dataclass

**Bases:** [OptimizerConfig](#nemo_automodel-components-optim-optimizer-OptimizerConfig)

`flashoptim.FlashAdamW`.

**`betas`** `tuple[float, float] = (0.9, 0.999)`

---

**`eps`** `float = 1e-08`

---

**`lr`** `float = 0.0001`

---

**`master_weight_bits`** `int = 24`

---

**`weight_decay`** `float = 0.01`

---

```python
nemo_automodel.components.optim.optimizer.FlashAdamWConfig._build_optimizer(
    params,
    foreach: bool | None = None
) -> torch.optim.Optimizer
```

```python
class nemo_automodel.components.optim.optimizer.FusedAdamConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0001,
    weight_decay: float = 0.01,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-08,
    adam_w_mode: bool = True,
    bias_correction: bool = True,
    master_weights: bool = True,
    master_weight_dtype: str | None = None
)
```

Dataclass

**Bases:** [OptimizerConfig](#nemo_automodel-components-optim-optimizer-OptimizerConfig)

`transformer_engine.pytorch.optimizers.FusedAdam`.

**`adam_w_mode`** `bool = True`

---

**`betas`** `tuple[float, float] = (0.9, 0.999)`

---

**`bias_correction`** `bool = True`

---

**`eps`** `float = 1e-08`

---

**`lr`** `float = 0.0001`

---

**`master_weight_dtype`** `str | None = None`

---

**`master_weights`** `bool = True`

---

**`weight_decay`** `float = 0.01`

---

```python
nemo_automodel.components.optim.optimizer.FusedAdamConfig._build_optimizer(
    params,
    foreach: bool | None = None
) -> torch.optim.Optimizer
```

```python
class nemo_automodel.components.optim.optimizer.LRSchedulerConfig(
    lr_warmup_steps: int | None = None,
    lr_decay_steps: int | None = None,
    lr_decay_style: str = 'cosine',
    init_lr: float | None = None,
    max_lr: float | None = None,
    min_lr: float | None = None,
    start_wd: float | None = None,
    end_wd: float | None = None,
    wd_incr_steps: int | None = None,
    wd_incr_style: str = 'constant',
    use_checkpoint_opt_param_scheduler: bool = True,
    override_opt_param_scheduler: bool = False,
    wsd_decay_steps: int | None = None,
    lr_wsd_decay_style: str | None = None
)
```

Dataclass

LR scheduler configuration.  `None` fields are computed by
`build` from the training schedule (total steps, optimizer base LR/WD).

**`end_wd`** `float | None = None`

---

**`init_lr`** `float | None = None`

---

**`lr_decay_steps`** `int | None = None`

---

**`lr_decay_style`** `str = 'cosine'`

---

**`lr_warmup_steps`** `int | None = None`

---

**`lr_wsd_decay_style`** `str | None = None`

---

**`max_lr`** `float | None = None`

---

**`min_lr`** `float | None = None`

---

**`override_opt_param_scheduler`** `bool = False`

---

**`start_wd`** `float | None = None`

---

**`use_checkpoint_opt_param_scheduler`** `bool = True`

---

**`wd_incr_steps`** `int | None = None`

---

**`wd_incr_style`** `str = 'constant'`

---

**`wsd_decay_steps`** `int | None = None`

---

```python
nemo_automodel.components.optim.optimizer.LRSchedulerConfig.build(
    optimizer: list[torch.optim.Optimizer] | torch.optim.Optimizer,
    step_scheduler: nemo_automodel.components.training.step_scheduler.StepScheduler
) -> list[nemo_automodel.components.optim.scheduler.OptimizerParamScheduler]
```

Build one LR scheduler per optimizer.

`None` fields are filled from the training schedule and each
optimizer's base LR/WD.

**Parameters:**

**`optimizer`** `list[torch.optim.Optimizer] | torch.optim.Optimizer`

The optimizer(s) to schedule.

---

**`step_scheduler`** `StepScheduler`

The step scheduler, used to derive total steps.

---

**Returns:** `list[OptimizerParamScheduler]`

class:`OptimizerParamScheduler` per optimizer.

```python
class nemo_automodel.components.optim.optimizer.MuonConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0005,
    weight_decay: float = 0.0,
    scalar_opt: str = 'adamw',
    scalar_betas: tuple[float, float] = (0.9, 0.999),
    scalar_eps: float = 1e-08,
    scalar_lr: float | None = None,
    embed_lr: float | None = None,
    lm_head_lr: float | None = None,
    mu: float = 0.95,
    betas: tuple[float, float] = (0.9, 0.95),
    epsilon: float = 1e-08,
    adjust_lr: str = 'spectral_norm',
    nesterov: bool = False,
    flatten: bool = False,
    use_triton: bool = False
)
```

Dataclass

**Bases:** [\_DionConfigBase](#nemo_automodel-components-optim-optimizer-_DionConfigBase)

`dion.Muon` — matrix-aware update for 2D+ params, scalar fallback for 1D.

**`adjust_lr`** `str = 'spectral_norm'`

---

**`betas`** `tuple[float, float] = (0.9, 0.95)`

---

**`epsilon`** `float = 1e-08`

---

**`flatten`** `bool = False`

---

**`mu`** `float = 0.95`

---

**`nesterov`** `bool = False`

---

**`use_triton`** `bool = False`

---

```python
nemo_automodel.components.optim.optimizer.MuonConfig._make_optimizer(
    param_groups,
    ctor_kwargs
)
```

```python
class nemo_automodel.components.optim.optimizer.NorMuonConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    lr: float = 0.0005,
    weight_decay: float = 0.0,
    scalar_opt: str = 'adamw',
    scalar_betas: tuple[float, float] = (0.9, 0.999),
    scalar_eps: float = 1e-08,
    scalar_lr: float | None = None,
    embed_lr: float | None = None,
    lm_head_lr: float | None = None,
    mu: float = 0.95,
    muon_beta2: float = 0.95,
    betas: tuple[float, float] = (0.9, 0.95),
    epsilon: float = 1e-08,
    adjust_lr: str = 'spectral_norm'
)
```

Dataclass

**Bases:** [\_DionConfigBase](#nemo_automodel-components-optim-optimizer-_DionConfigBase)

`dion.NorMuon` — Muon variant with neuron-wise normalization.

**`adjust_lr`** `str = 'spectral_norm'`

---

**`betas`** `tuple[float, float] = (0.9, 0.95)`

---

**`epsilon`** `float = 1e-08`

---

**`mu`** `float = 0.95`

---

**`muon_beta2`** `float = 0.95`

---

```python
nemo_automodel.components.optim.optimizer.NorMuonConfig._make_optimizer(
    param_groups,
    ctor_kwargs
)
```

```python
class nemo_automodel.components.optim.optimizer.OptimizerConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list()
)
```

Dataclass

Base optimizer config.

Subclasses expose their full field surface and implement
`_build_optimizer`, the per-part hook that constructs a single
optimizer from a list of parameters.  `build` owns the shared
orchestration (per-part loop, TP `foreach`, per-group LR overrides) and is
rarely overridden — only by configs whose construction does not fit the
`parameters -> optimizer` shape (e.g. `MuonConfig`).  Megatron-FSDP
optimizer sharding is no longer applied here; the recipe layer re-applies it
via `shard_optimizers_for_megatron_fsdp(...)`.

**`param_group_overrides`** `list[ParamGroupOverride] = field(default_factory=list)`

---

**`supports_megatron_fsdp_sharding`** `bool = True`

---

```python
nemo_automodel.components.optim.optimizer.OptimizerConfig.__post_init__() -> None
```

```python
nemo_automodel.components.optim.optimizer.OptimizerConfig._build_optimizer(
    params,
    foreach: bool | None = None
) -> torch.optim.Optimizer
```

Construct a single optimizer for `params` (one model part).

```python
nemo_automodel.components.optim.optimizer.OptimizerConfig._constructor_kwargs() -> dict[str, typing.Any]
```

`asdict(self)` with non-constructor (grouping) fields removed.

```python
nemo_automodel.components.optim.optimizer.OptimizerConfig.build(
    model: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    is_peft: bool = False
) -> list[torch.optim.Optimizer]
```

Build one optimizer per `model.parts` (or `[model]`).

Applies the shared per-part concerns (TP `foreach` disabling, per-group
LR/WD overrides) and delegates the actual optimizer instantiation to
`_build_optimizer`. Megatron-FSDP optimizer sharding is applied by
the recipe layer, not here.

**Parameters:**

**`model`** `torch.nn.Module`

Model (or model with `.parts`) to optimize.

---

**`device_mesh`** `DeviceMesh | None` — default: None

Device mesh used for tensor/data parallelism.

---

**`is_peft`** `bool` — default: False

Whether the model is being trained with PEFT (suppresses the
bf16 torch-Adam precision warning).

---

**Returns:** `list[torch.optim.Optimizer]`

One optimizer per model part.

```python
nemo_automodel.components.optim.optimizer.OptimizerConfig.build_from_param_groups(
    param_groups: list[dict[str, typing.Any]],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> torch.optim.Optimizer
```

Build one optimizer from caller-defined parameter groups.

```python
class nemo_automodel.components.optim.optimizer.OptimizerFromFactoryConfig(
    param_group_overrides: list[nemo_automodel.components.optim.optimizer.ParamGroupOverride] = list(),
    factory: collections.abc.Callable[..., torch.optim.Optimizer] | None = None,
    kwargs: dict[str, typing.Any] = dict()
)
```

Dataclass

**Bases:** [OptimizerConfig](#nemo_automodel-components-optim-optimizer-OptimizerConfig)

Build an optimizer from an arbitrary factory callable plus kwargs.

The integration escape hatch (e.g. veRL): rather than exposing typed fields,
it wraps an optimizer class/callable and the `**kwargs` to construct it.
This keeps the factory path on the same `config.build(model, ...)` contract
as the typed configs, so `build_optimizer` never has to special-case it.

Hyperparameters live in `kwargs`; the inherited `lr`/`weight_decay`
fields are unused.  The factory is called as `factory(params=..., **kwargs)`;
Dion-family optimizers (which need parameter grouping) should use the typed
`MuonConfig` instead.  A `param_group_overrides` entry in
`kwargs` is consumed here (not forwarded to the factory) to drive
per-group LR/WD, matching the typed-config behavior.

**`factory`** `Callable[..., Optimizer] | None = None`

---

**`kwargs`** `dict[str, Any] = field(default_factory=dict)`

---

```python
nemo_automodel.components.optim.optimizer.OptimizerFromFactoryConfig.build(
    model: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    is_peft: bool = False
) -> list[torch.optim.Optimizer]
```

```python
nemo_automodel.components.optim.optimizer.OptimizerFromFactoryConfig.build_from_param_groups(
    param_groups: list[dict[str, typing.Any]],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> torch.optim.Optimizer
```

```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
class nemo_automodel.components.optim.optimizer.ParamGroupOverride(
    pattern: str,
    lr_mult: float = 1.0,
    wd_mult: float = 1.0
)
```

Dataclass

Per-parameter-group learning-rate / weight-decay override.

Parameters whose (module-qualified) name matches `pattern` are placed
in their own optimizer parameter group carrying `lr_mult` /
`wd_mult`, which the LR scheduler multiplies into the group's learning
rate and weight decay every step (see
`OptimizerParamScheduler.step`). This mirrors Megatron-LM's per-group
`lr_mult` scaling of `max_lr` / `min_lr`.

**`lr_mult`** `float = 1.0`

Multiplier applied to this group's learning rate.

---

**`pattern`** `str`

Python regular expression matched against each parameter name
with `re.search` (so a plain substring like `"router"` also
works).

---

**`wd_mult`** `float = 1.0`

Multiplier applied to this group's weight decay.

---

```python
nemo_automodel.components.optim.dion.build_dion_optimizer(
    config: '_DionFamilyConfig',
    model: torch.nn.Module,
    device_mesh: typing.Any | None = None,
    mesh_kwarg: str | None = 'distributed_mesh'
) -> tuple[list[dict[str, typing.Any]], dict[str, typing.Any]]
```

Build the parameter groups and resolve the device mesh for a Dion-family
optimizer.

This does not instantiate the optimizer; it returns `(param_groups,
mesh_kwargs)` so the caller (a typed config in
`nemo_automodel.components.optim.optimizer`) can assemble its own
constructor kwargs and instantiate the optimizer itself.  `mesh_kwargs` is a
dict that maps `mesh_kwarg` to the resolved mesh (or is empty when there is
no mesh), ready to splat into the optimizer constructor.

The parameter-grouping settings are read off `config`: `lr`,
`weight_decay`, `scalar_opt`, `scalar_betas`, `scalar_eps`
(required), and the optional `scalar_lr`, `embed_lr`, `lm_head_lr` and
`no_compile`.

**Parameters:**

**`config`** `'_DionFamilyConfig'`

The dion-family config (see `_DionFamilyConfig`) to read settings from.

---

**`model`** `nn.Module`

Model whose parameters are to be optimized.

---

**`device_mesh`** `Any | None` — default: None

Optional DeviceMesh for FSDP/TP. When non-empty it is
resolved to a 1-D Dion submesh.

---

**`mesh_kwarg`** `str | None` — default: 'distributed\_mesh'

Name of the constructor argument that receives the resolved
mesh (`"distributed_mesh"` for Muon/Dion2/NorMuon,
`"outer_shard_mesh"` for legacy Dion).  Set to `None` to never
include the mesh.

---

**Returns:** `list[dict[str, Any]]`

A `(param_groups, mesh_kwargs)` tuple: the per-group parameter dicts and

```python
nemo_automodel.components.optim.optimizer.build_optimizer(
    model: torch.nn.Module,
    config: nemo_automodel.components.optim.optimizer.OptimizerConfig | tuple[str, dict[str, typing.Any]],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> list[torch.optim.Optimizer]
```

Build one optimizer per `model.parts` (or `[model]`).

Thin dispatcher: it normalizes `config` to an `OptimizerConfig` and
returns `config.build(model, ...)`.  Per-part concerns (TP `foreach`,
Dion param grouping) live on the config.  Megatron-FSDP optimizer sharding is
re-applied separately by the recipe layer.

`config` is one of:

* a typed `OptimizerConfig` instance — the Automodel-native path.
* a `(name_or_path, kwargs)` tuple, where `name_or_path` is a short
  registry name (see `OPTIMIZER_CONFIG_REGISTRY`, e.g. `"adamw"`) or a
  dotted import path (e.g. `"torch.optim.AdamW"`), and `kwargs` are the
  constructor arguments.  A registry/import-path that resolves to an
  `OptimizerConfig` subclass is built from its typed fields; any other
  callable is wrapped in an `OptimizerFromFactoryConfig` (the escape
  hatch for external integrations, e.g. veRL).

**Parameters:**

**`model`** `torch.nn.Module`

Model (or model with `.parts`) to optimize.

---

**`config`** `OptimizerConfig | tuple[str, dict[str, Any]]`

An `OptimizerConfig` instance or a `(name_or_path,
kwargs)` tuple.

---

**`device_mesh`** `DeviceMesh | None` — default: None

Device mesh used for tensor/data parallelism.

---

**Returns:** `list[torch.optim.Optimizer]`

One optimizer per model part.

```python
nemo_automodel.components.optim.optimizer.build_optimizer_config(
    target: nemo_automodel.components.optim.optimizer.OptimizerConfig | str | type[nemo_automodel.components.optim.optimizer.OptimizerConfig] | collections.abc.Callable[..., torch.optim.Optimizer],
    kwargs: dict[str, typing.Any] | None = None
) -> nemo_automodel.components.optim.optimizer.OptimizerConfig
```

Normalize an optimizer `target` plus `kwargs` into an `OptimizerConfig`.

This is the single normalization entry point shared by the recipe layer
(which resolves a YAML `_target_` to a Python object) and
`build_optimizer` (which accepts `(name_or_path, kwargs)` tuples).

`target` is one of:

* an `OptimizerConfig` instance — returned as-is (`kwargs` ignored,
  since the instance already carries its typed fields).
* an `OptimizerConfig` subclass — instantiated from its typed fields
  with `**kwargs`.
* a string — a registry short name (see `OPTIMIZER_CONFIG_REGISTRY`,
  e.g. `"adamw"`) or a dotted import path (e.g. `"torch.optim.AdamW"`);
  it is resolved and then handled as a subclass or callable.
* any other optimizer callable/class — wrapped in an
  `OptimizerFromFactoryConfig` (the escape hatch for external
  integrations, e.g. veRL).

**Parameters:**

**`target`** `OptimizerConfig | str | type[OptimizerConfig] | Callable[..., torch.optim.Optimizer]`

The optimizer config instance/subclass, registry name or import
path, or optimizer callable to normalize.

---

**`kwargs`** `dict[str, Any] | None` — default: None

Constructor arguments for the resolved config/callable.

---

**Returns:** `OptimizerConfig`

class:`OptimizerConfig` ready to `build(...)`.

```python
nemo_automodel.components.optim.dion.is_dion_optimizer(
    optimizer_factory: typing.Any
) -> bool
```

Return whether an optimizer factory targets a Dion-family optimizer.

```python
nemo_automodel.components.optim.optimizer.OPTIMIZER_CONFIG_REGISTRY: dict[str, type[OptimizerConfig]] = {'adam': AdamConfig, 'adamw': AdamWConfig, 'fused_adam': FusedAdamConfig, 'flash...
```