George Lin's Blog

在上一篇文章中，我们深入剖析了 Tutel 的自适应并行切换与流水线优化、FasterMoE 的影子专家与拓扑感知调度、以及 DeepSpeed-MoE 的层次化存储与多维度并行。这些工作的共同特征是在不改变 MoE 基本计算范式的前提下，通过系统层面的优化来缓解通信瓶颈和负载不均。

本篇文章将转向三条更为激进的技术路线。它们不再满足于在现有范式下做工程优化，而是试图从更根本的层面重新思考 MoE 训练系统的设计：MegaBlocks 将 MoE 计算重新定义为块稀疏矩阵乘法，从根本上消除了 Token Drop 问题；FlexMoE 引入 vExpert 抽象和动态专家放置机制，让系统得以实时响应负载变化；AMD 的 Fused Comm-Comp 则打破了 GPU 计算与通信之间泾渭分明的传统边界，将通信融入计算内核的执行过程之中。

这三条路径分别代表了稀疏化、动态放置和通信融合三种技术哲学，它们从不同角度出发，却指向同一个核心目标——让 MoE 训练更加高效且不牺牲模型质量。理解这三种方法的利弊取舍，对于构建下一代 MoE 训练系统具有重要的参考价值。

一、MegaBlocks：将 MoE 重新定义为块稀疏矩阵乘法

1.1 Token Drop 的代价

在传统的 MoE 实现中，每个专家都有一个固定的容量（capacity）。Token 通过门控网络（Gate Network）被分配到不同的专家，但如果某个专家收到的 Token 数量超过了其容量上限，多余的 Token 就会被直接丢弃——这就是所谓的 Token Dropping。

这个看似简单的工程权衡，实际上对模型质量有着显著的影响。让我们先看一组来自 MegaBlocks 论文的关键数据：

在使用容量因子（capacity factor）为 1.0 时，MoE 模型相对于 Dense Baseline 的验证损失（validation loss）提升仅为 0.15；而当模型采用 Dropless MoE（完全不丢弃 Token）时，同样的提升幅度达到了 0.26。这意味着，Token Drop 使得 MoE 相对于 Dense 模型的收益打了将近 4 折——容量因子 1.0 下的 MoE 只实现了 Dropless MoE 约 58% 的质量收益。

这是一个非常尴尬的局面。容量因子的选择陷入了一个三难困境：

容量因子 1.0：计算效率最高，不会浪费算力，但大量 Token 被丢弃，模型质量受损严重。
容量因子 2.0：丢弃的 Token 显著减少，但实际只有少数热门专家需要这么大的容量，大量冷门专家的容量白白浪费。这意味着每轮训练中有相当比例的浮点运算实际上在做无效计算（padding zeros）。
动态容量因子（Tutel 的方案）：根据实时负载动态调整每个专家的容量，避免了 Token Drop，但容量峰值可能飙升至平均值的 11 倍，导致单轮迭代的计算量剧烈波动，对分布式训练的同步效率造成极大冲击。

更麻烦的是，容量因子本身变成了一个需要精心调校的超参数。在实际训练中，工程师不仅需要搜索学习率、批次大小、优化器参数等常规超参数，还需要为容量因子找到一个”恰到好处”的值——太大了浪费算力，太小了损害模型质量。在 MoE 模型已经引入了专家数量、Top-K 值、负载均衡损失权重等额外超参数的情况下，容量因子的调参负担无疑是雪上加霜。

1.2 块稀疏矩阵乘法的核心洞察

MegaBlocks 的核心洞察在于：如果我们能用一个数学上等价的稀疏矩阵乘法来替代传统的 batched matrix multiplication，那么所有专家都可以处理任意数量的 Token，无需任何容量限制，自然也就不存在 Token Drop。

让我们来仔细拆解这个转换过程。

在传统的 MoE 实现中，每个专家的前向传播通常被实现为一个 batched GEMM：将所有分配给同一个专家的 Token 收集起来，打包成一个批次，然后在专家的权重矩阵上执行矩阵乘法。假设我们有 E 个专家，第 i 个专家分配到 T_i 个 Token，输入维度为 d_model，隐藏层维度为 d_ff（对每个专家而言是 d_ff / E）。那么 MoE 的第一层计算可以表述为 E 个并行的矩阵乘法：

Output_i = Input_i @ W_i^1    (其中 Input_i ∈ R^{T_i × d_model}, W_i^1 ∈ R^{d_model × d_ff/E})

这 E 个矩阵乘法通过批处理方式执行——要么用 cuBLAS 的 batched GEMM API，要么在同一个 GEMM 调用中将不同专家的计算拼接成一个大矩阵。无论哪种方式，都要求每个专家的 Token 数量是预先确定的，这就是容量限制的来源。

MegaBlocks 的突破性想法是：为什么不把所有专家的计算放在一个巨大的稀疏矩阵乘法中完成？

具体来说，如果我们将所有 Token 按专家分配关系排列，输入矩阵可以看作一个由 E 个块组成的块对角矩阵（每个块的大小为 T_i × d_model），而权重矩阵同样由 E 个块组成（每个块大小为 d_model × d_ff/E）。这样，MoE 的第一层计算就变成了一个块稀疏矩阵乘法。

但这里有一个关键的区别：在传统 MoE 中，每个块的大小是相同的（由容量决定）；在 MegaBlocks 中，每个块的大小 T_i 可以不同，完全由实际的 Token 分配决定。这就是”块稀疏”（Block-Sparse）的由来——稀疏矩阵中的非零元素以块的形式组织，但每个块的大小可以不同。

这种计算模式在数值线性代数中有一个经典的名字：Sampled Dense-Dense matrix multiplication，简称 SDD。对于一个两层的专家 MLP（中间有激活函数），计算流程变为 SDD（第一层）→ 激活函数 → DSD（第二层，权重矩阵的转置形式）。对于更多层的专家网络，模式自然延伸为 SDD → DSD → SDD → DSD ……

值得注意的是，DSD（Dense-Sparse-Dense）操作涉及的是权重矩阵的转置。这对于块稀疏矩阵来说并非天然支持——我们将在 1.3 节中看到 MegaBlocks 如何通过”转置索引”来解决这个问题。

1.3 块稀疏 GPU 内核设计

将块稀疏矩阵乘法的理论转化为高效的 GPU 实现，是 MegaBlocks 最具技术含量的部分。面临的挑战包括：稀疏格式的选择、块大小的确定、转置操作的支持、以及与现有 GPU 线性代数库的集成。

块大小选择：128×128 图块

MegaBlocks 选择 128×128 作为块大小，这是在 A100 GPU 上经过系统实验得出的结论。实验对比了 64×64、128×128 和 256×256 三种配置，结果表明 128×128 在吞吐量上始终持平或优于其他配置，且在端到端训练中表现最佳。

这个选择背后有充分的硬件依据。A100 的张量核心（Tensor Core）在执行 128×128 的矩阵乘法时可以达到接近峰值的利用率。更小的 64×64 块虽然提供更细粒度的稀疏性，但每个块的启动开销占比过大；更大的 256×256 块则粗化了稀疏模式，可能导致更多不必要的零填充。128×128 在两者之间取得了最优的平衡。

