DeepSeek V4的tid2eid是怎么来的?

原文链接 | 2026-05-15 | 苏剑林

训过MoE的同学都知道,如果把整个模型的MLP部分都换成MoE,那么靠近Embedding的前几层MoE往往很难实现负载均衡。对此,DeepSeek V3,包括我们的Kimi K2,采取的应对策略是“first_k_dense”——顾名思义,就是前\(k\)层不用MoE而用常规的Dense型GLU。到了DeepSeek V4,这个策略改成了“first_k_hash”。

这里的“hash”指的是“Hash Routing”,提出自《Hash Layers For Large Sparse Models》,它通过一张事先确定的映射表tid2eid来为每个Token分配Expert。本文来探讨一下如何生成这个tid2eid。

问题背景#

首先要指出的是,对于前几层MoE,如果用笔者在《MoE环游记:6、最优分配促均衡》和《MoE环游记:7、动态激活极简解》所提的Quantile Balancing,其实已经能较大程度上缓解它们的不均衡问题。

至于V4采用Hash Routing,可以说是同一问题在不同方向的尝试,其想法是对于前几层MoE来说,上下文信息并不多,所以基于它的输入向量来选Expert,跟直接基于没有任何上下文信息的Token Id来选,也许并无太大区别。既然如此,不如干脆基于Token Id来预先写死前几层每个Token该激活的Expert,这便是“tid2eid (Token Id to Expert Id)”的由来。

那怎么生成tid2eid这个表格呢?完全随机可以吗?看上去不大行,因为每个Token的出现频率并不一样,完全随机的话反而会导致不均衡。DeepSeek并没有公布这个表格的生成细节,我们只能按照自己的思路去猜测一下。

数学描述#

假设一共有\(m\)个不同的Token,第\(i\)个Token的出现频率是\(p_i\),不失一般性,假设它们已经按从大到小排列,即\(p_i\geq p_{i+1}\)。我们用\(x_{i,j}\in\{0, 1\}\)表示Token \(i\)是否应该激活Expert \(j\)(\(0\)表示不激活,\(1\)表示激活),共有\(n\)个Expert,每个Token选\(k\)个Expert,那么我们实际上要求解方程组

\[\begin{equation}x_{i,j}\in\{0, 1\},\qquad\sum_{j=1}^n x_{i,j} = k,\qquad \sum_{i=1}^m p_i x_{i,j}\approx \frac{k}{n}\end{equation}\]

注意最后一个等式我们用了约等于\(\approx\),这是因为如果改为\(=\),方程组严格来讲未必有解。所以我们保留\(\approx\)的记号,其含义是两端应该尽可能接近。也可以考虑写成优化形式

\[\begin{equation}\min_{x_{i,j}\in\{0, 1\}} \sum_{j=1}^n \left(\sum_{i=1}^m p_i x_{i,j} - \frac{k}{n}\right)^2 \qquad\text{s.t.}\qquad\sum_{j=1}^n x_{i,j} = k\end{equation}\]

该问题的一个比较朴素的解法是贪心算法:

按频率从高到低逐Token处理,先统计各Expert的已分配负载,然后选择负载最轻的\(k\)个Expert作为当前Token的激活Expert,随后更新Expert的负载分布。

参考实现#

贪心算法虽然原理上是贪心的,但它大多数情况下都能得到一个近乎最优的解,一个参考实现如下:

import numpy as np

# 模拟分布
m, n, k = 80000, 128, 4
p = 1 / (10 + np.arange(m))
p /= p.sum()

# 贪心处理
x, f = np.zeros((m, k), dtype='int32'), np.zeros(n)
for i in range(m):
    j = f.argsort()[:k]  # 选择负载最轻的k个
    f[j] += p[i] / k  # 更新负载分布
    x[i] = j  # 记录到tid2eid中

# 评估均衡
max_vio, min_vio = f.max() * n - 1, f.min() * n - 1

注意,该代码没有任何随机的地方,原则上是一个确定性的算法。那如果我们前几层想要不同的tid2eid呢?可以考虑加入一些随机性,比如

range(m)

换成

np.random.permutation(m)

,即不按照频率从高到低的顺序执行,这样依然能够得到可用的解,只不过均衡性相对差点。我们可以重复跑多次,然后挑max_vio最低的几个解。

极端情形#

现在我们考虑一种极端的情况:\(p_1 \gg k / n\),即某个Token的频率已经远超均衡水平,此时无论tid2eid怎么编排,都不可能实现负载均衡。换句话说,只根据单个Token Id来决策,是无论如何也均衡不了的。这种情况下又该如何应对呢?

答案很简单,换一种依赖于更多输入信息的Hash方法。前面我们只根据当前Token Id来决策,实际上每个Token都处于某个序列中,它是有上下文的,如果说之前的做法是“1-gram to expert”,现在我们可以考虑跟它前面的Token一起,组成“2-gram to expert”、“3-gram to expert”等等。

以2-gram \((a, b)\)为例,一种朴素的Hash方式是:选大于\(m\)的质数\(q\),计算\((aq+b)\mod n\)作为激活的Expert,用不同的质数重复\(k\)次就得到\(k\)个Expert。由于增加了gram数,输入的组合数大大增加,然后映射到有限的\(n\)个数中,每个数大概率都能被“塞满”,所以只要Hash函数不是特别糟糕,那么结果几乎就是均衡的。

因此,它的优点是不用考虑频率,也不用事先记忆一个tid2eid。至于缺点,就是Hash函数的计算会稍微复杂一些了,而且同一个Token可能会选到重复的Expert,但这些都并不是特别严重的问题。

文章小结#

本文简单梳理了DeepSeek V4中的Hash Routing的基本思想,并重点探讨了它的tid2eid映射表的构造原理。