L大语言模型从理论到实践
阅读 PDF ↗
CHAPTER 36

混合专家模型

稠密前馈层对每个词元应用同一组参数。增大前馈宽度可以增加表达能力,但也使每个词元承担更多矩阵计算。混合专家模型试图部分分离参数容量与单次计算:保存多组可学习变换,由输入决定本次调用其中哪些变换。然而,减少参与计算的参数并不自动减少全部权重的存储,也不自动降低跨设备通信与执行尾部。

本章承接第3章《神经网络基础》的链式求导,从专家混合的函数形式出发,推导稀疏路由、分发与合并、容量约束以及路由器的梯度。随后用完整数值例题说明权重归一化、专家负载和溢出处理如何改变输出,最后讨论共享专家、软专家混合与专家并行。具体模型家族的结构差异在开源模型章节展开,集群拓扑和分布式训练将在相应章节继续细化。

36.1条件计算及专家结构

专家函数的定义

混合专家模型(Mixture of Experts,MoE)包含若干专家函数 \(F_e\),并由路由器(Router)或门控网络(Gating Network)确定各专家对当前输入的贡献。专家通常是结构相同、参数独立的前馈网络,但“结构相同”不要求其训练后的函数相同。名称中的专家描述条件化计算组件,不能据此将某一编号直接命名为数学、法律或编程专家。 1

在 Transformer 中,常见做法是用 MoE 替换部分逐位置前馈子层,注意力和残差主干仍保留。稀疏门控使一次输入只调用少量专家,是条件计算(Conditional Computation)的一种形式(Shazeer 等 2017)。它并不等于为每种任务部署一个完整语言模型,也不等于多个独立模型生成答案后再投票。

设一个专家采用门控前馈形式,行向量输入 \(x\in\Real^{1\times d}\) 的输出为

\[ F_e(x)=\left[\phi(xW_{e,g})\odot(xW_{e,u})\right]W_{e,d}, \tag{36.1}\]

其中 \(W_{e,g},W_{e,u}\in\Real^{d\times d_f}\)\(W_{e,d}\in\Real^{d_f\times d}\)\(\phi\) 是固定的激活函数。每个专家都输出宽度 \(d\),以便加权合并并与残差相加。普通两层前馈专家也可使用相同路由机制;专家内部结构与路由结构是两个可分别改变的维度。

统一符号及适用假设

将当前路由组内的有效词元表示展平为 \(X\in\Real^{N\times d}\)。这里 \(N\) 是实际参与路由的位置数,可以来自多个序列,但不把填充位置算入其中。路由组可能只覆盖一个本地微批次,也可能覆盖一个专家并行组;统计量必须明确采用哪一种范围。

符号 含义 形状或范围
\(N,d,d_f\) 有效词元数、模型宽度、专家中间宽度 正整数
\(E_r,k\) 路由专家数、每词元所选专家数 \(1\le k\le E_r\)
\(W_r,b_r\) 路由器权重和偏置 \(d\times E_r,\ E_r\)
\(G,P\) 路由分数与全专家概率 \(N\times E_r\)
\(S_t\) 词元 \(t\) 选中的专家集合 大小为 \(k\)
\(a_{te}\) 稀疏混合权重 非负标量
\(n_e,C_e\) 专家申请负载与容量 非负整数
\(D_e\) 词元到专家槽位的选择矩阵 \(C_e\times N\)
\(H_e,Y\) 专家输出与合并结果 \(C_e\times d,\ N\times d\)
\(f_e,\bar p_e\) 硬分配比例与平均路由概率 无量纲
\(P_{\mathrm{total}},P_{\mathrm{active}}\) 总参数量、每词元激活参数量 参数个数

本章用 \(E_r\) 表示专家数量,避免与嵌入矩阵混用。\(P\) 在路由部分表示专家概率矩阵,带下标的 \(P_{\mathrm{total}}\) 等表示参数计数。

基本稀疏路由模型

路由组含 \(N>0\) 个有效词元;各专家输入输出宽度相同;每个词元先选择 \(k\) 个不同专家;分数有限且并列选择规则明确。先讨论没有容量溢出的情况,再显式加入容量与拒绝规则。对梯度进行局部推导时,假定选择集合和接收集合保持不变。

稠密混合及稀疏混合

路由器先计算

\[ G=XW_r+\mathbf1_Nb_r^{\mathsf T},\qquad p_{te}=\frac{\exp(g_{te})}{\sum_{j=1}^{E_r}\exp(g_{tj})}. \tag{36.2}\]

Softmax 沿专家轴归一化。稠密专家混合(Dense Mixture of Experts)对每个词元执行所有专家,再计算

\[ y_t=\sum_{e=1}^{E_r}p_{te}F_e(x_t). \tag{36.3}\]

概率可以接近零,但除非执行过程真正跳过相应专家,其算术开销仍然存在。“权重很小”与“没有执行”不是同一个条件。

稀疏专家混合(Sparse Mixture of Experts)只执行选中集合 \(S_t\) 中的专家:

\[ S_t=\operatorname{TopK}(g_{t,:},k),\qquad y_t=\sum_{e\in S_t}a_{te}F_e(x_t). \tag{36.4}\]

Top-\(k\) 表示取最高的 \(k\) 个分数对应的不同编号。由于 Softmax 对同一行保持分数次序,在没有额外变换时,按分数或概率选择得到同一集合。但概率如何用于合并,仍有不同约定。

36.2路由权重及局部梯度