需要注意的是，128 的块大小意味着 Token 数量必须被填充到 128 的倍数。这引入了一定的填充开销——每个专家的 Token 数量可能需要最多 127 个 Token 的填充。然而，在实际的 MoE 训练中，每个专家通常分配到数千甚至数万个 Token，额外的填充 Token 占比微乎其微。MegaBlocks 将填充操作融合进了排列（permutation）内核中，进一步摊销了开销。

混合 BCSR-COO 格式

稀疏矩阵的存储格式选择对计算效率至关重要。MegaBlocks 需要同时支持两种操作：DSD（以行主序迭代非零块）和 SDD（需要随机访问任意位置的非零块）。这两种操作对稀疏格式的要求恰恰相反。

MegaBlocks 的解决方案是采用混合 BCSR-COO 格式：

BCSR（Blocked Compressed Sparse Row） 作为主存储格式。BCSR 非常适合 DSD 操作，因为它允许按行高效迭代非零块——每一行的非零块在内存中是连续存储的，column index 和 block pointer 数组提供了 O(1) 的行内定位能力。
在 BCSR 的基础上，MegaBlocks 额外物化了每个非零块的行索引。这使得格式同时具备了类似 BCOO（Blocked Coordinate）格式的随机访问能力，可以高效支持 SDD 操作。由于每个块大小为 128×128 = 16,384 个值，而只需要一个额外的整数索引，存储开销仅为 1/16,384——完全可以忽略不计。

这种混合格式的设计非常精巧：它不复制数据，只是在元数据层面提供了双重视图，使得同一份块稀疏矩阵可以高效地同时支持两种访问模式。

转置索引（Transpose Indices）

DSD 操作的一个麻烦之处在于需要以转置顺序访问权重矩阵。在 BCSR 格式中，原始的行主序迭代非常高效，但如果要按列访问（即转置访问），就需要在每一行中搜索对应的列索引——这会导致严重的随机内存访问。

MegaBlocks 的解决方案是构建一个转置索引数组：为每个非零块创建一个额外的索引条目，并按转置后的顺序存储。这个索引数组可以理解为数据库中的”二级索引”——它不复制数据本身，而是提供了从转置维度访问数据的高效路径。

具体的实现方式是：对于原始 BCSR 格式中的每个非零块（位于行 r、列 c），在转置索引中添加一个条目，表示该块在转置视图中位于行 c、列 r。转置索引按行（即原始矩阵的列）排序，使得 DSD 操作可以像 DSD 操作一样高效地迭代”行”（实际上是原始矩阵的列）。

唯一的性能代价出现在 DSTD（Dense-Sparse-Transpose-Dense，即权重梯度计算）操作中。由于转置索引中的数据块地址不是按原始内存布局顺序排列的，访问的空间局部性有所下降。不过，考虑到权重梯度计算在整个训练循环中的占比相对较小，这一代价是可以接受的。

与现有稀疏库的对比

MegaBlocks 的内核基于 CUTLASS 扩展实现，这需要回答一个自然的问题：为什么不直接使用现有的 GPU 稀疏矩阵库？

cuSPARSE 提供了 blocked-ELL 格式，但要求每一行的非零块数量相同。这恰恰违背了块稀疏的核心目的——我们正是希望每个专家可以处理不同数量的 Token，允许不同行有不同数量的非零块。强制对齐会使我们重新陷入容量限制的困境。

Triton 的 Blocksparse 模块通过预计算查找表来加速块稀疏操作，但这种查找表是针对静态稀疏模式优化的。在 MoE 训练中，Token 到专家的分配在每一步都可能不同，稀疏拓扑是动态变化的。频繁重建查找表的开销使得 Triton Blocksparse 不适合这种场景。

因此，MegaBlocks 需要在 CUTLASS 的基础上从头构建对动态块稀疏矩阵的支持，这正是其核心工程贡献所在。

1.4 高效路由与排列

在块稀疏计算之前，还需要完成 Token 到专家的路由和重新排列。MegaBlocks 在这方面的设计同样值得关注。

Token 排列（permutation）是 MoE 训练中的常见开销来源。在传统实现中，Token 需要从原始的批次顺序重新排列为按专家分组的顺序。MegaBlocks 将两个操作融合进一个定制的 CUDA 内核：

Token 重排：将 Token 从批次顺序重新排列为按专家分组的顺序。
填充：将每个专家组的 Token 数量填充到块大小（128）的最近倍数。

由于每个专家通常分配到数千到数万个 Token，最多 127 个填充 Token 带来的计算开销在实际中几乎可以忽略不计（通常低于 1%）。

此外，MegaBlocks 还需要构建稀疏拓扑元数据——即哪些（专家索引，Token 行范围）对应非零块的映射关系。这个构建过程以及转置元数据的生成被整合进了一个单独的 CUDA 内核，避免了多次内核启动的开销。

1.5 性能结果解码

MegaBlocks 的性能数据需要放在适当的上下文中解读。让我们逐一分析各项指标：

块稀疏内核的微观性能

在 A100 上，MegaBlocks 的块稀疏内核平均达到了 cuBLAS 密集 GEMM 吞吐量的 98.6%（标准差 4%），最高达到 104%，最低为 91%。这个数字意味着，块稀疏格式的额外开销平均仅为 1.4%——这在稀疏计算领域是一个相当突出的成绩。大多数通用稀疏矩阵乘法难以超过密集计算 80% 的吞吐量，MegaBlocks 能做到接近持平，得益于精心设计的块大小、混合存储格式和充分的 CUTLASS 底层优化。

值得玩味的是最高 104% 的情况。这意味着在特定的矩阵形状和稀疏模式下，块稀疏内核的吞吐量居然超过了 cuBLAS 的密集 GEMM。这并不是说稀疏计算本身比密集计算快，而是因为块稀疏内核可以跳过全零的块，在某些特定配置下，跳过这些块节省的时间可能超过稀疏格式管理的额外开销。

端到端训练与 Tutel dMoE 的对比

MegaBlocks 在端到端训练中与 Tutel 的 Dropless MoE（dMoE）进行了系统的对比。Tutel dMoE 是 Tutel 的去 Token-Drop 模式，通过动态容量因子来实现，但面临容量峰值可达 11 倍的问题。

在三组不同规模的模型上，MegaBlocks 分别实现了：

MoE-XS：1.38 倍加速
MoE-Small：2.0 倍加速
MoE-Medium：4.35 倍加速

这个加速比随模型规模递增的趋势非常值得分析。随着模型规模增大，Tutel 的动态容量方法引入了越来越多的填充开销。在 MoE-Medium 规模上，Tutel 的填充占用的额外内存导致最大微批次大小减少了 2 到 8 倍——这意味着在相同硬件上，Tutel 需要更多的梯度累积步骤来完成相同的有效批次大小，从而增加了训练时间。MegaBlocks 的块稀疏方法没有这个问题，因为填充仅发生在每个专家内部（最多 127 个 Token），不会随模型规模增长而恶化。

与密集基线的对比

相对于 Megatron-LM 的密集模型训练，MegaBlocks 的 MoE 训练实现了 1.8 到 2.4 倍的加速（在相同的验证损失下）。这意味着使用 MoE 架构配合 MegaBlocks 训练系统，可以在大致相同的模型质量下获得接近两倍的训练吞吐量提升。这对于大模型训练的成本控制具有直接的现实意义。

与有 Token Drop 的 Tutel 对比

