MoE Inference Communication

引言：训练赢了，推理怎么办？

如果你读完了这系列的前十三篇，你应该对 MoE 训练的通信问题有了相当清晰的认识：AllReduce + All-to-All 双重轰炸、跨节点带宽瓶颈、负载不均引发的尾延迟、还有 1:1 的计算-通信比。DeepSeek-V3 用 DualPipe、节点限制路由、warp specialization 等一系列工程技巧硬生生把这个瓶颈”吃”掉了——但这一切都是训练场景的设计。

把同样的 MoE 模型部署到推理场景，你会发现面前不是同一张地图。

训练和推理在通信上的差异，比大多数工程师意识到的要大得多。用一个简单的比喻：训练像工厂流水线——你关心的是每天能生产多少产品（throughput），一条产线的短暂停顿可以由其他产线补上。推理像餐厅上菜——每道菜必须在规定时间内送到顾客桌上（latency），任何延迟都直接影响用户体验。

在这个比喻下，MoE 推理的通信问题变得异常尖锐。我们来看看具体的数字：

• 通信模式变了：训练有 AllReduce（Dense 部分的梯度同步）+ All-to-All（专家路由），而推理只有 All-to-All。但这不意味着推理更简单——少了一种通信模式，却让剩余的那种变成了唯一且无可分摊的瓶颈。没有梯度同步可以做计算-通信重叠，没有 micro-batch 的流水线可以隐藏延迟。
• 优化目标变了：训练优化 throughput（每秒处理多少 token），推理优化 latency（每个请求几毫秒内响应）。All-to-All 对 throughput 的伤害可以通过增加 batch size 摊薄，但对 latency 的伤害几乎是线性的——多一次转发就多几毫秒，没有商量的余地。
• batch 语义变了：训练时的 batch 可以是数百甚至数千个 sequence 同时推进，巨大的 batch 能有效摊薄 All-to-All 的启动开销。推理时——特别是在线服务场景——batch=1 或 batch=4 是常态。batch 越小，All-to-All 的单位通信开销越大。当 batch=1 时，每个 GPU 在每次 All-to-All 中可能只发送寥寥几个 token 的 hidden state，但通信建立的开销（握手、协商、内存 pinning）一分不少。
• 路由不均衡极端化：训练时可以通过 auxiliary loss 强制负载均衡，推理时不能。用户请求的 token 分布是无法控制的——可能某个时间段内所有 query 都落入同一个专家，导致一块 GPU 被挤爆而其余 GPU 几乎空闲。这种极端不均衡在 batch=1 时影响尤为严重：单个 GPU 被塞满，整个推理的延迟取决于这条最长路径。

这些差异构成了 MoE 推理通信优化的核心挑战。本文选择三篇对这个挑战给出了不同方向回答的关键工作进行深入分析：

• ExFlow（IPDPS’24）：从模型结构本身挖潜——发现跨层专家亲和性（inter-layer expert affinity），把推理中的 All-to-All 从每层 2 次减少到单次，不需要任何 fine-tuning。
• Aurora（2024）：首次从理论层面推导 MoE 推理的最下界（theoretical minimum），在四个集群场景中给出最优或接近最优的专家放置、GPU 分配和通信调度方案。
• DeepSpeed-MoE 推理系统（ICML’22）：从系统和模型压缩两个维度同时推进，实现了万亿参数 MoE 模型 <25ms 延迟、比等质量 Dense 模型便宜 9 倍的推理性能。

这三篇工作的共同特点是：它们都意识到 MoE 推理的通信优化是一场不同于训练的战争。每个团队选择了不同的主攻方向，三家路线互为补充，共同构成了当前 MoE 推理通信优化的全貌。

系统剖析：MoE 推理的三条优化路线

在深入每项工作之前，我们先用一个统一的框架来审视 MoE 推理的通信优化空间。MoE 推理的一层前向传播可以分解为以下步骤：

Gate → All-to-All dispatch → Expert FFN → All-to-All combine → Aggregation

在这个流程中，两个 All-to-All 是通信开销的全部来源。优化可以从三个维度切入：

路线一：通信消除（Communication Elimination）——减少或完全消除 All-to-All 操作本身。ExFlow 代表这条路线：通过发掘模型内在的专家亲和性模式，让 token 尽可能停留在当前 GPU 上，从而减少需要跨 GPU 路由的 token 数量。极端情况下，一个 token 的整个跨层路径都不需要离开当前 GPU——All-to-All 虽然还在，但它的有效通信量降到了接近零。这条路线的本质是”如果模型本身告诉你 token 大概率去哪，为什么不让专家提前站好位置？”

路线二：部署调度（Deployment Scheduling）——在给定通信需求的前提下，通过优化专家在 GPU 上的放置顺序和 token 的传输时序来最小化总延迟。Aurora 代表这条路线：它接受 All-to-All 必须发生这个事实，但在”如何发生”这个层面上做精细的控制——哪些专家放在同一块 GPU？哪个 token 先发？发送的时序如何避免带宽争用？这条路线的本质是”既然通信躲不掉，就让它完成得最快。”

路线三：系统压缩（System Compression）——从模型层面减少需要通信的参数总量。DeepSpeed-MoE 代表这条路线：通过 PR-MoE（Pyramid-Residual MoE）架构设计和 MoS（Mixture-of-Students）知识蒸馏，在保持模型质量的前提下把 MoE 模型的总参数减少 3.7 倍。参数减少直接转化为更少的专家、更少的 All-to-All 参与方、更低的通信开销。这条路线的本质是”从源头上让需要通信的东西变少。”

三条路线的选择不是互斥的，实践中可以叠加。但在研究中，每项工作各自深入了其中一条路线，把它们的研究路径和关键发现梳理清楚，才能理解 MoE 推理通信的全貌。

三条路线的核心区别

