nemo_automodel.components.moe.experts

Bases: Module

Sparse MoE implementation using all-gather/reduce-scatter primitives.

Supports two compute backends:

• Per-expert loop with gather/scatter (default)
• torch._grouped_mm with argsort-based permutation (backend.experts=“torch_mm”; “gmm” remains a deprecated compatibility alias)

nemo_automodel.components.moe.experts.GroupedExperts._forward_grouped_mm(
    x,
    token_mask,
    weights,
    indices,
    gate_and_up_projs,
    down_projs,
    gate_up_proj_bias,
    down_proj_bias,
    n_local_experts,
    experts_start_idx
)

nemo_automodel.components.moe.experts.GroupedExperts._forward_loop(
    x,
    weights,
    indices,
    token_mask,
    gate_and_up_projs,
    down_projs,
    gate_up_proj_bias,
    down_proj_bias,
    n_local_experts,
    experts_start_idx,
    experts_end_idx
)

nemo_automodel.components.moe.experts.GroupedExperts._gather_ep_inputs(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor,
    ep_group: torch.distributed.ProcessGroup,
    ep_size: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[int]]

nemo_automodel.components.moe.experts.GroupedExperts.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.moe.experts.GroupedExperts.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Bases: Module

Sparse MoE implementation using grouped GEMM with DeepEP token dispatch.

Uses torch._grouped_mm for the canonical experts="torch_mm" backend. experts="gmm" remains a deprecated compatibility alias.

Once the experts for a particular token have been identified, this module is invoked to compute and average the output of the activated experts.