两种归一化约定

一种常见形式在选中专家内重新归一化:

\[ a_{te}=\begin{cases} \displaystyle\frac{\exp(g_{te})}{\sum_{j\in S_t}\exp(g_{tj})},&e\in S_t,\\ 0,&e\notin S_t. \end{cases} \tag{36.5}\]

此时每个词元保留的混合权重和为一。也可以先对全部专家归一化,再直接保留选中的概率,即 \(a_{te}=p_{te}\mathbf1[e\in S_t]\)。后者保留权重之和一般小于一;输出幅度因而也受到路由置信度影响。Switch 的单专家路由保留了被选专家的全专家门控概率(Fedus, Zoph, 和 Shazeer 2022),不能把它改写成选中集合内重新归一化后仍声称函数相同。

两种形式在 \(k=1\) 时差异尤其明显。选中集合内归一化会得到唯一权重为一,在固定选择集合的邻域内,任务损失不能再通过这个恒为一的权重更新路由分数。保留全专家概率则仍有连续梯度路径。路由器是否能从任务损失学习,不能只看代码里是否出现 Softmax,还要看后续归一化是否消去了其作用。

并列分数处需要一个明确的编号选择规则;选择边界附近,微小浮点扰动也可能改变专家集合。确定性并列规则可以使重复执行更可解释,但不消除硬选择本身的非光滑性。

专家及路由器梯度

暂时固定一个词元的选中集合,并省略词元下标。令专家输出 \(z_e=F_e(x)\),混合结果 \(y=\sum_ea_ez_e\),上游梯度为列向量 \(h=\frac{\partial\mathcal L}{\partial y}\)。定义标量敏感度 \(r_e=h^{\mathsf T}z_e^{\mathsf T}\),则

\[ \frac{\partial\mathcal L}{\partial a_e}=r_e,\qquad \frac{\partial\mathcal L}{\partial z_e}=a_eh^{\mathsf T}. \tag{36.6}\]

专家参数的梯度由其局部网络反向传播获得,并乘以该词元的混合权重。未选专家不接收这个词元经专家执行路径产生的梯度,但可能从同批其他词元、正则项或其他共享路径获得贡献。

对式(36.5),选中集合内的 Softmax 雅可比为 \(a_e(\mathbf1[e=j]-a_j)\),所以

\begin{align} \frac{\partial\mathcal L}{\partial g_j} &=\sum_{e\in S}r_ea_e(\mathbf1[e=j]-a_j)\tag{36.7}\\ &=a_j\left(r_j-\sum_{e\in S}a_er_e\right),\qquad j\in S. \tag{36.8}\end{align}

未选分数的这一局部梯度为零。式中括号比较某个专家的输出敏感度与当前混合平均值;因此路由器可以学习改变已选专家之间的连续权重。

若保留全专家概率,令 \(r_e=0\) 对所有未选专家成立,则同样的 Softmax 反向得到

\[ \frac{\partial\mathcal L}{\partial g_j} =p_j\left(r_j-\sum_{e\in S}p_er_e\right),\qquad j=1,\ldots,E_r. \tag{36.9}\]

未选分数也可因全局分母收到梯度,但这不意味着未选专家已经执行或获得同样的参数梯度。必须区分路由器参数与专家参数。

最后,若所有词元的分数梯度构成 \(G_G\),路由器线性层满足

\[ \frac{\partial\mathcal L}{\partial W_r}=X^{\mathsf T}G_G,\qquad \frac{\partial\mathcal L}{\partial b_r}=G_G^{\mathsf T}\mathbf1_N. \tag{36.10}\]

输入梯度还应加上所有已选专家的回传与路由器的回传。漏掉任一路径,都不再对应完整混合函数的导数。

硬选择的导数边界

Top-\(k\) 的编号集合在分数排序不变的小邻域中保持常数,在排序交换处发生跳变。通常的反向传播沿已选择路径求导,不会自动告诉模型“如果换成另一个专家,损失会怎样”。加噪声、辅助目标或其他选择机制可以改变探索和学习行为,但每种方法都需说明自己的前向与反向定义。

如果所有词元过早集中到少数专家,这些专家得到更多任务训练机会,其他专家则缺少有效更新,可能形成自我强化的负载差异。负载均衡机制因此不仅是运行调度,也影响训练过程中不同参数获得监督的分布。不过,强制每个专家处理完全相同数量、完全相同类型的输入也不是模型质量的充分条件。

36.3词元路由

词元轴的专家批次变换

MoE 的输出仍按原词元顺序排列,但专家计算希望把送往同一专家的词元放在同一个批次中。分发(Dispatch)将原词元重排或复制到专家局部缓冲,合并(Combine)则将处理后的结果按原位置加权累加。

专家 \(e\)\(C_e\) 个槽位,定义选择矩阵 \(D_e\in\{0,1\}^{C_e\times N}\)。如果原词元 \(t\) 被接收到槽位 \(c\),令 \((D_e)_{ct}=1\);一个槽位至多接收一个词元。于是

\[ X_e=D_eX,\qquad H_e=F_e(X_e), \tag{36.11}\]

其中 \(F_e\) 对每个有效槽位按行计算。再定义合并矩阵 \(A_e\in\Real^{N\times C_e}\),将对应位置设为接收后的权重,得到

\[ Y=\sum_{e=1}^{E_r}A_eH_e. \tag{36.12}\]

