4.8 MOE Series

4.8 MOE Series

4.8.1 Gshard

• Gate: The Gate module uses a linear layer to decide which expert a token should be routed to. In other MoE (Mixture of Experts) architectures, the Gate is sometimes referred to as the Router. The Gate has a shape of (M, E), where E is the number of experts. After passing input data of shape (S, M) through the Gate, we obtain a probability matrix of shape (S, E). This matrix represents, for each token, the probability of being routed to each expert. Since GShard uses the top2Expert routing mechanism, we only care about the top 2 experts with the highest probabilities for each token. In the illustration below, the dark color represents the expert with the highest routing probability, while the lighter color represents the second-highest. For example, for token 0, the probability of being routed to expert 0 is the highest, followed by expert 1.

• Expert Overflow Handling: To handle expert overflow, we define an expert buffer. Suppose we have 8 tokens and 4 experts, each expert should ideally receive 8 / 4 = 2 tokens. Considering we're using the top2Expert strategy, the upper bound of tokens routed to each expert would be (8 / 4) × 2 = 4. This value is also the buffer size shown in the diagram. The capacity can be defined as:

• Random Routing: The 1st expert is always selected deterministically. However, the 2nd expert selection includes some randomness. Here's how it works: Given the probability (or more precisely, logit) for all experts, we randomly sample 4 noises from a certain distribution. Then, noise is added to these 4 logits (each corresponds to an expert). After that, the logit corresponding to the already selected 1st expert is masked (excluded), and the maximum of the remaining 3 is selected as the 2nd Expert. Once we have selected the top 2 experts for a token, they can process this token. Finally, we retrieve the corresponding original (non-noised) probabilities, and normalize them as follows:
Here, and are treated as weights. After the token is processed independently by expert 0 and expert 1, each producing one output token, the final output token is a weighted combination of the two outputs using the and .

• Token Overflow: If only one expert overflows, the token's weight for another expert is set to 1, and then compute the weighted computation as usual. If both experts overflow, the token is not routed to any expert. Instead, it bypasses the expert layer and directly connects to the next layer (e.g., Attention) via a residual connection.

• Auxiliary Loss: In addition to the capacity and random routing mechanisms, GShard also introduces an auxiliary loss to encourage a balanced expert load. It is defined as:
Where: E is the number of experts. is the number of tokens currently assigned to an expert's buffer (typically tokens for which the expert is selected as the 1st expert). S is the total number of tokens. is the average weight of the tokens assigned to the expert (again, only counting tokens for which this expert was chosen as the 1st expert)

• Zero Padding: When there are unused buffer slots in an expert, they are filled with zeros. This is called Zero Padding.

• Drop Tokens: When overflow occurs, not all tokens are guaranteed to be processed by experts. This behavior is referred to as Drop Tokens.

4.8.2 Mistral

• Model Architecture of One Layer



After each token passes through the Attention layers and residual connections, it is routed—via the Gating/Router mechanism—into two expert FFNs. The outputs of these experts are then aggregated (weighted sum), followed by another residual connection to produce the final output of the current layer. This structure breaks a large model into multiple “experts,” each responsible for handling a specific aspect of the input data.
• Specifically, each expert specializes in processing a certain type of input. For example, one expert might focus on syntactic structure, while another specializes in understanding semantic content.
• The key component here is the gating mechanism, which determines which experts are selected for which inputs. This selection is made dynamically based on the characteristics of the input data.
• At each layer, for each token, the model chooses two experts (from a total of 8) using a gating network, and the outputs from those two “experts” are combined via a weighted sum.
• By using MoE, the Transformer model can efficiently scale to large datasets or complex tasks, since only the two most relevant experts are activated per token.

• Improvements
• Sentence-Level Routing: Routing is applied per sample.
• Token-Level Routing: Routing is applied per token within each sample.
• Task-Level Routing: Different experts are explicitly responsible for different tasks.
• Uses the same SwiGLU architecture similar to the expert function . With K = 2, each token is routed to two SwiGLU modules with different weights. Overall, the output y for input token x is computed as:

• MoE Parallel Training: Usually adopts EP (Expert Parallelism) + DP (Data Parallelism). Experimental setup: 16 GPUs. For non-MoE parts: use TP (Tensor Parallelism) + DP. For MoE parts: use EP + DP.
a. Distributed Initialization
Assume that each MoE layer contains some experts. We now want to distribute these experts across multiple GPUs. The most intuitive method is to first decide how many GPUs we want to use to hold one full set of experts (EP - Expert Parallelism). From that, we can determine how many such expert replicas exist in total (DP - Data Parallelism). The code below shows how to configure this expert and data parallel distribution:
As shown below, we have 4 GPUs for all experts, and 4 such replicas, in total 16 GPUs.



b. FWD and BWD Process
There are 8 mini-batches. Goal: After performing FWD (Forward) and BWD (Backward) passes across 16 GPUs, the results should be equivalent to the result of processing the combined large batch (composed of those 8 mini-batches) on a single GPU.