这张表揭示了一个重要的现象：三条路线的优化对象在系统栈上处于不同层次。ExFlow 在应用层面做模型内部结构的发现和利用，Aurora 在调度层面做形式化优化，DeepSpeed-MoE 在模型和系统层面做全局改造。这种层次分化意味着它们天然可以叠加——你可以先用 PR-MoE 压缩模型（路线三），再对压缩后的模型计算专家亲和性矩阵做 placement（路线一），最后用 Aurora 的调度算法优化 All-to-All 的传输时序（路线二）。

路线一深入：ExFlow —— 用模型的结构取代通信

核心发现：跨层专家亲和性（Inter-Layer Expert Affinity）

ExFlow 的起点是一个简单的经验问题：在一个预训练的 GPT MoE 模型中，如果在第 i 层某个 token 被路由到了专家 E_a，那么在第 i+1 层，这个 token 被路由到专家 E_b 的概率是多少？是完全随机的（在 E 个专家中等概率分布），还是存在某种稳定的模式？

Ohio State 团队的答案是：存在极强的模式，而且这种模式在一开始训练时就出现了，并随着训练趋于稳定。

下图是这个发现最直观的表达（原文 Figure 2）。在一个 12 层、每层 32 个专家的预训练 GPT 模型中，研究人员选取了四对连续层来绘制条件概率热力图。Y 轴是第 j 层的专家索引，X 轴是第 j+1 层的专家索引，每个点的颜色深度表示”token 在第 j 层选了专家 E_i，然后在第 j+1 层选了专家 E_j”的条件概率：

• 在 Layer 0→Layer 1 的热力图中，每行只有少数几个列呈现深红色——这意味着从任何一个专家出发，token 在下一层大概率只会去到少数几个目标专家。
• 这个模式在 Layer 3→Layer 4、Layer 7→Layer 8、甚至 Layer 11→Layer 12 的所有层对之间都持续存在——不是偶然现象，而是模型学习到的一种结构性特征。
• 不同层之间的亲和性矩阵不同（每对层有自己的”路由偏好图”），但它们的共同点是：每一行都是稀疏的——一个专家和下一层的绝大多数专家之间几乎没有亲和性。

ExFlow 将这个条件概率形式化地定义为：

P(E_{p,j+1} | E_{i,j}) = (1/N) * Σ_k P(E_{p,j+1} | E_{i,j}, T_k)

其中 N 是 token 总数，T_k 是第 k 个 token。这个公式计算的是：在所有 token 中，那些在第 j 层被路由到专家 E_i 的 token，在 j+1 层被路由到 E_p 的比例。

为什么这个亲和性对推理至关重要？

标准 MoE 专家并行在每层需要两次 All-to-All：

1. 第一次 All-to-All（dispatch）：每个 GPU 把 token 发送到目标专家所在的 GPU。
2. 第二次 All-to-All（combine）：专家计算完成后，每个 GPU 把结果传回 token 原始所在的 GPU——因为下一层的 attention 需要 token 和它的 context 在同一个 GPU 上。

第二次 All-to-All 的根因是数据并行（DP）：在标准专家并行中，不同 GPU 上的 token 拥有不同的 context（prompts + 之前生成的 token），attention 计算只能在自己 GPU 的 context 上进行。所以 token 做完专家的 FFN 后必须”回自己家”——它的 context 在家。

ExFlow 从两个层面切入这个问题：

层面一：上下文一致（Context Coherence）。在执行推理前，先用一次 AllGather 把所有 GPU 的 context 广播到所有 GPU。此后每个 GPU 都拥有全部请求的完整 context。这样 token 到了任何 GPU 都可以直接做 in-place attention——不需要再回到原始 GPU。

这个操作的效果是：All-to-All 从 2 次减少到 1 次。 代价是在每次迭代（生成每个新 token 后）做一次 AllGather 来同步新生成的 token，保持所有 GPU 的 context 最新。但 AllGather 的开销远小于第二层 All-to-All——尤其是当模型有很多层时，AllGather 被摊薄到整个迭代中微不足道。

层面二：专家放置优化。在上下文一致后，还有一次 All-to-All 无法消除——token 必须跨 GPU 找专家。但如果专家在不同的 GPU 上按照亲和性排列——把高亲和性专家放在同一块 GPU、中亲和性专家放在同一节点——那么 token 一旦到达某块 GPU，就有很大概率在后续层继续留在同一块 GPU 上（或同一节点内）。

ExFlow 的实验数据很好地说明了这个效果：

• 用 4 块 GPU 跑 MoE-64 模型（每 GPU 持有 16 个专家/层），有超过 50% 的 token 在整个推理过程中不需要离开初始 GPU。
• 用 8 块 GPU 跑 MoE-64 模型（每 GPU 持有 8 个专家/层），40% 的 token 留在本地。
• 即使扩展到 32 块 GPU（每 GPU 仅 2 个专家/层），仍有 28% 的 token 不需要跨 GPU 通信。

专家亲和性放置相比随机的 DeepSpeed 默认放置，跨 GPU 路由减少了 25% 到 40%。

形式化：整数线性规划建模

ExFlow 将专家放置问题形式化为一个整数线性规划（ILP）问题。其核心变量和目标函数如下：

决策变量：

• x_{i,j}^p ∈ {0, 1}：专家 E_{i,j}（第 j 层的第 i 个专家）是否放在 GPU p 上。
• R_{k,j} ∈ {0, 1}：token k 在第 j 层是否发生了跨 GPU（或跨节点）路由。

目标函数（最小化跨 GPU 路由总量）：

Minimize  Σ_{k=1}^{N} Σ_{j=1}^{L-1} R_{k,j}

约束条件：

1. 负载均衡：每块 GPU 在每一层持有的专家数量相等——

   Σ_{i=1}^{E} x_{i,j}^p = E/P,  ∀j, ∀p
   

   （E 为专家总数，P 为 GPU 数量）

2. 专家互斥：每个专家在每层只能放在一块 GPU 上——

   Σ_{p=1}^{P} x_{i,j}^p = 1,  ∀j, ∀i
   