没有拒绝且采用原权重时,\((A_e)_{tc}=a_{te}(D_e)_{ct}\)。这一写法用于说明映射关系,不要求运行时物化巨大稀疏矩阵;索引列表、前缀和与分组缓冲可以表达同样的关系。

36.1概括了分组与回收路径。一个词元选择两个专家时,分发阶段逻辑上产生两份输入引用,合并时必须将两份结果相加。只保留最后返回结果会丢失一条专家路径。空槽位即使因偏置产生非零专家输出,其合并权重也必须为零;否则填充缓冲会污染真实词元。

稀疏 MoE 的数据流。专家批次大小由路由决定;选择多个专家意味着多条返回路径。
图 36.1 稀疏 MoE 的数据流。专家批次大小由路由决定;选择多个专家意味着多条返回路径。

路由容量及溢出

令申请负载 \(n_e=\sum_t\mathbf1[e\in S_t]\),无拒绝时有 \(\sum_en_e=Nk\),平均负载为 \(\frac{Nk}{E_r}\)。一种容量规则是

\[ C=\left\lceil\gamma\frac{Nk}{E_r}\right\rceil, \tag{36.13}\]

其中 \(\gamma>0\)容量因子(Capacity Factor),所有专家暂取相同容量。Switch 中的单专家形式对应 \(k=1\)(Fedus, Zoph, 和 Shazeer 2022);其他系统可能按局部分组、不同专家宽度或不同对齐规则定义容量,不能只比较容量因子的数字。

即使总槽位 \(E_rC\ge Nk\),局部专家仍可能溢出。例如一个专家收到大量申请时,其他专家的空槽位不能自动代替它执行同一个函数。容量通常与静态缓冲大小、通信消息形状和最坏负载相关,提高容量可以减少拒绝,却可能增加填充计算和临时空间。

超过容量后的处理必须成为模型定义的一部分。拒绝某条专家路径,可以使该路径贡献为零;重新路由会调用另一组参数;等待或变长分组可以保留全部路径,但付出动态调度和尾延迟代价。被称为“不丢词元”的实现通常是保留全部路由计算,而非移除了负载和通信约束。

若一个词元的全部专家路径被拒绝,MoE 子层贡献可能为零,残差主干仍可保留该位置。这里的“丢词元”常指跳过该层的专家处理,不是从序列中删除编号,也不是将其监督标签自动移除。 2

接收规则可能引入批次依赖

容量竞争发生在路由组内。若采用先到先得,后部位置可能更容易被拒绝;若按门控权重排序,则低权重路径更容易失去计算。于是同一词元在不同批次组成中,可能因为其他词元的申请而获得不同输出。这与单纯逐词元前馈函数的性质不同。

对于严格自回归训练,还应警惕一种间接依赖:若先观察整段序列的路由权重,再按全局排名决定早期词元能否获得容量,未来词元可能影响早期位置的接收结果。避免这一问题需要容量方案和接收规则与因果定义相容,而不能仅在注意力层放置因果掩码。实际设计应区分前向内容读取、容量竞争以及训练辅助统计三类关系。

36.4负载均衡及路由稳定性

硬负载及软概率描述不同对象

定义容量截断前的分配比例和平均概率:

\[ f_e=\frac1{Nk}\sum_{t=1}^{N}\mathbf1[e\in S_t],\qquad \bar p_e=\frac1N\sum_{t=1}^{N}p_{te}. \tag{36.14}\]

二者各自之和为一。\(f_e\) 衡量实际申请的离散次数,\(\bar p_e\) 衡量路由器的连续概率质量。一个专家经常以第二名被选中,可能有较大的申请次数,却不拥有最大的平均概率,因此两种统计不能互换。

若只统计容量截断后的负载,被饱和容量裁平的结果可能看起来很均衡,却掩盖大量溢出。应同时保留申请数、接收数、拒绝数和执行耗时。若专家宽度不同,次数均衡也不等于算术量均衡。

辅助均衡目标及其梯度范围

Switch 使用硬分配比例与平均概率乘积构造辅助目标。按照式(36.14)的 top-\(k\) 归一化,本章讨论其一种相应扩展:

\[ \mathcal L_{\mathrm{bal}}=E_r\sum_{e=1}^{E_r}f_e\bar p_e,\qquad \mathcal L=\mathcal L_{\mathrm{task}}+\lambda\mathcal L_{\mathrm{bal}}. \tag{36.15}\]

其中 \(\lambda\ge0\) 控制任务目标与均衡目标的相对权重,top-1 时恢复相应统计定义。不同实现对 \(k\)、序列和设备的归一化可能不同,系数不能脱离公式移植。

在反向时通常将硬计数 \(f_e\) 视为当前已确定的常数,而通过 \(\bar p_e\) 对路由器求导。由 Softmax 雅可比,

\begin{align} \frac{\partial\mathcal L_{\mathrm{bal}}}{\partial g_{tj}} &=\frac{E_r}{N}\sum_ef_ep_{te}(\mathbf1[e=j]-p_{tj})\tag{36.16}\\ &=\frac{E_r}{N}p_{tj}\left(f_j-\sum_ef_ep_{te}\right). \tag{36.17}\end{align}

因此它可以为未被选择的分数提供连续训练信号,但没有对离散计数进行通常意义下的求导。它不是对实际最大执行时间的精确梯度,也不能保证每一步更新后负载立刻均匀。

