nemo_automodel.components.optim.muown_triton

Fused local Muown row operations; Newton-Schulz and collectives stay in Muown.

Only contiguous FP32 local tensors with complete input-feature axes use these kernels. Matrix batches and output-neuron shards are independent. Public DTensor ownership and autograd version counters remain the optimizer’s responsibility.

Module Contents

Functions

Name	Description
`_direction_kernel`	-
`_finish`	Apply direction/magnitude updates, normalize and decay local weights.
`_finish_kernel`	-
`_magnitude_grad_kernel`	-
`_prepare`	Prepare complete local rows and mutate momentum.
`_prepare_kernel`	-
`_supported`	Check local dtype/strides and the bounded row-reduction footprint.

Data

HAVE_TRITON

API

nemo_automodel.components.optim.muown_triton._direction_kernel(
    W,
    G,
    N,
    Grad,
    Mom,
    D,
    U,
    GG,
    SIZE: tl.constexpr,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    MU: tl.constexpr,
    EPS: tl.constexpr,
    NESTEROV: tl.constexpr,
    BLOCK: tl.constexpr
)

nemo_automodel.components.optim.muown_triton._finish(
    weight: torch.Tensor,
    magnitude: torch.Tensor,
    norm: torch.Tensor,
    first_moment: torch.Tensor,
    second_moment: torch.Tensor,
    direction: torch.Tensor,
    update: torch.Tensor,
    magnitude_grad: torch.Tensor,
    transposed: bool,
    lr: float,
    scale: float,
    beta1: float,
    beta2: float,
    step: int,
    epsilon: float,
    weight_decay: float
) -> None

Apply direction/magnitude updates, normalize and decay local weights.

Tensor inputs follow _prepare’s contiguous local-row contract. Update is BF16 with weight’s shape; direction is FP32 with that shape. All other tensors are FP32 with magnitude’s row shape. Weight, magnitude, norm and both moments are mutated in place; direction/update/gradients are preserved. Scale must be computed from the GLOBAL matrix shape, including on FSDP output shards. LR and bias corrections are runtime scalars.

nemo_automodel.components.optim.muown_triton._finish_kernel(
    W,
    G,
    N,
    M,
    V,
    D,
    U,
    GG,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    TRANSPOSED: tl.constexpr,
    LR,
    DIRECTION_LR,
    DECAY,
    B1: tl.constexpr,
    B2: tl.constexpr,
    BC1,
    BC2,
    EPS: tl.constexpr,
    WD: tl.constexpr,
    BR: tl.constexpr,
    BI: tl.constexpr
)

nemo_automodel.components.optim.muown_triton._magnitude_grad_kernel(
    W,
    G,
    Grad,
    GG,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    EPS: tl.constexpr,
    BR: tl.constexpr,
    BI: tl.constexpr
)

nemo_automodel.components.optim.muown_triton._prepare(
    weight: torch.Tensor,
    gradient: torch.Tensor,
    momentum: torch.Tensor,
    magnitude: torch.Tensor,
    norm: torch.Tensor,
    transposed: bool,
    mu: float,
    epsilon: float,
    nesterov: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Prepare complete local rows and mutate momentum.

Parameters:

weight

Tensor

Contiguous FP32 CUDA […, output, input], or […, input, output] when transposed. Leading dimensions batch independent matrices; the input axis must be complete on this rank.

gradient

Tensor

Same shape/dtype/device/layout as weight, preserved.

momentum

Tensor

Same shape/dtype/device/layout as weight, updated in place.

magnitude

Tensor

Contiguous FP32 tensor with the input axis reduced to 1.

norm

Tensor

Same shape/layout as magnitude; direction norms.

transposed

bool

Whether output neurons occupy the final physical axis.

float

Momentum coefficient.

epsilon

float

Signed magnitude division guard.

nesterov

bool

Whether to use Nesterov momentum.

Returns: Tensor

Independent direction (FP32), update (BF16), and magnitude-gradient

nemo_automodel.components.optim.muown_triton._prepare_kernel(
    W,
    G,
    N,
    Grad,
    Mom,
    D,
    U,
    GG,
    R: tl.constexpr,
    INPUTS: tl.constexpr,
    TRANSPOSED: tl.constexpr,
    MU: tl.constexpr,
    EPS: tl.constexpr,
    NESTEROV: tl.constexpr,
    BR: tl.constexpr,
    BI: tl.constexpr
)

nemo_automodel.components.optim.muown_triton._supported(
    tensors: tuple[torch.Tensor, ...],
    transposed: bool
) -> bool

Check local dtype/strides and the bounded row-reduction footprint.

nemo_automodel.components.optim.muown_triton.HAVE_TRITON = _HAS_TRITON and _HAS_LANGUAGE