nemo_automodel.components.optim.muown

Muown: Muon directions and Adam neuron magnitudes, using Dion’s FSDP2 all-to-all ownership scheme.

Algorithm reference: https://arxiv.org/abs/2605.10797. Numerical conventions follow the reference implementation’s ordinary Muown / newtonschulz5_torch backend.

Module Contents

Classes

Name	Description
`Muown`	Learn per-output-neuron magnitudes with Adam and matrix directions with Muon.

Functions

Name	Description
`_matrix_sharding`	Resolve the communication group for the last two matrix axes.
`_newton_schulz`	Orthogonalize complete local matrices with the reference quintic polynomial.
`_nonzero_magnitude`	Preserve signed neuron magnitudes while guarding division by zero.
`_prepare_direction`	Reconstruct a direction and apply the weight-normalization Jacobian.
`_recompose`	Compose effective weights with reference-compatible fused arithmetic.
`_row_norm`	Compute per-output-neuron norms without gathering the weight matrix.
`_row_sum`	Reduce an input-feature axis and resolve DTensor partial sums.

Data

_DION_MUON_HOOKS

_MuonBase

API

class nemo_automodel.components.optim.muown.Muown(
    params: torch.optim.optimizer.ParamsT,
    distributed_mesh: torch.distributed.tensor.DeviceMesh | torch.distributed.ProcessGroup | None = None,
    lr: float = 0.0003,
    mu: float = 0.95,
    betas: tuple[float, float] = (0.9, 0.95),
    weight_decay: float = 0.0,
    epsilon: float = 1e-08,
    nesterov: bool = True,
    ns_steps: int = 5,
    ns_epsilon: float = 1e-07,
    muon_update_scale: float = 1.0,
    use_triton: bool = False
)

Bases: _MuonBase

Learn per-output-neuron magnitudes with Adam and matrix directions with Muon.

Uses the existing optional Dion dependency. Leading dimensions of 3D+ weights are independent matrices (experts). Groups with matrix_transposed=True store matrices as […, input, output]; the default is […, output, input]. Embeddings, biases and heads should be placed in algorithm=“adamw” or “lion” groups by MuownConfig.

FSDP communication follows each parameter’s DTensor mesh, so dense and expert parameters can share an optimizer despite different DP/EP meshes. Only one evenly sharded matrix axis is supported. EP-only and replicated matrices are orthogonalized locally. FP32 parameter storage is recommended.

Parameters:

params

ParamsT

Parameters or parameter-group dictionaries. Matrix weights have global shape […, output, input], or […, input, output] for matrix_transposed groups; leading axes are independent experts. AdamW/Lion groups accept arbitrary parameter shapes. DTensor matrix sharding follows _matrix_sharding. Gradients must use matching placements. Parameters and optimizer state are updated in-place; caller-owned gradients are preserved.

distributed_mesh

DeviceMesh | dist.ProcessGroup | NoneDefaults to None

Compatibility argument for Dion; matrix collectives use each parameter’s actual DTensor mesh.

floatDefaults to 0.0003

Base magnitude LR; direction LR is scaled by matrix dimensions.

floatDefaults to 0.95

Direction momentum coefficient.

betas

tuple[float, float]Defaults to (0.9, 0.95)

Adam coefficients for learned magnitudes and scalar fallback.

weight_decay

floatDefaults to 0.0

Decoupled decay applied to the original weight.

epsilon

floatDefaults to 1e-08

Adam denominator and zero-norm stabilizer.

nesterov

boolDefaults to True

Enable Nesterov direction momentum.

ns_steps

intDefaults to 5

Newton-Schulz iteration count.

ns_epsilon

floatDefaults to 1e-07

Newton-Schulz normalization stabilizer.

muon_update_scale

floatDefaults to 1.0

Nonnegative multiplier for direction updates only. Magnitude Adam and auxiliary AdamW/Lion updates are unchanged. This run-level setting is read from config, including on resume.

use_triton

boolDefaults to False

Fuse local FP32 row updates on CUDA. Unsupported layouts, wide rows and input-feature shards retain the Torch implementation.

nemo_automodel.components.optim.muown.Muown._create_adamw_tasks(
    param_groups: list[dict[str, typing.Any]],
    algo_name: str = 'adamw'
) -> collections.abc.Generator[dion.opt_utils.AsyncTask, None, None]

Schedule PyTorch AdamW for the non-matrix parameter groups.

Parameters:

param_groups

list[dict[str, Any]]

Dion groups of arbitrary-shaped floating-point parameters. DTensors and their gradients/states share the same placements; each local shard is updated independently.

algo_name

strDefaults to 'adamw'

Dion’s scalar-algorithm name.

nemo_automodel.components.optim.muown.Muown._create_muon_tasks(
    param_groups: list[dict[str, typing.Any]],
    algo_name: str = 'muon'
) -> collections.abc.Generator[dion.opt_utils.AsyncTask, None, None]

Schedule groups with matrix layouts and ownership documented by Muown.

Parameters:

param_groups

list[dict[str, Any]]

Dion group dictionaries containing matrix parameters of global shape […, rows, columns] and Muown hyperparameters.

algo_name

strDefaults to 'muon'

Dion’s matrix-algorithm name.

nemo_automodel.components.optim.muown.Muown._update_adamw_group(
    group: dict[str, typing.Any]
) -> collections.abc.Generator[None, None, None]

Apply AdamW to local shards using PyTorch’s reference arithmetic.

Parameters:

group

dict[str, Any]