即使 Tutel 调到了最优的容量因子（允许一定的 Token Drop），MegaBlocks 仍然实现了 1.18 到 1.38 倍的加速——而且 MegaBlocks 不需要丢弃任何 Token。这意味着 MegaBlocks 同时做到了”更快”和”更好”：不需要在容量因子上做妥协，不需要承担 Token Drop 带来的质量损失，同时还能获得更快的训练速度。

消解容量因子超参数

对于实际使用者而言，MegaBlocks 最直接的实用价值可能不是那 40% 的速度提升，而是”不再需要调容量因子”这件事。在 MoE 训练中消除一个需要精细调节的超参数，意味着更少的实验轮次、更快的调参收敛、以及更可复现的训练流程。从工程效率的角度看，这个收益可能比绝对的训练速度提升更有长期价值。

1.6 MegaBlocks 的局限性与适用场景

尽管 MegaBlocks 在技术上非常优雅，它并非万能灵药。理解其局限性对于正确选择技术方案至关重要。

首先，块稀疏内核的开发维护成本很高。MegaBlocks 需要修改 CUTLASS 的底层来支持块稀疏矩阵，这意味着它绑定了特定的 GPU 架构和 CUDA 版本。当新的 GPU 架构发布时，需要大量的适配工作。

其次，128 的块大小虽然经过精心选择，但仍然是一个硬性的粒度限制。当 Batch Size 很小时（例如在推理场景或小规模微调中），每个专家可能只分配到几十个 Token，128 的块大小会导致显著的填充开销。MegaBlocks 的最佳应用场景是中等以上规模的 MoE 训练，其中每个专家的 Token 数量远超 128。

另外，块稀疏内核在极大规模（例如 GPT-4 级别的数千专家）下的可扩展性也值得进一步验证。虽然理论上块稀疏格式不受专家数量的限制，但稀疏拓扑元数据的构建和管理开销可能随着规模增长而成为新的瓶颈。

总体来说，MegaBlocks 最适合的场景是：中等至大规模 MoE 训练，内存效率是关键约束，模型质量不容妥协（不能有 Token Drop），且愿意承担定制内核的维护成本。

二、FlexMoE：动态专家放置与 vExpert 抽象

如果说 MegaBlocks 是通过重新定义计算范式来解决问题，那么 FlexMoE 则选择了一条更”务实”的路径：接受现有的 MoE 计算范式，但引入动态的专家放置机制来应对负载不均。FlexMoE 来自北京大学、CMU 和微软的联合团队，是 Hetu 分布式深度学习系统的一部分。

2.1 负载不均的深层分析

FlexMoE 的作者对 MoE 训练中的负载不均问题进行了相当细致的实证分析，揭示了几个关键的观察：

负载分布的极端偏斜性

在 GPT-MoE 模型（64 个专家）上，FlexMoE 发现了一个惊人的统计事实：Top-10 专家（即最热门的 10 个，仅占专家总数的 15.6%）接收了约 75% 的 Token。这意味着剩余的 54 个专家共同分担仅 25% 的 Token。这种极度偏斜的分布意味着，哪怕最精细的容量因子调节也无法避免严重的计算浪费——大多数专家的容量大部分时间处于闲置状态，而少数热门专家始终处于超负荷运转。

这种偏斜有其深刻的根源。在 Transformer 的 MoE 层中，专家的功能分化是自然形成的：某些专家在处理特定类型的语义模式时表现更优，而自然语言中某些模式的出现频率远高于其他模式。负载均衡损失（load balancing loss）可以在一定程度上抑制这种偏斜，但过度施加均衡约束会迫使 Token 被分配到不合适的专家，直接损害模型质量。

FlexMoE 的实验中有一个数据非常有说服力：在 Swin-MoE 模型上，施加负载均衡损失后，Top-5 准确率从 94.588%（无均衡约束）下降到 93.981%。近 0.6 个百分点的下降在图像分类任务中是非常显著的。这意味着负载均衡和模型质量之间存在真实的、不可忽视的权衡。

负载变化的连续性

FlexMoE 的另一个关键发现是专家负载在训练过程中呈现连续而缓慢的变化模式。具体来说，同一专家的负载可能从”热门”逐渐变为”冷门”（从多到少），也可能反向变化（从少到多），还可能呈现周期性的循环波动。但重要的是，这种变化是平滑的——不会出现某一训练步中某个专家的负载突然从 50 个 Token 暴增到 5000 个的情况。

这个”连续性”属性是 FlexMoE 动态放置策略得以工作的前提条件。如果负载变化是突变的、不可预测的，那么任何基于历史信息的调度决策都会失效。正是因为负载变化是渐进的，FlexMoE 才能基于对近期负载的观察来做出有效的、前瞻性的调度决策。

2.2 vExpert 抽象与三模块架构

FlexMoE 的核心创新在于引入了一个新的抽象层：vExpert（虚拟专家）。理解 vExpert 对于把握 FlexMoE 的整体设计至关重要。

vExpert 的定义与意义

在 FlexMoE 中，vExpert 是专家调度的最小单元。每个 vExpert 是某个真实专家的一个副本（replica），负责处理该专家的一部分 Token。同一专家在同一个 GPU 上的所有 vExpert 共享同一份权重参数（packed），同一个专家在不同 GPU 上的 vExpert 各自持有独立副本。

vExpert 的设计精妙之处在于它解耦了两个维度：专家的”逻辑身份”和”物理部署”。在传统 MoE 中，每个专家在一个 GPU 上最多有一个实例；在 FlexMoE 中，一个热门专家可以对应多个 vExpert，分布在一个或多个 GPU 上，每个 vExpert 处理一部分 Token。这种一对多的映射关系使得系统可以灵活地调整每个专家的计算资源分配。

更重要的是，vExpert 之间的负载分区是均匀的——对于同一个专家的所有 vExpert，输入 Token 被均匀地分配。这意味着如果一个热门专家有 4 个 vExpert 分布在 4 个 GPU 上，每个 vExpert 恰好处理 1/4 的 Token，从而精确地实现了负载均衡。

三模块架构

FlexMoE 的系统架构由六个组件构成，分为两个层次：

基础 MoE 组件（传统架构中也存在）：

Gate Network（门控网络）：决定每个 Token 的专家分配。
Router（路由器）：执行实际的 Token 到专家的路由。
Expert Network（专家网络）：专家的实际计算，支持多副本。

新增的调度组件（FlexMoE 的核心贡献）：

Scheduler（调度器）：持续监控每个 vExpert 的实时负载，收集运行时统计信息。
Policy Maker（策略制定器）：基于调度器的实时数据，决定何时以及如何调整 vExpert 的部署布局。
Expert Network（增强版）：支持 vExpert 的动态增删和跨 GPU 迁移。

这种架构设计将”监控”和”决策”分离——Scheduler 专注于数据采集，Policy Maker 专注于策略优化。这种分离使得两个组件可以独立演化和优化，也使得系统可以适应不同的调度策略而不需要修改监控基础设施。

2.3 三种修改原语：Expand、Shrink 与 Migrate

FlexMoE 对 vExpert 布局的调整通过三个原子操作完成。理解这三个原语的行为和成本对于评估 FlexMoE 的实用性至关重要。

Expand（扩展）

当一个专家过载时，Policy Maker 决定为其分配新的 vExpert。Expand 操作需要：