nemo_automodel.components.moe.experts.GroupedExpertsDeepEP._dispatch_tokens(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.moe.experts.GroupedExpertsDeepEP._init_deepep_buffer(
    ep_group: torch.distributed.ProcessGroup
) -> None

nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.get_pipeline_runtime_initializers() -> tuple[nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer, ...]

nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.init_token_dispatcher(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh
)

nemo_automodel.components.moe.experts.GroupedExpertsDeepEP.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Bases: Module

MoE experts using TE’s GroupedLinear module directly.

Uses TE’s native GroupedLinear for computation, providing:

• Optimized grouped GEMM kernels from TE

For expert parallelism, each rank creates GroupedLinear with num_local_experts = n_routed_experts / ep_size.

nemo_automodel.components.moe.experts.GroupedExpertsTE._get_stacked_bias(
    linear: transformer_engine.pytorch.GroupedLinear
) -> torch.Tensor | None

nemo_automodel.components.moe.experts.GroupedExpertsTE._get_stacked_weight(
    linear: transformer_engine.pytorch.GroupedLinear,
    transpose: bool = False
) -> torch.Tensor

nemo_automodel.components.moe.experts.GroupedExpertsTE._load_from_state_dict(
    state_dict: typing.Dict[str, typing.Any],
    prefix: str,
    local_metadata,
    strict,
    missing_keys,
    unexpected_keys,
    error_msgs
)

nemo_automodel.components.moe.experts.GroupedExpertsTE._normalize_moe_mesh(
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> torch.distributed.device_mesh.DeviceMesh | None

nemo_automodel.components.moe.experts.GroupedExpertsTE._set_stacked_bias(
    linear: transformer_engine.pytorch.GroupedLinear,
    stacked: torch.Tensor
)

nemo_automodel.components.moe.experts.GroupedExpertsTE._set_stacked_weight(
    linear: transformer_engine.pytorch.GroupedLinear,
    stacked: torch.Tensor,
    transpose: bool = False
)

nemo_automodel.components.moe.experts.GroupedExpertsTE._to_ep_dtensor(
    tensor: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.moe.experts.GroupedExpertsTE.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.moe.experts.GroupedExpertsTE.get_pipeline_runtime_initializers() -> tuple[nemo_automodel.components.moe.megatron.token_dispatcher.HybridEPPipelineRuntimeInitializer, ...]

nemo_automodel.components.moe.experts.GroupedExpertsTE.init_token_dispatcher(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
)

nemo_automodel.components.moe.experts.GroupedExpertsTE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

nemo_automodel.components.moe.experts.GroupedExpertsTE.set_moe_mesh(
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> None

nemo_automodel.components.moe.experts.GroupedExpertsTE.state_dict(
    args = (),
    destination = None,
    prefix = '',
    keep_vars = False,
    kwargs = {}
) -> typing.Dict[str, typing.Any]

Bases: Function

All-gather with variable local lengths and autograd-safe backward.

Backward uses all-reduce + local narrow instead of reduce-scatter to avoid monitoredBarrier deadlocks observed with mixed FSDP/EP backward collective ordering.

nemo_automodel.components.moe.experts._AllGatherConcatVarlenFn.backward(
    ctx,
    grad_output: torch.Tensor
)

nemo_automodel.components.moe.experts._AllGatherConcatVarlenFn.forward(
    ctx,
    local_tensor: torch.Tensor,
    group: torch.distributed.ProcessGroup,
    gathered_lens: list[int],
    max_len: int
)

Bases: Function

Add expert bias without materializing a full-size expanded bias.

nemo_automodel.components.moe.experts._ChunkedBiasAdd.backward(
    ctx,
    grad_output
)

nemo_automodel.components.moe.experts._ChunkedBiasAdd.forward(
    ctx,
    value,
    bias,
    token_counts,
    probs = None,
    reuse_input = False
)

Bases: Function

Expand expert biases while reducing their gradients deterministically.

nemo_automodel.components.moe.experts._DeterministicBiasRepeatInterleave.backward(
    ctx,
    grad_output
)

nemo_automodel.components.moe.experts._DeterministicBiasRepeatInterleave.forward(
    ctx,
    bias,
    token_counts,
    output_size
)

Allocate device workspace required by Triton tensor descriptors.

Parameters:

Returns: torch.Tensor

One-dimensional CUDA byte tensor with size elements.

Apply per-expert bias to grouped GEMM output.

NOTE: torch._grouped_mm accepts a bias kwarg in its schema but raises “RuntimeError: Bias not supported yet” as of PyTorch 2.9.0. Additionally, down projection bias needs weighting by routing probs (bias * permuted_probs) which native bias support wouldn’t handle.

Parameters:

Returns:

Grouped GEMM output with per-expert bias applied, shape

Run a grouped expert MLP in bounded row chunks.

Each chunk is independently checkpointed, so backward recomputes its two grouped GEMMs instead of retaining full [tokens, 2 * inter_dim] and [tokens, inter_dim] activations. The final [tokens, hidden] tensor is allocated once and populated by slice assignment; no concatenation or stateful scratch buffer is used.

The checkpoint callback captures the input dtype rather than the routed tensor so an enclosing activation checkpoint can release that tensor.

This helper is for the plain torch._grouped_mm path. MXFP8 uses a different quantization contract and remains on its existing whole-dispatch implementation.

Parameters:

Returns: torch.Tensor

Tensor of shape [tokens, hidden] in hidden_states.dtype with independent storage.

Reduce one expert’s contiguous token gradients without atomic writes.

grad_output_ptr is a contiguous row-major [tokens, hidden] allocation and probs_ptr is a contiguous FP32 [tokens, 1] allocation when USE_PROBS is true. The caller ensures the gradient row stride is descriptor-aligned and token offsets fit in int32. Each program exclusively owns one contiguous FP32 [expert, hidden block] tile in grad_bias_ptr, so the reduction order is deterministic.

Permute tokens by expert assignment and compute offs for torch._grouped_mm.

Takes the raw router outputs and produces sorted token IDs, routing weights, tokens_per_expert counts, and cumulative offsets ready for grouped GEMM.

Returns:

Token indices sorted by expert assignment.

Reduce bias gradients without full-size intermediates.

Parameters:

Returns: torch.Tensor

Tensor of shape [experts, hidden] with one weighted gradient sum per expert.

Reduce grouped token gradients without requiring Triton.

Parameters:

Returns: torch.Tensor

Tensor of shape [experts, hidden] with one gradient sum per expert.

Reduce grouped CUDA token gradients with deterministic compensated FP32 sums.

Parameters:

Returns: torch.Tensor

CUDA tensor of shape [experts, hidden] and output_dtype when provided,

Resolve the per-expert split sizes for TE’s grouped GEMM.

Parameters:

Returns: list

The split sizes as a Python list. Under static routing with a warm cache the

Reduce contiguous expert segments into an accumulation-dtype tensor.

Parameters:

Returns: torch.Tensor

Tensor of shape [experts, hidden] in accumulation_dtype.

Keep empty expert probabilities stable across checkpoint recomputation.

DeepEP may return an empty probability tensor in either the router dtype or the activation dtype. With non-reentrant activation checkpointing, a later recomputation can choose the other representation, and PyTorch checkpointing rejects the resulting saved-tensor metadata change even though both tensors contain zero elements. Empty probabilities have no numerical contribution, so canonicalize only that case to the dispatched activation dtype. Non-empty probabilities retain their original (normally FP32) router precision.

Clamped SwiGLU (DeepSeek V4 style) for DeepEP.

Gate is clamped at max=limit and up at (-limit, +limit) in FP32 before silu(gate) * up; the result is multiplied by the permuted routing probs and cast back. Matches the official V4 Expert.forward::

gate = self.w1(x).float() up = self.w3(x).float() if self.swiglu_limit > 0: up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit) gate = torch.clamp(gate, max=swiglu_limit) y = F.silu(gate) * up

x has shape [..., 2 * inter_dim] with gate in the first half and up in the second half (same layout as weighted_bias_swiglu_impl).

Add weighted expert bias without expanded bias or FP32 activation buffers.

Parameters:

Reduce each row’s bias-weighted gradient without expanded bias buffers.

Parameters:

Return the DeepEP expert activation function selected by the MoE config.

Check if activation requires gating (gate_proj + up_proj).

Gated activations (SwiGLU, Quick-GEGLU) use both gate_proj and up_proj, requiring gate_and_up_projs tensor with shape [n_experts, dim, 2*inter_dim].

Non-gated activations (ReLU²) only use up_proj, requiring up_projs tensor with shape [n_experts, dim, inter_dim] - 50% memory savings.

Apply DeepEP Quick-GEGLU activation and routing probabilities.

ReLU² activation for DeepEP: relu(x)^2

For DeepEP with ReLU², x is the output of the up projection (already computed). x already has shape […, inter_dim] from efficient up_proj.

Clamped SwiGLU (DeepSeek V4 style) for DeepEP: Triton kernels on CUDA, the compiled function elsewhere.

See nemo_automodel.components.moe.swiglu_triton for the kernels and the layout guard.

SwiGLU-OAI (GPT-OSS / MiniMax-M3) activation for grouped experts.

Computes gate * sigmoid(alpha * gate) * (up + 1) in fp32 with gate clamped max=limit and up clamped +/-limit (when limit > 0).

Unlike quick_geglu_deepep (which expects an interleaved gate/up layout, x[..., ::2] / x[..., 1::2]), this reads the concatenated [gate | up] layout produced by MoESplitExpertsStateDictMixin (torch.cat([gate_t, up_t], dim=-1)), matching sglang’s swiglu_no_interleaved_with_alpha_and_limit.