3. 路由判定：如果专家 E_i 在第 j 层放在 GPU p，但在第 j+1 层不在 GPU p，则发生了跨 GPU 路由——

   R_{k,j} ≥ x_{i,j}^p - x_{i,j+1}^p
   R_{k,j} ≥ x_{i,j+1}^p - x_{i,j}^p
   

这个 ILP 的巧妙之处在于：它完全不依赖任何额外的超参数（如 affinity 的强弱系数），而是通过目标函数直接捕获问题的核心——最小化那些”token 在连续两层的专家位于不同 GPU”的情况。

ExFlow 进一步将其拓展为分阶段优化（Staged Expert Affinity）：

• 第一阶段：以节点为单元，最小化跨节点的路由（因为 Node→Node 通信最贵）。
• 第二阶段：在完成跨节点优化的基础上，再做节点内 GPU 之间的放置优化（因为 NVLink 带宽远高于 IB）。

这个两阶段的 ILP 带来的实际收益是：

• 在 16 节点配置下，token 留在同一节点的概率提升了 2 倍。
• 跨节点通信减少了 15% 到 35%。

专家亲和性的时间演化

ExFlow 做了另一个有价值的分析：专家亲和性是什么时候出现的？

通过追踪训练过程中每层的专家路由分布，他们发现：

• 训练初期（前数百步）：专家选择高度不均衡——少数几个专家接收了绝大多数 token。这与 FasterMoE 等工作的观察一致（MoE 训练早期的”路由坍缩”现象）。
• 随着训练推进：由于 auxiliary load balancing loss 的作用，专家选择逐渐趋于均匀分布。但亲和性的模式并没有消失——在均衡的表面下，每个专家仍然”偏爱”下一层的特定几个专家。
• 亲和性稳定后：在训练的后期，专家亲和性矩阵变得稳定。这意味着你可以在训练的中间阶段就提取亲和性矩阵，用于后续的离线放置优化。

这个发现的重要工程含义是：ExFlow 不需要任何 fine-tuning。 你可以拿到一个预训练好的 GPT MoE 模型，跑一批采样 token 通过它来收集每层的路由决策，求解 ILP 得出最优的专家放置方案，然后按这个方案加载模型开始推理。整个过程对模型本身是零侵入的。

实验结果

ExFlow 在由 4 个 A100-SXM4-80GB GPU 组成的节点（NVLink 互联）+ 双轨 HDR200 InfiniBand 互联的 Wilkes3 集群上做了全面验证。关键数据：

一个值得注意的规律是：当每块 GPU 持有的专家数量越多，ExFlow 的效果越显著。 在 MoE-64 用 8 块 GPU 时（每 GPU 8 个专家），性能提升最大——因为 GPU 内部的专家间亲和性发挥空间最大。当每 GPU 只有 1 个专家时（如 MoE-16 用 16 个 GPU），ExFlow 只能利用节点内的 GPU 间亲和性，改善幅度有限——跨节点通信的开销压倒了节省下来的节点内通信。

这揭示了 ExFlow 设计的一个关键适用条件：它在”专家密度”较高的配置下最有效——即每块 GPU 持有多于 1 个专家。当专家数等于或小于 GPU 数时（每 GPU 最多 1 个专家），亲和性只能在节点内层面发挥作用，优化空间被大幅压缩。

另一个发现是 All-to-All 在总推理时间中的占比随规模急剧增长：

• 1 节点：All-to-All 仅占 15.3%——NVLink 足够快，计算主导。
• 2 节点：All-to-All 占比飙升至 62.5%——跨节点 IB 通信成为主要瓶颈。
• 4 节点：70.2%。
• 8 节点：76.0%。

在后三种情况下，ExFlow 能通过减少需跨节点路由的 token 数量实现对总延迟的显著改善。

路线二深入：Aurora —— 理论最优的推理调度

问题的重新定义：推理时间的理论下界

Aurora 的工作方式和 ExFlow 完全不同。ExFlow 走的是”发现和利用模型结构”的路径，Aurora 走的则是”形式化建模→理论推导→算法设计”的经典计算机系统研究路径。

Aurora 团队提出：MoE 推理时间的理论最小化还没有被形式化地研究过。 现有的 MoE 推理优化都是经验性的——尝试不同的放置策略、调整 batch size、调节容量因子——但没有一个统一的数学框架来回答”在这种情况下，推理时间的最优值是多少？”

这个问题的形式化定义如下：

对于一个 MoE 模型在 GPU 集群上的推理，任一层的前向时间被分解为：

Inference_time = max_i(|G_i| + |N_i|) + max_j(|F_j| + |C_j|) + max_k(|A_k|)

其中：

• G_i	：GPU i 上 Gate 网络的计算时间

• N_i	：GPU i 上第一次 All-to-All 的通信时间

• F_j	：GPU j 上 FFN（专家计算）的时间

• C_j	：GPU j 上第二次 All-to-All 的通信时间

• A_k	：GPU k 上 Aggregation 的计算时间

公式中的 max 操作是最关键的部分：由于 All-to-All 是同步的（所有 GPU 必须等通信完成才能继续），推理时间由最慢的 GPU 决定。这是 MoE 推理和 Dense 推理的根本差异——在 Dense 推理的 Tensor 并行中，AllReduce 也是同步的，但数据分布是均匀的；在 MoE 推理中，All-to-All 的数据分布高度不均衡，某块 GPU 可能被”淹没”。

四种场景的分类与理论分析

Aurora 将 MoE 推理的所有实际部署场景归纳为两个维度、四个象限：

每个场景对应不同的优化变量组合：

• 独占 + 同构（最简单）：只需优化通信时序（transmission order）。所有 GPU 一样快，专家放哪不重要，关键是 All-to-All 中 token 的传输顺序。
• 独占 + 异构：需要优化 GPU 分配（哪个专家放哪种 GPU）和通信时序。
• 共置 + 同构：需要优化专家配对（哪两个专家放在同一 GPU）和通信时序。
• 共置 + 异构（最复杂）：需要同时优化专家配对、GPU 分配、通信时序——被证明是 NP-hard 的 3 维匹配问题。