AdamW group with arbitrary-shaped parameters; DTensor parameters, gradients and moments have matching placements. Parameters, moments and per-parameter steps are mutated in place.

nemo_automodel.components.optim.muown.Muown._update_batch(
    params: list[torch.Tensor],
    group: dict[str, typing.Any],
    shard_dim: int | None,
    process_group: torch.distributed.ProcessGroup | None,
    world_size: int
) -> collections.abc.Generator[None, None, None]

Update a batch while keeping temporary direction buffers batch-local.

Parameters:

params

list[Tensor]

Parameters sharing global shape […, rows, columns], placements, dtype, device, and mesh. Last-two-axis ordering is given by group[“matrix_transposed”]. Gradients are preserved.

group

dict[str, Any]

Dion hyperparameters and the Muown matrix layout.

shard_dim

int | None

Evenly sharded matrix axis, or None for complete matrices.

process_group

dist.ProcessGroup | None

Group owning the matrix shards, or None.

world_size

int

Size of process_group, or 1 for local matrices.

nemo_automodel.components.optim.muown._matrix_sharding(
    param: torch.Tensor
) -> tuple[int | None, torch.distributed.ProcessGroup | None]

Resolve the communication group for the last two matrix axes.

Parameters:

param

Tensor

Parameter of global shape […, rows, columns]. Leading axes are matrix batches; DTensors may shard those axes and at most one matrix axis. Matrix shards must be even and Partial is unsupported.

Returns: int | None

Matrix tensor axis and its process group, or (None, None) when each

nemo_automodel.components.optim.muown._newton_schulz(
    update: torch.Tensor,
    epsilon: float | torch.Tensor,
    matrix_transposed: bool,
    steps: int
) -> torch.Tensor

Orthogonalize complete local matrices with the reference quintic polynomial.

Parameters:

update

Tensor

Local tensor of shape […, output, input], or […, input, output] when matrix_transposed is true. Leading axes are independent matrices; this function does not accept matrix-sharded DTensors.

epsilon

float | Tensor

Stabilizer for the matrix Frobenius norm. Dion supplies a zero-dimensional CPU tensor; Python scalars are also accepted.

matrix_transposed

bool

Whether the last axis contains output neurons.

steps

int

Number of polynomial iterations.

Returns: Tensor

Independent contiguous BF16 tensor matching update’s shape and axis

nemo_automodel.components.optim.muown._nonzero_magnitude(
    value: torch.Tensor,
    epsilon: float
) -> torch.Tensor

Preserve signed neuron magnitudes while guarding division by zero.

Parameters:

value

Tensor

Magnitude tensor of shape […, output, 1] or […, 1, output], with arbitrary leading batch axes; DTensor placements are preserved.

epsilon

float

Minimum absolute magnitude.

Returns: Tensor

Independent tensor with the same shape, dtype, device and placements.

nemo_automodel.components.optim.muown._prepare_direction(
    weight: torch.Tensor,
    magnitude: torch.Tensor,
    norm: torch.Tensor,
    gradient: torch.Tensor,
    axis: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Reconstruct a direction and apply the weight-normalization Jacobian.

Parameters:

weight

Tensor

Effective weights […, output, input], or […, input, output] when axis is -2. DTensor shapes are global.

magnitude

Tensor

Signed neuron magnitudes with axis reduced to size 1.

norm

Tensor

Direction norms with the same shape and placements as magnitude.

gradient

Tensor

Weight gradients matching weight’s shape and placements.

axis

int

Physical input-feature axis; partial row sums are resolved.

Returns: Tensor

Tuple (direction, magnitude_gradient, direction_gradient). Direction

nemo_automodel.components.optim.muown._recompose(
    weight: torch.Tensor,
    magnitude: torch.Tensor,
    direction: torch.Tensor,
    axis: int
) -> torch.Tensor

Compose effective weights with reference-compatible fused arithmetic.

Parameters:

weight

Tensor

Destination weights […, output, input], or […, input, output] for axis -2. Mutated in place; DTensor shapes are global.

magnitude

Tensor

Signed magnitudes with the input-feature axis reduced to 1.

direction

Tensor

Updated directions matching weight’s shape and placements.

axis

int

Physical input-feature axis.

Returns: Tensor

Independent row norms with magnitude’s shape and placements. The caller

nemo_automodel.components.optim.muown._row_norm(
    value: torch.Tensor,
    axis: int,
    epsilon: float
) -> torch.Tensor

Compute per-output-neuron norms without gathering the weight matrix.

Parameters:

value

Tensor

Tensor of shape […, output, input] or […, input, output], including globally shaped DTensors as described by _row_sum.

axis

int

Physical input-feature axis, -1 or -2.

epsilon

float

Minimum returned norm.

Returns: Tensor

Norm tensor in value’s dtype, with the input-feature axis of size 1

nemo_automodel.components.optim.muown._row_sum(
    value: torch.Tensor,
    axis: int
) -> torch.Tensor

Reduce an input-feature axis and resolve DTensor partial sums.

Parameters:

value

Tensor

Tensor of shape […, output, input] or […, input, output], with arbitrary leading matrix-batch dimensions. DTensors may shard batch or matrix axes; shapes are global.

axis

int

Physical input-feature axis, -1 or -2.

Returns: Tensor

Tensor with the input-feature axis reduced to size 1. Partial sums are

nemo_automodel.components.optim.muown._DION_MUON_HOOKS = ('_create_muon_tasks', '_create_adamw_tasks', '_get_or_initialize_state')

nemo_automodel.components.optim.muown._MuonBase = _dion.Muon if _HAS_DION else Optimizer