在目标 GPU 上分配内存（权重参数 + 优化器状态）。
通过 P2P（Peer-to-Peer）通信将专家参数从已有副本拷贝到目标 GPU。
注册新的 vExpert，使其可以接收 Token。

Expand 的主要开销在于参数拷贝。对于典型的 MoE 专家（例如 d_model = 1024, d_ff/E = 4096），权重矩阵的大小约为 8MB（含 FP16 权重和 FP32 优化器状态）。通过 NVLink 的 P2P 带宽（A100 上约 600 GB/s），这个拷贝可以在几十微秒内完成。更大规模模型上的拷贝时间会相应增加，但仍远小于一个训练步的时间（通常在数百毫秒级别）。

Shrink（缩减）

当一个 vExpert 利用率不足时，系统可以将其释放。Shrink 操作仅通过修改标签（tag）来执行——不需要任何数据传输或通信。具体来说，Policy Maker 只需将被标记为”释放”的 vExpert 从活跃列表中移除。该 vExpert 的 GPU 内存在后续的 Expand 操作中被复用。

零通信开销是 Shrink 的一大优势，但这并不意味着 Shrink 完全无代价。被缩减的 vExpert 之前接收的 Token 需要被重新路由到该专家的其他副本，这可能增加那些副本的负载。Policy Maker 在做 Shrink 决策时需要考虑这种负载转移的影响。

Migrate（迁移）

Migrate 操作用于交换两个 vExpert 的模型状态——将一个冷门专家的副本从一个 GPU 移动到另一个 GPU，同时将一个热门专家的副本反向移动。Migrate 的目标是减少持有某个专家副本的 GPU 总数，从而降低 AllReduce 同步（需要在该专家的所有副本之间同步权重梯度）的成本。

Migrate 的开销与 Expand 类似——需要传输权重参数和优化器状态。但 Migrate 的净收益可能更大，因为它同时优化了两个专家的部署布局，而且直接减少了同步成本。

2.4 调度算法与成本模型

Balance Ratio：直接瞄准最慢 GPU

FlexMoE 提出了一个简洁而精确的负载不均度量指标——Balance Ratio（均衡比）：

Balance Ratio = max_gpu(Σ_{vExpert on GPU} load(vExpert)) / (Σ_{all GPUs} load(all vExperts) / num_GPUs)

这个公式的含义是：负载最重的 GPU 的负载，除以所有 GPU 的平均负载。当所有 GPU 负载完全均衡时，这个值为 1.0；负载越不均衡，比值越大。

选择 max_gpu 而非方差或基尼系数等传统不均度量，体现了 FlexMoE 的工程思维：在同步训练中，整体训练速度由最慢的 GPU 决定（木桶效应）。无论其他 GPU 多么空闲，只要有一个 GPU 负载过重，整个训练步就会被拖慢。因此，优化目标应该直接针对”最慢 GPU”的负载，而非”平均不均度”。

当 Balance Ratio 超过预设阈值时，Policy Maker 被激活，开始寻找可以改善负载均衡的 vExpert 布局调整。阈值的选择是一个工程权衡：阈值太低，Policy Maker 频繁介入，调整开销抵消甚至超过收益；阈值太高，系统对负载不均的响应不够及时。FlexMoE 建议使用 1.05 到 1.10 之间的阈值，在实践中表现良好。

Policy Maker 的贪心搜索算法

Policy Maker 的核心算法（Algorithm 2）采用了贪心搜索策略：

在所有专家中找出当前负载最重的专家（记为 id_0）和负载最轻的专家（记为 id_1）。
模拟一对操作的组合：对 id_0 执行 Expand（增加一个 vExpert） + 对 id_1 执行 Shrink（释放一个 vExpert）。
使用成本模型估算修改后的训练步时间。如果估算的修改后时间小于当前时间，则执行这组修改。
重复上述过程，直到找不到可以改善性能的修改组合。

贪心策略不能保证全局最优，但在实践中效果很好，原因在于负载变化的连续性（如 2.1 节所述）。由于负载变化缓慢，每一步只需要小幅调整，贪心策略足以跟踪负载变化的最陡下降方向。

在 Expand/Shrink 循环结束后，后端还会执行 Migrate 操作来优化专家副本的物理布局——减少持有同一专家副本的 GPU 数量，从而降低梯度同步的 AllReduce 开销。

成本模型的三维分解

FlexMoE 的成本模型（Equation 5）将训练步时间分解为三个组成部分：

T_step = T_comp + T_a2a + T_sync

T_comp：计算时间，与分配的 Token 数量和 GPU 的计算吞吐量（TPS）相关。计算公式为 tokens / TPS，其中 TPS 通过对不同批次大小和专家配置的实际运行进行 profiling 获得。
T_a2a：All-to-All 通信时间。FlexMoE 采用拓扑感知的带宽模型，区分节点内（intra-node，NVLink/NVSwitch）和节点间（inter-node，InfiniBand/RoCE）通信。实际的 Token 传输量乘以相应的带宽参数，得到通信时间的估计。
T_sync：同步时间。当同一个专家在不同 GPU 上有多个副本时，反向传播阶段需要对这些副本的梯度进行 AllReduce。这部分时间通过对不同设备组进行 AllReduce 的 profiling 来建模。

FlexMoE 报告成本模型的平均预测误差低于 3%，这个精度对于调度决策来说完全足够。值得注意的是，成本模型不需要完美——它只需要正确排序不同配置的相对性能，而非精确预测绝对时间。只要模型能准确判断”配置 A 比配置 B 快”，Policy Maker 就能做出正确的调度决策。

2.5 Token 路由策略

FlexMoE 的 Token 路由策略（Algorithm 3）采用了”优先局部、比例远程”的贪心策略。

局部优先原则

对于每个 Token，路由器首先检查本地 GPU 上是否有目标专家的活跃 vExpert。如果有，就优先将 Token 路由到本地副本。这个策略的最大好处是避免了跨 GPU 的 All-to-All 通信——本地处理的 Token 不需要任何网络传输。

局部优先的第二个好处是减少了后续梯度同步的通信量。如果专家的所有活跃 vExpert 都集中在少数几个 GPU 上（甚至只有一个 GPU），梯度 AllReduce 的范围就相应缩小，同步成本大幅降低。

远程分配的比例原则

当本地没有目标专家的 vExpert 时，Token 需要被路由到远程 GPU。FlexMoE 按照各 GPU 上目标专家的可用容量比例进行分配——可用容量越大的 GPU 分到越多的 Token。

这个比例分配策略有助于维持负载均衡：它自然地倾向于将更多 Token 分配到负载较轻的 GPU，避免已过载的 GPU 雪上加霜。可用容量的计算考虑了 GPU 上所有 vExpert 的当前负载情况，而非仅考虑目标专家。

高效布局转换

路由器需要将 Token 从批次顺序重新排列为专家-vExpert 分组顺序（然后反过来，在专家计算完成后恢复为批次顺序）。FlexMoE 实现了一种高效的专家级布局转换（expert-wise layout transformation），能够以最小的内存和延迟开销完成这个排列操作。

2.6 实现细节中的工程智慧

FlexMoE 的论文中包含了若干值得关注的实现细节，它们体现了系统设计中的工程智慧。

并行操作合并