场景 1：独占+同构 —— 最小化推理时间等于最小化通信时间

在这个最简场景中，Aurora 证明了两个定理：

定理 4.1：在独占+同构场景中，最小化推理时间等价于最小化通信时间。

证明是直观的：同构 GPU 意味着每块 GPU 的计算时间相同（Gate 和 Aggregation 与 token 数无关，FFN 时间与 token 数线性相关）。由于两块 All-to-All 是镜像的（第一次从 GPU a 到 b 的数据量与第二次从 b 到 a 的完全相同），最慢的发送 GPU 和最慢的接收 GPU 是同一块。因此：

Min(Inference_time) = |G| + max(|F_i|) + min(max(|N_i|)) + |A|

唯一可优化的项是 max(|N_i|)——All-to-All 通信的最大单 GPU 时间。

定理 4.2：最小通信时间由流量矩阵 D 的最大行和或列和决定：

τ_min = max( max_i Σ_j D_{i,j}, max_j Σ_i D_{i,j} ) / B

其中 D_{i,j} 是从 GPU i 发往 GPU j 的 token 数量，B 是每块 GPU 的带宽。

这个定理的工程意义是：All-to-All 的最优延迟只取决于”最忙的那块 GPU 的发送量或接收量”这个瓶颈值——一旦确定了这个值，你可以通过合理安排 token 的传输顺序来达到这个下界。

Aurora 给出了一个具体的调度算法（Algorithm 1）：

1. 找到瓶颈 GPU（发送/接收量最大的那个）。
2. 瓶颈 GPU 的传输顺序随机（因为它每一刻都在收发，顺序不影响完成时间）。
3. 对其他 GPU，按流量从大到小排序。
4. 依次为每块 GPU 安排传输顺序，确保任何时刻没有两块 GPU 同时向同一目标发送——避免接收端的带宽争用。
5. 重复直到所有流量被调度完毕。

这个算法的核心直觉是：收端不争用。确保每块 GPU 在任何时刻只从一块 GPU 接收数据，则所有 GPU 可以以全带宽持续收发，总通信时间等于瓶颈 GPU 的传输时间。

在同构集群上，Aurora 的调度相比最短作业优先（SJF）调度快了最高 1.38x。

场景 2：独占+异构 —— GPU 分配顺序的重要性

异构集群引入了新的变量：不同的 GPU 有不同的计算能力和网络带宽。Aurora 需要回答：如何分配专家到不同类型的 GPU？

定理 5.1：最优的 GPU 分配是按专家负载（处理的 token 数量）从高到低排序，然后分配给性能从高到低的 GPU。

这个定理的直觉同样清晰：流行的专家需要大量的计算和通信，应该放在最快的 GPU 上。Aurora 的证明使用了交换论证（exchange argument）——任何偏离这个顺序的分配方案，都可以通过交换两个专家来严格减少或不增加推理时间。

定理 5.2：在同构场景中推导出的传输顺序（定理 4.2 的 Algorithm 1）在异构环境中仍然是最优的。唯一的修改是将带宽 B 替换为每块 GPU 各自的带宽 B_i。

在异构集群上，Aurora 的分配方案相比随机 GPU 分配（RGA）实现了 1.36x 到 1.81x 的加速。

场景 3：共置+同构 —— 多模型混部的优化

这是 Aurora 最独特的贡献。在实际生产环境中，多模型混部（multi-model colocation）是常见的——同一块 GPU 同时服务不同的 MoE 模型，以提高 GPU 利用率（工业 GPU 集群的利用率通常不足 50%）。

关键洞察：传统的多模型混部方案——如 Lina——将同一模型的热门专家和冷门专家放在同一 GPU 上。但这样做的效率有限，因为同一模型的两个专家受制于同步 All-to-All——一个模型在通信时，另一个模型即使有计算任务也在等待。

Aurora 提出：将不同模型的专家配对放在同一 GPU 上，让它们在计算和通信上形成互补。 一个模型的 All-to-All 通信与另一个模型的专家计算重叠，实现资源的时间复用。

定理 6.1：在共置+同构场景中，最小化两个模型聚合的通信时间等于最小化推理时间。

Theorem 6.2（Case I：收发对称）：将模型 A 的专家按负载从小到大排序，模型 B 的专家按负载从大到小排序，顺序配对——即在每块 GPU 上，将模型 A 的一个冷门专家与模型 B 的一个热门专家放在一起。

这种”交错配对”的效果是：在每块 GPU 上，两个模型的通信需求互补——一个模型在”猛发”时另一个在”少收”，两者的聚合通信时间不会超过单一热门专家的通信时间。

Case II（收发不对称）：当专家的发送量和接收量不对等时，问题被建模为二分图上的瓶颈匹配（bottleneck matching）。图的左侧节点是模型 A 的专家，右侧节点是模型 B 的专家，边的权重是”如果这两个专家在同一个 GPU 上，该 GPU 的收发量最大值”。目标是在所有完美匹配中，最小化最大边权。这可以用二分搜索 + Hopcroft-Karp 算法在 O(E^2.5) 时间内求解。

在这个场景中，Aurora 实现了：

• 相比 Lina 的推理时间加速：1.25x 到 2.38x（同构集群）。
• GPU 利用率提升：相比独占模式 1.57x-1.72x，相比 Lina 方案 1.28x-1.50x。

场景 4：共置+异构 —— NP-hard 的最复杂战场

当 GPU 既异构（不同性能）又混部（多模型）时，优化问题变为 3 维匹配：

• 维度 1：模型 A 的专家
• 维度 2：模型 B 的专家
• 维度 3：GPU 类型

每个超边（hyperedge）连接一个模型 A 的专家、一个模型 B 的专家、和一个 GPU，超边的权重是该 GPU 上运行这两个专家的推理时间。目标是在所有 3 维完美匹配中最小化最大超边权重。

