> 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 |
| ---------------------------------------------------------------------- | ----------- |
| [`_init_weights`](#nemo_automodel-components-moe-layers-_init_weights) | -           |

### 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`**

---

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

---

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

---

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

---

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

---

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

---

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

---

**`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_func`** `= config.score_func`

---

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

---

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

---

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

---

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

---

**`weight`**

---

```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`**

---

**`gate_proj`**

---

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

---

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

---

**`up_proj`**

---

```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`** `= config.dim`

---

**`experts`** `= GroupedExpertsMoK(config, backend)`

---

**`fc1_latent_proj`**

---

**`fc2_latent_proj`**

---

**`gate`**

---

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

---

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

---

**`shared_expert_gate`**

---

**`shared_experts`**

---

```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._init_weights(
    module,
    buffer_device: torch.device,
    init_std: float = 0.02
)
```