当 Policy Maker 决定执行多个 Expand/Shrink/Migrate 操作时，FlexMoE 尽可能并行化这些操作。例如，多个 Expand 操作可以同时通过不同的 P2P 链路拷贝参数，多个 GPU 之间的 Migrate 可以组成环状传输以最大化带宽利用率。这种并行化显著降低了调度调整的时间开销。

逻辑专家 ID 与 NCCL 死锁避免

在 AllReduce 同步时，FlexMoE 使用逻辑专家 ID 而非物理 vExpert ID 来组建 AllReduce 通信组。这个看似微小的设计选择实际上很关键：物理 vExpert 在 GPU 之间的分布是动态变化的，如果直接使用物理 ID 组建 NCCL 通信组，频繁的组创建和销毁可能导致 NCCL 内部状态不一致，甚至死锁。逻辑专家 ID 提供了一层稳定的抽象——无论物理 vExpert 如何在 GPU 间迁移，同一个逻辑专家的 AllReduce 组成员可以保持不变。

最佳努力（Best-Effort）调度

Policy Maker 的调度操作在一个独立的 CUDA 流上以非阻塞模式执行。这意味着：如果调度操作没有在下一个训练步开始前完成，训练不会被阻塞等待调度完成；当前的 vExpert 布局继续使用，调度调整的效果在下一个训练步才生效。

这种”最佳努力”的设计哲学是务实的选择。它避免了调度开销直接影响训练吞吐量，同时利用负载变化的连续性——延迟一个步生效的调度调整在大多数情况下不会导致显著的负载失衡。

NCCL 通信组管理的 LRU 缓存

随着 vExpert 在 GPU 之间的动态迁移，NCCL 通信组（用于梯度同步的 AllReduce）需要频繁地创建和销毁。NCCL 通信组的创建是一项开销不小的操作（涉及多 GPU 的握手和内存注册）。FlexMoE 使用 LRU（最近最少使用）缓存来管理 NCCL 通信组——最近使用过的通信组被保留在缓存中，当相同的 GPU 组合再次需要同步时可以直接复用，避免了重复创建的延迟。

2.7 性能与质量评估

FlexMoE 在多个维度上与竞品进行了对比，我们需要仔细解读这些数据。

系统效率对比

在 64-GPU A100 集群上，FlexMoE 相对于 DeepSpeed-MoE 实现了平均 1.70 倍加速（最高 2.10 倍），相对于 FasterMoE 实现了平均 1.30 倍加速（最高 1.45 倍）。

这个性能优势有三个主要来源：

零 Token Drop：DeepSpeed 存在 Token Drop，这虽然不直接影响吞吐量，但意味着 DeepSpeed 需要更多的训练步才能达到相同的有效 Token 处理量。换句话说，在相同的吞吐量下，FlexMoE 的每个训练步都在处理更多的有效信息。
更好的负载均衡：FasterMoE 实现了负载均衡，但依赖的是静态的影子专家机制。FlexMoE 的动态 vExpert 放置能够更精确地匹配实际负载分布，减少了 GPU 的空闲等待时间。
优化的通信拓扑：FlexMoE 的局部优先路由策略和 Migrate 优化减少了跨节点 All-to-All 通信量。在 64-GPU 集群中（通常跨越多台服务器），节点间通信是主要的性能瓶颈，减少这部分通信直接转化为训练速度的提升。

模型质量对比

在 BERT、GPT 和 Swin 三个不同架构上，FlexMoE 在困惑度（PPL）和准确率（Accuracy）指标上均一致地优于 DeepSpeed-MoE。这意味着动态专家放置不仅没有损害模型质量，反而——由于避免了 Token Drop——帮助模型实现了更好的收敛效果。

特别值得注意的是 Swin-MoE 上的结果。图像模型对 Token Drop 更为敏感（因为图像 Patch 之间的语义关联比文本 Token 更强），FlexMoE 的 Dropless 特性在这个场景下带来了特别显著的质量优势。

可扩展性

FlexMoE 在 8 GPU 到 64 GPU 的范围内展现了良好的可扩展性。随着 GPU 数量的增加，FlexMoE 相对于 DeepSpeed 和 FasterMoE 的性能优势基本保持甚至略有扩大。这表明动态放置的开销不会随规模增长而失控——Policy Maker 的计算复杂度与 vExpert 数量相关，而 vExpert 数量通常远小于 Token 数量，因此在大规模场景下调度开销仍然是可控的。

2.8 FlexMoE 的适用场景与局限性

FlexMoE 最大的优势在于它的实用性。它构建在 PyTorch 之上，不需要修改底层的 GPU 内核，与现有的分布式训练基础设施兼容。对于大多数工业级 MoE 训练任务，FlexMoE 提供了一个”即插即用”的效率提升方案。

然而，FlexMoE 也有其局限性：

Profiling 开销：成本模型需要对不同配置的计算和通信进行 profiling，这在模型或硬件发生变化时需要重新执行。
调度延迟：虽然负载变化是连续的，但 Policy Maker 的决策-执行-生效链路仍有一定延迟（通常一个训练步）。在负载变化速度较快的场景下，这个延迟可能导致调度落后于实际情况。
vExpert 粒度：vExpert 作为调度的最小单元，其粒度选择影响调度效果。太粗则无法精确匹配负载，太细则管理开销过大。FlexMoE 需要为不同的模型和硬件配置选择合适的 vExpert 粒度。
节点间带宽假设：成本模型中的 All-to-All 通信估计依赖于对网络拓扑和带宽的准确建模。在实际的云环境中，网络带宽可能因多租户竞争而波动，影响成本模型的预测精度。

最适合 FlexMoE 的场景是：生产环境中的大规模 MoE 训练，负载分布随时间变化，模型质量不能妥协，且用户不希望修改底层 GPU 内核。

三、Fused Comm-Comp：AMD 的内核级通信-计算融合

前两种方法——MegaBlocks 的重构计算和 FlexMoE 的动态放置——都是在现有通信-计算分离架构下做优化。AMD 在 SC’24 上提出的 Fused Comm-Comp 则挑战了这个基本架构本身：如果计算和通信不再是两个相继执行的阶段，而是在 GPU 内核执行过程中交织进行，会发生什么？

3.1 核心创新：打破内核边界的通信

传统分布式训练中，计算和通信之间存在明确的边界。一个训练步的典型流程是：前向计算（一系列 GPU 内核）→ All-to-All 通信（由 CPU 触发，通过 NCCL/RCCL 执行）→ 专家计算（一系列 GPU 内核）→ All-to-All 通信 → 后续计算。通信操作在整个 GPU 内核执行完毕之后才开始，而通信期间 GPU 的计算单元处于闲置状态——反之亦然。

AMD Fused Comm-Comp 的核心思想是：在 GPU 工作组（Workgroup，WG）级别实现计算和通信的重叠。当一个工作组完成其分配的计算任务后，立即将结果通过 GPU 网络发送到目标 GPU，而同一内核中的其他工作组仍在继续执行它们的计算任务。这种方式将传统上”整个内核完成后才开始通信”的粗粒度流水线，细化为”单个工作组完成后即可开始通信”的微流水线。

支撑这个愿景的是两个关键技术支柱：

GPU 发起的通信（GPU-Initiated Communication）：传统分布式训练中，通信操作由 CPU 发起——CPU 调用 NCCL/RCCL API，后者在 GPU 上执行集合通信。CPU 发起的通信有一个根本性的限制：它只能在 GPU 内核的边界处触发（内核启动和完成是 CPU 可见的事件），无法在内核执行过程中动态触发。

