1. MoE 简介
MoE(Mixed Expert Models),混合专家模型。在 Transformer 的 FFN 中,有一个重要的观察是,其计算过程中的神经元激活是非常稀疏的,在一次计算中只有 90%的输入激活不到 5%的神经元,省略计算未激活的神经元就可以大大降低冗余计算。因此,通过训练可以将相关神经元有效组织在一起,这样就形成了 MoE 整体框架,它是一个整体参数较大,而推理时参数规模较小的模型。
1.1 密集模型与稀疏模型
密集模型一般指 Transformer 模型中密集的前馈神经网络 (Feedforward Neural Network, FFNN) ,下图为标准的 decoder-only Transformer 模型中 FFNN
- 整体 Decoder-only Transformer 结构(图一):每个输入 token 会被转换为一个向量(embedding),并加上位置编码,Decoder Blocks 通常是堆叠的 Transformer 层,每个 decoder block 都包括:Masked Self-Attention 掩码自注意力、LayerNorm 归一化层、FFNN
- Encoder 使用标准 Self-Attention 编码时需要全局上下文信息,Decoder 解码时按 token 顺序生成,不能看到“未来”token,在计算权重时屏蔽掉当前 token 之后的位置
- 每一层 decoder block 有两个子层(图二):Masked Self-Attention,通过 attention 机制捕捉 token 间的依赖关系,加入 LayerNorm + 残差连接(Residual);FFNN 子层也是在 LayerNorm 和残差连接后使用,输入是一个长度为 512 的embedding,所有连接都是 密集的,即每个输入维度都连接到每个输出神经元
- Dense FFNN 是 Transformer block 中计算最重的部分之一
- 所谓 “密集模型” 指的是:所有 token 的每层 FFNN 都计算完整网络,所有参数(包括权重和偏置项)都会被激活,所有参数都被用于计算输出,不跳过任何部分
- 相比之下,MoE 模型 会将 FFNN 替换成多个“专家”,并且只激活其中一部分,从而提升参数规模但保持计算量不变(稀疏激活)
稀疏模型(Sparse Models)仅激活总参数中的一部分,可以将密集模型分解为多个部分(称为专家):
- 其核心思想是,每个专家在训练过程中学习不同的信息。而在推理时,仅使用与当前任务最相关的特定专家
- MoE 有两种类型:稀疏专家混合模型(Sparse Mixture of Experts)和密集专家混合模型(Dense Mixture of Experts),两者都使用路由器来选择专家,但稀疏 MoE 只选择少数几个专家,而密集 MoE 则选择全部专家,但可能会以不同的分布进行选择。
- 目前的 LLM 中 MoE 通常指的是稀疏 MoE,MoE 的一个显著优势是能够在远少于 Dense 模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下,可以显著扩大模型或数据集的规模。特别是在预训练阶段,与稠密模型相比,混合专家模型通常能够更快地达到相同的质量水平
稠密和稀疏模型小结:以人类分工为例,稠密就是类似手工业时代的生产模式,在这种模式下,每个工人(即神经元)都需要参与处理所有类型的任务,就像手工业时代的工匠需要精通产品制作的各个环节,掌握所有的生产技能,这种方法虽然直观且易于实现,但在面对复杂多变的任务时,往往效率低下且难以扩展;稀疏 MoE 则是工业革命之后的分工模式:每个岗位(专家)只需要完成一部分生产任务,这种分工的方式大大提高了生产效率
1.2 MoE 架构
MoE 基于 Transformer 架构,主要由两部分组成:
- 稀疏 MoE 层:这些层代替了传统 Transformer 模型中密集的前馈网络 (FFN) 层,每个 MoE 层包含一定数量的 expert,每个 expert 都是一个神经网络,每次只有其中部分 expert 参与计算。从这个角度看,整体上的计算是稀疏的,例如 DeepSeek V3 的 expert 总数为 256,推理 expert 数为 8 个。
- DeepSeek‐V3 的 MoE 模块中,主要包含两类专家:
- 路由专家(Routed Experts):每个 MoE 层包含 256 个路由专家,这些专家主要负责处理输入中某些特定、专业化的特征
- 共享专家(Shared Expert):每个 MoE 层中还有 1 个共享专家,用于捕捉通用的、全局性的知识,为所有输入提供基本的特征提取支持
- 注意,专家不是按领域划分,而是按 token 类型或上下文中的语义功能划分的。这些“专家”并不像人类领域中的“心理学”或“生物学”专家那样在特定学科上表现出高度专业化。实际上,它们更多是从词汇层次上学习句法信息(这种多样性是通过端到端的训练自发形成的),例如:Expert 1 专注于处理标点符号,Expert 2 专注于处理动词,Expert 3 处理连词,Expert 4 处理视觉描述相关的词汇,这些分类只是示意说明,实际训练过程中,专家会根据 token 使用频率、上下文分布等自动学习其“擅长”的 token 类型。
- 门控网络或路由(Gate Network 或 Router):这个部分用于决定哪些 token 被发送到哪个专家。一个 token 甚至可以被发送到多个专家。
- 门控网络一开始的路由决策是近似随机的,但随着专家逐步积累专长,门控网络也会调整其路由策略:
- 正反馈循环:如果某个专家因早期获得较多特定类型数据而表现出色,门控网络便倾向于将更多此类数据路由给它
- 专家和路由网络的协同演化:专家因接收到较多特定数据而“专长”,而门控网络根据反馈不断更新参数,使得路由更加精准
- 门控网络一开始的路由决策是近似随机的,但随着专家逐步积累专长,门控网络也会调整其路由策略:
- MoE 整体框架,它是一个整体参数较大,而推理时参数规模较小的模型,例如 DeepSeek V3 总参数量 671B,推理参数 37B,推理参数占比 5.5%
1.3 专家架构
在 MoE 架构中,专家是指训练好的子网络(神经网络或层),它们专门处理特定的数据或任务,专家和门控机制都通过梯度下降与其他网络参数一起进行联合训练。
一般来说,MoE 中的每个专家都是具有相同架构的前馈神经网络 FFNN,替换掉稠密模型中的FFN(图一),专家本身就是完整的FFNN:
- FFN 层被分解为多个专家,每个专家实际上是 FFN 参数的一个子集。专家并不是对FFN的平均切分,实际上可以任意指定每个expert的大小,每个expert甚至可以大于原来单个FFN层,这并不会改变MoE的核心思想:对于一个token,部分专家的计算量要小于所有专家的计算量
- 不同的专家会专注于不同的主题(图二),每个专家模型可以专注于处理自己接受到的输入数据,学习数据中的一种特定模式或特征
大多数 LLM 有多个 Decoder Blocks,因此在生成文本之前通过多个专家,不同 tokens 被选中的专家可能各不相同,这导致了不同的“路径”被选择:
- 专家学习了什么?MoE 里的“专家”是一种拟人的形象化的说法,专家学习的信息比整个领域的信息更加细粒度,并不专攻“心理学”或“生物学”等特定领域,最多只是在单词层面学习句法信息:更具体地说,它们擅长于在特定上下文中处理特定的 tokens
- 下图为 Mixtral 8x7B 论文中,每个 token 都被其选择的第一个专家进行了着色。在代码中,缩进标记总是分配给相同的专家:
1.4 门控网络
路由器(或门控网络)本身也是一个 FFNN,它根据特定的输入选择专家。路由器会输出概率值,并利用这些概率来选择最匹配的专家,专家层返回被选定专家的输出,并乘以门控值(选择概率),路由器和专家(其中仅选择少部分)共同构成了 MoE 层:
门控网络可以说是 MoE 中最重要的组件,因为它不仅决定了推理时要选择哪些专家,还决定了训练时的选择。
- 最基本的形式是,将输入(x)与路由器的权重矩阵(W)相乘,然后对输出应用 SoftMax 操作,为每个专家创建一个概率分布 G(x),路由器利用这个概率分布来为给定的输入选择最匹配的专家,最后,将每个路由器的输出与各自选择的专家输出相乘,并将结果相加
然而,这个简单的功能往往会导致路由器总是选择相同的专家,因为某些专家可能比其他专家学习得更快:这不仅会导致专家选择的不均匀分布,还会导致某些专家几乎没有被训练过。这会在训练和推理过程中引发问题。因此,希望在训练和推理期间,各个专家的使用具有同等的重要性,这就是所谓的负载平衡。某种程度上,这是为了防止模型在同一组专家上过拟合。
1.5 负载均衡
因为 MoE 的结构由多个 expert 构成,每次只选择部分执行,如果 Expert 的分配不均衡,就可能出现如下局面:某些 Expert 几乎一直闲置,浪费算力;某些Expert要处理的Token太多,根本忙不过来,只能Token Drop(即放弃处理部分Token)。这样的话,既浪费了显存,又使模型没有达到预期的参数量。
因此负载均衡就显得非常重要,MoE 的负载均衡问题主要体现在两个层面:
- expert 的负载均衡
- GPU 的计算均衡
KeepTopK
一种对路由器进行负载平衡的方法是使用一个简单的扩展策略,称为 KeepTopK。
1. 通过引入可训练的(高斯)噪声,可以防止总是选择相同的专家:
除了希望被激活的前 k 个专家(例如 2 个)以外的所有专家权重都将被设为 -∞,将这些权重设为 -∞ 时,SoftMax 操作后的输出概率将变为 0:
2. 不使用额外的噪声,KeepTopK 策略会将每个 token 路由到若干选定的专家。这种方法被称为 Token 选择策略(Token Choice),它允许一个给定的 token 被路由到一个专家(top-1 路由)或者被路由到多个专家(top-k 路由):
辅助损失
为了在训练过程中实现专家之间的均衡使用,除了常规损失函数之外,还引入了一个辅助损失(也称为负载均衡损失)。这个辅助损失起到了约束作用,使模型在训练时对所有专家赋予相似的重要性,避免某些专家被频繁调用,而其他专家很少参与计算,鼓励所有专家被均匀使用,避免“热门专家”和“闲置专家”的极端情况。
1. 计算每个专家的重要性(Importance)
- 在一个 batch 中,每个 token 都会被路由到少数几个专家
- 将整个 batch 中每个专家被选中的概率值累加,计算每个专家在整个 batch 中被路由器选中的概率总和,得到这个专家的 Importance(重要性)总和,这个值反映了各个专家在当前批次中的使用频率
2. 引入 Coefficient of Variation(CV,变异系数),衡量专家使用不均的程度,表示各个专家的重要性得分之间的差异程度(标准差除以均值):
例如,如果重要性得分之间的差异较大,那么 CV 值就会较高,如果所有专家的得分都相似,则 CV 值较低(期望的情况),例如左边柱状图中,重要性差异大(0.8 vs. 0.05),CV = 1.11 ➜ 不理想;右边柱状图中,专家得分接近(0.3 ~ 0.2),CV = 0.19 ➜ 理想情况:
3. 构造辅助损失项(Auxiliary Loss)
通过使用这个 CV 得分,可以在训练过程中更新辅助损失,使其尽可能降低 CV 得分,其中 w 是一个常数系数(scaling factor),CV 值越大,辅助损失越高,模型就会被惩罚更多,进而优化路由机制,使专家分布趋于均衡:
辅助损失将作为一个独立的损失项(是训练总 loss 的一部分,但单独计算),参与训练优化。
专家容量
专家的不平衡不仅体现在被选中的专家上,还体现在分配给这些专家的 token 分布上。
例如,如果输入 token 被不成比例地分配到某些专家上,而不是平均分配,这可能导致某些专家的训练不足(图一);解决这个问题的方法是限制每个专家能够处理的 token 数量,即专家容量(Expert Capacity)。当一个专家达到其容量时,多余的 token 将被分配到下一个专家(图二);如果两个专家都达到了其容量,token 将不会被任何专家处理,而是直接传递到下一层。这种情况被称为 token 溢出(token overflow)(图三)。
以上为MoE的基本原理(门控函数等都没有具体展开),现在主流的模型架构中基本都有 MoE 变体,也有很多工作对MoE模型做了优化,此处不再一一赘述。
MoE 整体计算流程
1. Routing:本质是对输入进行分类的过程,选择最适合处理输入的专家模型。在语言模型的应用中,当输入 token 通过 MoE 层时,Token 通过和 Router 的权重矩阵相乘得到一个 Expert Indices(决策矩阵)和一个概率张量,即索引和概率:
- Expert indices 是 expert-to-token 映射,用于指示每个 token 被分配给了哪个 expert,即张量中第 i 个值代表本 token 应该分配到第 i 个专家
- Probabilities 张量是分配置信度的概率,其中第 i 个值代表这个专家对于该 token 最终结果的权重
2. Permutation(排列/置换):根据路由决策(expert-to-token映射)将 Token 分配给对应的专家,中间可能会有drop操作
3. Computation:每个专家网络并行处理其分配到的 token,计算输出。把输入矩阵 𝑥 与专家网络的权重矩阵相乘 𝑦 = 𝑥 × 𝑊
4. Un-Permutation:收集专家的计算结果。这是 Permutation 的逆运算,将从各个 experts 收集到的处理后的 tokens 组合成一个完整的序列,这个序列保持了原始 tokens 的顺序。即,将每个专家网络的输出根据原始的 token 顺序重新排列。接着使用 Routing 步骤生成的分配置信度概率对结果进行加权求和,以得到最终的模型输出,然后将这个结果继续向下游处理
Permutation 的主要作用:
- 分发 token。Permutation 会依据 Expert Indices 构建本地的置换 Token 位置后的临时矩阵(将输入Token根据路由结果重新排列),这样可以把属于每个专家的 token 分别放在一起,然后把 tokens 送给对应的专家。比如上图中,“the” 和 “jumped” 应该分配给专家1,所以就把它们放在一起,“quick” 和 “fox" 都应该被发送给专家2,所以把它们也放在一起
- 维持 token 和 expert 的顺序。因为一个batch里有很多token,将这些token发往不同的expert做计算后,专家输出结果的顺序肯定是打乱的,所以需要通过一种方式追踪顺序,把 token permute 回正常的位置再输入下一层网络。通过构建的矩阵,Permutation 在计算时,就可以维护住这种顺序
- 负载均衡。Permutation 可以实现输入数据在不同专家之间的合理分配,平衡各个专家的计算负载。不同的输入样本可能对不同专家的计算资源需求不同。通过对输入样本进行置换,使得每个专家能够相对均匀地接收到需要处理的样本,避免某些专家过度使用而其他专家闲置的情况。
- 增加多样性。Permutation 可以增加模型对输入数据处理的多样性。因为不同的置换顺序可能会导致不同专家组合对数据进行处理,从而挖掘出数据的不同特征。
从代数的角度来看,MoE 计算实际上是对 token 进行一次置换群的操作,那么MoE实现的本质问题是:基于 Permutation 矩阵后构建的稀疏矩阵乘法如何进行并行。
2. EP 原理
专家并行(Expert Parallelism+, EP)的思路是将 MoE 层中不同的专家分配到不同的计算设备上,每个设备负责存储和计算部分专家,而所有非专家层则在设备间复制。
如图所示,一个包含 6 个专家的 MoE 模型在 EP=2 时的专家分布情况:
可以把专家并行理解成模型并行的一种形态(模型被切分成多份),但是输入的数据又是不同的(DP),因此 token 在经过 Router 之后,可能会选中别的卡上的专家,此时就需要将这些 token 发送过去,即 EP 进程组内需要通过 AlI2AII 通信交换 token。
2.1 EP 推理流程
如图为在一个 MoE 层的推理流程,非专家层(如 Input、Attention、Gate)在所有 GPU 上复制部署,每个 GPU 卡上只包含一部分 expert,token 在一个 MoE 层处理完成后,输入到下一层进行处理。专家并行的推理流程包括以下几个顺序操作:门控路由、输入编码、All-to-All dispatch、专家计算、All-to-All combine和输出解码
- 门控路由:每个 GPU 上会有一个门控模块(Gate),对输入 token 做出判断:这个 token 应该被哪个 expert 处理(expert 是稀疏选择,通常只选择 top-1 或 top-2)
- 输入编码与 Permutation(排列):每一张卡的输入是一个完整的任务,包含多个 batch 与序列(sequence),被路由到同一个 expert 的所有 token,排列在内存中的一块连续空间里,这个“排列”过程本质是把原始 token 的顺序打乱,把属于同一个 expert 的 token 聚集在一起,便于发送和批量计算
- All-to-All dispatch(分发):每个设备根据之前的排列,把 token 按照 expert 所在的设备位置,通过 All-to-All 通信机制发送过去。比如:GPU 0 上的某些 token,路由结果是 expert 2,这个 expert 在 GPU 1,那么这些 token 就通过 All-to-All 发送到 GPU 1。
- 专家计算:每个 GPU 上含有多个 expert,所有 token 被传输到目标 expert 所在卡之后,expert 在本地计算,每个 expert 仅处理被分配给自己的 token
- All-to-All Combine 与输出解码(Output Decode):expert 完成计算后,每个 token 的结果通过 All-to-All 返回到原来的 GPU 卡。然后根据最初记录的排列顺序,进行反排列(de-permutation),将 token 恢复成原始顺序,拼接成完整的输出。最终,token 在一个 MoE 层处理完成后,继续进入下一层
2.2 EP 与 DP、TP、PP
下图展示了 四种 MoE(Mixture of Experts)中的并行策略组合示意图:
(a) Data + Expert Parallelism(数据并行 + 专家并行):输入 token X1 X2 分配在两个 GPU 上,每张卡上的 gate 决定当前 token 应该被分发给哪个 expert,使用 All-to-All Dispatch 把 token 发往 expert 所在的卡,Encode 准备要跨卡传输的数据,加上位置等信息,All-to-All Combine 再把结果聚合回来, Decode 将合并后的结果还原到原 token 顺序、格式,最后继续执行后续 Add + Normalize,输出 Y1 / Y2
(b) Data + Expert + Tensor Parallelism(数据并行 + 专家并行 + 张量并行):加入了 Tensor 并行:模型参数本身被张量切分,不同 GPU 协同完成一个 FFN。Self-Attention 本身也用了 Tensor 并行:比如一个 Attention 权重被横向切为两半 → GPU1 和 GPU2 分别算一部分;每个 expert(比如 FFN1)也被横向拆开,两张卡各算一半;Gate 还是各自独立计算(每张卡计算本地 token);没有 All-to-All,因为 Tensor 并行下,每张卡都需要协同计算同一个 expert 的一部分;所有 GPU 对参数是 切片共享的,非副本;最后通过 All-Reduce 聚合 tensor 切片,才能得到完整输出
(c) Data + Expert + Pipeline Parallelism(数据并行 + 专家并行 + 流水线并行):把 Transformer 模型切成多个 stage,每张卡负责不同层(层间流水线),并且专家分布在不同卡上,每张 GPU 执行模型的一部分。GPU1 和 GPU2 执行 MoE block 的第 1 层(SA1 → FFN11/12);GPU3 和 GPU4 执行第 2 层(SA2 → FFN21/22);每层内依然有 Expert 并行,每个 FFN expert 只存在于某一张卡上,每一阶段处理完后,将结果传到下一阶段卡上,Pipeline 并行让每张 GPU 专注于一部分层的计算
(d) Expert + Tensor Parallelism(专家并行 + 张量并行):专家部分内部用 Tensor 并行,但不做数据并行或流水线分层。一张卡内有多个专家 FFN,这些 experts 是分别独立的,但内部每个 expert 的 FFN 由多个 GPU 协同计算,Gate 分配 token 后,在每张卡上 local route,不需要 All-to-All,每张卡只执行本地 expert 负责的 token 部分,Tensor 并行用于分担一个 expert 的模型计算负担
关系总结:
- DP是“复制模型,切分数据”
- PP是“切分模型层,数据流式处理”
- TP是“切分层内算子(权重矩阵)”
- EP是“切分MoE层内的专家”
下图是从模型参数切分和数据切分的角度(只考虑 FFN 层)来比较几种并行策略:(https://arxiv.org/pdf/2101.03961)
每个 4×4 的虚线网格代表 16 个 Device,阴影方块是该 Device 上包含的数据(模型权重或token)。第一行:模型权重如何在 Device 之间分配,阴影方块的大小为不同的权重矩阵大小,每种颜色都标识了一个权重矩阵。第二行:数据batch如何在 Device 之间分割,不同颜色代表不同 token。
- 第一列:数据并行,上方所有设备(1-16)都有相同、全部的模型参数;下方每个设备只有一个数据分片,共 16 个数据片。
- 第二列:模型并行,上方所有设备(1-16)都只有模型参数的一部分,共 16 个分片;下方所有设备使用共同的一份数据。
- 第三列:模型并行+数据并行,设备分为 4 组。上方表明每组都有完整的模型参数副本,但是每组内的设备只有参数的一部分;下方表明数据分为 4 个切片,每组对应一个数据切片。
- 第四列:专家并行+数据并行,设备分为 16 组(1-16)。上方表明每一个设备都有不同的专家,共 16 个专家。下方表明每个设备都有不同的数据分片(Token),共 16 个数据分片,一个专家对应一个分片。
- 第五列:专家并行+模型并行+数据并行,有 4 组设备(1-4,5-8,9-12,13-16)。上方表明有 4 个专家,每个专家分布在对应的 4 个设备上,比如绿色专家分布在 5,6,7,8 设备上。下方表明有 4 个数据分片,每组设备(每个专家)对应一个数据分片,一组里的 4 个设备共享一份数据分片。
2.3 AlltoAll 通信
AlltoAll 是集合通信库(NCCL)中一种常见的通信机制,用于多个设备之间进行数据交换。AlltoAlll 操作允许每个参与的设备将其本地数据分发到其他设备,同时从其他设备接收数据。AlltoAll 的作用相当于分布式转置Transpose 操作。具体如下图所示,可以看到,GPU0 把自己收到的 4 个绿色的块分配给了全部 4 个 GPU 。
结合MoE来看,需要通过 AlltoAll 将 token 发去指定的 expert 做计算,再通过 AlltoAll 将计算结果返回。假设我们有一个 4 张卡的 GPU 集群。下图步骤 1 表示首次做 AlltoAll(All-to-All Dispatch)的过程,这个过程的目的是将 token 发去对应的 expert 上进行计算,对比一下左侧图和中间图的数据块排布,会发现 AlltoAll 就相当于做了一次矩阵转置。因此通过 AlltoAll,我们就让数据块去到了它对应的位置:A0、B0、C0和D0去GPU0,A1、B1、C1和D1去GPU1,以此类推。而为了实现这种转置,我们必须提前把 token 做分块排序,让它按照要去的专家位置排好。步骤 2 表示第二次做 AlltoAll(All-to-All Combine)的过程,这个过程的目的是将MoE算完的token再返回给各卡,原理和上述一致。
在 All-to-All Dispatch 操作之前准备好 All-to-All 输入的过程叫输入编码,即需要对本 GPU 上的 local token 按照路由结果进行重新排列(permute 或 group),将发往同一个专家的 token 进行分组,让属于同一个专家的 token 在 memory 中邻近,方便批量发送。随后,这些 token 通过 All-to-All Dispatch 通信发送到对应的 expert rank(每个 EP rank 上只包含一部分 expert)。在 All-to-All Combine 操作之后需要把这些 token 的结果需要送回原先对应的输入位置,还原顺序,这叫输出解码。大多数流行的 DL 框架利用 NCCL 的点对点(P2P)通信来实现线性 All-to-All 算法。
这个伪代码实现的是:在 n 个 GPU 间进行 Linear All-to-All 通信,也就是每张卡将它本地输入的第 r 块发给第 r 张卡,同时接收来自第 r 张卡的一块,是一种 线性复杂度的 All-to-All 实现(每 GPU 发/收 n 次)。这段代码完成的操作是:每个 GPU 把自己的输入 input 分成 n 块,每块分别发送到 n 个 GPU 的对应位置,同时接收从 n 个 GPU 发来的块,实现一种简化版的 All-to-All:每块数据从 sender 到对应的 receiver,用的是 NCCL 的点对点通信接口(而不是全局的 AllToAll)
2.4 MoE 分布式计算过程
下图为 TP、DP 和 EP 互相结合的 MoE 分布式计算过程,其中并行划分为:Transformer layer 在四张 GPU 卡上并行,TensorParallel=2,DataParallel=2,ExpertParallel=2。张量并行组(GPU0 和 GPU1)在处理token 1和 token 2(w1、w2),而张量并行组(GPU 2 和 GPU 3)处理 token 3和 token 4(w3、w4)。
整体计算过程分为7步:
- 每个 GPU 首先计算各自自注意力块的分区。{GPU0,GPU1} 与 {GPU2,GPU3} 之间是数据并行;
- 在每个自注意力块的张量并行组内,每个 GPU 执行 All-Reduce 聚合它们各自 token 的完整输出(a1、a2、a3和a4),聚合自注意力块的张量并行输出;
- 每个 GPU 对于自己的本地 token 执行 MoE 路由功能;
- 根据路由结果,将 token 发送到对应的专家。假设路由函数将 a1 和 a3 映射到专家1,且将 a2 和 a4 映射到专家2。
- 在 All-to-All 通信操作完成之后,每个专家分块在自己所在的 GPU 进行计算;
- 在 FFN 的张量并行组内执行 All-reduce 操作来聚合完整输出。这一步是聚合FFN的张量并行输出;
- 执行 All-to-All 通信操作,将 token 发回它们的原始所在 GPU
2.5 通信复杂度
如上图左边所示,TP 的基本思路是将模型参数切分到多个 GPU 进行计算。面对参数量大幅增大但计算量不变(每个 token 只激活一小部分专家)的 MoE 架构,TP 方案暴露出两大核心问题:通信会成为瓶颈,内存也会逐渐成为瓶颈。
1. 通讯会成为瓶颈
- 假设每次推理一个 batch 里一共有 S 个 token,hidden dimension 是 D,每个 GPU 的通信量 = 2 × S × D
- 一次 AllReduce 实际上是一个 scatter + gather,两次通信:一次发送(send),一次接收(recv)
- 和 TP size(并行粒度)无关 —— 也就是说你分得再细,也改变不了 hidden dim 的维度,也就改变不了通信量,通信开销不会因为增加 TP size 而减少,TP 的通信开销与 TP 粒度无关
- TP 的部署方式是每个 GPU 上都需要 AllReduce 来聚合所有 input tokens 的激活值,无论划分多少份,都是不会改变结果矩阵的大小,TP 中的 AllReduce 的通信量会随着整个部署实例 batch size 的增大而增大,并且即使增大 TP 并行度,通信量也不会变小——这是因为 TP 是源自稠密模型的设计范式
2. 内存会成为瓶颈
- 由于 TP 划分的是权重,对同一个实例里的每一张卡都有相同的输入,TP 中每个 GPU 上处理的都是 相同的 input token(全 batch),所以 TP 的整体通信量是和当前实例的 batch-size 成正比的,且不受 TP划分粒度的增加而减少。整个实例增大 batchsize,那么对于实例内的每一张卡都需要增加 batchsize,极大限制了推理的batch size(会被单卡显存限制住)
- 而 MoE 的计算需要足够多的 token 才能实现有效的 Expert 负载均衡,如果 batch size 小,那每个 expert 分到的 token 可能很少,甚至有的 expert 不被激活,大大降低了GPU的利用率。然而,TP 方案要求所有 GPU 加载全部专家的参数,即使某些专家未参与计算,也会占用显存资源
3. EP的作用
EP 方案为大规模 MoE 推理提供了一种全新的并行思路,能够有效解决 TP 方案的两大核心问题。
- 在通信开销方面,EP 采用 All-to-All 进行数据交换,通信成本更低:在 EP size 增大的情况下,EP 能大幅降低计算相同数量 token 的情况下单个 GPU 的通讯开销
- 假设一个 batch 共有 S 个 token,每个 token 会路由到 top-k 个专家,Expert 并行度(EP size)为 M,每个专家的 hidden dim 仍为 D
- 每个 GPU 通信量 ≈ 2 × K × S / M × D,这是因为每张卡只需要处理自己对应专家的 token 子集,token 会被 dispatch 到目标专家(combine 过程同理),如果 K/M < 1,即只有极少部分专家被激活,那么通信量远低于 TP
- 更好的内存使用 & batch size 支持:每个 GPU 只处理自己专家负责的 token,所以每张卡只需要存储被分配到的 token,而不是整个 batch,这就意味着 可以增大整个 batch size(在多卡上分发),而不会受单卡内存限制,所有专家都有可能被激活,也能分到更多 token,提高 GPU 利用率
小结:TP 更适合稠密模型;EP 是为稀疏 MoE 模型量身定做的更高效并行方案,能缓解通信瓶颈、显存瓶颈并扩展推理能力。
在专家并行中,每个 MoE 层在前向和反向传播阶段中,一共需要进行四次 All-to-All 通信,这会产生显著的开销,甚至成为效率的主要制约因素。这种通信的效率取决于多个因素,包括通道带宽的异质性、网络拓扑结构和集体通信算法。此外,MoE 固有的负载不均衡可能通过引发同步延迟来加剧这些低效。为了优化节点内高带宽和节点间低带宽的使用,研究人员做了很多努力,比如:
-
最小化网络流量并利用高带宽连接。比如引入分层 All-to-All 、拓扑感知的路由策略、利用专家亲和性来进行分配等;
-
考虑到通信和计算的并发性,把流水线并行和专家并行集成,以此协调 All-to-All 通信和专家计算的重叠。也有研究人员利用GPU的大规模并行性和GPU发起的通信,将计算与依赖的集合通信进行融合。或者将通信依赖关系进行解耦来通信与计算之间的重叠。
3. DeepEP
项目地址:https://github.com/deepseek-ai/DeepEP
DeepEP 是 DeepSeek 专为混合专家模型(MoE)设计的专家并行(EP)通信库,解决了 Token 在 GPU 之间分发和聚合过程中的效率问题。该库提供了高吞吐量、低延迟的 All-to-All 的 GPU 内核(通常称为 MoE 分发与聚合内核),支持 FP8 低精度计算,并通过通信-计算重叠技术实现资源高效利用,内节点通信性能可达到 153-158 GB/s 的 NVLink 带宽,而跨节点通信可达 43-47 GB/s 的 RDMA 带宽。
简单来说,DeepEP 是一个库,它的目的是帮助加速和改进 GPU 在处理复杂的机器学习任务时的沟通效率,特别是那些涉及到混合专家模型(MoE)的任务,这些模型会使用多个“专家”来处理一个问题的不同部分,而 DeepEP 的作用就是确保数据在这些“专家”之间快速、高效地传递。DeepEP 就像机器学习系统中的一个智能交通管理员,确保所有“专家”都能及时收到数据,并协同工作,不会出现延误,从而让整个系统更高效、更快速。
3.1 核心功能
DeepEP 的核心功能包括:
- 高效的 All-to-All 通信:通过软硬件协同优化,DeepEP 实现了专家之间的高速数据传递,显著提升了训练和推理效率
- 支持 NVLink 和 RDMA:无论是节点内还是节点间通信,DeepEP 都能充分利用 NVLink 和 RDMA 的高带宽和低延迟特性(针对Hopper架构)
- 高吞吐量与低延迟内核:在训练和推理预填充阶段,DeepEP 提供了高吞吐率计算核;而在推理解码阶段,则提供了低延迟计算核,满足不同场景的需求
- 原生 FP8 支持:DeepEP 支持包括 FP8 在内的低精度运算,进一步优化了计算资源的利用
- 灵活的 GPU 资源控制:通过计算与通信的重叠处理,DeepEP 实现了资源的高效调度,提升了整体性能
NVLink 与 RDMA
DeepEP 优化了两种不同通信域之间的数据传输:
- NVLink 域:NVIDIA GPU 之间的高速直连通道,带宽可达约 160 GB/s,适合服务器内部 GPU 间通信
- RDMA 域:远程直接内存访问技术,允许计算机不经过操作系统直接访问远程内存,带宽约 50 GB/s,适合服务器间通信
DeepEP 专门为非对称域带宽转发场景提供了优化核心。
3.2 优化原理
DeepEP 提供了两种主要类型的内核,以满足不同的操作需求:
1. 常规内核(Normal kernels):这些内核针对需要高吞吐量的场景进行了优化,例如在推理或训练的预填充阶段。它们利用 NVLink 和 RDMA 网络技术,在 GPU 之间高效地转发数据。测试显示,在 Hopper GPU 上,节点内通信的吞吐量约为 153 GB/s,而使用 CX7 InfiniBand(带宽约为 50 GB/s)的节点间通信性能稳定在 43–47 GB/s 之间。通过最大化可用带宽,这些内核减少了在 token 分发和结果合并过程中的通信开销。
- 上图为示意图:IB chunk 和 NVL chunk :表示数据块来源不同 —— IB chunk 来自节点间(通过 IB),NVL chunk 来自节点内(通过 NVLink)
- 整体流程:CPU 启动一个 GPU 的通知内核(Notify)后开始等待GPU结果,GPU 执行Notify 内核启动后返回,CPU 获得通知后,得知“tensor 大小”——提前知道数据规模,为下一步 tensor 分配做准备;Tensor allocation:根据 Notify 阶段得到的信息,CPU 分配好 tensor 空间;Launch dispatch:启动 GPU 的 dispatch 内核;Launch computation:启动计算任务,然后 GPU 执行 Dispatch:读取 IB chunks 和 NVL chunks,Computation kernels:执行计算;在计算完成后,CPU 启动 GPU 的 combine 操作;GPU 执行Combine,将 IB chunk 和 NVL chunk 的结果合并起来,准备好最终输出给下游任务
- 通过 notify 提前获得 tensor 大小,加快后续分配,优化延迟;带宽利用最大化:节点内用 NVLink,节点间用 InfiniBand;CPU 启动内核、GPU 并行执行,避免空闲
2. 低延迟内核(Low-latency kernels):DeepEP 提供了仅依赖 RDMA 的低延迟内核,专为处理小批量数据而设计。此外,设计中还引入了一种基于hook(在每个 kernel 结束后判断是否需要触发通信操作)的通信与计算重叠技术,使数据传输可以与计算同时进行,所有 RDMA 通信在后台异步完成,而不占用 GPU 的流式多处理器。
- 上半部分:传统通信重叠方式(需要 communication SMs),Stream 0 和 Stream 1 分别表示两个并发 CUDA 流
- 流程如下(每条流):[Attention] → [Dispatch] → [MoE] → [Combine] → [下一步 Attention]
- Dispatch、MoE、Combine 等通信与分发相关的操作在不同 stream 上执行,虽然可以部分重叠,但它们仍占用 SM 资源
- 下半部分:DeepEP 的优化方式 ——“无通信 SM 重叠”,所有任务集中在 Stream 0,计算和调度在一个 stream 中,避免多流同步和切换
- Step 1:在 Stream 0 执行 Attention 0:正常运行的注意力计算模块,没有通信,纯计算,完成后触发 Dispatch 0 issue,发出路由请求,将token发给专家
- Step 2: Attention 1 的计算照常进行,同时,RDMA 正在后台执行 Dispatch 0 的数据路由,将 token 发往专家(RDMA 不用 SM,这两个操作真正重叠了),Attention 1 计算完成后,触发 Dispatch 0 接收 + Dispatch 1 issue
- Step 3: MoE 0 计算,RDMA 正在后台执行 Dispatch 1,Dispatch 0 的远端数据已到达 GPU,本地 SM 可开始计算,下一步:Combine 0 issue、MoE 1 issue
- Step 5:MoE 1 执行时,Combine 0 开始从远端收集 MoE 输出结果,Combine 0 的结果将是下一步 Attention 所需输入
- Step 6: 又回到了 Attention 阶段,所有前序 Combine 操作都在后台通过 RDMA 执行
3.3 性能表现
使用 NVLink 和 RDMA 转发的常规内核
DeepSeek 在 H800 GPU 上测试了常规内核(NVLink 最大带宽约为 160 GB/s),每块 GPU 连接一张 CX7 InfiniBand 400 Gb/s RDMA 网卡(最大带宽约为 50 GB/s)。此外,遵循 DeepSeek-V3/R1 的预训练设置,包括每批 4096 个 token、隐藏层维度 7168、Top-4 分组、Top-8 专家、FP8 分发以及 BF16 结果合并。性能测试的结果如下表所示:
Type | Dispatch #EP | Bottleneck bandwidth | Combine #EP | Bottleneck bandwidth |
Intranode | 8 | 153 GB/s (NVLink) | 8 | 158 GB/s (NVLink) |
Internode | 16 | 43 GB/s (RDMA) | 16 | 43 GB/s (RDMA) |
Internode | 32 | 58 GB/s (RDMA) | 32 | 57 GB/s (RDMA) |
Internode | 64 | 51 GB/s (RDMA) | 64 | 50 GB/s (RDMA) |
使用纯 RDMA 的低延迟内核
DeepSeek 在 H800 GPU 上测试了低延迟内核,每块 GPU 连接一张 CX7 InfiniBand 400 Gb/s RDMA 网卡(最大带宽约为 50 GB/s)。此外,遵循典型的 DeepSeek-V3/R1 生产环境设置,包括每批 128 个 token、隐藏层维度 7168、Top-8 专家、FP8 分发以及 BF16 结果合并。性能测试的结果如下表所示:
Dispatch #EP | Latency | RDMA bandwidth | Combine #EP | Latency | RDMA bandwidth |
8 | 77 us | 98 GB/s | 8 | 114 us | 127 GB/s |
16 | 118 us | 63 GB/s | 16 | 195 us | 74 GB/s |
32 | 155 us | 48 GB/s | 32 | 273 us | 53 GB/s |
64 | 173 us | 43 GB/s | 64 | 314 us | 46 GB/s |
128 | 192 us | 39 GB/s | 128 | 369 us | 39 GB/s |
256 | 194 us | 39 GB/s | 256 | 360 us | 40 GB/s |
3.4 通信方案
DeepEP 是专门针对 MoE 模型大规模专家并行场景进行优化的高效通信库,其依赖 NVIDIA 软件生态的 NVSHMEM GPU通信库、GDRCopy 低延时 GPU 显存拷贝库及 IBGDA 等核心技术,实现了经典的以 NCCL 为主的高效集合通信库的功能。官方 DeepSeek-R1 的推理系统即是依赖 DeepEP+NVSHMEM+GDRCopy+IBGDA 的方案替代了 NCCL 进行高效通信的。整个软件栈的架构图如下:
- 最底层:CUDA(NVIDIA GPU 驱动层):CUDA 是 NVIDIA 提供的通用 GPU 编程平台,为上层通信与计算库提供基础执行环境和硬件控制接口
- 第 2 层:通信加速库:NCCL —— 传统的 GPU 通信库;DeepEP + NVSHMEM + GDRCopy + IBGDA —— DeepSeek 通信方案
- 第 3 层:PyTorch 框架层
- 第 4 层:训练与推理部署框架
NVSHMEM 通信库
DeepEP 利用了 NVSHMEM 的能力进行高效通信。NVSHMEM(NVIDIA SHared MEMory)是 NVIDIA 专为多 GPU 环境下的高效通信与共享内存访问开发的一套通信库。它基于 OpenSHMEM 通信模型,结合了 GPU-Direct、RDMA 和 CUDA 技术,实现了 低延迟、无 CPU 参与、GPU 到 GPU 直接数据传输。
其核心思想是将所有 GPU 节点上的显存视为一个大的显存池来进行管理即分区全局地址空间(PGAS),所有参与通信的 GPU 都能 访问全局地址空间中的任意段内存,图中绿色框内的内存,其他 GPU 可直接读写,不依赖 CPU 参与。该库支持通过GPU共享内存直接进行数据访问,提供如 shmem_put、shmem_get 等可以进行细粒度数据传输的 API 接口。这些 API 都可以从 GPU kernel 内部直接调用,不需 host 配合,是 NVSHMEM 区别于 NCCL 的关键能力。
NCCL VS NVSHMEM
NVSHMEM通信库和经典的NCCL集合通信库的对比:
| NCCL | NVSHMEM |
设计目标 | GPU集合通信 | GPU共享内存通信 |
API设计 | 更高级别的API ncclAllReduce、ncclAllGather... API设计简洁 | 更低级别的API,对共享内存、同步、数据传输等进行细粒度控制 shmem_put、shmem_get… |
高性能模式 | 支持NVLink、GDR | 支持NVLink、GDR、IBGDA |
使用场景 | 大规模分布式数据并行或模型并行 | 应用程序需要直接对内存进行操作 |
GDRCopy 低延时库
GPU ↔ CPU 之间数据拷贝方式的原理和性能差异:cudaMemcpy 和 GDRCopy:
- cudaMemcpy 是 CUDA 提供的标准 API,用于在 GPU 与 CPU 之间传输数据。其工作机制是:触发 GPU 的 DMA 引擎(Direct Memory Access Engine),由它将数据从 GPU 显存拷贝到 CPU 内存,或反过来。
- GDRCopy 是基于 GPUDirect RDMA 技术的低延时 GPU 显存拷贝库,允许 CPU 直接访问 GPU 显存。它的关键机制是:使用 PCIe 的 BAR(Base Address Register)映射机制,将 GPU 显存映射到 CPU 地址空间。CPU 可以像访问普通内存一样,直接访问 GPU 显存,实现低延迟的数据读取与写入,不需要触发 DMA 引擎,也不需要上下文切换
从下图中可以看出,在使用了 GDRCopy 后,H2D 的链路缩短了,这优化了 H2D 的延时。NVIDIA 官方给出的性能测试结果如下:
- 测试环境: NVIDIA DGX-1V 服务器,8 个 Tesla V100 GPU,CUDA 10.1,Driver 418
- 在小数据量下实现极低的主机与设备之间的数据传输延迟:主机到设备(Host-to-Device)拷贝的延迟大约为 1 微秒,而使用 cudaMemcpy 的延迟约为 7 微秒
- 高带宽的主机到设备内存拷贝:通过 写合并(write-combining)(受 NUMA 结构影响)或在部分基于 POWER9 架构的平台上通过 缓存映射(cached mappings) 实现更高带宽
在一个 MoE 路由推理中:每个 token 会分发到多个专家 → 分发数据是 小块 token embedding,使用 cudaMemcpy,每次 token 分发都需要 DMA 调度,延迟高;使用 GDRCopy,CPU 或其他 GPU 线程能 直接写入目标 GPU 显存中的 expert buffer,快速且高效。
InfiniBand GPUDirect Async技术
IBGDA 与 IBRC
- InfiniBand GPUDirect Async 简称 IBGDA,是 NVIDIA 推出的基于 InfiniBand GPUDirect RDMA(简称GDR)技术进一步优化的高效通信技术。
- IBRC(InfiniBand Reliable Connection)是 InfiniBand 协议栈中定义的一种传输模式,属于 InfiniBand 协议自身的通信语义。
技术细节对比:
| IBRC | IBGDA |
通信控制 | 由 CPU 代理线程 发起和管理 | 由 GPU 内核线程(CTA) 直接发起 |
RDMA 支持 | 使用 GPUDirect RDMA,但依赖 CPU 协调 | 基于 GPUDirect RDMA 的进一步优化版本 |
小消息处理 | CPU 线程串行化带来瓶颈,带宽低、延迟高 | GPU 线程并发发起,带宽高、延迟低 |
大消息传输 | 通过 RDMA pipeline 实现高带宽,效果好 | 仍可用,但对大消息不如 IBRC 稳定 |
通信并行性 | 受限于代理线程调度 | 支持 成百上千个 CTA 并发通信 |
自动合并功能 | 无 | 有:在地址连续的 warp 内自动合并消息,减少消息数量,降低开销 |
适配场景 | Prefill(传输大 token embeddings) | Decode(高频小 token 广播) |
- 打个比方:IBRC 像是让“前台客服”用一个电话一个电话拨号传消息(适合讲清楚一个大段内容),慢但稳定;BGDA 像是让“广播系统”同时对几十人喊话,适合快速同步短内容
- IBGDA 主要用于LLM 推理阶段的 Decode 中,每生成一个 token,需要将它广播到多个 GPU(或者多个进程),每条消息的大小很小(通常是几个 byte 到十几 KB),通信非常频繁,每一步 decode 都要同步,IBGDA 是 GPU 多线程并发、无代理线程,能提供更高消息速率,吞吐和延迟表现远好于 IBRC。
通信流程
在介绍 IBGDA 之前,先看下 NVSHMEM 在引入 IBGDA 之前、使用 CPU 代理线程 + IBRC 传输模式 下通信流程:
整个流程如下:① GPU 的 SM 通过 CUDA kernel 在本地 GPU memory 中生成需要通信的数据;②GPU 使用 nvshmem_put 等接口,把通信请求写入主机内存的 Proxy Buffer(给CPU代理线程读取的任务描述区域);③ CPU代理线程检测 Proxy Buffer 中的新任务,发现新数据后,开始准备通信;④ CPU 将任务写入 WQ(Work Queue),WQ 位于主机内存中,由 NIC 网卡读取;⑤ CPU 写入 DBR(Doorbell Record),写入它是为了容错,如果 NIC 忽略了真实门铃(DB),可以从 DBR 中恢复;⑥ CPU 通过 PCie 向 NIC 的 Doorbell 寄存器发信号(写DB);⑦ NIC 读取 WQ 中的工作请求;⑧ NIC 用 GDR 从 GPU memory 直接读取数据,不再走 CPU 中转,直接从 GPU memory 中拷贝要发送的数据;⑨ NIC 通过 InfiniBand 通道将数据发送到其他节点(GPU 或 Host);通信完成后,⑩ NIC 写入 Completion Queue(CQ)表示操作完成,通知 CPU 当前 RDMA 操作完成;⑪ CPU 代理线程周期性轮询 CQ,检测哪些操作已经完成;⑫ CPU 通知 GPU 内存,告诉它数据是否成功发送(使用了 GDRCopy,CPU 直接把状态写到 GPU memory,否则写回 Proxy Buffer,让 GPU kernel 自己去 poll)
总结:GPU发起通信请求写入 Proxy Buffer,代理线程读取请求,准备并通知 NIC,NIC 从 GPU memory 读取数据并发送,通信完成后,NIC 通知 CPU,CPU 通知 GPU
从上述流程可以看出,多跳通信路径 GPU → Host 内存 → CPU → NIC,有比较多的非应用数据传输的步骤需要 CPU 的参与,由于 GPU 和 Mellanox 高性能网卡的数据处理能力都在快速增长,且远远超过 CPU 的处理能力,因此在对延时有极高要求的场景下,经典的 GDR 技术在 CPU 侧会成为瓶颈,CPU 不可能每秒处理数十亿次通信请求,这也正是 IBGDA 设计要解决的核心问题:让 GPU 自己直接驱动 NIC,绕过 CPU 全部流程!
为进一步优化通信效率,NVIDIA在GDR的基础上推出了IBGDA,如下图所示:
整个流程如下:① CUDA Kernel 在 GPU 的 SM 上运行,生成数据并写入 GPU memory 中的 Data 缓冲区,准备要传输的数据;② SM 直接写入 GPU memory 中的 WQ,这个步骤是传统上由 CPU 代理完成的,现在变成了 GPU 自主完成;③ SM 更新 DBR 缓冲区,这个缓冲区用于回退或辅助通知 NIC。和 WQ 一样,也放在 GPU memory 中,便于 GPU 快速访问;④ SM 直接写入 NIC 的 DB(doorbell)寄存器,通知 NIC 去读取 WQ,开始执行通信操作(⚠️ 这是一个关键突破点:GPU 不再需要 CPU 中转,能像 CPU 一样直接下发网络操作!);⑤ NIC 收到 DB 通知后,从 GPU memory 中读取 SM 写入的 WQ 描述符(使用GDR),了解该做什么操作。⑥ NIC 根据 WQ 中描述的地址,从 GPU memory 中读取实际数据(也是通过GDR),准备发送;⑦ 将数据通过 InfiniBand 网络发送到其他节点的 GPU/Host;⑧ 数据传输完成后,NIC 使用 GDR 向 GPU memory 中的 CQ(Completion Queue) 写入事件,告知 GPU 操作完成。
从上述流程可以看出,IBGDA 将在 CPU 上进行的相关操作全部放到 GPU 中,整个过程完全不需要 CPU 的参与,进一步减少了通信链路,提高了通信效率。
小结:
| 传统 IBRC + CPU Proxy | IBGDA(Kernel-Initiated) |
控制流程 | GPU → Host Proxy Buffer → CPU 代理线程 → NIC | GPU SM → NIC(直接) |
WQ/DBR/CQ 位置 | Host Memory | GPU Memory |
Doorbell 发起者 | CPU | GPU |
数据路径 | GDR 从 GPU memory 读取 | GDR 从 GPU memory 读取(一样) |
通信延迟 | 高,受 CPU 调度限制 | 低,无中间人 |
并发能力 | 限制在 CPU 上 | 支持数千个 GPU CTA 并发 |
适合场景 | Prefill、大消息批量传输 | Decode、细粒度通信 |
❓ 什么是 Doorbell?
- Doorbell(门铃) 是网络接口卡(NIC)中的一个 特殊寄存器或地址空间,它的作用是通知 NIC 有新的通信请求需要处理
- 发起方(CPU 或 GPU)写入 Doorbell 寄存器 → 相当于“按了一下门铃”
- 接收方(NIC)收到这个写操作后 → 就知道:“哦,有新任务来了,我要去工作队列看看”
- Doorbell 不传输数据,也不是控制命令的详细内容,它只是一个“轻量级触发器”
❓ Doorbell 工作在哪一层?
- Doorbell 是 RDMA / InfiniBand 等 通信协议栈的底层机制之一,用于将 通信请求从主机控制单元(CPU/GPU)提交给 NIC 执行
- 在通信流程中,每次构造了一个 RDMA 操作,将操作信息写入 Work Queue(WQ),必须写 Doorbell,告诉 NIC “现在有新任务了,请去看 WQ”,NIC 就会去解析 WQ 中的描述符并发起通信操作
❓ 什么是 DBR(Doorbell Record)
- Doorbell Record,是一个位于主机或 GPU 内存中的缓冲区,它是 Doorbell(DB)写操作的副本或备份,用于支持:容错、重放(replay)、以及 NIC 对 Doorbell 操作的冗余备份读取。维护一份 DBR,NIC 可以定期从 DBR 中拉取数据,确保不会漏任务
- DBR 是一种冗余机制,作为 Doorbell 写入的内存记录副本,用于容错和冗余恢复,确保 NIC 不会错过任务
NVIDIA 官方基于 IBGDA 技术在 All-to-All 场景下的延时测试如下:
- IBRC 在大消息传输时能够达到高带宽,但对于小消息存在性能瓶颈,更适合大消息的批量传输。
- IBGDA 在小消息传输中表现出色,能够实现更高的吞吐量和更低的延迟,更适合涉及大量小消息的细粒度通信。
通信模式对比:
通信方式 | 是否需 CPU | 延迟 | GPU 直达 | 应用场景 |
NCCL All-to-All | ✅ 需要 | 高 | ❌ | 通用集合通信 |
CUDA Memcpy | ❌ 可绕过 | 中 | ⚠️ 仅同节点 | 拷贝数据用 |
NVSHMEM | ❌ 不需要 | 低 | ✅ | Token 精细调度、MoE 通信 |
GDRCopy | ❌ 不需要 | 低 | ✅ | 高速内存拷贝 |
MPI | ✅ 需要 | 高 | ❌ | CPU 主导并行 |
4. EPLB
EPLB(Expert Parallel Load Balancing,专家并行负载均衡器)是 DeepSeek 为解决 MoE 模型在分布式训练和推理中的负载不平衡问题而开发的优化策略。
平衡每张 GPU 上的工作负载:
- 问题:在专家并行中,不同的专家会被分配到不同的 GPU 上运行,导致每张 GPU 上的专家负载不均,部分 GPU 空转或拥堵
- 解决方案:EPLB 采用了 “冗余专家” 策略,即复制高负载的专家,然后将这些复制后的专家合理分配到各个 GPU 上,以实现负载均衡
在 DeepSeek-V3 的训练中,EPLB 将 GPU 利用率提升超 30%,减少因负载不均导致的节点空转现象,通过优化专家任务的本地化计算,跨节点通信量减少约25%,尤其在高稀疏性MoE模型(如每层仅激活8/256个专家)中效果显著。
基本原理
- 冗余专家策略:实时监控每个专家的计算负载,创建高负载专家副本,从而分散计算压力
- 通信优化:将关联性强的专家模块部署在同一计算节点,降低跨节点通信的数据量,从而减少网络带宽占用和延迟
- 无辅助损失的负载均衡:区别于传统方法需引入额外的负载均衡损失函数,EPLB通过可学习的偏置项(Bias Term),动态调整专家激活概率(降低高负载专家的激活权重,提升低负载专家的优先级),不增加训练目标复杂性的情况下实现均衡
- 跨并行模式的协同优化:EPLB支持与数据并行(DP)、专家并行(EP)等多种并行策略协同工作,通过全局调度器协调不同并行模式下的资源分配,避免单一策略的局限性
4.1 负载均衡策略
EPLB 负载均衡算法提供了两种策略,适用于不同的场景:
- 分层负载均衡(Hierarchical Load Balancing):当服务器节点数量能够整除专家组数量时,使用该策略来利用分组限制专家路由。具体步骤为:首先将专家组均匀分配到各个节点,确保不同节点的负载平衡;然后在每个节点内复制专家;最后将复制后的专家分配到各个 GPU 上,确保不同 GPU 的负载平衡。该策略适用于预填充阶段,专家并行规模较小的情况。
- 全局负载均衡全局负载均衡全局负载均衡(Global Load Balancing):在全局范围内复制专家,不考虑专家组的划分,然后将复制后的专家分配到各个 GPU 上。该策略适用于解码阶段,专家并行规模较大的情况。
4.2 DeepSeek 负载均衡方案
Prefill 阶段:路由专家 EP32、MLA 和共享专家 DP32,一个部署单元是 4 节点,32 个冗余路由专家,每张卡 9 个路由专家和 1 个共享专家
- Prefill 是生成任务中输入 token 的处理阶段,特点是并行度高(多个 token 同时处理),计算密集,容易产生专家负载不均
- EP32(路由专家):每层用 32 个路由专家副本参与计算(实际上是从 256 个专家中按需复制出来的)
- DP32(共享专家):每层有 1 个共享专家,复制到每个 GPU 上,采用 Data Parallel
- 冗余路由专家:32 个 —— 这些是对 负载较高的专家 进行复制,以分担负载
- GPU 数:部署单元 = 4 节点 × 每节点 8 卡 = 32 卡;专家副本数量:256 + 32 = 288(路由专家副本总数);平均每张卡分配:288 / 32 = 9 个路由专家/卡
- 专家组:32 个专家组,每组 9 个专家(288 / 32 = 9),每个专家组就部署在一个 GPU 上
Decode 阶段:路由专家 EP144、MLA 和共享专家 DP144,一个部署单元是 18 节点,32 个冗余路由专家,每张卡 2 个路由专家和 1 个共享专家
- Decode 是生成 token 时逐个处理的阶段,token 数量少,但对延迟要求高
- EP144(路由专家):每层有 144 个路由专家副本
- DP144(共享专家):每个 GPU 上也部署一个共享专家
- 冗余专家:32 个副本(和 Prefill 类似,拷贝负载高的专家)
- 节点数 = 18 × 每节点 8 卡 = 144 张卡,所以 144 张卡 × 2 = 288 个路由专家副本
小结:
阶段 | 路由专家配置 | 节点数 | 每卡专家数(路由 + 共享) | 冗余专家 | 特点说明 |
Prefill | EP32 + DP32 | 4 | 9 路由 + 1 共享 | 32 | token 多,需高吞吐,专家负载不均,复制高负载专家 |
Decode | EP144 + DP144 | 18 | 2 路由 + 1 共享 | 32 | token 少,需低延迟,资源分得更分散 |
DeepSeek-R1 为了解决 MoE 推理阶段负载不均问题采用了静态 + 局部冗余复制(32 个专家副本),根据阶段切换部署策略(Prefill 更集中,Decode 更分散),每卡部署有限个专家,避免内存溢出,同时提升负载均衡,兼顾了吞吐量(Prefill)与延迟(Decode)。
4.3 源码分析
项目地址:https://github.com/deepseek-ai/eplb
代码就一个eplb.py,它实现了负载均衡算法的核心代码,包含以下主要函数:
- replicate_experts:将逻辑专家复制为物理专家,以最小化所有副本的最大负载
- balanced_packing:将加权对象打包到指定数量的包中,使每个包的权重尽可能平衡
- rebalance_experts_hierarchical:实现分层负载均衡策略
- rebalance_experts:负载均衡算法的入口函数,根据情况选择合适的负载均衡策略
- inverse:用于计算排列的逆排列
from typing import Tuple # 导入元组类型注解
import torch # 导入 PyTorch 库def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor, torch.Tensor]:"""采用贪心算法:将 n 个带权重的对象打包到 m 个包中,使得每个包恰好包含 n/m 个对象,并且所有包的权重尽可能平衡参数:- weight: 一个形状为 [X, n] 的 PyTorch 张量,代表权重- num_packs: 包的数量返回: pack_index: [X, n],每个物品所在的包的索引rank_in_pack: [X, n],物品在包中的排名"""# 1.平均分组,从而实现包的权重平衡(R1 的专家数量 256 个比较多,并不是一个一个去分配,首先会进行一个合并,合并成 num_groups 组,然后将 group 分配到不同节点上)num_layers, num_groups = weight.shape # 获取权重张量的层数和组数assert num_groups % num_packs == 0 # 确保组数能被包数整除groups_per_pack = num_groups // num_packs # 计算每个包中的组数# 2. 特殊情况处理:如果每个包只有一个组,计算每个组的包索引和排名并返回,从而为后续的负载均衡提供信息if groups_per_pack == 1: # 生成每个组的包索引pack_index = torch.arange(weight.size(-1), dtype=torch.int64, device=weight.device).expand(weight.shape)rank_in_pack = torch.zeros_like(weight, dtype=torch.int64) # 生成每个组在包中的排名return pack_index, rank_in_pack # 返回包索引和排名# 3. 处理每个包含有多个分组的情况indices = weight.float().sort(-1, descending=True).indices.cpu()# 对权重张量进行降序排序并获取索引,从而优先处理权重较大的组pack_index = torch.full_like(weight, fill_value=-1, dtype=torch.int64, device='cpu')# 初始化包索引张量,初始值为-1表示未分配rank_in_pack = torch.full_like(pack_index, fill_value=-1) # 初始化每个组在包中的排名张量,初始值为 -1for i in range(num_layers): # 外层循环遍历每个包pack_weights = [0] * num_packs # 初始化每个包的权重pack_items = [0] * num_packs # 初始化每个包中的物品数量for group in indices[i]: # 内层循环遍历当前层的每个组# 核心:对于每个组,找到当前权重最小且物品数量未达到上限的包,将该组分配到这个包中。pack = min((i for i in range(num_packs) if pack_items[i] < groups_per_pack), key=pack_weights.__getitem__)assert pack_items[pack] < groups_per_pack # 确保所选包中的物品数量未达到上限pack_index[i, group] = pack # 记录当前组所在的包的索引rank_in_pack[i, group] = pack_items[pack] # 记录当前组在包中的排名pack_weights[pack] += weight[i, group] # 更新所选包的权重pack_items[pack] += 1 # 更新所选包中的物品数量return pack_index, rank_in_pack # 返回包索引和排名def replicate_experts(weight: torch.Tensor, num_phy: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:"""创建冗余专家,具体是将 num_log 个逻辑专家复制为 num_phy 个物理专家副本,遍历冗余专家并将每个冗余专家映射到当前负载最大的逻辑专家参数:weight: [X, num_log] 形状的张量,代表每个逻辑专家的负载num_phy: 复制后专家的总数返回:phy2log: [X, num_phy] 形状的张量,phy2log[i][j] 表示在第 i 个场景下,第 j 个物理专家对应的逻辑专家的 IDrank: [X, num_phy] 形状的张量,rank[i][j] 表示在第 i 个场景下,第 j 个物理专家是其对应逻辑专家的第几个副本logcnt: [X, num_log] 形状的张量,logcnt[i][k] 表示在第 i 个场景下,第 k 个逻辑专家有多少个物理专家副本"""# 1. 获取专家的相关信息n, num_log = weight.shape # 获取权重张量的行数和逻辑专家数量num_redundant = num_phy - num_log # 计算冗余专家的数量assert num_redundant >= 0 # 确保冗余专家数量非负device = weight.device # 获取权重张量所在的设备# 2. 初始化映射和phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1) # 初始化物理专家到逻辑专家的映射rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device) # 表示每个物理专家在其对应的逻辑专家副本中的次序logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device) # 初始化每个逻辑专家的物理专家副本数量arangen = torch.arange(n, dtype=torch.int64, device=device) # 生成一个从 0 到 n-1 的张量,n 代表组的数量# 3. 复制逻辑专家为物理专家副本for i in range(num_log, num_phy): # 遍历冗余专家(逻辑专家num_log到物理专家num_phy间为冗余的)redundant_indices = (weight / logcnt).max(dim=-1).indices # 找到负载最大的逻辑专家的索引phy2log[:, i] = redundant_indices # 将当前物理专家映射到负载最大的逻辑专家,更新 phy2log 张量rank[:, i] = logcnt[arangen, redundant_indices] # 记录当前物理专家的副本排名,更新 rank 张量logcnt[arangen, redundant_indices] += 1 # 更新负载最大的逻辑专家的副本数量,更新 logcnt 张量。# 返回物理专家到逻辑专家的映射、副本排名和逻辑专家的副本数量return phy2log, rank, logcntdef rebalance_experts_hierarchical(weight: torch.Tensor, num_physical_experts: int, num_groups: int, num_nodes: int, num_gpus: int):"""分层负载均衡策略,将逻辑专家复制为物理专家副本,并在节点和 GPU 层面进行负载均衡参数:weight: [num_moe_layers, num_logical_experts] 形状的张量,代表每个逻辑专家的负载num_physical_experts: 复制后物理专家的总数num_groups: 专家组的数量num_nodes: 服务器节点的数量,节点内网络(如 NVLink)更快num_gpus: GPU 的数量,必须是 num_nodes 的倍数返回: physical_to_logical_map: [num_moe_layers, num_physical_experts] 形状的张量,每个物理专家对应的逻辑专家的 IDlogical_to_physical_map: [num_moe_layers, num_logical_experts, X] 形状的张量,每个逻辑专家对应的物理专家的索引logical_count: [num_moe_layers, num_logical_experts] 形状的张量,每个逻辑专家的物理专家副本数量"""# 1. 输入检查与初始化num_layers, num_logical_experts = weight.shape # 获取权重张量的层数和逻辑专家数量assert num_logical_experts % num_groups == 0 # 确保逻辑专家数量能被专家组数量整除group_size = num_logical_experts // num_groups # 计算每个专家组中的专家数量assert num_groups % num_nodes == 0 # 确保专家组数量能被节点数量整除groups_per_node = num_groups // num_nodes # 计算每个节点中的专家组数量assert num_gpus % num_nodes == 0 # 确保 GPU 数量能被节点数量整除assert num_physical_experts % num_gpus == 0 # 确保物理专家数量能被 GPU 数量整除phy_experts_per_gpu = num_physical_experts // num_gpus # 计算每个 GPU 上的物理专家数量# 2. 定义逆排列计算辅助函数def inverse(perm: torch.Tensor) -> torch.Tensor:"""计算排列的逆排列,接受一个排列张量 perm 作为输入,返回其逆排列"""inv = torch.empty_like(perm) # 初始化逆排列张量inv.scatter_(1, perm, torch.arange(perm.size(1), dtype=torch.int64, device=perm.device).expand(perm.shape)) # 计算逆排列return inv # 返回逆排列# 3. 阶段的负载均衡# 3.1 将专家组分配到节点上tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1) # 计算每个专家组的总权重group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes) # 使用 balanced_packing 函数将专家组分配到节点上# log2mlog:表示逻辑专家在节点上的位置# mlog2log:表示节点上的位置对应逻辑专家的位置log2mlog = (((group_pack_index * groups_per_node + group_rank_in_pack) * group_size).unsqueeze(-1) + torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device)).flatten(-2) # 计算逻辑专家到中间逻辑专家的映射mlog2log = inverse(log2mlog) # 计算中间逻辑专家到逻辑专家的映射# 3.2 在节点内创建冗余专家tokens_per_mlog = weight.gather(-1, mlog2log).view(-1, num_logical_experts // num_nodes) # 计算每个中间逻辑专家的总权重phy2mlog, phyrank, mlogcnt = replicate_experts(tokens_per_mlog, num_physical_experts // num_nodes) # 在节点内创建冗余专家# 3.3 将物理专家分配到 GPU 上,借助中间物理专家这个中间层,能分阶段处理负载均衡问题tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog) # 计算每个物理专家的总权重pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes) # 使用 balanced_packing 函数将物理专家分配到 GPU 上phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack # 计算物理专家到中间物理专家的映射pphy2phy = inverse(phy2pphy) # 计算中间物理专家到物理专家的映射pphy2mlog = phy2mlog.gather(-1, pphy2phy) # 计算中间物理专家到中间逻辑专家的映射pphy2mlog = (pphy2mlog.view(num_layers, num_nodes, -1) + torch.arange(0, num_logical_experts, num_logical_experts // num_nodes).view(1, -1, 1)).flatten(-2)pphy2log = mlog2log.gather(-1, pphy2mlog) # 计算物理专家到逻辑专家的映射pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1) # 计算每个物理专家的副本排名logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog) # 计算每个逻辑专家的副本数量return pphy2log, pphyrank, logcnt # 返回物理专家到逻辑专家的映射、副本排名和逻辑专家的副本数量def rebalance_experts(weight: torch.Tensor, num_replicas: int, num_groups: int,num_nodes: int, num_gpus: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:"""专家并行负载均衡器的入口点。参数:weight: [layers, num_logical_experts],所有逻辑专家的负载统计信息num_replicas: 物理专家的数量,必须是 `num_gpus` 的倍数num_groups: 专家组的数量num_nodes: 服务器节点的数量,节点内网络(例如,NVLink)更快num_gpus: GPU 的数量,必须是 `num_nodes` 的倍数返回: physical_to_logical_map: [layers, num_replicas],每个副本对应的专家索引logical_to_physical_map: [layers, num_logical_experts, X],每个专家对应的副本索引expert_count: [layers, num_logical_experts],每个逻辑专家的物理副本数量"""# 1. 数据预处理num_layers, num_logical_experts = weight.shape # 获取权重张量的层数和逻辑专家数量weight = weight.float().cpu() # 将权重张量转换为 float 类型并移动到 CPU 上# 2. 策略选择:如果专家组数量能被节点数量整除使用分层负载均衡,如果不能使用全局负载均衡if num_groups % num_nodes == 0:phy2log, phyrank, logcnt = rebalance_experts_hierarchical(weight, num_replicas, num_groups, num_nodes, num_gpus)else:phy2log, phyrank, logcnt = replicate_experts(weight, num_replicas)# 计算逻辑专家到物理专家的映射maxlogcnt = logcnt.max().item() # 获取最大的逻辑专家副本数量# 初始化逻辑专家到物理专家的映射log2phy: torch.Tensor = torch.full((num_layers, num_logical_experts, maxlogcnt), -1, dtype=torch.int64, device=logcnt.device)# 计算逻辑专家到物理专家的映射log2phy.view(num_layers, -1).scatter_(-1, phy2log * maxlogcnt + phyrank, torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(num_layers, -1))# 返回物理专家到逻辑专家的映射、逻辑专家到物理专家的映射和逻辑专家的副本数量return phy2log, log2phy, logcnt# 定义模块的公共接口
__all__ = ['rebalance_experts']
4.4 使用示例
pip install eplb
EPLB 支持灵活的专家复制和分配,能够适配不同规模的模型和硬件配置。
例如,在一个包含 2 个节点、每个节点 4 个 GPU 的集群上,EPLB 可以动态规划专家副本的分配。两层 MoE 模型,每层包含 12 名专家,每层引入 4 个冗余专家,总共 16 个副本放置在 2 个节点上:
# 使用 eplb.rebalance_experts 这个函数为一个两层 MoE 模型进行专家负载均衡
import torch
import eplbweight = torch.tensor([[ 90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],[ 20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27]]) # [2, 12]的张量表示两层MoE模型中每个专家的当前负载(token调度次数)num_replicas = 16 # 每层总共有 16 个专家副本(原始12 + 冗余4)
num_groups = 4 # 每个 node 拥有的专家组数
num_nodes = 2 # 两个物理节点(machines)
num_gpus = 8 # 总共 8 块 GPU(每节点4块)phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
print(phy2log)# Output:
# tensor([[ 5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1], # 第 1 层的 16 个副本对应的逻辑专家编号
# [ 7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1]]) # 第 2 层的 16 个副本对应的逻辑专家编号
-
第 0 层:16 个物理副本分别负责专家 5, 6, 5, 7, ... 等(例如专家 5 被分配了两个副本,专家 1 被分配了三个副本,等等)
style="display: none !important;">