\(f_e=\bar p_e=\frac{1}{E_r}\) 时,该辅助项为1,形成一个便于比较的均匀参考值。但它不是均衡偏差的范数,也不能只根据某个标量接近1就认定系统均衡。特别是硬分配与软概率的关系、局部分组及有限样本都会影响这个乘积;应联合观察各专家分布和最大负载。

统计范围及任务权重

如果各设备本地批量大小不同,先计算各自平均概率再无权平均,会改变全局统计。应先合并概率和与有效词元数,再形成所需范围内的均值。另一方面,全局均衡可能掩盖单台设备上的局部热点;设计目标必须说明是专家层级、设备层级还是通信组层级。

均衡压力过大时,路由可能优先满足平均分配而牺牲任务有效性;压力过小时,少数专家可能成为瓶颈并获得不成比例的更新。除了辅助损失,也可以通过容量分配、专家复制或路由偏置控制负载,但这些机制改变的对象不同:有的改变训练目标,有的改变离散选择,有的仅改变相同专家副本的执行位置。

路由分数的数值范围同样影响稳定性。稳定 Softmax 通过减去行最大值避免指数溢出,但不能防止分数排序因低精度舍入而变化。路由计算与专家矩阵乘法可以采用不同精度;是否合适应结合选择差异、梯度和运行成本判断,而不是要求全部运算使用同一种数据类型。

\(N=d=4\)\(E_r=3\),每词元选择 \(k=2\) 个专家,使用选中集合内重新归一化。令输入 \(X=I_4\),路由器无偏置,并设

\[ W_r=\begin{bmatrix} \log0.7&\log0.2&\log0.1\\ \log0.6&\log0.3&\log0.1\\ \log0.55&\log0.35&\log0.1\\ \log0.1&\log0.2&\log0.7 \end{bmatrix}. \tag{36.18}\]

因为每行概率之和为一,所以 \(G=XW_r=W_r\),Softmax 后恰好得到矩阵中的对数自变量。为使全部专家结果可手算,本例用三个真实线性映射 \(F_1(x)=x\)\(F_2(x)=2x\)\(F_3(x)=3x\);这些映射仅用于构造算例,不声称来自训练或具备专家分工。

词元 全专家概率 选中集合 选中后权重
1 \((0.7,0.2,0.1)\) \(\{1,2\}\) \((\frac{7}{9},\frac{2}{9})\)
2 \((0.6,0.3,0.1)\) \(\{1,2\}\) \((\frac{2}{3},\frac{1}{3})\)
3 \((0.55,0.35,0.1)\) \(\{1,2\}\) \((\frac{11}{18},\frac{7}{18})\)
4 \((0.1,0.2,0.7)\) \(\{2,3\}\) \((\frac{2}{9},\frac{7}{9})\)

最后一行按专家编号列权重,而不是按分数从高到低排列。保存路由记录时必须同时保存编号与权重,不能假定两种排序始终一致。

没有容量拒绝时,四个输出分别为

\begin{align} y_1&=(\frac{7}{9})x_1+(\frac{2}{9})(2x_1)=(\frac{11}{9})x_1,\tag{36.19}\\ y_2&=(\frac{2}{3})x_2+(\frac{1}{3})(2x_2)=(\frac{4}{3})x_2,\tag{36.20}\\ y_3&=(\frac{11}{18})x_3+(\frac{7}{18})(2x_3)=(\frac{25}{18})x_3,\tag{36.21}\\ y_4&=(\frac{2}{9})(2x_4)+(\frac{7}{9})(3x_4)=(\frac{25}{9})x_4. \tag{36.22}\end{align}

由于 \(X=I_4\)\(Y\) 就是以这些系数为对角元素的矩阵。路由权重和为一,但输出系数可以大于一,因为专家输出本身并不被限制在单位范围。

专家申请数为 \((n_1,n_2,n_3)=(3,4,1)\),总计8条路径。取容量因子 \(\gamma=1\),由式(36.13)得到每专家容量 \(C=3\)。虽然总槽位数为9,大于申请总数8,专家2仍会溢出。

设按词元顺序接收,专家2接收词元1、2、3,拒绝词元4。若拒绝路径直接置零且不重新归一化,词元4输出成为

\[ y_4^{\mathrm{drop}}=(\frac{7}{9})(3x_4)=(\frac{7}{3})x_4. \tag{36.23}\]

若改为对剩余接收路径重新归一化,唯一专家3的权重变为一,输出则为 \(3x_4\)。两个结果均不同于原来的 \((\frac{25}{9})x_4\),而且彼此也不相等。溢出策略因此是明确的函数差异,不能只作为不影响结果的内存优化。

本例拒绝1条路径,路径拒绝率为 \(\frac{1}{8}\);有1个词元失去部分路径,词元受影响比例为 \(\frac{1}{4}\);没有词元失去全部路径,完全跳过专家层的比例为零。这三个数字分别回答不同问题。

截断前的归一化分配比例为

\[ f=(\frac{3}{8},\frac{1}{2},\frac{1}{8}),\qquad \bar p=(0.4875,0.2625,0.25). \tag{36.24}\]

代入式(36.15)得到

\[ \mathcal L_{\mathrm{bal}} =3\left[(\frac{3}{8})(0.4875)+(\frac{1}{2})(0.2625)+(\frac{1}{8})(0.25)\right] =1.0359375. \tag{36.25}\]

专家2的申请次数最多,却不是平均概率最高的专家,因为它在多个词元上作为第二选择进入集合。这解释了为什么仅检查平均门控概率不足以诊断实际执行负载。

