# How to Implement Mixture of Experts (MoE) with Expert Parallelism in Nanotron

> Learn to implement Mixture of Experts MoE with expert parallelism in Nanotron. Configure expert parallelism, create parallel context, and utilize Qwen2MoELayer for efficient MoE implementation.

- Repository: [Hugging Face/nanotron](https://github.com/huggingface/nanotron)
- Tags: how-to-guide
- Published: 2026-03-03

---

**To implement mixture of experts with expert parallelism in Nanotron, configure `expert_parallel_size` in your parallelism arguments, instantiate a `ParallelContext` to create the expert parallel group (`ep_pg`), and use the `Qwen2MoELayer` class which orchestrates token routing, fused expert computation via `GroupedMLP`, and cross-rank communication through `ops.permute` and `ops.unpermute` primitives.**

Nanotron is Hugging Face’s open-source framework for pre-training large language models using up to five dimensions of parallelism. When implementing mixture of experts (MoE) architectures that exceed single-GPU memory capacity, expert parallelism becomes essential to distribute expert parameters across devices while maintaining efficient token routing and communication.

## Configure the Expert Parallel Context

The foundation of expert parallelism in Nanotron is the `ParallelContext` class defined in [`src/nanotron/parallel/context.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/context.py). During initialization, lines 60-82 construct a 5-D rank matrix (`world_ranks_to_pg`) where the first axis represents the expert dimension. The code creates the **expert parallel group** (`ep_pg`) by reshaping the world ranks:

```python
ranks = np.arange(0, self.world_size).reshape(
    (self.expert_parallel_size,
     self.pipeline_parallel_size,
     self.data_parallel_size,
     self.context_parallel_size,
     self.tensor_parallel_size),
)
self.ep_pg = self.create_new_group(
    ranks.transpose((1, 2, 3, 4, 0)).reshape((-1, self.expert_parallel_size))
)

```

All MoE modules access `parallel_config.expert_parallel_size` to calculate **local expert counts**. The constraint `num_experts % expert_parallel_size == 0` must hold, ensuring each rank manages `num_local_experts = num_experts // expert_parallel_size` experts.

## Define MoE Architecture with MoEConfig

Located in [`src/nanotron/config/models_config.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/config/models_config.py), the `MoEConfig` dataclass (lines 41-48) stores hyperparameters that control the sparse expert layer:

```python
@dataclass
class MoEConfig:
    num_experts: int = 8               # Total experts in the model

    top_k: int = 2                      # Experts activated per token

    moe_intermediate_size: int = 1408   # Hidden dimension inside each expert

    enable_shared_expert: bool = False  # Optional shared expert pathway

```

These fields integrate into `Qwen2Config` and determine how `Qwen2MoELayer` partitions weights. When `enable_shared_expert` is true, the layer instantiates an additional **shared MLP** that processes every token alongside the routed experts, blending outputs via a learned `shared_gate` scalar.

## Route Tokens Using the Router Class

The `Router` class in [`src/nanotron/nn/moe.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/nn/moe.py) (lines 26-64) computes the probability distribution over experts. It stores a **float32** weight matrix of shape `[num_experts, hidden_dim]` and executes three operations during the forward pass:

1. **Gating**: `F.linear(hidden_states, self.weight)` produces logits for every expert.
2. **Normalization**: `F.softmax` converts logits to routing probabilities.
3. **Selection**: `torch.topk` selects the `k` highest-scoring experts per token (lines 55-58).

The router returns `int32` indices because the **grouped-gemm** kernel expects that data type for the subsequent permute operations.

## Execute Expert Computation with GroupedMLP

The `GroupedMLP` class (lines 66-78 in [`src/nanotron/nn/moe.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/nn/moe.py)) stores all local experts in two fused weight tensors:

- `merged_gate_up_proj`: `[num_local_experts, hidden_size, 2 * moe_intermediate_size]`
- `merged_down_proj`: `[num_local_experts, moe_intermediate_size, hidden_size]`

During the forward pass, the layer uses **grouped matrix multiplication** (`ops.gmm`) to compute all expert transformations in a single fused kernel:

```python

# Move token counts to CPU as required by the grouped-gemm backend

num_tokens_per_expert = num_tokens_per_expert.to("cpu")

merged_states = ops.gmm(hidden_states, self.merged_gate_up_proj, num_tokens_per_expert, trans_b=False)
gate_states, up_states = torch.split(merged_states, merged_states.shape[-1] // 2, dim=-1)
hidden_states = self.act(gate_states) * up_states
hidden_states = ops.gmm(hidden_states, self.merged_down_proj, num_tokens_per_expert, trans_b=False)

```

This approach eliminates separate loops over experts and maximizes GPU utilization through coalesced memory access.

## Assemble the Complete MoE Layer

`Qwen2MoELayer` (lines 101-199 in [`src/nanotron/nn/moe.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/nn/moe.py)) composes the routing and computation primitives into a unified block. The forward pass follows a strict dispatch-combine pattern:

```python
routing_weights, routing_indices = self.router(hidden_states)

# Dispatch: reorder tokens so each expert's tokens are contiguous

dispatched_inputs, inverse_permute_mapping, num_tokens_per_expert = self._dispatch_tokens(
    hidden_states, routing_indices
)

# Compute: run grouped MLP on local experts only

expert_outputs = self.experts(dispatched_inputs, num_tokens_per_expert)

# Combine: restore original token order and apply routing weights

output = self._combine_expert_outputs(
    expert_outputs["hidden_states"], inverse_permute_mapping, routing_weights
)

```

The `_dispatch_tokens` method internally calls `ops.permute`, which performs an **all-to-all** communication across `self.ep_pg` to send tokens to the ranks hosting their target experts. After computation, `ops.unpermute` reverses the permutation and scales outputs by the routing softmax weights.

## Launch Training with Expert Parallelism

To train an MoE model across multiple GPUs, use `torchrun` with a configuration file that specifies both expert and data parallelism. The repository provides [`examples/config_qwen_with_moe.yaml`](https://github.com/huggingface/nanotron/blob/main/examples/config_qwen_with_moe.yaml) as a reference:

```bash
torchrun --nproc_per_node=4 \
    examples/moe/train_moe.py \
    --config-file examples/config_qwen_with_moe.yaml \
    --master-port 12345

```

The YAML configuration must declare compatible parallelism dimensions:

```yaml
parallelism:
  expert_parallel_size: 4    # Must divide model_config.moe_num_experts

  tensor_parallel_size: 1
  pipeline_parallel_size: 1
  data_parallel_size: 2

model_config:
  moe_num_experts: 8
  moe_top_k: 2

```

When scaling across nodes, ensure the total world size equals the product of all parallel dimensions: `expert_parallel_size × tensor_parallel_size × pipeline_parallel_size × data_parallel_size × context_parallel_size`.

## Summary

- **5-D parallelism**: Nanotron treats expert parallelism as a first-class citizen alongside tensor, pipeline, data, and context parallelism through `ParallelContext` in [`src/nanotron/parallel/context.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/context.py).
- **Configuration**: Set `expert_parallel_size` to divide `num_experts` evenly; each rank stores `num_local_experts` parameters via `MoEConfig`.
- **Routing**: The `Router` class in [`src/nanotron/nn/moe.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/nn/moe.py) applies softmax and top-k selection to determine expert assignment for each token.
- **Computation**: `GroupedMLP` uses `ops.gmm` for fused expert matrix multiplication, requiring `num_tokens_per_expert` on CPU.
- **Communication**: `Qwen2MoELayer` handles cross-rank token movement through `ops.permute` (dispatch) and `ops.unpermute` (combine) across the `ep_pg` process group.

## Frequently Asked Questions

### What is expert parallelism in Nanotron and how does it differ from tensor parallelism?

Expert parallelism distributes distinct expert parameters across different GPUs, whereas tensor parallelism shards individual weight matrices within a layer. According to the source code in [`src/nanotron/parallel/context.py`](https://github.com/huggingface/nanotron/blob/main/src/nanotron/parallel/context.py), expert parallelism creates a dedicated process group (`ep_pg`) that manages **all-to-all** communication of tokens between ranks, allowing each GPU to compute only its assigned experts via `GroupedMLP` without storing full model weights.

### How do I choose the correct values for `num_experts` and `expert_parallel_size`?

The `num_experts` parameter in `MoEConfig` must be divisible by `expert_parallel_size` specified in `ParallelismArgs`. For example, with 8 total experts and `expert_parallel_size=2`, each rank hosts 4 local experts. This divisibility constraint ensures balanced load distribution and prevents uneven memory consumption across the expert parallel group.

### What is the role of `ops.permute` and `ops.unpermute` in the MoE implementation?

These primitives handle the **token dispatch** and **collection** phases of expert parallelism. In `Qwen2MoELayer`, `ops.permute` reorders the token tensor so that all tokens destined for the same expert are contiguous in memory before the grouped matrix multiplication, while `ops.unpermute` restores the original sequence order and applies the routing weights after expert processing completes.

### Can expert parallelism be combined with pipeline and tensor parallelism?

Yes, Nanotron supports **5-D parallelism** where expert parallelism operates orthogonally to other dimensions. The `ParallelContext` manages separate process groups for each dimension, allowing configurations such as `expert_parallel_size=4`, `tensor_parallel_size=2`, and `pipeline_parallel_size=2` to scale MoE models across hundreds of GPUs while maintaining efficient communication patterns.