Aurora 证明这是一个 NP-hard 问题（被归约到 3-Dimensional Matching Variant）。在实际中，Aurora 提供了一个多项式时间的子优算法：

1. 暂时忽略 GPU 分配维度，先用瓶颈匹配求解模型 A 和模型 B 的专家配对（将 3 维问题降为 2 维）。
2. 基于第 1 步的配对结果，再用一个瓶颈匹配将配对后的专家组合分配到合适的 GPU 类型上。

在评估中，这个子优方案的推理时间与暴力搜索的最优解相比仅多 1.07x。

鲁棒性：在不确定中保持性能

Aurora 做了一个重要的鲁棒性测试：假设优化方案是基于收集到的历史流量矩阵离线计算的，但实际推理时流量分布与历史不同（最高达 75% 的噪声）。结果：

• 在有 75% 噪声的情况下，独占+异构场景的加速从 1.90x 退化到 1.60x（仅 15.8% 的性能退化）。
• 共置+异构场景的加速从 ~2.0x 退化到 ~1.80x。

这个结果说明 Aurora 的优化方案在实际生产环境中——流量分布总有一些不可预测的变化——仍然保持了大部分收益。这不奇怪，因为 Aurora 的核心调度策略（避免收端争用）对具体的流量分布细节不敏感。

Aurora 的性能汇总

路线三深入：DeepSpeed-MoE 推理 —— 让万亿参数模型跑在 25ms 内

问题：MoE 推理的”参数效率悖论”

DeepSpeed-MoE 推理系统的设计动机来自一个令 MoE 部署者头疼的悖论：

• 最好情况：MoE 模型（Top-1 路由）中，每个 token 在每个 MoE 层只激活一个专家。如果这个专家的参数大小和 base dense 模型的一个 FFN 相同，那么从”路径上的计算量”来看，一个 52B 参数的 MoE 模型（1.3B base + 128 experts）只消耗约 1.3B 参数的计算。
• 最差情况：一个 batch 中的所有 token 可能激活不同的专家，如果 batch 足够大，最终覆盖了所有 128 个专家，那么 GPU 必须把这 52B 参数全部从显存中读取一遍——推理的瓶颈往往不在计算，而在内存带宽。52B 参数的读取时间是 1.3B 的 40 倍。

这个悖论意味着：MoE 在实际推理中的延迟不取决于”路径上计算了多少 FLOPs”，而取决于”需要从显存中读多少参数”。 如果专家参数量巨大，即使每次只激活一个专家，为了保证所有可能的专家都处于可访问状态，整个推理系统的模型加载时间仍然由总参数量决定。

DeepSpeed-MoE 的解决方案是双管齐下：

1. 模型层面：通过 PR-MoE 和 MoS 把总参数量从源头砍下来（最高 3.7x 压缩）。
2. 系统层面：通过专家并行、Tensor 切片、层级化 All-to-All、并行协调通信等多个维度的系统优化，最大化聚合内存带宽。

PR-MoE：从”所有层一样多专家”到”哪里需要哪里多”

PR-MoE（Pyramid-Residual MoE）基于两个经验观察：

现象一（Phenomenon-I）：不是所有层的 MoE 都一样有用。

DeepSpeed-MoE 团队训练了两个变体：First-Half-MoE（前 6 层用 MoE，后 6 层用 Dense）和 Second-Half-MoE（前 6 层用 Dense，后 6 层用 MoE）。结果很明显：Second-Half-MoE 的收敛远好于 First-Half-MoE。深层的 MoE 比浅层的 MoE 对模型质量贡献更大。 这个观察在 CV 领域（浅层学通用特征，深层学特定特征）有对应，但在 NLP 领域，尤其是 MoE 架构中被首次系统验证。

现象二（Phenomenon-II）：残差 MoE 可以替代 Top-2 路由。

标准的 Top-2 路由需要每个 token 发往两个专家——意味着加倍通信量。DeepSpeed-MoE 发现，如果将一个专家改成固定的 Dense MLP（所有 token 都经过），另一个专家仍然由 gate 路由决定，两者输出相加（Residual-MoE），模型的收敛质量与 Top-2 路由几乎相同，但通信量等同于 Top-1（因为固定的 Dense MLP 不需要跨 GPU 通信）。

PR-MoE 将两个现象合并：

• Pyramid：浅层用少量专家，深层用大量专家。例如 350M 的 PR-MoE 模型在前 10 层用 32 个专家，最后 2 层用 64 个专家。相比标准 MoE 的全部 12 层都用 128 个专家，总参数量减少到不到 1/3。
• Residual：每个 MoE 层包含一个共享的 Dense MLP（所有 token 都经过）+ 一个路由选择的专家，两者输出相加。

最终效果：

• 350M+PR-MoE-32/64 仅用 4B 参数，达到与 350M+MoE-128（13B 参数）相同的模型质量——参数量减少 3.25x，质量不降。
• 1.3B+PR-MoE-64/128 仅用 31B 参数，达到与 1.3B+MoE-128（52B 参数）相同的模型质量——参数量减少 1.68x。

MoS：蒸馏让模型更小

在 PR-MoE 的基础上，DeepSpeed-MoE 进一步使用知识蒸馏（Knowledge Distillation）技术来减少模型深度。他们将其称为 Mixture-of-Students（MoS），因为学生模型保留了与教师模型相同的稀疏 MoE 架构——而不是像传统的 MoE 蒸馏那样将稀疏模型蒸馏成稠密模型（这样就失去了 MoE 的推理优势）。

MoS 的一个关键创新是分阶段知识蒸馏（Staged KD）：

• 在训练的早期阶段（前 400K 步），学生模型同时优化标准语言模型损失和蒸馏损失（与教师输出的 KL 散度）。
• 在训练的中后期，停止蒸馏，让学生模型只优化标准语言模型损失。