Fused Comm-Comp 使用 ROC SHMEM（对标 NVIDIA 的 NVSHMEM）来实现 GPU 发起的通信。ROC SHMEM 允许 GPU 内核中的线程直接发起 Put、Get 等单边通信操作到远程 GPU 的内存，无需 CPU 介入。对于单机多卡场景（scale-up），GPU 线程甚至可以直接通过 Infinity Fabric 将数据 store 到对等 GPU 的内存（零拷贝），完全不经过任何通信库。

持久化内核（Persistent Kernel）：传统的 GPU 内核在完成所有任务后退出。Fused Comm-Comp 采用持久化内核设计——内核中的逻辑工作组在一个任务循环中持续执行：获取计算任务 → 执行计算 → 发起通信（如果需要）→ 获取下一个任务，直到所有任务完成。

持久化内核避免了多次内核启动的开销，更重要的是使得内核内部的动态调度成为可能——工作组可以根据任务的实时状态决定下一步做什么。

3.2 三个融合算子的设计与实现

Fused Comm-Comp 论文展示了三个融合算子的设计和实现，分别针对不同的分布式 ML 场景。

3.2.1 Embedding + All-to-All（DLRM 场景）

DLRM（Deep Learning Recommendation Model）是 Meta 提出的推荐系统模型，其核心操作是 Embedding 查表和 All-to-All 通信的交替执行。在大规模 DLRM 中，Embedding 表分布在多个 GPU 上，每个 GPU 持有 Embedding 表的一部分切片。一个批次中的输入 ID 需要被路由到持有对应 Embedding 向量的 GPU 上完成查表，然后结果通过 All-to-All 返回。

Fused Comm-Comp 将这个过程的计算和通信进行了深度融合：

Scale-out（跨节点）设计：

内核中的逻辑工作组被组织为多个切片（slice），每个切片负责一部分 Embedding 查表计算。
工作组的调度顺序经过精心设计——那些计算结果需要发送到远程 GPU 的切片（remote-bound slices）被优先调度执行。这背后的逻辑是：远程传输的延迟更高，越早开始传输，整体完成时间越短。
每个切片完成后，最后一个完成的工作组负责将该切片的结果通过 ROC SHMEM 发出 Remote Put 操作。

簿记机制：

每个切片维护一个 WG Done 位掩码（bitmask），记录该切片中有哪些工作组已经完成计算。
切片就绪标志（sliceRdy）存放在对称内存（symmetric memory）中，所有 GPU 都可以访问。当一个切片的所有工作组都标记为完成后，该切片的数据即可被远程 GPU 消费。

Scale-up（单机多卡）设计：

在单机环境下（GPU 之间通过 Infinity Fabric 连接），Embedding 查表的结果可以直接写入对等 GPU 的内存：

GPU 线程使用标准的 store 指令将计算结果写入对等 GPU 的 HBM 地址空间。
这个 write 操作通过 Infinity Fabric 的缓存一致性协议自动传播到目标 GPU。
完全零拷贝 —— 不需要在本地 GPU 分配发送缓冲区，也不需要远程 GPU 上的接收缓冲区。

性能结果：Fused Embedding+All-to-All 在单机 8 GPU 配置下实现了 20% 的延迟缩减，在跨节点配置下实现了 31% 的延迟缩减。节点间收益更大的原因在于，传统的 CPU 触发的集合通信在跨节点场景下的开销更高（包括 RDMA 连接的建立和同步），GPU 直接发起的单边通信绕过了这些开销。

3.2.2 GEMV + AllReduce（Transformer 推理场景）

Transformer 推理中的关键操作是 GEMV（矩阵-向量乘法，权重矩阵乘以单个 Token 向量）和 AllReduce（在张量并行的 GPU 之间同步部分结果）。

Fused Comm-Comp 的实现集中在 scale-up 场景（推理通常部署在单机多卡上），使用两阶段直接 AllReduce：

两阶段直接 AllReduce：

Reduce-Scatter：每个 GPU 计算自己的权重分片与输入向量的 GEMV。结果不是写回本地内存，而是直接按分片写入不同 GPU 的内存。具体来说，如果输出向量的前 1/N 归 GPU 0 负责最终输出，那么所有 GPU 在计算完前 1/N 的部分后，直接将结果 store 到 GPU 0 的对应内存地址。
All-Gather：Reduce-Scatter 完成后，每个 GPU 持有输出向量的 1/N。然后每个 GPU 将自己的那部分结果 store 到所有其他 GPU 的对应内存位置。

整个过程完全通过 GPU store 指令完成，没有任何中间缓冲区，没有任何通信库调用。

调度优化：类似 DLRM 场景，那些计算结果需要归约到远程 GPU 的工作组被优先调度。在典型的张量并行配置中（例如 4 GPU），每个 GPU 的 3/4 的输出需要发送到其他 3 个 GPU。通过优先计算这些”远程绑定”的部分，通信可以与本地计算在时间上最大程度地重叠。

性能结果：在单机 4 GPU 配置下，Fused GEMV+AllReduce 平均实现了 13% 的延迟缩减（最高 22%）。推理场景的收益低于训练场景的主要原因是 GEMV 的计算量本来就很小（O(d_model²) vs 训练中 GEMM 的 O(batch_size × d_model²)），通信在总延迟中的占比更高，因此融合带来的相对收益更有限。

3.2.3 GEMM + All-to-All（MoE 场景）

这是与本系列文章最直接相关的融合算子。Fused Comm-Comp 将 MoE 专家计算（GEMM）与 All-to-All 通信进行融合。

Triton 实现：AMD 选择在 Triton 语言中实现这个融合算子，而非直接用 HIP/CUDA。这个选择有几个原因：

Triton 提供了更友好的开发体验，特别是对于需要频繁实验的融合算子设计。
Triton 自动处理很多底层优化（如共享内存管理、内存合并访问等），减少了手动调优的工作量。
AMD 为 Triton 添加了 Python 包装器，封装了 ROC SHMEM 和 scale-up store 通信 API，使得在 Triton 内核中发起通信成为可能。

融合逻辑：类似 DLRM 场景，GEMM 的结果分片在计算完成后立即发出，而非等待整个 GEMM 完成。对于 MoE 中典型的专家 GEMM（权重矩阵可能达到 d_model × d_ff/E = 1024 × 4096），这种微流水线化可以显著减少通信关键路径的长度。

性能结果：Fused GEMM+All-to-All 平均实现了 12% 的延迟缩减（最高 20%）。收益低于 Embedding 场景的原因在于 GEMM 的计算密度更高（计算/通信比更大），通信在总时间中的占比相对较小，因此融合的绝对收益也相对有限。

3.3 开销与权衡分析

Fused Comm-Comp 并非没有代价。论文中坦诚地分析了几个关键的性能权衡。

寄存器压力与 Occupancy 下降

GPU 发起的通信需要额外的硬件资源——用于存储网络操作的状态、目标地址、以及通信标志位的寄存器。Fused Comm-Comp 报告这导致 GPU 占用率（occupancy）下降了约 12.5%。较低的占用率意味着更少的工作组可以同时在 GPU 上执行，可能降低计算单元的利用率。