最大申请负载相对平均负载为 \(\frac{4}{\frac{8}{3}}=1.5\)。若只观察截断后的接收数 \((3,3,1)\),这个热点会显得减弱,但那是拒绝计算造成的结果。质量分析和性能分析都必须保留截断前后的统计。

算法36.1 带显式容量规则的稀疏专家前向

输入:\(X\in\Real^{N\times d}\),路由器 \(W_r,b_r\)\(E_r\) 个专家,选择数 \(k\),容量 \(C\);采用选中内归一化、按词元顺序接收、拒绝路径置零且不二次归一化。

输出:\(Y\in\Real^{N\times d}\),申请数、接收数、拒绝数和路由记录。状态:每专家槽位计数 \(c_e=0\),输出 \(Y=0\),以及接收记录 \((t,e,c,a_{te})\)

  1. \(N=0\),返回空输出和零计数;否则检查 \(1\le k\le E_r\)\(C\ge1\),计算路由分数及全专家概率。非有限分数作为输入失败处理,不进入排序。

  2. 对词元 \(t=1,\ldots,N\),按分数降序选取不同的 \(k\) 个专家;并列时按专家编号升序。计算选中集合内权重,记录所有申请。

  3. 对该词元每条选中路径,若 \(c_e<C\),将 \(x_t\) 写入专家 \(e\) 的下一个槽位,并记录原词元编号与权重;增加 \(c_e\)。否则记录拒绝,不产生专家输入。

  4. 每个专家仅对其接收的 \(c_e\) 个有效槽位计算 \(F_e\)。专家位于远端时,按同一槽位记录分发输入并回收输出。

  5. 对每条接收记录执行 \(Y_{t,:}\leftarrow Y_{t,:}+a_{te}(H_e)_{c,:}\)。全部接收记录合并完毕即停止。

不变量:每个槽位至多对应一个输入;每条接收记录仅合并一次;每专家接收数不超过容量;申请数等于接收数加拒绝数。若无拒绝,每词元合并权重和为一;有拒绝时本算法不再要求这一行和为一。

路由线性投影计算量为 \(O(NdE_r)\);直接完整排序的上界为 \(O(NE_r\log E_r)\),选择算法可避免完整排序。路由记录为 \(O(Nk)\),专家计算按实际接收数计。分发缓冲、网络代价和各专家矩阵大小需另行计入。

算法36.1完整定义了一种容量受限函数。修改排序、二次归一化或拒绝规则时,应先修改这个函数定义,再讨论底层实现是否保持相同结果。它并不要求高性能运行时逐词元串行执行,批量排序和索引操作可以在保持接收语义的前提下并行化。

36.5总参数、激活参数及执行成本

参数容量及词元计算

设非路由部分参数为 \(P_0\),每个路由专家有 \(P_e\) 个参数,路由器有 \(P_r\) 个参数,另有始终执行的共享专家参数 \(P_s\)。在各路由专家同尺寸、没有拒绝的单层计数下,

\[ P_{\mathrm{total}}=P_0+P_r+P_s+E_rP_e,\qquad P_{\mathrm{active}}=P_0+P_r+P_s+kP_e. \tag{36.26}\]

对于多层模型应逐层累加,不能假定每层都是同样的 MoE。每词元激活比例 \(\frac{k}{E_r}\) 仅对应可路由专家部分,不等于整个模型的计算比例。

门控前馈专家忽略偏置时约有 \(P_e=3dd_f\) 个参数。若取 \(d=1024\)\(d_f=4096\)\(E_r=8\)\(k=2\),则单专家参数为 \(12\,582\,912\),全部路由专家为 \(100\,663\,296\),每词元激活的路由专家参数为 \(25\,165\,824\)。这些是给定矩阵形状的算术计数,不是某个发布模型的测量结果。

以每权重两字节存储,全部路由专家权重为 \(201\,326\,592\) 字节,即192 MiB;每词元只执行其中两组专家,不能据此把常驻权重预算改成48 MiB。实际内存还包括路由器、共享层、缓存、激活和工作空间。

激活参数不等于存储容量

稀疏性减少每词元经过的专家路径;总权重与优化器状态仍由完整参数集合及其放置方式决定。专家分片或卸载可以降低单设备常驻量,但需计入通信、加载与峰值缓冲,不能用激活参数直接替代权重预算。

批次专家覆盖率

单个词元选择少量专家,不表示整个批次只读取少量权重。若多个词元选择不同专家,批次专家集合是 \(\bigcup_tS_t\),其大小可以接近 \(E_r\)。增大批量可能改善每个专家的矩阵尺寸,也可能增加本步触及的权重范围,两种效应需要同时分析。

在乘法与加法分别计一次浮点操作、专家主要由矩阵乘法组成的近似下,执行 \(N\) 个词元的专家前向成本约为 \(2NkP_e\)。共享部分、路由排序、激活函数、通信与填充不包含在这个式子中;训练的反向成本还取决于保存和重计算的中间状态。因此不能将稠密模型的 \(6PD\) 估算原样代入总 MoE 参数而忽略稀疏性,也不能只代入激活参数后宣布得到了精确时间。

负载不均及小矩阵效率

若每个专家处理一个局部矩阵乘法,理想均衡时每专家批量约为 \(\frac{Nk}{E_r}\)。专家数增加而总词元数不变,会减小这一批量。矩阵变小可能降低设备利用率,路由和内核启动开销占比则升高。更多专家不是无代价地增加容量。

