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

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

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, the MoEConfig dataclass (lines 41-48) stores hyperparameters that control the sparse expert layer:

@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 (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) 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:

# 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) composes the routing and computation primitives into a unified block. The forward pass follows a strict dispatch-combine pattern:

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 as a reference:

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:

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

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}