> This page is for version 0.5.0 · 26.06.
> 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.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.moe.layers

## Module Contents

### Classes

| Name                                                                         | Description                                                                     |
| ---------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`FakeBalancedGate`](#nemo_automodel-components-moe-layers-FakeBalancedGate) | Load balanced gate implementation, spreads tokens uniformly across all experts. |
| [`Gate`](#nemo_automodel-components-moe-layers-Gate)                         | Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.        |
| [`MLP`](#nemo_automodel-components-moe-layers-MLP)                           | Multi-Layer Perceptron (MLP) used as a feed-forward layer.                      |
| [`MoE`](#nemo_automodel-components-moe-layers-MoE)                           | Mixture-of-Experts (MoE) module.                                                |
| [`_GateRoutingCore`](#nemo_automodel-components-moe-layers-_GateRoutingCore) | Parameterless, fixed-shape portion of a learned MoE router.                     |

### Functions

| Name                                                                                   | Description                                                                                         |
| -------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------- |
| [`_clamped_swiglu_act`](#nemo_automodel-components-moe-layers-_clamped_swiglu_act)     | fp32 clamped SwiGLU on separate gate / up projections, cast back to the input dtype.                |
| [`_init_weights`](#nemo_automodel-components-moe-layers-_init_weights)                 | -                                                                                                   |
| [`_shared_expert_stream`](#nemo_automodel-components-moe-layers-_shared_expert_stream) | Return the per-device side stream used for shared-expert overlap (created lazily, one per process). |

### Data

[`_SHARED_EXPERT_STREAMS`](#nemo_automodel-components-moe-layers-_SHARED_EXPERT_STREAMS)

### API

```python
class nemo_automodel.components.moe.layers.FakeBalancedGate(
    config: nemo_automodel.components.moe.config.MoEConfig,
    skip_first_n_experts: int = 0,
    noise: float = 0.0
)
```

**Bases:** `Module`

Load balanced gate implementation, spreads tokens uniformly across all experts.
The rationale for this class is to do performance experiments to understand
how the load imbalance with real data is impacting end-to-end performance.

When `noise > 0`, random perturbation is added to mimic realistic routing
imbalance.  A noise value of 0.0 gives perfectly balanced assignment, while
1.0 gives fully random expert selection and non-uniform weights.

**`bias_update_factor`** `= 0.0`

---

**`n_activated_experts`** `= config.n_activated_experts`

---

**`n_routed_experts`** `= config.n_routed_experts`

---

```python
nemo_automodel.components.moe.layers.FakeBalancedGate.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Forward pass for the gating mechanism.

**Parameters:**

**`x`** `torch.Tensor`

Input tensor.

---

**`token_mask`** `torch.Tensor`

Boolean mask indicating valid tokens.

---

**`cp_mesh`** `Optional[DeviceMesh]`

Device mesh for context parallel computation.

---

**Returns:** `torch.Tensor`

Routing weights for the selected experts.

```python
nemo_automodel.components.moe.layers.FakeBalancedGate.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
nemo_automodel.components.moe.layers.FakeBalancedGate.update_bias() -> None
```

```python
class nemo_automodel.components.moe.layers.Gate(
    config: nemo_automodel.components.moe.config.MoEConfig,
    gate_precision: torch.dtype | None = None
)
```

**Bases:** `Module`

Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.

**`_cumulative_expert_load`** `Tensor | None = None`

---

**`_last_aux_loss`** `Tensor | None = None`

---

**`_last_expert_load`** `Tensor | None = None`

---

**`_track_load_balance`** `bool = False`

---

**`aux_loss_coeff`** `= config.aux_loss_coeff`

---

**`bias`** `Optional[torch.nn.Parameter]`

Optional bias term for the gate.

---

**`bias_update_factor`** `= config.gate_bias_update_factor`

---

**`dim`** `int = config.dim`

Dimensionality of input features.

---

**`n_experts`** `= config.n_routed_experts`

---

**`n_groups`** `int = config.n_expert_groups`

Number of groups for routing.

---

**`norm_topk_prob`** `= config.norm_topk_prob`

---

**`route_scale`** `float = config.route_scale`

Scaling factor for routing weights.

---

**`router_replay`** `RouterReplay | None`

---

**`router_weight_uses_score_correction_bias`** `= config.router_weight_uses_score_correction_bias`

---

**`router_weights_fp32`** `= config.router_weights_fp32`

---

**`routing_core`** `= _GateRoutingCore()`

---

**`score_dtype`**

---

**`score_func`** `str = config.score_func`

Scoring function ('softmax', 'sigmoid', 'sigmoid\_with\_bias',
'softmax\_with\_bias', or 'sqrtsoftplus').

---

**`softmax_before_topk`** `= config.softmax_before_topk`

---

**`topk`** `int = config.n_activated_experts`

Number of top experts activated for each input.

---

**`topk_groups`** `int = config.n_limited_groups`

Number of groups to route inputs to.

---

**`train_gate`** `= config.train_gate`

---

**`weight`** `torch.nn.Parameter`

Learnable weights for the gate.

---

```python
nemo_automodel.components.moe.layers.Gate._compute_aux_loss(
    original_scores: torch.Tensor,
    expert_load: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> torch.Tensor
```

Computes the auxiliary loss for load balancing.

**Warning**: Assumes batch size = 1, if batch size > 1, the aux\_loss will
be computed across multiple sequences.

**Parameters:**

**`original_scores`** `torch.Tensor`

Original scores from the gating mechanism.
Shape is \[num\_tokens, num\_experts].

---

**`expert_load`** `torch.Tensor`

Load of each expert (number of tokens routed to each expert).
Shape is \[num\_experts].

---

**`token_mask`** `torch.Tensor`

Boolean mask indicating valid tokens.
Shape is \[num\_tokens].

---

**`cp_mesh`** `Optional[DeviceMesh]`

Device mesh for context parallel computation.

---

**Returns:** `torch.Tensor`

torch.Tensor: Auxiliary loss for load balancing.
Shape is \[].

```python
nemo_automodel.components.moe.layers.Gate._compute_expert_load(
    indices: torch.Tensor,
    token_mask: torch.Tensor
) -> torch.Tensor
```

Computes the load of each expert based on the selected indices.
Args:
indices (torch.Tensor): Indices of the selected experts.
Shape is \[num\_tokens, num\_activated\_experts].
token\_mask (torch.Tensor): Boolean mask indicating valid tokens.
Shape is \[num\_tokens].

**Returns:** `torch.Tensor`

torch.Tensor: Load of each expert (number of tokens routed to each expert).
Shape is \[num\_local\_experts].

```python
nemo_automodel.components.moe.layers.Gate._load_from_state_dict(
    state_dict: dict[str, torch.Tensor],
    prefix: str,
    local_metadata: dict[str, object],
    strict: bool,
    missing_keys: list[str],
    unexpected_keys: list[str],
    error_msgs: list[str]
) -> None
```

Load gate state while preserving the routing-bias DTensor layout.

**Parameters:**

**`state_dict`** `dict[str, torch.Tensor]`

State mapping whose correction bias is either a full tensor
of shape `[experts]` or an already distributed tensor.

---

**`prefix`** `str`

Prefix applied to this module's state-dict keys.

---

**`local_metadata`** `dict[str, object]`

Metadata saved for this module.

---

**`strict`** `bool`

Whether state-dict keys must match exactly.

---

**`missing_keys`** `list[str]`

Missing state-dict keys collected during loading.

---

**`unexpected_keys`** `list[str]`

Unexpected state-dict keys collected during loading.

---

**`error_msgs`** `list[str]`

State-dict loading errors collected during loading.

---

**Raises:**

* `RuntimeError`: If a full correction-bias tensor targets an incompatible
  shape or a DTensor that is not replicated on every mesh dimension.

```python
nemo_automodel.components.moe.layers.Gate._local_score_correction_bias() -> torch.Tensor | None
```

Return the local tensor used to adjust this rank's routing scores.

**Returns:** `torch.Tensor | None`

Tensor of shape \[experts] replicated across the gate's DP/CP mesh, or

```python
nemo_automodel.components.moe.layers.Gate._route_scores(
    scores: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Apply fixed-shape expert selection and probability math to router logits.

**Parameters:**

**`scores`** `torch.Tensor`

Tensor of shape \[tokens, experts] containing local router logits.

---

**Returns:** `torch.Tensor`

Tuple containing routing weights of shape \[tokens, activated\_experts],

```python
nemo_automodel.components.moe.layers.Gate.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Forward pass for the gating mechanism.

**Parameters:**

**`x`** `torch.Tensor`

Input tensor.

---

**`token_mask`** `torch.Tensor`

Boolean mask indicating valid tokens.

---

**`cp_mesh`** `Optional[DeviceMesh]`

Device mesh for context parallel computation.

---

**Returns:** `torch.Tensor`

Routing weights for the selected experts.

```python
nemo_automodel.components.moe.layers.Gate.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
nemo_automodel.components.moe.layers.Gate.update_bias() -> None
```

Updates the correction bias used in the gate based on the popularity of experts.
This function is a NoOp if the gate is not trained.

To avoid routing collapse, and to promote better load balance of experts,
DeepSeek-V3 uses a correction mechanism to adjust the scores of experts using
a learned bias parameter. The bias parameter is updated based on the popularity
of experts, i.e., the number of tokens routed to each expert. If an expert is
more popular than the average, its bias term is decreased, and vice versa.
This encourages the model to route tokens to less popular experts, promoting
better load balance.

```python
class nemo_automodel.components.moe.layers.MLP(
    dim: int,
    inter_dim: int,
    backend: str,
    dtype: torch.dtype = torch.bfloat16,
    activation: str = 'swiglu',
    bias: bool = False,
    swiglu_limit: float = 0.0
)
```

**Bases:** `Module`

Multi-Layer Perceptron (MLP) used as a feed-forward layer.

Supports both gated activations (SwiGLU) and simple activations (ReLU²).

**`down_proj`** `nn.Module`

Linear layer for hidden-to-output transformation.

---

**`gate_proj`** `nn.Module`

Linear layer for gate in gated activations (or up\_proj for simple).

---

**`is_gated`** `= is_gated_activation(activation)`

---

**`swiglu_limit`** `= float(swiglu_limit)`

---

**`up_proj`** `nn.Module`

Additional linear layer for gated activations (None for simple).

---

```python
nemo_automodel.components.moe.layers.MLP.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Forward pass for the MLP layer.

**Parameters:**

**`x`** `torch.Tensor`

Input tensor.

---

**Returns:** `torch.Tensor`

torch.Tensor: Output tensor after MLP computation.

```python
nemo_automodel.components.moe.layers.MLP.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.moe.layers.MoE(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Mixture-of-Experts (MoE) module.

**`cp_mesh`** `DeviceMesh | None = None`

---

**`dim`** `int = config.dim`

Dimensionality of input features.

---

**`experts`** `nn.ModuleList = GroupedExpertsMoK(config, backend)`

List of expert modules.

---

**`fc1_latent_proj`**

---

**`fc2_latent_proj`**

---

**`gate`** `nn.Module`

Gating mechanism to route inputs to experts.

---

**`n_activated_experts`** `int = config.n_activated_experts`

Number of experts activated for each input.

---

**`n_local_experts`** `int`

Number of experts handled locally in distributed systems.

---

**`n_routed_experts`** `int = config.n_routed_experts`

Total number of experts in the model.

---

**`shared_expert_gate`**

---

**`shared_experts`** `nn.Module`

Shared experts applied to all inputs.

---

```python
nemo_automodel.components.moe.layers.MoE._join_shared_experts(
    z: torch.Tensor,
    stream: torch.cuda.Stream
) -> None
```

staticmethod

Make the current stream wait for the side-stream shared-expert output `z`.

```python
nemo_automodel.components.moe.layers.MoE._shared_expert_output(
    x: torch.Tensor
) -> torch.Tensor
```

Shared-expert output for `x`, gated by `shared_expert_gate` when the model has one.

```python
nemo_automodel.components.moe.layers.MoE._shared_experts_on_side_stream(
    x: torch.Tensor
) -> tuple[torch.Tensor, torch.cuda.Stream]
```

Run the shared experts for `x` on the per-device side stream.

Returns the shared-expert output and the stream it was computed on; the caller runs the routed
path on the current stream and then joins with `_join_shared_experts`. Subclasses with their
own `forward` (Kimi K3) reuse both helpers so the stream choreography lives in one place.

```python
nemo_automodel.components.moe.layers.MoE.forward(
    x: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> torch.Tensor
```

Route tokens through shared and routed experts.

**Parameters:**

**`x`** `torch.Tensor`

Input tensor of shape `[..., hidden]`.

---

**`padding_mask`** `torch.Tensor | None` — default: None

Boolean tensor matching `x.shape[:-1]` where true
entries are padding.

---

**`cp_mesh`** `DeviceMesh | None` — default: None

Optional context-parallel mesh used by the router.

---

**Returns:** `torch.Tensor`

Tensor with the same shape and dtype as `x`.

```python
nemo_automodel.components.moe.layers.MoE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.moe.layers._GateRoutingCore()
```

**Bases:** `Module`

Parameterless, fixed-shape portion of a learned MoE router.

The gate projection remains eager so FSDP can unshard its DTensor parameters.
The owning gate is passed as a non-tensor control so this child remains an
independent, parameterless CUDA graph boundary and survives model copies.

```python
nemo_automodel.components.moe.layers._GateRoutingCore.forward(
    scores: torch.Tensor,
    gate: nemo_automodel.components.moe.layers.Gate
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Select experts and compute routing probabilities from projected scores.

```python
nemo_automodel.components.moe.layers._clamped_swiglu_act(
    gate: torch.Tensor,
    up: torch.Tensor,
    limit: float
) -> torch.Tensor
```

fp32 clamped SwiGLU on separate gate / up projections, cast back to the input dtype.

Same arithmetic as the eager branch of `MLP.forward` (the DeepSeek-V4 reference Expert.forward); inductor
fuses it into one kernel per direction. Dynamo's default shape policy: the first token count compiles a static
graph, the second recompiles once with a symbolic row count and every later shape reuses it — `dynamic=False`
with `fullgraph=True` would instead raise `FailOnRecompileLimitHit` after `recompile_limit` distinct
token counts (variable-length batches). CPU inputs keep the eager chain.

```python
nemo_automodel.components.moe.layers._init_weights(
    module,
    buffer_device: torch.device,
    init_std: float = 0.02
)
```

```python
nemo_automodel.components.moe.layers._shared_expert_stream(
    device: torch.device
) -> torch.cuda.Stream
```

Return the per-device side stream used for shared-expert overlap (created lazily, one per process).

```python
nemo_automodel.components.moe.layers._SHARED_EXPERT_STREAMS: dict[int, Stream] = {}
```