但占用率并不是越高越好。Fused Comm-Comp 的实验揭示了一个反直觉的现象：在 25% 到 75% 的占用率范围内，性能提升了 46%，符合”更高占用率 = 更好性能”的常规认知。然而，当占用率从 75% 进一步提升到 87.5% 时，性能反而下降了 25%。

原因在于内存竞争（memory contention）。过高的占用率意味着过多的并发工作组同时竞争有限的内存带宽和缓存资源。每个工作组可用的缓存容量减少，导致缓存命中率下降，从而抵消了更高占用率带来的隐藏延迟收益。

这个发现暗示 Fused Comm-Comp 的 12.5% 占用率下降可能并非纯负面——如果融合内核本来就处于高占用率区域（例如 75% 以上），适当降低占用率反而可能改善性能。

同步开销

工作组的完成状态跟踪需要同步机制。Fused Comm-Comp 使用跨通道操作（cross-lane operations）而非跨工作组屏障（inter-WG barrier）来实现工作组的协调。跨通道操作（在 AMD GPU 上）允许同一工作组内的线程高效地交换数据，但不提供不同工作组之间的同步。

选择跨通道操作而非工作组屏障的设计选择反映了对 GPU 硬件特性的深刻理解。工作组屏障会强制所有参与的工作组在屏障点等待，这在异构执行场景下（不同工作组完成计算的时间不同）可能导致显著的等待时间。跨通道操作则允许更灵活的同步——工作组只在需要交互数据时进行同步，而非无条件等待。

通信发起延迟的摊销

每次通信操作（ROC SHMEM Put 或跨 GPU Store）都涉及固定的 API 调用延迟（通常是微秒级）。如果每个工作组的计算结果都单独触发通信，累积的 API 延迟可能抵消融合带来的收益。

Fused Comm-Comp 通过”分片级”（而非工作组级）的通信触发来摊销 API 延迟：不是每个工作组完成就发一次通信，而是等一个切片的所有工作组都完成后再统一发出。在典型的配置下，切片数量远少于工作组数量（例如每 GPU 16-64 个切片，每个切片包含多个工作组），将 API 调用的频率降低到可接受的水平。

3.4 与 ML 框架的集成

AMD 为 Fused Comm-Comp 设计了与 PyTorch 和 Triton 的集成方案，这对实际可用性至关重要。

PyTorch 集成：通过新增 API 来支持对称堆内存分配和融合算子启动。例如：

torch.symmetricHeapAlloc() 分配在多个 GPU 之间对称的内存空间
torch.embeddingAll2AllOp() 调用融合的 Embedding + All-to-All 算子
torch.gemmAll2AllOp() 调用融合的 GEMM + All-to-All 算子

这些 API 设计为 PyTorch 的 autograd 兼容，使得它们可以无缝地插入到现有的训练流程中，自动微分仍然可以正常传播梯度。

Triton 扩展：AMD 为 Triton 开发了 Python 级别的通信 API 包装器，使得开发者可以在 Triton 内核中调用 GPU 发起的通信操作。这大大降低了开发自定义融合内核的门槛——开发者不需要写 HIP/CUDA 代码，可以在 Python-like 的 Triton 语言中描述计算和通信的交织逻辑。

3.5 大规模仿真：128 节点 DLRM 训练

为了验证 Fused Comm-Comp 在大规模场景下的有效性，AMD 在 ASTRA-Sim 仿真器中进行了 128 节点 DLRM 训练的仿真实验。

整体收益：在 128 节点配置下，Fused Comm-Comp 实现了 21% 的整体训练时间缩减。这个数字表明融合的收益在大规模场景下并未被稀释——通信时间在多节点环境中通常更长，融合消除的通信等待时间反而更多。

通信感知调度的影响：与”无感知”调度（oblivious scheduling，不考虑通信拓扑的计算调度）相比，通信感知的 WG 调度（优先调度 remote-bound 切片）额外带来了 6% 的执行偏差缩减。执行偏差（execution skew）指的是不同 GPU 完成同一训练步的时间差异，更小的偏差意味着更少的同步等待时间。

这个 6% 的收益看起来不大，但它反映了一个深层优势：通信感知调度使得系统的执行时间更加可预测和稳定。在大规模训练中，执行时间的稳定性和可预测性有时比绝对的平均速度更重要——它影响着检查点保存策略、故障恢复设计、以及资源调度的效率。

3.6 Fused Comm-Comp 的前景与局限

Fused Comm-Comp 代表了一个有前景的技术方向，但当前阶段还有一些重要的限制：

硬件依赖：GPU 发起的通信需要特定的硬件支持——AMD 的 Infinity Fabric 用于 scale-up，ROC SHMEM 支持的网卡用于 scale-out。NVIDIA 的对应方案是 NVSHMEM + NVLink/NVSwitch。这些技术的可用性限制了 Fused Comm-Comp 的部署范围。不过，随着 GPU 互联技术的发展和 SHMEM 编程模型的普及，这个限制正在逐渐减弱。

实现复杂度：设计高效的融合内核需要对 GPU 微架构、网络协议、以及应用特性有深入理解。即使有 Triton 扩展支持，开发一个正确且高效的融合内核仍然是非平凡的工作。这与 MegaBlocks 面临的挑战类似——定制内核的开发和维护成本是不可忽视的。

收益的边界：Fused Comm-Comp 的收益很大程度上取决于通信在关键路径上的占比。对于计算密集型的操作（如大规模 GEMM），通信占比小，融合的绝对收益有限（12% 级别）。对于通信密集型的操作（如 Embedding 查表），融合的收益更为显著（20-31%）。这限制了 Fused Comm-Comp 的适用范围——它并非对所有 ML 操作都有相等的加速效果。

与现有框架的集成深度：虽然 AMD 提供了 PyTorch 和 Triton 的集成，但要充分利用 Fused Comm-Comp 的潜力，开发者可能需要改写现有的模型代码。这与 MegaBlocks 类似——不是即插即用的方案，需要一定的适配投入。

适合 Fused Comm-Comp 的场景是：通信在延迟中占比较高、对延迟敏感的应用（如在线推理、推荐系统训练）、以及有 GPU 发起的通信硬件支持的部署环境。

四、三条技术路径的比较与展望

4.1 稀疏化 vs 动态放置 vs 通信融合

在深入理解了三种方法的技术细节之后，我们现在可以从更高的维度进行比较。

MegaBlocks：改计算

MegaBlocks 的策略可以概括为”不改变通信模式，改变计算范式”。通过将 MoE 计算重写为块稀疏矩阵乘法，它从根本上消除了 Token Drop 的需求，进而消解了容量因子这个超参数。这是最”激进”的方案——它要求重写底层 GPU 内核，放弃了 cuBLAS 等成熟库的直接支持。

其核心优势在于零 Token Drop 和消解容量因子超参数，在内存效率方面表现突出。主要瓶颈有三：定制内核的开发维护成本高；128×128 的块大小是硬性粒度限制，小批次场景下填充开销上升；块稀疏元数据的构建和管理在极大规模下可能成为新瓶颈。

对于中等规模 MoE 训练、重视模型质量和内存效率、且有能力维护定制内核的团队，MegaBlocks 是最佳选择。

FlexMoE：改放置

FlexMoE 的策略是”不改变计算内核，改变资源分配”。通过 vExpert 抽象和动态放置机制，它在保持 100% Token 处理率和现有框架兼容性的前提下，实现了显著的性能提升。这是最”务实”的方案——它接受现有生态系统的约束，在系统层面做优化。

