> This page is for version Nightly (default).
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.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](https://arxiv.org/abs/2605.10797). Numerical conventions follow
the reference implementation's ordinary Muown / newtonschulz5_torch backend.

## Module Contents

### Classes

| Name                                                    | Description                                                                   |
| ------------------------------------------------------- | ----------------------------------------------------------------------------- |
| [`Muown`](#nemo_automodel-components-optim-muown-Muown) | Learn per-output-neuron magnitudes with Adam and matrix directions with Muon. |

### Functions

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

### Data

[`_DION_MUON_HOOKS`](#nemo_automodel-components-optim-muown-_DION_MUON_HOOKS)

[`_MuonBase`](#nemo_automodel-components-optim-muown-_MuonBase)

### API

```python
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](#nemo_automodel-components-optim-muown-_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 | None` — default: None

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

---

**`lr`** `float` — default: 0.0003

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

---

**`mu`** `float` — default: 0.95

Direction momentum coefficient.

---

**`betas`** `tuple[float, float]` — default: (0.9, 0.95)

Adam coefficients for learned magnitudes and scalar fallback.

---

**`weight_decay`** `float` — default: 0.0

Decoupled decay applied to the original weight.

---

**`epsilon`** `float` — default: 1e-08

Adam denominator and zero-norm stabilizer.

---

**`nesterov`** `bool` — default: True

Enable Nesterov direction momentum.

---

**`ns_steps`** `int` — default: 5

Newton-Schulz iteration count.

---

**`ns_epsilon`** `float` — default: 1e-07

Newton-Schulz normalization stabilizer.

---

**`muon_update_scale`** `float` — default: 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`** `bool` — default: False

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

---

```python
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`** `str` — default: 'adamw'

Dion's scalar-algorithm name.

---

```python
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`** `str` — default: 'muon'

Dion's matrix-algorithm name.

---

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

---

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

---

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

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

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

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

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

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

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

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

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