nemo_automodel.components.moe.optimized_ops

Memory-optimized MoE elementwise ops extracted from experts.py.

Chunked custom-autograd router-weight fp32 multiply: computes the identical fp32 math in row chunks and saves only low-precision inputs, removing the full-size fp32 intermediates that otherwise pin ~7 GiB blocks per MoE layer under activation checkpointing.

Module Contents

Classes

Name	Description
`_RouterWeightMulFunction`	Chunked fp32 router-weight multiply that saves only the raw inputs.

Functions

Name	Description
`_apply_router_weight_fp32`	Apply routing probabilities to expert outputs with fp32 arithmetic.

Data

_RW_CHUNK_ROWS

_RW_CHUNK_THRESHOLD

API

class nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction()

Bases: Function

Chunked fp32 router-weight multiply that saves only the raw inputs.

The plain (x.float() * probs.float()).to(dtype) lets autograd keep full-size fp32 [tokens, hidden] intermediates alive for backward (the upcast input, the product, and a recompute copy under activation checkpointing). This Function computes the same fp32 multiply in row chunks and saves only the low-precision inputs. The forward is bitwise-identical; the backward matches autograd’s fp32 chain (grad_x = (g_f32 * probs_f32).to(x.dtype), grad_probs = (g_f32 * x_f32).sum(-1, keepdim=True) cast to probs.dtype).

nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None]

staticmethod

Compute chunked fp32 gradients for the router-weight multiply.

Parameters:

ctx

Any

Autograd context holding the saved inputs.

grad_out

torch.Tensor

Upstream gradient of shape [tokens, hidden].

Returns: torch.Tensor | None

Tuple (grad_x, grad_probs, None, None) where grad_x has

nemo_automodel.components.moe.optimized_ops._RouterWeightMulFunction.forward(
    ctx: typing.Any,
    x: torch.Tensor,
    probs: torch.Tensor,
    out_dtype: torch.dtype,
    save_x: bool
) -> torch.Tensor

staticmethod

Multiply expert outputs by routing probabilities in fp32.

Parameters:

ctx

Any

Autograd context.

torch.Tensor

Expert outputs of shape [tokens, hidden].

probs

torch.Tensor

Routing probabilities of shape [tokens, 1].

out_dtype

torch.dtype

Output dtype (the dispatcher’s expected activation dtype, or fp32 for the scatter-add reduction path).

save_x

bool

Whether backward needs x. x is only consumed by the probs gradient; when probs carries no grad (e.g. FakeBalancedGate emits constant weights), saving it would pin a full-size [tokens, hidden] tensor per MoE layer across the activation-checkpointing backward window for nothing. Callers pass probs.requires_grad.

Returns: torch.Tensor

Tensor of shape [tokens, hidden] and dtype out_dtype.

nemo_automodel.components.moe.optimized_ops._apply_router_weight_fp32(
    output2: torch.Tensor,
    permuted_probs: torch.Tensor,
    compute_dtype: torch.dtype
) -> torch.Tensor

Apply routing probabilities to expert outputs with fp32 arithmetic.

Large 2-D row-aligned inputs go through the chunked custom Function so autograd does not retain full-size fp32 intermediates; every other shape keeps the plain eager multiply (bitwise-identical result).

Parameters:

output2

torch.Tensor

Expert down-projection outputs of shape [tokens, hidden].

permuted_probs

torch.Tensor

Routing probabilities broadcastable against output2, typically of shape [tokens, 1].

compute_dtype

torch.dtype

Output dtype.

Returns: torch.Tensor

(output2 * permuted_probs) computed in fp32 and cast to

nemo_automodel.components.moe.optimized_ops._RW_CHUNK_ROWS = 8192

nemo_automodel.components.moe.optimized_ops._RW_CHUNK_THRESHOLD = 12288