核心优势包括：100% Token 效率；保持模型质量不降低；与 PyTorch 和主流分布式框架兼容；3% 误差以内的成本模型使其性能可预测。主要局限性在于：需要预先 profiling 构建成本模型；调度决策-生效链路存在延迟；vExpert 粒度需要针对不同配置进行调优。

对于生产环境中的大规模 MoE 训练、不愿意修改底层内核、需要与现有训练流程无缝集成的团队，FlexMoE 是务实之选。

Fused Comm-Comp：改通信

Fused Comm-Comp 的策略是”不改变计算范式，改变通信的时机和方式”。通过在 GPU 工作组级别实现计算-通信重叠，从架构层面减少了通信关键路径的长度。这是最”前沿”的方案——它依赖于正在发展中的 GPU 发起的通信技术。

核心优势有：减少通信延迟对训练时间的影响；通信感知调度提升执行稳定性；适用于多种分布式 ML 架构（DLRM、Transformer、MoE）；Triton 扩展降低了定制开发的门槛。主要约束是：需要 GPU 发起的通信硬件支持；寄存器压力导致占用率下降；实现复杂度高；对计算密集型操作收益有限。

对于延迟敏感的推理和训练场景、拥有支持 GPU 发起通信的硬件环境、追求极致性能的团队，Fused Comm-Comp 是最前沿的选择。

4.2 可能的技术融合方向

这三条路径并非互斥。在未来的 MoE 训练系统中，我们可能会看到它们的技术要素被融合进统一的架构：

MegaBlocks + FlexMoE：使用块稀疏计算消除 Token Drop，同时使用动态放置优化专家到 GPU 的映射。块稀疏计算确保了计算资源的无浪费使用，动态放置确保了 GPU 之间的负载均衡。两者结合可以在模型质量和系统效率两个维度上同时达到最优。

FlexMoE + Fused Comm-Comp：在动态调整专家布局的同时，使用工作组级别的通信融合减少每次布局调整后的通信开销。特别是对于 Expand 和 Migrate 操作中涉及的大量 P2P 参数传输，融合通信可以显著降低调度开销。

MegaBlocks + Fused Comm-Comp：块稀疏计算的输出可以直接通过 GPU 发起的通信发送到目标 GPU，消除了”先写回 HBM、再通过 NCCL 发送”的两阶段开销。这需要扩展块稀疏内核以支持通信感知的调度和直接远程写入。

4.3 未解决的问题与开放挑战

在 MoE 训练系统的研究中，仍有一些重要问题尚未得到充分解决：

超大规模下的容错与弹性：当前的 MoE 训练系统研究主要关注性能和效率，对训练过程中的 GPU 故障、网络瞬断、以及动态资源变化考虑不足。当训练规模扩展到数千 GPU 时，故障从”异常事件”变为”常态”，现有的检查和恢复机制将成为严重瓶颈。

MoE 特定的自动并行化：虽然 Tutel 的自适应并行切换是一个良好开端，但 MoE 训练的并行化策略空间远未被充分探索。专家并行度、数据并行度、张量并行度、流水线并行度之间的最优配比与模型架构、硬件拓扑、批次大小等因素高度相关，且最优配比可能随训练阶段而变化。自动化这个决策过程仍然是一个开放挑战。

异构硬件环境下的 MoE 训练：现有工作几乎都假设同构的 GPU 集群。然而，在实际的云环境和算力中心中，不同 GPU 型号、不同网络拓扑、不同内存容量的异构组合越来越常见。如何在异构环境下高效运行 MoE 训练，使得慢速设备不拖累整体效率，是一个极具现实意义的问题。

MoE 推理系统的协同设计：当前的 MoE 训练系统研究大多独立于推理系统。然而，训练阶段的架构选择（如专家数量、容量策略、路由机制）直接影响推理时的部署效率和延迟特征。训练-推理协同设计——在训练阶段就考虑推理时的部署约束——是 MoE 系统研究的重要方向。

4.4 系列回顾与展望

本文是本系列关于工业 MoE 通信研究的收官之作。在这个系列中，我们从通信优化的不同维度，系统地审视了从 DeepSpeed-MoE、Tutel、FasterMoE，到 MegaBlocks、FlexMoE、Fused Comm-Comp 六项代表性工作。

回顾全局，我们可以看到一条清晰的技术演进脉络：

第一代系统（DeepSpeed-MoE、GShard）解决了”能跑起来”的问题——如何在多 GPU 上基本可工作地运行 MoE 训练。它们的核心贡献在于建立了 MoE 分布式训练的基本范式（All-to-All 通信 + 容量限制 + 负载均衡损失）。

第二代系统（Tutel、FasterMoE）解决了”跑得更快”的问题——通过自适应并行切换、阴影专家、拓扑感知调度等手段显著提升了训练效率。它们意识到静态配置无法适应动态变化的负载，开始引入自适应性。

第三代系统（MegaBlocks、FlexMoE、Fused Comm-Comp）开始挑战”范式本身”——MegaBlocks 重写了计算范式，FlexMoE 打破了专家放置的静态性，Fused Comm-Comp 模糊了计算和通信的边界。这些工作标志着 MoE 训练系统研究从”工程优化”迈向”架构创新”。

展望未来，随着 MoE 架构在 GPT-4、Gemini、Llama 4 等前沿模型中的广泛采用，以及模型规模的持续增长，MoE 训练系统的重要性只会越来越高。我们预计下一代系统可能具备以下特征：

训练-推理一体化设计：在训练阶段就针对推理部署进行优化，避免训练时的高效架构在推理时变得低效。
全自动自适应：从并行策略到专家放置，从通信调度到容错恢复，整个训练系统的行为完全由数据驱动，无需人工干预。
跨硬件平台的统一抽象：将 GPU、TPU、ASIC 等不同加速器的 MoE 训练整合在统一的编程模型和调度框架下。
Fault-Tolerant by Design：在系统架构层面内置容错能力，而非将容错作为外挂式补丁。

MoE 训练系统的研究正处于一个激动人心的阶段。我们期待看到学术界和工业界继续推进这些方向的探索，为更大规模、更高质量、更高效率的 AI 模型训练奠定坚实的系统基础。

参考文献

Gale, T., Narayanan, D., Young, C., & Zaharia, M. (2022). MegaBlocks: Efficient Sparse Training with Mixture-of-Experts. Proceedings of MLSys 2023. arXiv:2211.15841.
Nie, X., Zhao, X., Miao, X., et al. (2023). FlexMoE: Dynamic Expert Placement for Mixture-of-Experts Models. Proceedings of IEEE HPCA 2024. arXiv:2304.03946.
Mahdieh, S., et al. (2023). Fused Comm-Comp: Integrated Communication-Computation for Distributed Machine Learning at Scale. Proceedings of SC’24. arXiv:2305.06942.
Hwang, C., et al. (2023). Tutel: Adaptive Mixture-of-Experts at Scale. Proceedings of MLSys 2023. arXiv:2206.03382.
He, J., et al. (2022). FasterMoE: Modeling and Optimizing Training of Large-Scale Dynamic Pre-Trained Models. Proceedings of PPoPP 2022.
Rajbhandari, S., Li, C., Yao, Z., et al. (2022). DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale. Proceedings of ICML 2022. arXiv:2201.05596.