各设备必须等待相关专家结果返回,执行尾部受最慢专家或最慢设备影响。即使平均专家耗时很低,单个热点也可能决定整个层的完成时间。应关注最大负载、最大接收延迟和跨设备流量,而不只是所有专家平均的利用率。

36.6共享专家、细粒度专家及软混合

(选修)

共享路径及细粒度组合

共享专家(Shared Expert)始终处理每个词元,路由专家则条件执行。一种输出形式为

\[ y_t=F_{\mathrm{shared}}(x_t)+\sum_{e\in S_t}a_{te}F_e(x_t). \tag{36.27}\]

共享项没有被路由权重缩小,其参数与算术开销都应每次计入。让公共变换与条件变换并存,是结构选择;不能仅凭该公式证明某条路径已经学到公共知识。

细粒度专家(Fine-Grained Expert)将较宽专家改为更多较窄专家,同时增加被选专家数。例如将中间宽度除以 \(m\)、专家数乘以 \(m\)、选择数也乘以 \(m\),在忽略路由和其他开销时,专家总参数与每词元专家矩阵计算可以近似保持不变。这增加了组合自由度,却使路由分数、分组和通信的组织更加复杂。DeepSeekMoE 将细粒度划分与共享专家作为两项结构设计(Dai 等 2024);具体型号仍应按精确配置和权重核对。

专家划分不是将一个已训练稠密矩阵任意切块后就自动得到等价 MoE。激活函数、门控组合以及条件选择会改变函数,需要相应训练或经过论证的转换过程。

软专家混合改变专家输入

软专家混合(Soft Mixture of Experts,Soft MoE)并不只是给所有专家保留一个很小的概率。其一种构造先把词元连续混合成有限个专家槽位,再由专家处理槽位,最后将槽位输出混合回词元(Puigcerver 等 2024)

设总槽位数 \(S=E_rp\),每专家有 \(p\) 个槽位,槽位打分矩阵为 \(Z=X\Phi\in\Real^{N\times S}\)。分发权重对每个槽位沿词元轴归一化,合并权重对每个词元沿槽位轴归一化:

\[ D_{ts}=\frac{e^{Z_{ts}}}{\sum_{u=1}^{N}e^{Z_{us}}},\qquad C_{ts}=\frac{e^{Z_{ts}}}{\sum_{v=1}^{S}e^{Z_{tv}}}. \tag{36.28}\]

于是槽位输入 \(\widetilde X=D^{\mathsf T}X\in\Real^{S\times d}\),各槽位按归属专家计算 \(\widetilde Y_s=F_{e(s)}(\widetilde X_s)\),最后 \(Y=C\widetilde Y\)。这里的 \(D\) 是连续混合权重,区别于前文硬选择矩阵 \(D_e\)

每个专家处理的是若干词元的加权组合,不是对每个原始词元都独立执行一次专家。因此它不同于前文的稠密混合。两次归一化也不能交换:分发列和为一,合并行和为一,各自保证不同的混合对象。

这类全局词元混合若直接用于自回归解码,会把未来位置混入早期位置的专家输入。它不能不经修改就替换因果语言模型中的逐位置 MoE;需要另行定义因果槽位构造并分析代价。论文中的结构与实验条件不能自动外推为任意解码器的兼容性证明。

路由方式的约束差异

方式 专家处理对象 离散选择 主要约束
稠密混合 每个词元 无需截断 每词元执行全部专家
词元选择的稀疏混合 选中词元 每词元 top-\(k\) 专家负载不固定
专家选择的稀疏混合 专家挑选的词元 每专家挑选输入 每词元路径数可变
软专家混合 连续混合的槽位 不必使用硬 top-\(k\) 槽位数与混合可见性

专家选择表示按专家分数列挑选有限词元,能够直接约束每专家的接收量,但一个词元可能得到多个专家,也可能没有专家处理。若需要固定每词元覆盖,还要增加额外约束。比较方法时,应分别说明固定的是词元路径数、专家容量还是槽位数量,而不能只用“稀疏”一词代替这些不同定义。

36.7专家并行的通信及训练推理差异

参数放置及词元移动

专家并行(Expert Parallelism,EP)把不同专家参数放在不同设备。路由完成后,词元必须送到相应设备;专家输出再送回原来的词元归属设备。全互换通信(All-to-All)可以表达这种多发送端到多接收端的数据交换,但并不意味着每对设备每步都拥有相同流量。

设远端执行的接收路径比例为 \(\rho\),每激活元素占 \(s_a\) 字节,输入输出宽度均为 \(d\)。忽略元数据和协议开销,前向分发与返回的通信量级为

\[ V_{\mathrm{comm}}\approx2\rho Nkd s_a. \tag{36.29}\]

该式按每条远端路径的单向有效载荷计数,并将发送输入与返回输出各计一次;不同统计工具若同时累计发送和接收端会得到不同口径。张量并行、专家复制、压缩传输与容量拒绝都会改变实际数值。

理想传输时间至少受有效带宽与启动延迟限制。若共有 \(m\) 轮不可忽略的启动,延迟为 \(\ell\)、有效带宽为 \(b_{\mathrm{eff}}\),可以用

\[ t_{\mathrm{comm}}\gtrsim m\ell+\frac{V_{\mathrm{comm}}}{b_{\mathrm{eff}}} \tag{36.30}\]

理解瓶颈;拥塞、同步和拓扑争用还会增加时间。计算与通信只有在依赖允许时才能重叠,不能把所有理论算力与所有网络带宽同时按满负荷利用。

