A few days ago, I was suddenly assigned a task to help a few other teams optimize the inference for the full 671B version of DeepSeek-R1. At this stage, we're mainly making some simple tweaks on SGlang and vLLM, since resources for the H20 cards in China are quite limited.
Recently, I've noticed that both the SGlang and vLLM teams are competing fiercely. vLLM 0.7.2 has introduced optimizations for Triton MLA and FusedMoE, plus it has the advantage of pp parallelism over SGlang. Of course, SGlang is also developing pp parallelism, while vLLM is working on MTP. Over the next two to three weeks, there's room for further performance improvements in both frameworks.
After a few days of inference work, I've gained a better understanding of DeepSeek MoE. Thanks to suggestions from the DeepSeek team, I realized my earlier understanding of MoE Group Limit had some errors, and now I'm writing a detailed note on MoE.
Additionally, getting hands-on experience reveals many ingenious engineering details that one can miss by just reading papers. Many might think, "It's just MoE, I know about it too," but by not diving deeper, they miss out on these intricate details.
One important thing to note is that with DeepSeek's fine-grained MoE processing, the actual activation for the 671B model is only 37B, which has led to some interesting approaches in distributed inference systems. The extreme resource elasticity in training and inference integration could eventually lead to life-time learning/training, potentially disheartening other players in the field.
TOC as below:
1. Optimization Space of Transformer Model
2. Basic Working Principle of Sparse MoE
3. Expert Load Balancing and AuxLoss
4. DeepSeek-V1 MoE
4.1 Fine-Grained Expert Segmentation
4.2 Shared Expert Isolation
4.3 Expert Load Balancing
5. DeepSeek-V2
5.1 Device-Limited Routing
5.2 Communication Load Balancing Loss
5.3 Token Drop Strategy
6. DeepSeek-V3
6.1 Gating Function Using Sigmoid
6.2 Expert Grouping
6.3 Load Balancing Without Auxiliary Loss Function
6.4 No Need for Device-Limited Routing and Token Drop
6.5 Modifications to AlltoAll Infrastructure
6.6 MoE Handling During Inference
6.7 Suggestions for Infrastructure Improvement
7. On the Evolution of MoE
1. Optimization Space of Transformer Model
From the earliest Transformer architecture, the computational complexity for the Attention Block is $O(N^2d)$, while for the MLP Block it is $O(Nd^2)$. Thus, algorithmic optimizations for larger model scales naturally focus on these two blocks, such as MHA for the Attention Block by DeepSeek MLA and Stepfun MFA. Many optimizations initially target long-context over N, which I will discuss separately in another document.
As for MoE optimizations, the open-source ecosystem primarily began with Mistral's Mixtral 8x7B. However, unfortunately, several major companies initially opted for Dense MLP. A few members of the DeepSpeed team moved from Microsoft to Snowflake and developed DenseMoE, which attempted to solve some communication issues by parallelizing MoE block and Attention Block. Unfortunately, it doesn't seem to have gained much traction afterwards.
Looking back at the titles of the three papers from DeepSeek now, there is a clear continuity.
- 《DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models》
- 《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》
- 《DeepSeek-V3 Technical Report》
The first paper, "Towards the Ultimate Specialization of Expert MoE Models," initially chose the sparse MoE model approach, aiming for maximum performance while maintaining smaller activation parameters.
However, existing work typically involves a small number of experts, and during training, the expertise of these experts can be overwhelmed by an excessive number of tokens. Figuratively speaking, a large amount of knowledge results in experts learning very diverse content, which compromises the density of information they can hold. This affects the Expert Specialization. From this perspective, the approach introduced Fine-Grained Expert Segmentation and utilized Shared Expert Isolation to absorb common knowledge, thereby reducing parameter redundancy in other experts. The final model is constructed as shown in the diagram below.
The V1 model was constructed with 64 experts, with each token activating 6 experts, and included 2 shared experts.
In the second iteration, DeepSeek-V2, the model was further expanded and strengthened by increasing the total number of routing experts to 160, still selecting 6 experts per token activation and maintaining 2 shared experts. Load balancing was further optimized, and MLA was introduced to achieve "Economical and Efficient" performance.
The third version, DeepSeek-V3, represents a culmination, expanding the number of experts to 256 with 8 selected per activation. There were further optimizations in load balancing and numerous infrastructure-based communication optimizations, such as DualPipe during training and offloading in MoE communication. This effectively enabled cross-node all-to-all (A2A) communication, as described in section 3.2.2 of the paper, where the model design considers the bandwidth ratio of NVLink to IB. It also uses PTX and auto-tunes chunk-size in communication to reduce L2Cache usage and minimize impact on other SMs.
It's important to stress that a model's distinction doesn't lie solely in whether MoE is used; there are many nuances in MoE details. These small differences result in significant cumulative changes. In hindsight, although no single step might seem remarkable, there's a fitting saying in tribute to DeepSeek: "A good player has no dazzling moves across the board."
2. Basic Working Principle of Sparse MoE
The first MoE model in the open-source ecosystem is likely Mixtral's 8x7B. The model structure is as follows:
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
print(model)
MixtralForCausalLM(
(model): MixtralModel(
(embed_tokens): Embedding(32000, 4096)
(layers): ModuleList(
(0-31): 32 x MixtralDecoderLayer(
(self_attn): MixtralAttention(
...
)
(block_sparse_moe): MixtralSparseMoeBlock(
(gate): Linear(in_features=4096, out_features=8, bias=False)
(experts): ModuleList(
(0-7): 8 x MixtralBLockSparseTop2MLP(
(w1): Linear(in_features=4096, out_features=14336, bias=False)
(w2): Linear(in_features=14336, out_features=4096, bias=False)
(w3): Linear(in_features=4096, out_features=14336, bias=False)
(act_fn): SiLU()
)
)
)
...Each Expert block is essentially a standard structure.
class Expert(nn.Module):
def __init__(self, dim: int, inter_dim: int):
super().__init__()
self.w1 = nn.Linear(dim, inter_dim)
self.w2 = nn.Linear(inter_dim, dim)
self.w3 = nn.Linear(dim, inter_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.w2(F.silu(self.w1(x)) * self.w3(x))sMoE implementation as follow:
class SMoE(nn.Module):
def __init__(self, args):
super().__init__()
self.hidden_dim = args.dim
self.ffn_dim = args.moe_inter_dim
self.num_experts = args.n_routed_experts
self.top_k = args.n_activated_experts
# gating
self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
self.experts = nn.ModuleList([Expert(self.hidden_dim,self.ffn_dim ) \
for _ in range(self.num_experts)])Since we need to analyze the Gating function, we'll discuss the specific forward function in SMoE a bit later. This establishes the structure of sMoE.
args = ModelArgs()
args.dim = 4096
args.inter_dim = 14336
args.n_routed_experts = 8
args.n_activated_experts =2
smoe =SMoE(args)
print(smoe)
MoE(
(gate): Linear(in_features=4096, out_features=8, bias=False)
(experts): ModuleList(
(0-7): 8 x Expert(
(w1): Linear(in_features=4096, out_features=512, bias=True)
(w2): Linear(in_features=512, out_features=4096, bias=True)
(w3): Linear(in_features=4096, out_features=512, bias=True)
)
)
)Its Gating function computation is as follows: map through a linear layer to N experts, then apply softmax and select the top K.
tokens = 13
x = torch.randn(1, tokens, args.dim)
scores = F.softmax(smoe.gate(x.view(-1, args.dim)), dim=-1)
weights, indices = torch.topk(scores, smoe.top_k, dim=-1)
We could plot softmax score as below:
plt.plot(scores.detach().to('cpu')[0].numpy())
plt.plot(scores.detach().to('cpu')[1].numpy())
plt.plot(scores.detach().to('cpu')[2].numpy())
Indices represent the selected expert numbers. For example, Token 0 selects experts 4 and 1, Token 1 selects experts 0 and 6, and Token 2 selects experts 2 and 3.
indices.T
tensor([[4, 0, 2, 1, 4, 6, 2, 3, 7, 3, 3, 7, 5],
[1, 6, 3, 0, 2, 3, 1, 0, 2, 0, 0, 6, 7]])
#Plot as below
r = list()
for i in range(3):
r.append(np.zeros(args.n_routed_experts))
for item in indices[i].numpy():
r[i][item] = 1
plt.plot(r[i])
The subsequent computation is as shown in the diagram: tokens are routed to the experts based on the output indices, computed, multiplied by weights, and then summed together.
3. Expert Load Balancing and AuxLoss
However, there can be significant load imbalance among experts, which can lead to some experts being overloaded with information while others are undertrained, causing expert routing to collapse.
For example, we scaled up to 256 experts, selecting 8 for a test.
args.n_routed_experts = 256
args.n_activated_experts = 8
smoe1 =SMoE(args)
tokens = 1024
x = torch.randn(1, tokens, args.dim)
scores = F.softmax(smoe1.gate(x.view(-1, args.dim)), dim=-1)
weights, indices = torch.topk(scores, smoe.top_k, dim=-1)
counts = torch.bincount(indices.flatten(), minlength=args.n_routed_experts)
plt.plot(counts.detach().to('cpu').numpy())
Therefore, a simple idea is to define an auxiliary loss function (Aux_Loss). For example, given $S$ tokens and $E$ experts, the loss is defined as the sum of the variances of the tokens received by each expert.
$$ \mathcal l_{aux}=\frac{1}{E}\sum_{e=1}^E (\frac{c_e}{S})^2$$
avg_counts = counts/tokens
loss = (avg_counts * avg_counts).sum()/args.n_routed_experts
However, such an auxiliary loss function does not include the parameters of the Gating function, making it impossible to train through gradient updates. To address this, Google introduced an approach in GShard by replacing one component of the squared term with the mean of the Gating softmax.
$$\mathcal g_{S,E}= softmax(wg \cdot x_S) \quad
\mathcal m_E = \frac{1}{S}\sum_{s=1}^S \mathcal g_{S,E}
$$
$$ \mathcal l_{aux}=\frac{1}{E}\sum_{e=1}^E \frac{c_e}{S}\cdot m_e$$
m = scores.mean(dim=0)
avg_counts = counts/tokens
loss_aux = (m * avg_counts).mean()
4. DeepSeek-V1 MoE
Although there are some load balancing algorithms, when the number of experts is limited, the tokens assigned to a specific expert might cover different types of knowledge. The designated expert will tend to learn vastly different types of knowledge within its parameters, which are difficult to utilize effectively at the same time. If each token can be routed to more experts, different types of knowledge could be decomposed and learned by different experts. In this case, each expert can maintain a high level of specialization, promoting more concentrated distribution of knowledge among the experts.
The work of DeepSeek-V1 MoE primarily focuses on making experts more fine-grained and specialized (Towards Ultimate Expert Specialization), and it includes several aspects of work:
4.1. Fine-Grained Expert Segmentation
By partitioning experts into finer-grained segments while keeping the same amount of expert parameters and computational cost, the activation of expert combinations becomes more flexible and adaptable. For example, dividing each expert FFN in a MoE into $m$ smaller experts and reducing the FFN's inter-dim to $1/m$ of the original, the number of activated experts can increase by $m$ times at the same computational cost. Originally, there might be $C_{16}^2= 120$ combinations of expert selections, but with $m=4$, it can expand to $C_{64}^8=4,426,165,368$ combinations. This significantly enhances the accuracy and targeted knowledge acquisition capabilities.
4.2 Shared Expert Isolation
In traditional routing strategies, tokens assigned to different experts might require some common knowledge or information, causing different experts to learn shared knowledge within their parameters, leading to redundancy. If there are dedicated shared experts responsible for capturing and integrating shared knowledge from different contexts, the parameter redundancy among other routing experts would be alleviated, resulting in a more parameter-efficient model with clearer expert specialization.
To achieve this, in addition to the fine-grained expert segmentation strategy, DeepSeek further isolates $Ks$ experts as shared experts. Regardless of how the routing module assigns tokens, each token is deterministically assigned to these shared experts. To maintain constant computational cost, the number of activated experts among other routing experts is reduced by $Ks$. With the shared expert isolation strategy, the complete DeepSeekMoE architecture is as follows:
$$h_t^t=\sum_{i=1}^{Ks}FFN_i(u_t^l)+\sum_{i=Ks+1}^{mN}(g_{i,t}FFN_i(u_t^l))+u_t^l$$
$$g_{i,t}=
\begin{cases}
s_{i,t}& s_{i,t} \in TopK(s_{j,t}|Ks+1 \leq j \leq mN, mK-Ks) \\
0& Otherwise
\end{cases}
$$
$$s_{i,t} = Softmax_i({u_t^l}^Te_i^l)$$
Shared Expert sum : $\sum_{i=1}^{Ks}FFN_i(u_t^l)$
Routed Expert sum : $\sum_{i=Ks+1}^{mN}(g_{i,t}FFN_i(u_t^l))$
4.3 Expert Load Balancing
If routing strategies are fully automatically constructed through learning, load imbalance issues may arise. On one hand, there is a risk of expert routing collapse, where the model consistently selects only a few experts, leaving others undertrained. On the other hand, if experts are distributed across multiple devices, this can lead to computational load imbalance, further impacting the Model Fully Utilized (MFU) rate of the entire training cluster. Therefore, DeepSeek undertook two main actions.
4.3.1 Expert-Level Balance AuxLoss
To prevent routing collapse, an expert-level load balancing loss function is defined, where $N' = mN-Ks$ is the number of routing experts, $K'=mK-Ks$ is the number of activated fine-grained experts, and $T$ is the number of tokens that need to be processed by the experts. $\alpha_1$ is a hyperparameter.
$$\mathcal L_{ExpBalance}=\alpha_1 \sum_{i=1}^{N'} f_iP_i$$
$$f_i = \frac{N'}{K'T}\sum_{t=1}^T \mathbb 1 (Token t selects Expert i)$$
$$P_i = \frac{1}{T}\sum_{t=1}^Ts_{i,t}$$
Compared to the auxiliary loss function in Gshard, a correction factor of $N'/K'$ is applied when calculating $f_i$.Since the probability of selecting each expert is $K'/N'$ when uniformly distributing across $N'$ experts choosing $K'$,this correction makes the overall loss function independent of the expert selection strategy.
4.3.2 Device-Level Balance AuxLoss
On the other hand, device-level load balancing is performed to ensure experts are evenly routed to multiple devices, keeping the computational load relatively balanced and avoiding long-tail effects. Therefore, experts are divided into $D$ group ${\mathcal E_1,\mathcal E_2,...,\mathcal E_D}$ with each group deployed on a single device. The loss is calculated as follows.
$$\mathcal L_{DevBalance}=\alpha_2 \sum_{i=1}^{D} f_i'P_i'$$
$$f_i' = \frac{1}{|\mathcal E_i|}\sum_{j\in\mathcal E_i}f_j$$
$$P_i' = \sum_{j\in\mathcal E_i}P_j$$
4.3.3 AuxLoss hyperparameter
In the design of the loss hyperparameters $\alpha_1$ and $\alpha_2$ , the expert load balancing loss parameter is set relatively small, while the device load balancing loss hyperparameter is set larger to better balance the load across devices.
4.3.4 Code overview
On HuggingFace, there is a DeepseekMoE function. We noticed that the implementation of the MoEGate function differs somewhat from the paper. The number of tokens $T$ is used to calculate the load balancing loss for all tokens within a batch. Additionally, in the 16B model, the cross-device loss function is not executed.
import torch.nn.init as init
import math
batch_size = 5
tokens = 1024
x = torch.randn(batch_size, tokens, args.dim)
gate_weight = nn.Parameter(torch.rand(args.n_routed_experts, args.dim))
init.kaiming_uniform_(gate_weight, a=math.sqrt(5))
### calculate over the entire batch
bsz, seq_len, h = x.shape
hidden_states = x.view(-1, h)
logits = F.linear(hidden_states, gate_weight, None)
scores = logits.softmax(dim=-1)
### TopK and normalization
topk_weight, topk_idx = torch.topk(scores, k=args.n_activated_experts,dim=-1, sorted=False)
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
topk_weight = topk_weight / denominator
The computation of the Expert-Level auxiliary loss can be configured in two ways, with the default being calculated for each sequence within the batch.
scores_for_aux = scores
aux_topk = args.n_activated_experts
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
alpha=0.001
if seq_aux: ### based on seq calculate score
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
ce = torch.zeros(bsz, args.n_routed_experts)
ce.scatter_add_(1, topk_idx_for_aux_loss, torch.ones(bsz, seq_len * aux_topk)).div_(seq_len * aux_topk / args.n_routed_experts)
aux_loss = (ce * scores_for_seq_aux.mean(dim = 1)).sum(dim = 1).mean() * alphaAnother point to note is that the current DeepSeek code allows for the use of MLP-Dense in the initial layers. This approach benefits model stability by avoiding some stability issues caused by early-stage Attention entering the MoE. However, in the publicly available DeepSeek-MoE-16B model, this feature is disabled, meaning every layer is MoE.
class DeepseekDecoderLayer(nn.Module):
def __init__(self, config: DeepseekConfig, layer_idx: int):
self.self_attn = Deepseek_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
self.mlp = DeepseekMoE(config) if (config.n_routed_experts is not None and \
layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0) \
else DeepseekMLP(config)5. DeepSeek-V2
DeepSeek-V2 further expands the fine-grained expert selection by adopting a scheme of routing 160 experts, selecting 6 + 2 shared experts. Compared to DeepSeek-V1-MoE, it adds two new communication-related constraints.
Note: Some skepticism remains about this approach, primarily due to the impact of Mellanox network devices, where a large QP-Scale and the incast caused by all-to-all lead to increased communication costs. If this issue is resolved, fewer restrictions may be necessary. But, DeepSeek tests have shown that as long as the number of grouped devices reaches a certain threshold $M \geq 3$ , there is no performance difference compared to the original TopK.
5.1 Device-Limited Routing
The first approach is device-constrained routing to limit MoE-related communication costs. This mainly involves ensuring that experts routed for parallel processing are distributed across multiple devices. When the number of experts is particularly large, token communication within a single batch will span numerous devices, increasing the communication cost for expert parallelism. Thus, DeepSeek-V2 introduces a constraint that limits each token to be routed to a maximum of $M$ devices. Specifically, for each token, the top $M$ devices containing experts with the highest affinity scores are selected. Then, top-K selection is performed among the experts on these $M$ devices. In practice, it is found that when
$ M \geq 3$, device-constrained routing achieves performance roughly equivalent to unconstrained top-K routing.
For a specific code implementation, you can check the source code on HuggingFace. The calculation of the gating function is unchanged, still using softmax across the entire batch, but FP32 is used for computational precision.
import torch.nn.init as init
import math
batch_size = 5
tokens = 1024
x = torch.randn(batch_size, tokens, args.dim)
gate_weight = nn.Parameter(torch.rand(args.n_routed_experts, args.dim))
init.kaiming_uniform_(gate_weight, a=math.sqrt(5))
### calculate for entire batch
bsz, seq_len, h = x.shape
hidden_states = x.view(-1, h)
### gating function with FP32
logits = F.linear(hidden_states.type(torch.float32), gate_weight.type(torch.float32), None)
scores = logits.softmax(dim=-1, dtype=torch.float32)
Before performing TopK and normalization selection, MoE Groups are computed, dividing the experts into 8 groups in total. Then, 3 groups are selected using TopK. For each group, the maximum softmax value is calculated as the group's score, and then $M$ groups are selected from these.
n_group = 8
topk_group = 3
### based on each token, calculate max score in each group
group_scores = (
scores.view(bsz * seq_len, n_group, -1).max(dim=-1).values
) # [n, n_group]
### select M Group
group_idx = torch.topk(
group_scores, k=topk_group, dim=-1, sorted=False
)[
1
] # [n, top_k_group]Then build gourp mask , after mask the score, execute topK
group_mask = torch.zeros_like(group_scores) # [n, n_group]
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
score_mask = (
group_mask.unsqueeze(-1)
.expand(
bsz * seq_len, n_group, args.n_routed_experts // n_group
)
.reshape(bsz * seq_len, -1)
) # [n, e]
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
topk_weight, topk_idx = torch.topk(
tmp_scores, k=args.n_activated_experts, dim=-1, sorted=False
)By examining the distribution of scores and tmp_scores, it can be seen that the softmax values for other groups are masked to 0. The blue represents the original softmax, while the yellow shows the values after applying the group mask.
plt.plot(scores.detach().to('cpu')[1].numpy())
plt.plot(tmp_scores.detach().to('cpu')[1].numpy())
Then run topK function with norm.
topk_weight, topk_idx = torch.topk(
tmp_scores, k=args.n_activated_experts, dim=-1, sorted=False
)
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
topk_weight = topk_weight / denominator
5.2 Communication Load Balancing Loss
Although Device-Limit Routing can address the communication domain issue and reduce the communication fanout, some receiving devices may still experience concentration of activations on a few experts, leading to communication bottlenecks. Therefore, a communication load balancing loss is introduced, where $D$ presents the number of devices, $T$ represents the number of tokens,$M$ represents the number of selected groups, and $\mathcal E_i$ represents the i-th expert.
$$\mathcal L_{CommBalance}=\alpha_3 \sum_{i=1}^{D} f_i''P_i''$$
$$f_i'' = \frac{D}{MT}\sum_{t=1}^T\mathbb 1 \text{(Token t is sent to Device i)}$$
$$P_i'' = \sum_{j\in\mathcal E_i}P_j$$
$f_i''$ apply a factor $D/M$ to make sure the loss is independent with $M$ group number and $D$ device number
5.3 Token Drop Strategy
Although two new load balancing strategies have been introduced, to ensure training efficiency and avoid long-tail effects, the average number of tokens per device is calculated from the total number of tokens in a batch, serving as the device capacity limit. When the actual capacity exceeds this limit, the softmax gating scores are sorted in descending order, and the computations for the MoE expert networks are skipped for the portion exceeding the capacity, directly moving to the combine phase.
6. DeepSeek-V3
DeepSeek-V3 continues with the original architecture of fine-grained Experts and Shared Experts, further increasing the number of experts. The V3 config file indicates that there are three model sizes.
| Model |
RouteExperts(Activate Experts) |
shared Expert |
expert_grp |
grp_limit |
dim |
MOE inter_dim |
| 16B |
64(6) |
2 |
- |
- |
2048 |
1408 |
| 236B |
160(6) |
2 |
8 |
3 |
5120 |
1536 |
| 671B |
256(8) |
1 |
8 |
4 |
7168 |
2048 |
6.1 Gating Function Using Sigmoid
Additionally, according to the configuration file, in the 671B model, the gating function has been changed from softmax to sigmoid, with normalization applied.
$$h_t'=\sum_{i=1}^{Ns}FFN_i^{(s)}(u_t)+\sum_{i=1}^{Nr}(g_{i,t}FFN_i^{(r)}(u_t))+u_t$$
$$g_{i,t} = \frac{g_{i,t}'}{\sum_{j=1}^{Nr}g_{j,t}'}$$
$$g_{i,t}'=
\begin{cases}
s_{i,t}& s_{i,t} \in TopK(s_{j,t}|1 \leq j \leq Nr, Kr) \\
0& Otherwise
\end{cases}
$$
$$s_{i,t} = Sigmoid({u_t}^Te_i)$$
Here, I personally speculate that when the number of experts is larger, using Sigmoid instead of Softmax might be intended to widen the range of expert scores. This is because, in DeepSeek-V2, FP32 precision was already used, and calculating softmax with more experts would result in very small values, increasing both the distinction and computational errors in gating scores, leading to significant errors in expert selection. Therefore, the Sigmoid function is used, along with a bias added to the gating mechanism. Refer to the code on GitHub.
class Gate(nn.Module):
def __init__(self, args: ModelArgs):
...
self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
### dim 7168 is for 671B model only, it add Bias
self.bias = nn.Parameter(torch.empty(args.n_routed_experts)) if self.dim == 7168 else NoneThe score function as below:
args = ModelArgs()
args.n_routed_experts = 256
args.n_activated_experts = 8
args.n_expert_groups = 8
args.n_limited_groups = 4
args.route_scale = 2.5
args.dim = 6178
batch_size = 7
tokens = 1024
x = torch.randn(batch_size, tokens, args.dim)
### in class MoE, it flatten the input token.
x = x.view(-1, args.dim)
gate_weight = nn.Parameter(torch.rand(args.n_routed_experts, args.dim))
gate_bias = nn.Parameter(torch.rand(args.n_routed_experts))
init.kaiming_uniform_(gate_weight, a=math.sqrt(5))
To verify the previously speculation , we calculate softmax here to compare with sigmoid
softmax_score = scores.softmax(dim=-1, dtype=torch.float32)
plt.plot(softmax_score.detach().to('cpu')[0][0].numpy())
DeepSeek-V3 671B adopt sigmoid , the score is in range [0,1] even larger than softmax
scores = scores.sigmoid()
plt.plot(scores.detach().to('cpu')[0].numpy())
6.2 Expert Grouping
DeepSeek-V3 employs expert grouping, but unlike Device-Limited Routing, it is primarily used for the 3.2x bandwidth ratio of NVLINK and IB. The Group Score is calculated as follows:
### backup score for weight without Bias
original_scores = scores
### Bias is mainly used for aux-loss free loadblance, just used for modify routing decision.
scores = scores + gate_bias
scores = scores.view(x.size(0), args.n_expert_groups, -1)
group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
indices = group_scores.topk(args.n_limited_groups, dim=-1)[1]
mask = torch.zeros_like(scores[..., 0]).scatter_(1, indices, True)
scores = (scores * mask.unsqueeze(-1)).flatten(1)
indices = torch.topk(scores, args.n_activated_experts, dim=-1)[1]
### weights apply on original score , not include bias
weights = original_scores.gather(1, indices)
plt.plot(original_scores.detach().to('cpu')[0].numpy())
plt.plot(scores.detach().to('cpu')[0].numpy())It can be seen that with the Group constraint, experts are limited to 4 groups.
Finally, weights are normalized, and a scaling factor of route-scale=2.5 is applied.
weights /= weights.sum(dim=-1, keepdim=True)
weights *= self.route_scale
6.3 Load Balancing Without Auxiliary Loss Function
The biggest change in DeepSeek-V3 is the removal of auxiliary loss functions. Although AuxLoss was introduced for load balancing, excessive AuxLoss can further harm model performance. To better balance load and model performance, DeepSeek adopts a clever load balancing strategy without auxiliary loss by adding a bias term to each expert's gating score.
$$g_{i,t}'=
\begin{cases}
s_{i,t}& s_{i,t}+b_i \in TopK(s_{j,t}|1 \leq j \leq Nr, Kr) \\
0& Otherwise
\end{cases}
$$
It is important to note that the bias is only used for routing computation; the final weight still uses the original sigmoid (Original_score in the code). During training, load balancing for each batch is continuously monitored. At the end of each step, if an expert's load is too high, its corresponding bias term is reduced by γ; if an expert's load is insufficient, its corresponding bias term is increased by γ, where γ is a hyperparameter known as the bias update speed. Through this dynamic adjustment, DeepSeek-V3 maintains balanced expert loads during training and achieves better performance compared to models that rely solely on auxiliary loss for load balancing.
On the other hand, to prevent extreme imbalance within a single sequence, a sequence-level auxiliary loss compensation is also introduced.
$$\mathcal L_{SeqBalance}=\alpha \sum_{i=1}^{Nr} f_iP_i$$
$$f_i = \frac{Nr}{KrT}\sum_{t=1}^T \mathbb 1 (s_{i,t} \in TopK({s_{j,t}|1 \leq j \leq Nr},Kr))$$
$$s_{i,t}' = \frac{s_{i,t}}{\sum_{j=1}^{Nr}g_{j,t}}$$
$$P_i = \frac{1}{T}\sum_{t=1}^Ts_{i,t}'$$
6.4 No Need for Device-Limited Routing and Token Drop
Since the load balancing without auxiliary loss functions is already effective, Token-Drop from DeepSeek-V2 is not used. For Device-Limit Routing, DualPipe and other overlap strategies should suffice.
Therefore, it is incorrect to view DeepSeek-V3's Group concept based on the Device-Limit Routing group concept of DeepSeek-V2.
6.5 Modifications to AlltoAll Infrastructure
To ensure high computational performance of DualPipe, DeepSeek has developed a high-performance cross-node all-to-all communication kernel (including Dispatch and Combine) to reduce the number of SMs dedicated to communication. This kernel is co-designed with the MoE gating algorithm and the network topology of the cluster. For inter-machine networking, DeepSeek uses 400Gbps IB interconnects, and within a node, the NVLINK bandwidth of H800 is 160GB/s (unidirectional, though theoretically rated at 200GB/s, actual collective communication via NCCL runs at around 334GB/s bidirectional, approximately 160GB/s unidirectional). Therefore, the NVLINK to IB bandwidth ratio is approximately 3.2x.
MoE class defined as below:
world_size = N
rank = M
class MoE(nn.Module):
self.n_routed_experts = args.n_routed_experts
self.n_local_experts = args.n_routed_experts // world_size
self.n_activated_experts = args.n_activated_experts
### calculate local expert range
self.experts_start_idx = rank * self.n_local_experts
self.experts_end_idx = self.experts_start_idx + self.n_local_experts
self.gate = Gate(args)
### only create local experts
self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None
for i in range(self.n_routed_experts)])
### Shared Expert
self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)Then, in the forward phase, dispatch is performed based on the indices calculated by the gating function. However, the code provided on GitHub does not include details related to dispatch.
def forward(self, x: torch.Tensor) -> torch.Tensor:
shape = x.size()
x = x.view(-1, self.dim)
weights, indices = self.gate(x)
y = torch.zeros_like(x)
counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
### only calculate the local expert, seems dispatch function is missing
for i in range(self.experts_start_idx, self.experts_end_idx):
if counts[i] == 0:
continue
expert = self.experts[i]
idx, top = torch.where(indices == i)
y[idx] += expert(x[idx]) * weights[idx, top, None]
z = self.shared_experts(x)
### combine stage for allreduce.
if world_size > 1:
dist.all_reduce(y)
return (y + z).view(shape)However, the paper mentions that to effectively utilize the different bandwidths of IB and NVLink, each token's dispatch is limited to a maximum of 4 nodes, reducing IB traffic. Once a token's routing decision is determined, it is first transferred via IB to the GPU on the target node with the same intra-node index. Upon reaching the target node, efforts are made to ensure that it can be immediately forwarded through NVLink to the specific GPU hosting the target expert, without being blocked by subsequent tokens. This way, the communications over IB and NVLink completely overlap, allowing each token to efficiently select an average of 3.2 experts per node without additional NVLink overhead.
Although the current choice is 8 experts, it can be expanded to up to 13 experts (4 nodes × 3.2 experts/node) while maintaining the same communication cost. However, this strategy still requires 20 SMs for communication.
In the communication kernel implementation, Hopper's warpspecialization capability is utilized, dividing the 20 SMs into 10 channels. During dispatch, IB Send, IB to NVLINK, and NVLINK reception are handled by different warps, and the number of warps allocated to each communication task is dynamically adjusted based on the SM workload. Similarly, during the combine process, forwarding and reduce accumulation are merged into the operator and can also dynamically adjust workload.
On the other hand, there is overlap between the communication kernel and the computation kernel. Some PTX instructions and dynamic chunk size adjustments are used to reduce L2 cache usage and minimize impact on other SMs.
6.6 MoE Handling During Inference
During the inference phase, MoE communication is also noteworthy. The minimum deployment consists of 320 GPUs across 40 H800 units, with each GPU handling a single expert, requiring a total of 256 GPUs. The remaining 64 GPUs serve as redundant and shared experts. At this point, all-to-all communication is conducted directly over IB, utilizing IBGDA to reduce latency and enhance communication efficiency.
6.7 Suggestions for Infrastructure Improvement
In fact, DeepSeek and me share several views on certain issues. One is the disconnect between NVLINK and RDMA semantics, leading to high communication complexity. Additionally, using 20 out of 132 SMs on the H800 for communication is too costly, and there's a desire to offload this. When offloading, it is also hoped that the reduce operation can be performed during the combine phase.
This issue is very clear to me 4 years ago when I develop NetDAM at Cisco, Providing memory interface support for LD/ST semantics on the network card.
This creates a unified network, eliminating the need to distinguish between ScaleUp and ScaleOut.
Then apply in-networking computing
Just a rant, but I really don't understand why so many people are working on UEC/UAL and various ScaleUp standards. As for the congestion control issue in the Fabric during the MoE dispatch phase, it's been completely resolved in the last couple of years by AlibabaCloud eRDMA. Once the patents are out, everyone will understand. There are also some ideas on the co-design of the network and MoE gating, and a few patents were filed last year and the year before.
7. On the Evolution of MoE
Continuing with the derivation of the next-generation model, consider further expanding the dimensions due to hidden states and increasing the number of experts, for example, setting n_routed_experts=2048/4096 and activated_experts = 16/32/64. How would this affect grouping and communication domain limitations? Further scrutiny: what if domestic cards are used to replace the computation of experts at this point?
Let's assume a slightly extreme scenario: what if we build a cluster using a small number of H100 cards in combination with a large number of commercial cards like RTX4090/RTX5090 or other ASIC based cards for training?
At this point, the design of the gating function becomes quite interesting. You could continue with the expert group grouping, having the gating function output <grp_id, expert_id>, thereby constructing some degree of locality. This is similar to the approach mentioned recently in "A Conjecture on MoE."
When the computational power of these commercial cards is relatively low, large matrix multiplications for MLPs can still be challenging. How should we optimize for this? Additionally, as multi-modal and reasoning model contexts become longer, how can we better hide communication latency under this mode? Could we potentially split the hidden-state dimension, structure it into several smaller MLPs, and then concatenate the results?
I also found that Microsoft's Wei Furu and his team have been working on MultiHead MoE for quite some time.
Specifically, MH-MoE might offer more unexpected advantages in multi-modal scenarios, as illustrated in examples from another MH-MoE paper.
Conclusion: Once we resolve certain communication efficiency issues (which, of course, don't require any ScaleUp network—just effectively manage tail latency and overlap communication delay, which has been done ), we can integrate the idea of a two-level gating with the assumed MoE Group. Perform one gating for the entire token, dispatch it to a group, and then perform another MultiHead MoE gating within the group. This ensures both data locality and hierarchy.
At the end of the day , we could remove mainframe(NVL72) out, enjoy by commercial grade distributed system :)
A few days ago, I was suddenly assigned a task to help a few other teams optimize the inference for the full 671B version of DeepSeek-R1. At this stage, we're mainly making some simple tweaks on SGlang and vLLM, since resources for the H20 cards in China are quite limited.
Recently, I've noticed that both the SGlang and vLLM teams are competing fiercely. vLLM 0.7.2 has introduced optimizations for Triton MLA and FusedMoE, plus it has the advantage of pp parallelism over SGlang. Of course, SGlang is also developing pp parallelism, while vLLM is working on MTP. Over the next two to three weeks, there's room for further performance improvements in both frameworks.
After a few days of inference work, I've gained a better understanding of DeepSeek MoE. Thanks to suggestions from the DeepSeek team, I realized my earlier understanding of MoE Group Limit had some errors, and now I'm writing a detailed note on MoE.
Additionally, getting hands-on experience reveals many ingenious engineering details that one can miss by just reading papers. Many might think, "It's just MoE, I know about it too," but by not diving deeper, they miss out on these intricate details.
One important thing to note is that with DeepSeek's fine-grained MoE processing, the actual activation for the 671B model is only 37B, which has led to some interesting approaches in distributed inference systems. The extreme resource elasticity in training and inference integration could eventually lead to life-time learning/training, potentially disheartening other players in the field.
TOC as below:
1. Optimization Space of Transformer Model
From the earliest Transformer architecture, the computational complexity for the Attention Block is$O(N^2d)$ , while for the MLP Block it is $O(Nd^2)$ . Thus, algorithmic optimizations for larger model scales naturally focus on these two blocks, such as MHA for the Attention Block by DeepSeek MLA and Stepfun MFA. Many optimizations initially target long-context over N, which I will discuss separately in another document.
As for MoE optimizations, the open-source ecosystem primarily began with Mistral's Mixtral 8x7B. However, unfortunately, several major companies initially opted for Dense MLP. A few members of the DeepSpeed team moved from Microsoft to Snowflake and developed DenseMoE, which attempted to solve some communication issues by parallelizing MoE block and Attention Block. Unfortunately, it doesn't seem to have gained much traction afterwards.
Looking back at the titles of the three papers from DeepSeek now, there is a clear continuity.
The first paper, "Towards the Ultimate Specialization of Expert MoE Models," initially chose the sparse MoE model approach, aiming for maximum performance while maintaining smaller activation parameters.
However, existing work typically involves a small number of experts, and during training, the expertise of these experts can be overwhelmed by an excessive number of tokens. Figuratively speaking, a large amount of knowledge results in experts learning very diverse content, which compromises the density of information they can hold. This affects the Expert Specialization. From this perspective, the approach introduced Fine-Grained Expert Segmentation and utilized Shared Expert Isolation to absorb common knowledge, thereby reducing parameter redundancy in other experts. The final model is constructed as shown in the diagram below.
The V1 model was constructed with 64 experts, with each token activating 6 experts, and included 2 shared experts.
In the second iteration, DeepSeek-V2, the model was further expanded and strengthened by increasing the total number of routing experts to 160, still selecting 6 experts per token activation and maintaining 2 shared experts. Load balancing was further optimized, and MLA was introduced to achieve "Economical and Efficient" performance.
The third version, DeepSeek-V3, represents a culmination, expanding the number of experts to 256 with 8 selected per activation. There were further optimizations in load balancing and numerous infrastructure-based communication optimizations, such as DualPipe during training and offloading in MoE communication. This effectively enabled cross-node all-to-all (A2A) communication, as described in section 3.2.2 of the paper, where the model design considers the bandwidth ratio of NVLink to IB. It also uses PTX and auto-tunes chunk-size in communication to reduce L2Cache usage and minimize impact on other SMs.
2. Basic Working Principle of Sparse MoE
The first MoE model in the open-source ecosystem is likely Mixtral's 8x7B. The model structure is as follows:
Each Expert block is essentially a standard structure.
sMoE implementation as follow:
Since we need to analyze the Gating function, we'll discuss the specific
forward functionin SMoE a bit later. This establishes the structure of sMoE.Its Gating function computation is as follows: map through a linear layer to N experts, then apply softmax and select the top K.
We could plot softmax score as below:
Indicesrepresent the selected expert numbers. For example, Token 0 selects experts 4 and 1, Token 1 selects experts 0 and 6, and Token 2 selects experts 2 and 3.The subsequent computation is as shown in the diagram: tokens are routed to the experts based on the output indices, computed, multiplied by weights, and then summed together.
3. Expert Load Balancing and AuxLoss
However, there can be significant load imbalance among experts, which can lead to some experts being overloaded with information while others are undertrained, causing expert routing to collapse.
For example, we scaled up to 256 experts, selecting 8 for a test.
Therefore, a simple idea is to define an auxiliary loss function (Aux_Loss). For example, given$S$ tokens and $E$ experts, the loss is defined as the sum of the variances of the tokens received by each expert.
However, such an auxiliary loss function does not include the parameters of the Gating function, making it impossible to train through gradient updates. To address this, Google introduced an approach in GShard by replacing one component of the squared term with the mean of the Gating softmax.
4. DeepSeek-V1 MoE
Although there are some load balancing algorithms, when the number of experts is limited, the tokens assigned to a specific expert might cover different types of knowledge. The designated expert will tend to learn vastly different types of knowledge within its parameters, which are difficult to utilize effectively at the same time. If each token can be routed to more experts, different types of knowledge could be decomposed and learned by different experts. In this case, each expert can maintain a high level of specialization, promoting more concentrated distribution of knowledge among the experts.
The work of DeepSeek-V1 MoE primarily focuses on making experts more fine-grained and specialized (Towards Ultimate Expert Specialization), and it includes several aspects of work:
4.1. Fine-Grained Expert Segmentation
By partitioning experts into finer-grained segments while keeping the same amount of expert parameters and computational cost, the activation of expert combinations becomes more flexible and adaptable. For example, dividing each expert FFN in a MoE into$m$ smaller experts and reducing the FFN's inter-dim to $1/m$ of the original, the number of activated experts can increase by $m$ times at the same computational cost. Originally, there might be $C_{16}^2= 120$ combinations of expert selections, but with $m=4$ , it can expand to $C_{64}^8=4,426,165,368$ combinations. This significantly enhances the accuracy and targeted knowledge acquisition capabilities.
4.2 Shared Expert Isolation
In traditional routing strategies, tokens assigned to different experts might require some common knowledge or information, causing different experts to learn shared knowledge within their parameters, leading to redundancy. If there are dedicated shared experts responsible for capturing and integrating shared knowledge from different contexts, the parameter redundancy among other routing experts would be alleviated, resulting in a more parameter-efficient model with clearer expert specialization.
To achieve this, in addition to the fine-grained expert segmentation strategy, DeepSeek further isolates$Ks$ experts as shared experts. Regardless of how the routing module assigns tokens, each token is deterministically assigned to these shared experts. To maintain constant computational cost, the number of activated experts among other routing experts is reduced by $Ks$ . With the shared expert isolation strategy, the complete DeepSeekMoE architecture is as follows:
Shared Expert sum :$\sum_{i=1}^{Ks}FFN_i(u_t^l)$ $\sum_{i=Ks+1}^{mN}(g_{i,t}FFN_i(u_t^l))$
Routed Expert sum :
4.3 Expert Load Balancing
If routing strategies are fully automatically constructed through learning, load imbalance issues may arise. On one hand, there is a risk of expert routing collapse, where the model consistently selects only a few experts, leaving others undertrained. On the other hand, if experts are distributed across multiple devices, this can lead to computational load imbalance, further impacting the Model Fully Utilized (MFU) rate of the entire training cluster. Therefore, DeepSeek undertook two main actions.
4.3.1 Expert-Level Balance AuxLoss
To prevent routing collapse, an expert-level load balancing loss function is defined, where$N' = mN-Ks$ is the number of routing experts, $K'=mK-Ks$ is the number of activated fine-grained experts, and $T$ is the number of tokens that need to be processed by the experts. $\alpha_1$ is a hyperparameter.
Compared to the auxiliary loss function in Gshard, a correction factor of$N'/K'$ is applied when calculating $f_i$ .Since the probability of selecting each expert is $K'/N'$ when uniformly distributing across $N'$ experts choosing $K'$ ,this correction makes the overall loss function independent of the expert selection strategy.
4.3.2 Device-Level Balance AuxLoss
On the other hand, device-level load balancing is performed to ensure experts are evenly routed to multiple devices, keeping the computational load relatively balanced and avoiding long-tail effects. Therefore, experts are divided into$D$ group ${\mathcal E_1,\mathcal E_2,...,\mathcal E_D}$ with each group deployed on a single device. The loss is calculated as follows.
4.3.3 AuxLoss hyperparameter
In the design of the loss hyperparameters$\alpha_1$ and $\alpha_2$ , the expert load balancing loss parameter is set relatively small, while the device load balancing loss hyperparameter is set larger to better balance the load across devices.
4.3.4 Code overview
On HuggingFace, there is a DeepseekMoE function. We noticed that the implementation of the MoEGate function differs somewhat from the paper. The number of tokens$T$ is used to calculate the load balancing loss for all tokens within a batch. Additionally, in the 16B model, the cross-device loss function is not executed.
The computation of the Expert-Level auxiliary loss can be configured in two ways, with the default being calculated for each sequence within the batch.
Another point to note is that the current DeepSeek code allows for the use of MLP-Dense in the initial layers. This approach benefits model stability by avoiding some stability issues caused by early-stage Attention entering the MoE. However, in the publicly available DeepSeek-MoE-16B model, this feature is disabled, meaning every layer is MoE.
5. DeepSeek-V2
DeepSeek-V2 further expands the fine-grained expert selection by adopting a scheme of routing 160 experts, selecting 6 + 2 shared experts. Compared to DeepSeek-V1-MoE, it adds two new communication-related constraints.
5.1 Device-Limited Routing
The first approach is device-constrained routing to limit MoE-related communication costs. This mainly involves ensuring that experts routed for parallel processing are distributed across multiple devices. When the number of experts is particularly large, token communication within a single batch will span numerous devices, increasing the communication cost for expert parallelism. Thus, DeepSeek-V2 introduces a constraint that limits each token to be routed to a maximum of$M$ devices. Specifically, for each token, the top $M$ devices containing experts with the highest affinity scores are selected. Then, top-K selection is performed among the experts on these $M$ devices. In practice, it is found that when
$ M \geq 3$, device-constrained routing achieves performance roughly equivalent to unconstrained top-K routing.
For a specific code implementation, you can check the source code on HuggingFace. The calculation of the gating function is unchanged, still using softmax across the entire batch, but FP32 is used for computational precision.
Before performing TopK and normalization selection, MoE Groups are computed, dividing the experts into 8 groups in total. Then, 3 groups are selected using TopK. For each group, the maximum softmax value is calculated as the group's score, and then$M$ groups are selected from these.
Then build gourp mask , after mask the score, execute topK
By examining the distribution of scores and tmp_scores, it can be seen that the softmax values for other groups are masked to 0. The blue represents the original softmax, while the yellow shows the values after applying the group mask.
Then run topK function with norm.
5.2 Communication Load Balancing Loss
Although Device-Limit Routing can address the communication domain issue and reduce the communication fanout, some receiving devices may still experience concentration of activations on a few experts, leading to communication bottlenecks. Therefore, a communication load balancing loss is introduced, where$D$ presents the number of devices, $T$ represents the number of tokens,$M$ represents the number of selected groups, and $\mathcal E_i$ represents the i-th expert.
5.3 Token Drop Strategy
Although two new load balancing strategies have been introduced, to ensure training efficiency and avoid long-tail effects, the average number of tokens per device is calculated from the total number of tokens in a batch, serving as the device capacity limit. When the actual capacity exceeds this limit, the softmax gating scores are sorted in descending order, and the computations for the MoE expert networks are skipped for the portion exceeding the capacity, directly moving to the combine phase.
6. DeepSeek-V3
DeepSeek-V3 continues with the original architecture of fine-grained Experts and Shared Experts, further increasing the number of experts. The V3 config file indicates that there are three model sizes.
6.1 Gating Function Using Sigmoid
Additionally, according to the configuration file, in the 671B model, the gating function has been changed from softmax to sigmoid, with normalization applied.
The score function as below:
To verify the previously speculation , we calculate softmax here to compare with sigmoid
DeepSeek-V3 671B adopt sigmoid , the score is in range [0,1] even larger than softmax
6.2 Expert Grouping
DeepSeek-V3 employs expert grouping, but unlike Device-Limited Routing, it is primarily used for the 3.2x bandwidth ratio of NVLINK and IB. The Group Score is calculated as follows:
It can be seen that with the Group constraint, experts are limited to 4 groups.
Finally, weights are normalized, and a scaling factor of route-scale=2.5 is applied.
6.3 Load Balancing Without Auxiliary Loss Function
The biggest change in DeepSeek-V3 is the removal of auxiliary loss functions. Although AuxLoss was introduced for load balancing, excessive AuxLoss can further harm model performance. To better balance load and model performance, DeepSeek adopts a clever load balancing strategy without auxiliary loss by adding a bias term to each expert's gating score.
It is important to note that the bias is only used for routing computation; the final weight still uses the original sigmoid (Original_score in the code). During training, load balancing for each batch is continuously monitored. At the end of each step, if an expert's load is too high, its corresponding bias term is reduced by γ; if an expert's load is insufficient, its corresponding bias term is increased by γ, where γ is a hyperparameter known as the bias update speed. Through this dynamic adjustment, DeepSeek-V3 maintains balanced expert loads during training and achieves better performance compared to models that rely solely on auxiliary loss for load balancing.
On the other hand, to prevent extreme imbalance within a single sequence, a sequence-level auxiliary loss compensation is also introduced.
6.4 No Need for Device-Limited Routing and Token Drop
Since the load balancing without auxiliary loss functions is already effective, Token-Drop from DeepSeek-V2 is not used. For Device-Limit Routing, DualPipe and other overlap strategies should suffice.
Therefore, it is incorrect to view DeepSeek-V3's Group concept based on the Device-Limit Routing group concept of DeepSeek-V2.
6.5 Modifications to AlltoAll Infrastructure
To ensure high computational performance of DualPipe, DeepSeek has developed a high-performance cross-node all-to-all communication kernel (including Dispatch and Combine) to reduce the number of SMs dedicated to communication. This kernel is co-designed with the MoE gating algorithm and the network topology of the cluster. For inter-machine networking, DeepSeek uses 400Gbps IB interconnects, and within a node, the NVLINK bandwidth of H800 is 160GB/s (unidirectional, though theoretically rated at 200GB/s, actual collective communication via NCCL runs at around 334GB/s bidirectional, approximately 160GB/s unidirectional). Therefore, the NVLINK to IB bandwidth ratio is approximately 3.2x.
MoE class defined as below:
Then, in the forward phase, dispatch is performed based on the indices calculated by the gating function. However, the code provided on GitHub does not include details related to dispatch.
However, the paper mentions that to effectively utilize the different bandwidths of IB and NVLink, each token's dispatch is limited to a maximum of 4 nodes, reducing IB traffic. Once a token's routing decision is determined, it is first transferred via IB to the GPU on the target node with the same intra-node index. Upon reaching the target node, efforts are made to ensure that it can be immediately forwarded through NVLink to the specific GPU hosting the target expert, without being blocked by subsequent tokens. This way, the communications over IB and NVLink completely overlap, allowing each token to efficiently select an average of 3.2 experts per node without additional NVLink overhead.
Although the current choice is 8 experts, it can be expanded to up to 13 experts (4 nodes × 3.2 experts/node) while maintaining the same communication cost. However, this strategy still requires 20 SMs for communication.
In the communication kernel implementation, Hopper's warpspecialization capability is utilized, dividing the 20 SMs into 10 channels. During dispatch, IB Send, IB to NVLINK, and NVLINK reception are handled by different warps, and the number of warps allocated to each communication task is dynamically adjusted based on the SM workload. Similarly, during the combine process, forwarding and reduce accumulation are merged into the operator and can also dynamically adjust workload.
On the other hand, there is overlap between the communication kernel and the computation kernel. Some PTX instructions and dynamic chunk size adjustments are used to reduce L2 cache usage and minimize impact on other SMs.
6.6 MoE Handling During Inference
During the inference phase, MoE communication is also noteworthy. The minimum deployment consists of 320 GPUs across 40 H800 units, with each GPU handling a single expert, requiring a total of 256 GPUs. The remaining 64 GPUs serve as redundant and shared experts. At this point, all-to-all communication is conducted directly over IB, utilizing IBGDA to reduce latency and enhance communication efficiency.
6.7 Suggestions for Infrastructure Improvement
In fact, DeepSeek and me share several views on certain issues. One is the disconnect between NVLINK and RDMA semantics, leading to high communication complexity. Additionally, using 20 out of 132 SMs on the H800 for communication is too costly, and there's a desire to offload this. When offloading, it is also hoped that the reduce operation can be performed during the combine phase.
This issue is very clear to me 4 years ago when I develop NetDAM at Cisco, Providing memory interface support for LD/ST semantics on the network card.
This creates a unified network, eliminating the need to distinguish between ScaleUp and ScaleOut.
Then apply in-networking computing
7. On the Evolution of MoE
Continuing with the derivation of the next-generation model, consider further expanding the dimensions due to hidden states and increasing the number of experts, for example, setting
n_routed_experts=2048/4096andactivated_experts = 16/32/64. How would this affect grouping and communication domain limitations? Further scrutiny: what if domestic cards are used to replace the computation of experts at this point?Let's assume a slightly extreme scenario: what if we build a cluster using a small number of H100 cards in combination with a large number of commercial cards like RTX4090/RTX5090 or other ASIC based cards for training?
At this point, the design of the gating function becomes quite interesting. You could continue with the expert group grouping, having the gating function output
<grp_id, expert_id>, thereby constructing some degree of locality. This is similar to the approach mentioned recently in "A Conjecture on MoE."When the computational power of these commercial cards is relatively low, large matrix multiplications for MLPs can still be challenging. How should we optimize for this? Additionally, as multi-modal and reasoning model contexts become longer, how can we better hide communication latency under this mode? Could we potentially split the hidden-state dimension, structure it into several smaller MLPs, and then concatenate the results?
I also found that Microsoft's Wei Furu and his team have been working on MultiHead MoE for quite some time.
Specifically, MH-MoE might offer more unexpected advantages in multi-modal scenarios, as illustrated in examples from another MH-MoE paper.
Conclusion: Once we resolve certain communication efficiency issues (which, of course, don't require any ScaleUp network—just effectively manage tail latency and overlap communication delay, which has been done ), we can integrate the idea of a two-level gating with the assumed MoE Group. Perform one gating for the entire token, dispatch it to a group, and then perform another MultiHead MoE gating within the group. This ensures both data locality and hierarchy.
At the end of the day , we could remove mainframe(NVL72) out, enjoy by commercial grade distributed system :)