• In the FWD (forward) phase, the data first passes through the non-MoE (Attention) layers. Since each TP group produces identical outputs for each GPU shard, the triangle-colored batches remain unchanged in terms of distribution. We move these triangle indicators to the corresponding non-MoE submodules to indicate that these tokens have already been seen by that submodule during the FWD pass.
• Continuing with FWD, now the data arrives at the MoE layer. Assuming the token-to-expert mapping is already computed, we only need to send each token to the correct expert via all2all communication within the corresponding ep_group—this is the role of the ep_group.
We’ll cover the technical details of all2all later, but for now, just remember that after this communication, each expert (e.g., e0) may have received data from multiple batches (e.g., both blue and orange). Each expert computes its own output, and then another all2all operation is performed to return the processed token results to the original GPUs within the same ep_group. After this, FWD continues with the usual non-MoE → MoE pipeline until FWD finishes.
• After FWD, we move on to BWD (backward). Let’s again take expert e0 as an example. According to distributed training strategy, we must allreduce the gradients computed from the 8 batches before updating weights of e0. The 8 batches are located in the same ep_dp_group. So during BWD, we perform an allreduce over the gradients for expert e0 within the ep_dp_group, and then use the result to update e0’s parameters. Now the role of ep_group and ep_dp_group should be clearer.

• MOE Inference Optimization
• Sliding Window Attention: The reason cache requires more memory is that our Attention is in causal decoder form—that is, each token attends to all previous tokens. As a result, the amount of data stored in the cache grows linearly with seq_len. But what if we change the way we think? Suppose each token only attends to the previous W tokens (including itself). In that case, the cache size could remain constant at W. Intuitively, this also makes sense: for the current token, tokens that are farther away contribute less useful information, so it may not be worth spending resources attending to those distant tokens.



For layer3 token9, although at each individual layer, it can only "see" part of the preceding tokens in the sequence (limited to a local window), as long as the model is deep enough, it will eventually be able to see all previous tokens at some layer. This is similar to the concept of a receptive field in CNNs.
• Rolling Buffer Cache



It’s easy to observe that the storage position of the i-th token in the prompt within the KV cache is: i % W
• Chunking: Long prompts have significant memory pressure. An intuitive solution is to split the prompt into multiple chunks, feed the model one chunk at a time, and update the KV Cache once per chunk. Typically, we set: chunk_size = cache_window = sliding_window = W. This means the size of each chunk and the size of the cache both match the sliding window size W.

4.8.3 Switch Transformer

• Innovations:
• MoE to Dense: MoE models that are trained to perform well can be distilled into dense models. Remarkably, even after reducing 99% of the MoE model's parameters, the distilled dense model outperforms a dense model trained from scratch.
• Training and Fine-tuning Techniques:
• For the first time, it successfully trains MoE models using bfloat16 (bf16) precision.
• Improved model initialization tailored to the MoE structure.
• Increased expert regularization that improves sparse model fine-tuning and multi-task training.
• Multilingual: When trained on multilingual datasets, Switch Transformer showed consistent performance improvements across 101 languages.

• Model Architecture Design
The overall architecture is like replacing the Feed-Forward Network (FFN) in each Transformer layer with a Mixture of Experts (MoE) layer. In Switch Transformer, the gating mechanism is simplified to select only one expert per forward pass, i.e., k = 1. This type of MoE layer is referred to as a Switch layer.

• Load Balancing
• If a token is assigned to an expert that has already reached its capacity, an overflow occurs. In that case, the token is not processed by the current layer but instead passed through via a residual connection to the next layer. This is similar to how GShard handles overflows.
• In Switch Transformer, the number of tokens each expert can handle is controlled by a capacity factor. The formula is:
• Note: A larger capacity factor means each expert can process more tokens, which reduces the chance of overflow. However, it also increases computational and communication load. Therefore, this factor must be carefully balanced.
• Setting Expert Capacity:
Empirically, maintaining a low token drop rate is crucial for model scaling. To train large-scale models, this issue must be addressed. By using a load balancing loss, the model can maintain good balance, minimizing overflow even with a relatively small capacity factor — thereby achieving a good tradeoff between performance and efficiency.
• The key issue is how to design the load balancing loss. Given N experts and a batch B containing T tokens, the load balancing loss is calculated as:
Where: is the fraction of tokens assigned to the i-th expert (not differentiable):
is the total probability assigned to the i-th expert over the batch (differentiable):
This loss function is designed the same way as in GShard.
• When token assignments are perfectly balanced, the vectors f and P align, and the load balancing loss reaches its minimum.
• The coefficient is experimented between 1e-5 and 1e-1. Empirically, a value like 1e-2 is sufficient to maintain load balance without hindering model convergence.
• Notably, under uniform token distribution, we have . This is why the loss is scaled by a factor of N, which ensures the loss remains consistent regardless of how many experts are used.