原因很有趣：在 MoE 模型中，由于 PR-MoE 已经减少了参数容量，进一步减少层数后，学生模型的容量可能不足。到训练后期，学生模型无法同时最小化两个目标（CE loss 和 KD loss），强制两者都优化反而会导致模型”分心”——在一个目标上变好，在另一个上变差。提前停止 KD 让学生模型专注在核心目标上，产生了更好的效果。

最终结果：

• 350M+PR-MoE+L21+MoS（3.5B 参数）保留了教师模型（4B）99.5% 的零样本评估性能。
• 1.3B+PR-MoE+L21+MoS（27B 参数）保留了教师模型（31B）99.1% 的零样本评估性能。
• 层数减少 12.5%（24 层→21 层）。
• PR-MoE + MoS 合计将标准 MoE 模型的大小压缩了最多 3.7 倍。

推理系统：多维并行 + 通信优化的交响乐

DeepSpeed-MoE 推理系统是整个 DeepSpeed 推理栈的 MoE 分支。它的核心设计哲学是：没有单一的并行策略适用于 MoE 模型的不同部分——专家参数和非专家参数需要不同的处理方式。

专家参数的处理：专家并行 + 专家切片

• 专家并行（Expert Parallelism）：将 E 个专家分布到 P 个 GPU 上，每个 GPU 持有 E/P 个专家。当 EP 度设置为 128 时，一个 1.3B+MoE-128 模型（52B 总参数）在每个 GPU 上的关键路径只有 1.3B 参数——比等质量的 6.7B Dense 模型小 5 倍。这意味着，在理论上，这个 MoE 模型可以比等质量 Dense 模型快 5 倍——前提是通信开销为零。
• 专家切片（Expert Slicing）：当 GPU 数量超过专家数时（追求更低的延迟），进一步将每个专家内部的参数矩阵做 Tensor 切片。这增加了通信（切片间需要 AllReduce），但可以进一步减少每个 GPU 上的参数加载量。

非专家参数的处理：Tensor 切片 + 数据并行

• 对于非专家部分（Attention、LayerNorm、Gate），使用节点内 Tensor 切片——将参数矩阵切到多个 GPU，通过 AllReduce 汇总结果。节点内 NVLink 带宽足够支持 Tensor 切片的通信需求。
• 对于跨节点的扩展，使用数据并行——不同节点处理不同的 batch，不需要通信（除了最终的结果同步）。

这种”专家参数和非专家参数分别用不同的并行模式”的设计，是 MoE 推理和 Dense 推理最大的工程差异。

层级化 All-to-All（Hierarchical All-to-All）

当专家分布在大量 GPU 上时，标准的 All-to-All 延迟随 GPU 数量线性增长——O(p)，其中 p 是 GPU 总数。DeepSpeed-MoE 实现了层级化 All-to-All，将延迟从 O(p) 降到 O(G + p/G)（G 是每节点的 GPU 数）：

1. 先在每个节点内部做数据布局变换。
2. 节点内 All-to-All。
3. 第二次数据布局变换。
4. 跨节点 All-to-All（p/G 个并发）。

虽然这使总通信量翻倍（因为数据经历了两次 All-to-All），但在小 batch size 推理场景下——消息大小小，通信是延迟敏感的而非带宽敏感的——减少通信跳数带来的延迟收益超过了额外数据拷贝的代价。

并行协调通信优化（Parallelism-Coordinated Communication）

当同时使用 Tensor 切片和专家并行时，两者之间的交互可以产生严重的通信浪费。标准的做法是将 Tensor 切片的 AllReduce 和专家并行的 All-to-All 作为独立的黑箱操作，但 DeepSpeed-MoE 利用了一个关键特性：

Tensor 切片的 AllReduce 会在所有参与 GPU 之间复制数据。这意味着在 AllReduce 之后，同一个 Tensor 切片组的 GPU 拥有相同的数据。DeepSpeed-MoE 利用这个特性，让 All-to-All 只在具有不同数据的 GPU 之间进行（跨 Tensor 切片组），而不是在所有 GPU 之间进行。

具体地，如果有 T 路 Tensor 切片、E 路专家并行：标准 All-to-All 的延迟是 O(T × E)。并行协调后，延迟降为 O(E/T) + O(T)（跨切片组的 All-to-All + 切片组内的 AllGather）。

在一个 128 GPU、8 路 Tensor 切片、128 路专家并行的配置中，这将 All-to-All 的延迟从 O(128) 降到 O(16) + O(8)，约为 6-7 倍的通信延迟减少。

MoE 计算核优化

DeepSpeed-MoE 还对 MoE 特有的计算操作（gate 函数、token-to-expert 映射、expert-to-token 反向映射）做了深度优化。传统实现使用稀疏张量（sparse einsum），计算复杂度为 O(S × E × M)——其中 E（专家数）通常很大，且大部分操作涉及与零的乘法和加法。

DeepSpeed-MoE 的优化：

• 将 Gate 函数中的所有操作（top-k、cumsum、scatter）融合到一个 CUDA kernel 中，减少 kernel launch 开销。
• 用稠密的 token-to-expert 映射表（mapping table）取代稀疏 einsum——复杂度从 O(S × E × M) 降到 O(S × M)，去掉了 E 因素。
• 使用 Blelloch 扫描算法在 GPU 上并行化 cumsum 操作。

这些优化合计带来了超过 6 倍的 MoE kernel 延迟减少。

推理性能：从不敢相信到不得不信

DeepSpeed-MoE 的最终性能数据展示了 MoE 推理在充分优化后的潜力：

延迟与吞吐的同步优化：

• 对于 52B 的 MoE 模型（1.3B+MoE-128），从 8 GPU 扩展到 64 GPU，每 GPU 的吞吐不降反升（与 Dense 模型的 sub-linear scaling 形成鲜明对比）。这是因为更多 GPU 意味着每 GPU 持有更少的专家，减少了每 GPU 的参数加载量，而 DeepSpeed 的通信优化使增加的 All-to-All 开销被控制在可接受范围内。
• Dense 模型做不到这一点——Dense 模型的 Tensor 切片在增加 GPU 数时会引入更多 AllReduce 通信，通常导致每 GPU 吞吐下降。MoE 模型却可以让延迟和吞吐同时改善，这是 MoE 结构在推理中的一个独特优势。

