nemo_automodel.components.optim.optimizer

Bases: OptimizerConfig

torch.optim.Adam.

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

Bases: OptimizerConfig

torch.optim.AdamW.

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

Bases: _DionConfigBase

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

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

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.

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

Bases: OptimizerConfig

flashoptim.FlashAdamW.

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

Bases: OptimizerConfig

transformer_engine.pytorch.optimizers.FusedAdam.

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

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

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]

Bases: _DionConfigBase

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

Grouped MoE experts are trained as batches of matrices. Their [expert, in, out] storage is orthogonalized as is (Newton–Schulz is transpose-equivariant), but dion.Muon’s spectral-norm LR scale reads the last two dims as [out, in], so _make_optimizer folds the correct per-shape scale into those groups.

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

nemo_automodel.components.optim.optimizer.MuonConfig._spectral_norm_groups(
    group: nemo_automodel.components.optim.optimizer._DionParamGroup
) -> list[nemo_automodel.components.optim.optimizer._DionParamGroup]

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

Bases: _DionConfigBase

Muown: per-neuron Adam magnitudes and Dion Muon directions, including MoE.

nemo_automodel.components.optim.optimizer.MuownConfig._make_optimizer(
    param_groups: list[dict[str, typing.Any]],
    ctor_kwargs: dict[str, typing.Any]
) -> torch.optim.Optimizer

Bases: _DionConfigBase

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

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

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

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

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

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]

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

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.

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

Per-parameter-group learning-rate and 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.

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

Bases: OptimizerConfig

Shared base for the dion-family typed configs (Muon / Muown / NorMuon / Dion2 / Dion).

Dion optimizers need Dion’s parameter grouping (built from the model) and the device mesh rather than a flat parameter list, so build runs grouping per model part. Grouping-only fields (scalar_* / *_lr and adamw_param_patterns) are stripped from the constructor kwargs. adamw_param_patterns selects trainable parameters for auxiliary AdamW by regex on canonical parameter names; an empty list keeps default routing. LR/WD overrides apply after this selection. Dion is incompatible with Megatron-FSDP optimizer sharding; this is enforced at the recipe layer (supports_megatron_fsdp_sharding = False drives an allow=False sharding call that asserts rather than silently returning an unsharded optimizer).

nemo_automodel.components.optim.optimizer._DionConfigBase._apply_adamw_param_patterns(
    model: torch.nn.Module,
    groups: list[dict[str, typing.Any]]
) -> list[dict[str, typing.Any]]

nemo_automodel.components.optim.optimizer._DionConfigBase._apply_param_group_overrides(
    model: torch.nn.Module,
    groups: list[dict[str, typing.Any]]
) -> list[dict[str, typing.Any]]

nemo_automodel.components.optim.optimizer._DionConfigBase._make_optimizer(
    param_groups: typing.Any,
    ctor_kwargs: dict[str, typing.Any]
) -> torch.optim.Optimizer

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

Bases: typing.TypedDict

Parameter group handed to a Dion-family optimizer.

params is always present. Layout groups built by _separate_param_groups carry matrix_transposed; scalar groups carry algorithm with their own lr, weight_decay and Adam/Lion hyperparameters. _DionConfigBase._apply_param_group_overrides adds lr_mult / wd_mult and MuonConfig._spectral_norm_groups sets adjust_lr.

Partition named_params into optimizer groups by name-pattern overrides.

Each parameter joins the group of the first override whose pattern matches its name; unmatched parameters form the default group. Override groups carry only lr_mult / wd_mult (read by OptimizerParamScheduler.step), not a pre-scaled lr / weight_decay: every group therefore inherits the optimizer’s base lr / weight_decay at construction, and the scheduler (which calls step(0) on init) applies the multipliers. Keeping the stored lr unscaled is what lets LRSchedulerConfig.build read an accurate base LR from param_groups[0] even when the default group is empty. Empty groups (a pattern matching nothing) are dropped with a warning.

Parameters:

Returns: list[dict[str, Any]]

A list of parameter-group dicts suitable for a torch optimizer, default

Coerce a list of dicts (as delivered by YAML) or objects into ParamGroupOverride.

Drop parameters whose rank-local shard holds zero elements.

FSDP2 shards every parameter along dim-0 across the shard group, so any parameter with dim-0 smaller than the group — e.g. the biases, norm weights, or class/position embeddings of a small dense vision tower sharded over a wide mesh — leaves zero-numel local shards on the tail ranks. TransformerEngine FusedAdam’s multi_tensor_apply kernel has no empty-tensor guard and faults (CUDA misaligned address / illegal memory access) at the first optimizer step. Dropping locally-empty shards is exact, not an approximation: every element of those parameters lives on other ranks, whose optimizers update them; this rank has nothing to do.

Parameters:

Returns: list[Any]

params with zero-numel local shards removed. Param groups keep

Raises:

• ValueError: If every parameter of the flat list — or of any single param group — is locally empty on this rank. Neither outcome has a safe representation: dropping a whole group makes optimizer.param_groups rank-asymmetric (LR/WD schedulers address groups positionally, e.g. param_groups[0], so ranks would silently schedule different lr/wd for shards other ranks own), while keeping an empty group breaks torch DCP’s flattened optimizer-state load, which indexes the first param of each group.

Return True if factory accepts a foreach kwarg.

torch.optim optimizers take foreach; external factories such as TE FusedAdam do not, so passing it would raise TypeError.

Return False when TP > 1 (foreach is unsupported), else None.

Import an object from a dotted path, e.g. "torch.optim.AdamW".

Return True if factory is TransformerEngine’s FusedAdam (or a subclass).

functools.partial wrappers are unwrapped (iteratively, for nested partials) before the identity check, so partial(FusedAdam, ...) factories are recognized. Other wrapper callables (closures, custom factory functions) are opaque and are NOT recognized; such factories must guard against zero-numel local shards themselves.

Identity-based and import-free: TE is an optional dependency, so this never imports it. If TE has not been imported yet, factory cannot be TE’s FusedAdam class and the check is trivially False.

Number of elements of param owned by this rank.

Parameters:

Returns: int

Element count of the rank-local shard; 0 when this rank owns no slice.

Return one model part’s trainable params, grouped by overrides when set.

Without overrides this is the flat trainable-parameter list (unchanged behavior); with overrides it is the list of parameter-group dicts from _build_param_groups.

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:

Returns: list[torch.optim.Optimizer]

One optimizer per model part.

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:

Returns: OptimizerConfig

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