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”)

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

Supports two GEMM backends via BackendConfig.experts:

• grouped_gemm.ops.gmm (experts=“gmm”, default)
• torch._grouped_mm (experts=“torch_mm”, no external dependency)

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._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.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
) -> typing.Optional[torch.Tensor]

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: typing.Optional[torch.distributed.device_mesh.DeviceMesh]
) -> typing.Optional[torch.distributed.device_mesh.DeviceMesh]

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.init_token_dispatcher(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh,
    moe_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = 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: typing.Optional[torch.distributed.device_mesh.DeviceMesh]
) -> 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

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

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

grad_output_ptr is a contiguous row-major [tokens, hidden] allocation. The caller ensures its row stride is descriptor-aligned and its token offsets fit in int32. Each program exclusively owns one [expert, hidden block] output tile, 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 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 the same dtype as grad_output.

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.

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).

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 :func: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.