nemo_automodel.components.optim

Submodules

Package Contents

Classes

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

Functions

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

Data

OPTIMIZER_CONFIG_REGISTRY

API

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

torch.optim.Adam.

amsgrad

bool = False

betas

tuple[float, float] = (0.9, 0.999)

eps

float = 1e-08

float = 0.0001

weight_decay

float = 0.01

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

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

torch.optim.AdamW.

amsgrad

bool = False

betas

tuple[float, float] = (0.9, 0.999)

eps

float = 1e-08

fused

bool = False

float = 0.0001

weight_decay

float = 0.01

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

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

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

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

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

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

float = 0.95

power_iters

int = 1

qr_method

str = 'rcqr'

rank_fraction

float = 1.0

rank_multiple_of

int = 1

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

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

flashoptim.FlashAdamW.

betas

tuple[float, float] = (0.9, 0.999)

eps

float = 1e-08

float = 0.0001

master_weight_bits

int = 24

weight_decay

float = 0.01

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

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

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

float = 0.0001

master_weight_dtype

str | None = None

master_weights

bool = True

weight_decay

float = 0.01

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

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

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.

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

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

float = 0.95

nesterov

bool = False

use_triton

bool = False

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

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

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

float = 0.95

muon_beta2

float = 0.95

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

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

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

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

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

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

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 | NoneDefaults to None

Device mesh used for tensor/data parallelism.

is_peft

boolDefaults to 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.

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.

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

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)

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]

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

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

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.

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 | NoneDefaults to None

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

mesh_kwarg

str | NoneDefaults to '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

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 | NoneDefaults to None

Device mesh used for tensor/data parallelism.

Returns: list[torch.optim.Optimizer]

One optimizer per model part.

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] | NoneDefaults to None

Constructor arguments for the resolved config/callable.

Returns: OptimizerConfig

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

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

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

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