专家并行及数据并行

数据并行复制同一模型参数并处理不同数据,专家并行则让一次前向调用不同设备上的不同专家参数。专家参数只在拥有同一专家副本的数据并行组内同步梯度,路由器与共享参数又可能采用不同复制范围。组定义错误会造成不同参数被错误求和,或应保持一致的副本没有同步。

同一设备可容纳多个专家,单个专家也可进一步进行张量切分。专家编号、物理设备编号和张量分片编号不是一个索引。模型配置记录逻辑专家结构,运行时放置表记录执行位置,检查点必须能从分片还原这些关系。

训练及推理张量形状

训练往往同时处理较多词元,专家批次相对较大,但需保存反向状态并执行反向通信。各专家梯度还可能十分稀疏或不均匀,优化器状态则不能仅按当前激活量估计。激活重计算需要重现路由选择、随机噪声与容量接收结果,否则反向不再对应原前向路径。

预填充可以并行处理一段输入,而逐词元解码每个请求每步通常只新增一个位置。并发较小时,各专家收到的矩阵可能很小,路由、权重读取和通信启动的比例增大。增加请求合批可能提高专家利用率,却引入等待时间及更复杂的容量竞争。

训练中接受的路径拒绝比例不应直接移植到服务端。一个被拒绝的关键词元可能改变后续生成,且输出依赖批次组成会使复现和回归更困难。推理可以采用变长无拒绝执行、扩大容量或其他明确方案,但每项变化都需要评估其函数、质量和延迟影响。

36.8可执行结构还原

模型名称中出现 MoE 或激活参数数字,只能提供结构线索。需要按每层核对路由专家数、共享专家数、每词元选择数、中间宽度、门控归一化、路由激活、容量策略及哪些层使用 MoE。再将这些配置与路由矩阵、专家参数张量和前向定义相互对应。

例如路由投影在本章约定中为 \(d\times E_r\);某些库按输出乘输入布局保存为 \(E_r\times d\),这属于转置约定,不应仅凭形状方向不同就断言错误。量化参数可能按组打包,存储元素数也不能直接替代原始参数计数。对于共享专家,既要确认参数存在,也要确认前向确实每次调用,而不是仅凭张量名称推断。

路由热力图说明某次输入如何被分配,不证明专家专业化;权重可加载说明运行时接受了张量,不证明合并规则与原模型一致;总参数和激活参数相符,也不证明设备拓扑能达到预期吞吐。完整判断需要把结构事实、数学语义、质量证据与系统测量分开记录。


  1. 专家编号是模型内部的参数索引。对专家参数及路由输出通道作一致置换,可以表示同一个函数,因此编号本身不具有稳定的外部语义。专家是否形成可解释分工,需要独立的行为与干预证据。↩︎

  2. 不同论文与运行时对 token dropping 的统计口径可能不同:拒绝的专家分配数、至少失去一条路径的词元数、以及失去全部路径的词元数,应分别报告。↩︎

WORKBOOK / 习题

配套习题与解析

先独立作答,再展开参考解析。选修题保留原书标记。

习题 36.1

\(E_r=4,k=2\)、分数 \((\log0.4,\log0.3,\log0.2,\log0.1)\),计算选中内归一化和保留全专家概率的两种权重。设已选专家输出为 \((1,0)\)\((0,2)\),比较最终输出。

展开参考解析

选中专家1、2。选中内权重为 \((\frac{4}{7},\frac{3}{7})\),输出 \((\frac{4}{7},\frac{6}{7})\);保留全专家概率时为(.4,.3),输出(.4,.6)。后者权重总和.7,缺少的.3没有被重分配,所以两种输出不同;选择集合相同不代表函数相同。

习题 36.2

从式36.8证明:固定选中集合且 \(k=1\) 时,选中内归一化的路由任务梯度为零。再解释保留全专家概率为何不必为零。

展开参考解析

固定集合、选中归一化时 \(y=\sum_iw_iF_i\),对选中logit的梯度为 \(w_jg^{\mathsf T}(F_j-y)\)\(k=1\)\(w_j=1,y=F_j\),因此为零。若保留全专家Softmax概率,\(y=p_jF_j\)\(p_j\) 仍受所有logit影响,故通常有非零梯度。选择边界本身不可微,局部零梯度不表示改变排名永远不改变输出。

习题 36.3选修

对某个词元,两个专家输出相同。分析选中内归一化时路由权重的任务梯度,并说明这是否意味着专家参数梯度也为零。

展开参考解析

两个选中专家均输出v时,归一化混合始终为v,\(F_j-y=0\),路由任务梯度为零。但专家输出的上游梯度分别是 \(w_jg\),专家参数通过其雅可比仍可得到非零梯度。路由选择无差别不等于专家停止学习;辅助均衡也可能继续更新路由。

习题 36.4

构造 \(E_rC\ge Nk\) 但仍存在容量溢出的例子。分别计算路径拒绝率、受影响词元比例和完全失去专家处理的比例。

展开参考解析

\(N=2\)\(k=1\)\(E=2\)、每专家容量\(C=1\),则 \(EC=Nk=2\)。若两个词元都选择专家1,按词元编号接收一个、拒绝另一个。路径拒绝率\(\frac{1}{2}\),受影响词元率\(\frac{1}{2}\),完全失去专家处理率\(\frac{1}{2}\)。专家2的空闲槽位不会自动把已选路径变为合法替代;重新路由是另一策略。