万亿参数级别的推理：

• 1 万亿参数 MoE 模型（24B+MoE-128）在 128 GPU 上的推理延迟 < 25ms（DeepSpeed），而 PyTorch baseline 是 > 70ms——7.3 倍的延迟 + 吞吐同时改善。
• 2 万亿参数 MoE 模型（47B+MoE-128）在 256 GPU 上的推理延迟 ~25ms（DeepSpeed），PyTorch baseline ~180ms。

与等质量 Dense 模型的对比：

• 52B MoE 模型（等质量于 6.7B Dense）在 DeepSpeed 上比 6.7B Dense 在 PyTorch 上推理快 2.4 倍、便宜 2.4 倍。
• 1.5 万亿 MoE 模型（等质量于 175B Dense，如 GPT-3 规模）在 DeepSpeed 上比 175B Dense 在 PyTorch 上推理快 4.5 倍、便宜 9 倍。
• MoE 的优势随规模增长而放大：因为更大规模的 Dense 模型需要更高度的 Tensor 切片，产生更多 AllReduce；而 MoE 模型通过专家并行 + 专家切片处理主要参数，AllReduce 的压力更小。

PR-MoE + MoS 的叠加效果：

• 在 349B 规模的 MoE 模型上（8B+MoE-128），PR-MoE+MoS 使得推理仅需 16 GPU（而非标准的 32 GPU）——资源需求减半。
• 延迟进一步降低约 30%-40%，因为每 GPU 需要加载的参数量更少。

横向对比：训练 vs 推理 —— 同一模型，两场战争

我们把训练和推理在通信维度的区别做一个系统性对比。这个对比不仅有助于理解为什么需要不同的优化策略，还能帮助工程师在设计 MoE 系统时做出正确的架构选择。

一、根本差异总览

二、通信量级的差异

训练和推理在通信量级上的差异比对数值本身更能说明问题。考虑一个具体例子：

模型：1.3B+MoE-128（52B 总参数，12 MoE 层，d_model=2048，128 experts/layer） 训练：batch_size=256，seq_len=2048 推理：batch_size=1，seq_len=2048（prompt）+ 生成 512 tokens

训练的单层 All-to-All 通信量：

dispatch: 256 × 2048 × 2048 × sizeof(fp16) = 256 × 8.4 MB = 2.1 GB
combine: 相同 = 2.1 GB
每层合计: 4.2 GB

这块 4.2 GB 的数据要在 128 块 GPU 之间通过 All-to-All 交换。在 50 GB/s 的 InfiniBand 上，理论最小传输时间是 84ms。但实际上，由于 All-to-All 的拓扑限制和消息建立开销，实际时间往往是 2-3x。

推理的单层 All-to-All 通信量（batch=1, prompt phase）：

dispatch: 1 × 2048 × 2048 × sizeof(fp16) = 8.4 MB
combine: 相同 = 8.4 MB
每层合计: 16.8 MB

数据量小了 250 倍！但这不意味着通信问题消失了。原因：

1. 消息太小：16.8 MB 在所有 GPU 间分摊，平均每个 GPU 对之间只有 16.8 / (128 × 127) ≈ 0.001 MB = 1 KB。这种大小的消息在 IB 上的 RDMA 传输效率极低——建立 RDMA 连接的固定开销可能比实际数据传输时间还大。
2. 没有摊薄空间：16.8 MB 不能像训练那样通过更大的 batch 摊到更多 token 上。
3. 12 层的累积：虽然每层只有 16.8 MB，但 12 层的串行 All-to-All 的延迟累积（加上每层的 gate、expert FFN、attention）构成了总推理延迟的主要部分。

推理生成阶段（decode, batch=1, 每次只生成 1 个新 token）：

dispatch: 1 × 1 × 2048 × sizeof(fp16) = 4 KB
combine: 4 KB
每层合计: 8 KB

这种情况下，All-to-All 几乎完全是建立开销。每层 8 KB 的数据需要 128 路 All-to-All——这意味着大量的 GPU 需要交换极少的数据。这就是为什么 ExFlow 的优化在生成阶段特别有价值（减少需要参与 All-to-All 的 token 数量），也是为什么层次化 All-to-All（DeepSpeed）在小消息场景下有效（减少通信跳数，避免全网 All-to-All）。

三、为什么负载均衡在推理中更棘手

训练中的负载均衡可以通过 auxiliary loss 强制实现。但推理中无法使用这个机制——你总不能因为某个专家的负载太高就修改推理请求的路由结果。

推理的负载不均问题在三个层面表现出来：

1. 请求分布不均：用户请求的语义内容不是均匀分布的。在某个时间段内，可能有大量请求都是”写 Python 代码”——这些请求的 token 可能会倾向于某些特定的专家。如果这些专家聚集在少数 GPU 上，这些 GPU 会被挤爆。
2. Prompt 长度不均：不同用户的 prompt 长度可能相差几个数量级（从一句话到一篇长文）。更长的 prompt 意味着更多的 token 需要路由。
3. “热点专家”问题：即使 auxiliary loss 在训练阶段强制了全局均匀路由，模型在特定领域（如编程、数学）的推理中仍然可能对某些专家产生倾向。这些”热点专家”在推理时成为单点瓶颈——它们所在的 GPU 延迟决定了整个推理的延迟。

四、优化策略的对比

这个对比表清楚地说明：训练和推理虽然都在处理 MoE 的通信问题，但它们最优策略集几乎没有交集。 训练靠大量 tricks 和调度自由度覆盖通信瓶颈，推理靠减少通信需求本身和精细的离线优化。

总结与清单：MoE 推理通信优化的决策框架

五条关键洞察

1. 推理的通信问题不是训练通信问题的子集。 它是一个独立的问题，有自己的约束和优化空间。把训练中的优化方案直接搬到推理中——比如 topology-aware gating——在大多数场景下是无效的。推理需要自己的技术栈。

2. 减少通信需求比优化通信执行更根本。 ExFlow 从每层 2 次 All-to-All 减到 1 次，比 Aurora 优化传输顺序的收益（最高 1.38x）更大。DeepSpeed-MoE 把总参数砍 3.7 倍，直接减少 All-to-All 的参与方和通信量。当推理 batch=1 时，通信建立开销占总通信时间的比重极高，所以”少通信”比”快通信”更重要。三条路线中，路线一（通信消除）和路线三（系统压缩）在极端推理场景下通常比路线二（部署调度）带来更大的绝对收益。

3. 专家在 GPU 上的”摆放”是一门值得投入的科学。 ExFlow 的 ILP 优化和 Aurora 的瓶颈匹配——两者都用数学规划来求解专家放置问题——说明这不仅是工程直觉问题，而是有形式化最优解的结构化问题。有充足理由将专家放置纳入 MoE 推理部署的标准流程。

4. 模型设计可以服务于推理效率。 PR-MoE 的 Pyramid 设计（浅层少专家、深层多专家）和 Residual 设计（替代 Top-2 路由）展示了”从推理通信约束出发设计模型架构”的潜力。MoS 的分阶段蒸馏——在训练早期使用教师信号、后期让学生独立——也表明训练过程中的决策可以影响最终推理系统的效率。

5. 异构和多模型的叠加是推理优化的新前沿。 Aurora 处理的异构 GPU 集群和多模型混部场景，是实际生产环境的真实写照。随着 MoE 模型多样化和推理集群生态复杂化，异构 + 混部的优化将越来越重要。Aurora 的 1.07x 子优解为这个 NP-hard 问题提供了实用方案。

MoE 推理通信优化决策清单

以下是一个面向工程师的决策清单，按照”部署前”和”部署中”两个阶段组织：

部署前阶段：

• 分析目标 MoE 模型的专家数量、每层激活模式、路由分布。确认是否存在跨层专家亲和性（ExFlow 的 critical prerequisite）。
• 如果存在专家亲和性且 GPU 资源允许每 GPU 持有多于 1 个专家：求解 ExFlow 的 ILP，产出最优专家放置方案。预期收益：Cross-GPU 路由减少 25%-67%，具体取决于每 GPU 专家数。
• 无论是否存在亲和性：按 Aurora 的算法对 All-to-All 通信矩阵做传输顺序优化（收端不争用策略）。预期收益：同构集群最高 1.38x 延迟改善。
• 评估 PR-MoE 是否适用于目标模型：如果 model quality 对参数效率敏感（即当前模型有”参数冗余”空间），考虑使用 PR-MoE 架构在相同质量下减少总参数量。
• 如果模型已经训练完毕且质量有冗余空间：考虑 MoS 蒸馏来压缩模型深度，进一步降低总参数和通信需求。

部署中阶段：

• 如果 GPU 集群是异构的（不同 GPU 型号混部）：按 Aurora 的负载-性能排序做 GPU 分配。
• 如果有多个 MoE 模型需要混部在同一集群：使用 Aurora 的 bottleneck matching 来配对跨模型专家。目标：计算密集型专家的通信阶段与通信密集型专家的计算阶段共享 GPU 时间。
• 在推理系统的实现层面：采用层次化 All-to-All（DeepSpeed 的设计）。如果同时使用 Tensor 切片和专家并行，务必实现并行协调优化——避免 Tensor 切片组内的冗余 All-to-All。
• 对于 MoE 的 gate 和 token mapping kernel：用稠密映射表取代稀疏 einsum。将 gate、cumsum、scatter 融合为单 kernel。预期收益：MoE kernel 延迟 6x 减少。
• 建立监控：持续跟踪推理时的 token-to-expert 分布。如果分布与离线优化时的假设产生显著偏差，重新运行专家放置优化（ExFlow 的重优化成本很低——只需重新求解 ILP）。

优先级建议：

如果资源有限，按以下顺序推进：

1. 首先做模型压缩（PR-MoE + MoS 如果可用）——这是最大杠杆的措施，因为减少参数总量从根源上降低所有推理成本。
2. 然后做专家放置优化（ExFlow ILP）——这是成本最低、收益显著的优化，不需要任何模型修改。
3. 再做通信调度优化（Aurora 传输顺序）——实现复杂度低，对现有系统的侵入性最小。
4. 最后实现系统级优化（层次化 All-to-All, kernel fusion, parallelism-coordinated communication）——这些需要更深的工程投入，但在前三项完成后仍然能提供额外的加速。

MoE 推理的通信优化是一场年轻的战争。ExFlow 和 Aurora 分别发表于 2024 年，说明学术界在最近一两年才真正开始正视这个问题。DeepSpeed-MoE 的推理系统虽然更早（2022 年），但更多是从工业工程的角度在”做”。随着 MoE 模型从实验室走向大规模在线推理服务，这条优化链条的每一个环节都会面临更严格的实用检验。

有一个尚未充分探索的方向值得关注：将 ExFlow 的专家亲和性思想与 Aurora 的形式化优化框架结合——用 ILP 确定专家放置，然后用 Aurora 的调度算法优化剩余 All-to-All 的传输时序。这两项工作在数学上天然互补（ILP 决定”谁在哪”，Aurora 决定”谁先发”），但至今没有人将它们整合到一个统一的推理优化框架中。这可能是 MoE 推理通信优化的下一个突破点。

下一篇预告：第十五篇将聚焦 MoE 通信中的容错与弹性训练——当 GPU 故障成为常态，如何在不中断训练的情况下动态调整专家分布和通信拓扑。

• Previous
  Full System Power, Thermal and Parallelism
• Next
  Specialized Chips and Future Outlook