习题 36.5选修

在本章数值例题中,把专家2接收规则改为按该专家的路由权重降序选择,并列时按词元编号升序接收。确定哪条路径被拒绝,再比较输出变化。说明规则不同为何不能只当成排序实现细节。

展开参考解析

专家2的四个权重为 \(\frac{2}{9},\frac{1}{3},\frac{7}{18},\frac{2}{9}\)。按降序为词元3、2,然后词元1与4并列;题设补充并列按词元编号升序,故仍拒绝词元4,输出与正文的同序接收例相同。若并列优先较大编号,则拒绝词元1:不重归一化时 \(y_1=(\frac{7}{9})x_1\),其余恢复原无拒绝输出。该例中的是否改变依赖明确的并列规则;不能凭方法名称声称一定改善。

习题 36.6

推导式36.17。如果辅助统计覆盖两个大小不同的本地批次,写出正确的全局平均概率,解释局部均值直接平均的问题。

展开参考解析

\(L=E\sum_if_i\bar p_i\)\(\bar p_i=N^{-1}\sum_tp_{ti}\),并对离散统计f停止梯度。由Softmax导数得 \(\frac{\partial L}{\partial z_{tj}}=(\frac{E}{N})p_{tj}(f_j-\sum_if_ip_{ti})\)。两个批次大小 \(N_1,N_2\) 时,正确全局概率为 \(\frac{N_1\bar p^{(1)}+N_2\bar p^{(2)}}{N_1+N_2}\);直接各占一半将小批次词元过度加权,路由计数也须在同一统计范围聚合。

习题 36.7

\(d=512,d_f=2048,E_r=16,k=2\),忽略偏置,计算门控专家的总参数与每词元激活参数。加入两个始终执行的同尺寸共享专家后,两个计数分别如何变化?

展开参考解析

一个无偏置门控专家含三矩阵,共 \(3dd_f=3145728\) 参数。16个路由专家总计50331648,每词元选2个执行6291456。新增2个同尺寸共享专家后,总量56623104,激活量12582912。这里只数专家矩阵;若计入路由器还需加 \(dE=8192\),注意力、归一化及输出头另计。

习题 36.8选修

保持每词元专家矩阵计算近似不变,将专家中间宽度减半、专家数与选择数加倍。哪些成本不能随专家矩阵算术量保持不变?

展开参考解析

\(d_f\)减半、E和k加倍保持 \(3kdd_f\) 及专家总矩阵参数近似不变,但路由logit数量翻倍,选择/分发元数据和路径数增多。每专家矩阵更窄、每路径批次更小,内核利用率可能下降;通信消息数量、负载尾部与合并开销也变化。需同硬件同token预算测量,不能由乘加次数守恒推出时延守恒。

习题 36.9选修

对 Soft MoE 分发与合并矩阵,分别验证列和与行和。举例说明若先混合全序列词元再调用专家,为什么早期输出可能读取未来输入。

展开参考解析

若分发权重 \(D_{ts}=\frac{\exp z_{ts}}{\sum_u\exp z_{us}}\),则每个槽位列和为1;合并权重 \(C_{ts}=\frac{\exp z_{ts}}{\sum_r\exp z_{tr}}\),则每个词元行和为1。两词元一个槽位、权重各\(\frac{1}{2}\)时,槽输入 \(\frac{x_1+x_2}{2}\),若专家为恒等,则第一个输出直接含x2。仅在最终合并处加因果mask不能消除已经混入的未来信息。

习题 36.10

某模型在训练大批次中专家利用率较高,在低并发解码中却很慢。用矩阵形状、权重读取、通信延迟和路由负载解释至少三种可能机制。

展开参考解析

低并发使每专家输入行数很少,矩阵乘法退化为低利用率的小矩阵;大量专家权重可能仍需从显存读取,计算量低但带宽占主导。跨设备每轮小消息暴露固定通信延迟,路由偏斜又使少数专家成为串行尾部。应记录逐专家批量、读取字节、通信耗时和请求尾延迟,再决定合批、放置或卸载策略。

习题 36.11

某报告只给出平均专家负载和激活参数量,便断言单卡内存与吞吐满足上线需求。说明尚缺哪些可支持结论的结构信息和运行证据。

展开参考解析

需要逐层专家数/容量/共享布局、总驻留权重精度、路由和KV状态、峰值工作区、设备分片及卸载方案;还需请求长度并发、批次分布、溢出与重试、内核和通信数据。以实际候选模型在目标硬件测吞吐、首词延迟和尾延迟,并验证任务质量。平均负载会掩盖热点,激活参数也不等于显存。

REFERENCES

参考文献

Dai, Damai, Chengqi Deng, Chenggang Zhao, 等. 2024. 《DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models》. 2024年1月11日. https://arxiv.org/abs/2401.06066.
Fedus, William, Barret Zoph, 和 Noam Shazeer. 2022. 《Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity》. Journal of Machine Learning Research 23 (120): 1–39. https://jmlr.org/papers/v23/21-0998.html.
Puigcerver, Joan, Carlos Riquelme, Basil Mustafa, 和 Neil Houlsby. 2024. 《From Sparse to Soft Mixtures of Experts》. 2024年. https://arxiv.org/abs/2308.00951.
Shazeer, Noam, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, 和 Jeff Dean. 2017. 《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》. 2017年1月23日. https://arxiv.org/abs/1701.06538.

搜索全书

搜索全书正文、习题与解析