Megatron-LM-Moe 论文阅读笔记

可扩展 Moe 模型在 MegatronLM 核心上的训练

原论文请点击：Scalable training of Mixture-of-Experts Models with Megatron Core

笔者注：最近感觉一直在挖坑，而且还根本没时间填。自己水平不够没有什么创造性的工作，总是搬一些源码解读或者读书笔记之类的东西也感觉跟偷窃一样。什么时候我才能做出比较好的工作呢...
总之先把这篇看完一遍再去分析其他的代码...

目录

0 摘要

Moe模型训练带来了稠密模型中不存在的新挑战，主要存在于内存，通信和计算三个方面。集中于解决其中一方必然会导致另一方的瓶颈问题，因此整体设计与解决非常重要。

团队通过集成多种跨内存/通信/计算的优化于一起解决这些挑战。主要有：

  • 内存：（细粒度重计算，卸载等）
  • 计算：（优化后的dispatcher，重叠等）
  • 通信：（分组Gemm，融合，cuda图等）。

这个框架也提供了：

  • 平行折叠 -> 针对灵活的多维度并行；
  • 低精度训练支持 -> 针对FP8/NVFP4；
  • 高效的长上下文训练。

优化后表现如下：

  • GB300/GB200 -> 1233/1048 TFLOPS/GPU，Dpsk V3 685B
  • GB300/GB200 -> 974/919 TFLOPS/GPU， Qwen3-235

接下来我们将详细看看这些技巧时如何工作的，他们都实现了哪些权衡，以及他们在系统级别的交互。

Megatron Framework

1 介绍

Moe架构已经被普遍证明，在指数级降低训练开销的情况下，其效果和等参数量大小的稠密模型性能相似甚至更好。然而有关Moe训练引入了稠密模型中不存在的挑战。本文主要是面向介绍Megatron-Core中的Moe，其架构/并行策略和系统优化是如何实现提升兆级别参数的吞吐的。

1.1 Moe

Moe其实是标准Transformers的一个变体：

  • 引入了一个轻量级的router（或门网络），动态地选择一个token应当传入的专家；
  • 同时将FFN网络更换成对应的多个专家组成的FFN子网络，每个token根据门网络选择的一个专家的小子集(top-k)来选择将token赋予对应的专家。

下图来自于论文：(ISCA 2025) Chimera: Communication Fusion for Hybrid Parallelism in Large Language Models

img

一个token通过门控网络计算出得分：

\[\begin{aligned} p(x) = \rm{softmax}(W, x) \end{aligned} \]

并且通过选择分数最高的K个专家得到最终的输出结果：

\[\begin{aligned} \rm{MoE}(x)=\sum\limits_{i\in\text{TopK}(p(x))}p_{i}(x)\cdot E_{i}(x) \end{aligned} \]

\(E_i\) 代表第i个专家网络。

Moe架构提供了三个关键优势：

  1. （可扩展容量） 模型容量可以独立于计算开销增长：通过加更多的专家！
  2. （计算高效性） 只有一部分模型的参数是激活的，减少了计算量。
  3. （特化性） 不同的专家可以专门负责不同的输入类型。

1.2 挑战

一切的挑战的根源来源于Moe的稀疏性，进一步说明包括：

  • 参数/计算不匹配；（参数远超活跃计算量）
  • 稠密/稀疏不匹配；（注意力机制和MoE层并行策略的不同配置）

稀疏性引入不对称问题

原先稠密模型每个参数都参与运算，因此直接切分模型实现跨GPU运算保证了每个GPU都有运算工作，通信开销不明显。

然而Moe的稀疏性打破了这样的美好幻想：E个专家中每个token只会使用到K个专家，并且 \(K << E\). Deepseek V3中，685B的整体参数只有37B是每个token能够激活的，接近18x的差距。

这样传统的切分方法（如TP）必然会导致某些GPU的计算量不足，降低他们的效率。

因为每个专家都是独立的网络，因此传统策略就是EP并行：将不同的专家放在不同的GPU上，保留全大小的专家GEMM，同时引入了all-to-all运算来实现路由不同的token到对应的专家上。这样的参数-计算不匹配，以及EP并行的通信需求带来了“三堵墙”：内存墙，计算墙和通信墙。

  • 内存墙：所有专家的参数/梯度/优化器状态必须在训练过程中呆在内存中，即使每个token仅激活K个专家。这样就会导致严重的内存压力。
    • 将参数分布到更多设备上->引来了通信带宽的开销；
    • 重计算激活值而不是储存他们->引来了额外的计算；
    • 卸载他们到主机内存上->引来了PCIe开销；
    • 动态路由token到每个专家上->负载不均造成内存激增/某些专家的负载过大
  • 通信墙：EP要求all-to-all通信来分发token，收集结果。我们可以假设 \(T\) 为本地token数目，\(K\) 为topk的数目，\(h\) 为隐藏层维度，这样每个GPU的通信量(send)就是 \(TKh\dfrac{E - 1}{E}\).
    • 随着专家数量不断增长，这个值会趋向饱和，但是通信将逐渐从高带宽的节点内带宽向低带宽额节点间带宽转移，这样可用带宽将严重下降。
    • 同时，稀疏激活值的模式将会导致计算和通信的重叠度不足。
  • 简单计算一下，T个token经过门控后形成的大小为 \(T\cdot h\), 同时从E个专家中选择k个专家的期望计算如下:
  • \(X = K - Y\) (Y: k次选择有没有选到自己)
  • \(E[X] = K - E[Y] = K - K\dfrac{1}{E}\)

  • 计算墙：Moe引入了稠密模型中不存在的如下情况：

    • 小GEMM：细粒度的专家引入了许多小的矩阵乘法，GPU利用率低。在Llama-3中占到了70%左右执行时间，Dpsk-V3中仅有不到50%。剩余时间的消耗在与张量数量相关而非浮点计算相关的操作上。

    这些操作主要就是量化过程中的scale-up/scale-down，因为量化需要寻找最大/最小值，取决于张量元素数量。以及门控计算中的softmax计算，也是浮点运算占用较小的部分。

    • 路由和排列开销。增加了～9%的开销。
    • 负载不均。有些专家很忙，有些专家很闲。浪费算力。
    • 主机开销。由于稀疏性和路由性，Moe会唤起更多的keernels。每次唤起kernel就会导致固定的主机侧开销，这些加起来会导致GPU在kernel之间的空闲。在无丢弃Moe中，动态的tensor形状更导致了主机/设备之间的同步开销。

    曾经有模型通过设置每个专家的capacity来限制专家过载。这样每次当前专家获得的张量就是固定大小的，也不会造成严重的负载不均问题，但是模型的质量可能会略微下降。

曾经的一些优化模式都无法达到很好的tradeoff：

  • 增大BatchSize会提升GEMM利用度，但是内存压力和通信量也会增加。
  • Cuda图消灭了主机调用kernel之间的开销，但是张量的大小必须固定，与无丢弃路由相互矛盾。
  • 跨专家将token聚合成一个组提升了计算效率单负载均衡复杂度增加。

1.3 Megatron-Core

Megatron通过以下方式同时减轻三堵墙问题：

  • 多维度的并行。EP与TP,PP,SP和DP集成。
    • 并行折叠将注意力和Moe层的配置相互解耦，打破了EP<=DP的限制，允许针对不同模型架构/硬件拓扑进行配置。
  • 内存优化。以下几种方式都可以减少内存占用的同时不牺牲模型的吞吐能力。同时，全面支持FP8/FP4降精度训练的专家GEMM/激活/通信进一步降低了激活值存储，同时通过选择性精度策略维持收敛。
    • 细粒度的激活重运算
    • 内存高效型排列
    • 精度注意型优化器
    • 激活值卸载
  • 通信优化。高性能的token分发最大化使用带宽、计算-通信交叠将all-to-all延迟隐藏在专家计算之后。
  • 计算优化。分组GEMM，kernel融合，cuda图，无同步执行流程解决了细粒度Moe架构中潜在的计算分片问题。
  • 其他特性：负载均衡策略，通过丢弃token控制容量，分布式优化器/FSDP支持，分布式ckpt与灵活分片；模型升级回收技术将稠密模型直接部署成专家。

2 Megatron-Core Moe 架构

主要分为内外两个部分介绍：第一部分是Moe层的设计，第二部分是外部设计：(1) 并行进程组如何组织以支持分布式执行；(2) 优化器如何处理与稠密层不同的专家参数。

2.1 Moe层的架构，前向传播

在Megatron-LM中，我们主要有四个部分来部署Moe模型。

  • TopKRouter：将token传输给expert；
  • TokenDispatcher：用于跨GPU通信；
  • Expert：计算激活值；
  • TokenCombiner：将计算好的值通过跨GPU通信传输回去给对应的GPU。

分离几个模块的用途是为了实现独立的优化。

  • router可以被融入CUDA图中，不影响分发逻辑。
  • dispatcher分发器可以在不进行修改专家计算方式的情况下在all-to-all和deepep模式间切换。
  • 专家部署可以使用不同的GEMM后端。

2.1.1 前向传播部分：路由Route，分发Dispatch，计算Compute，聚合Combine

Moe层的处理经过了如下的4个阶段：

img

  • 1.路由。TopK路由将决定当前token由哪一个专家进行处理。
    • 对于一个批次的 \((T, h)\) token，在门函数中通过简单线性变换变成 \((T, E)\)，并通过softmax/sigmoid打分函数得到 \((T, E)\) logits. 随后，通过top-k选择器选择logits中最高分数的 K 个专家，得到最终的路由表，也就是 \((T, K)\) 张量 probs。同时，路由器生成一个 \((T, E)\) 布尔张量来表明我们的路由目的地，也就是 routing_map。
    • 为了保证数据的稳定性，路由器可以通过FP32来运行。--moe-router-dtype fp32
  • 2.分发。分发器将会准备tokens用于跨GPU通信。
    • 首先它需要排列token，这样发送给同一个专家的token就是连续的，这样的排列对于专家侧的稠密高效GEMM非常重要。
    • 分发器通过3个后端的其中之一来将token移动到拥有他们对应专家的GPU上：AllGather(简单，内存要求高)；all-to-all(标准NCCL支持)；Flex（支持优化后的后端比如DeepEP和HybridEP）。
  • 3.专家计算：每个GPU将会对其收到的token执行本地专家计算。
    • 所有的本地专家将会通过TEGroupedMLP在一个简单的分组GEMM中运行，即使每个专家的工作负载很小，也可以最大化GPU利用率。
  • 4.合并：反向通信将会把处理好的token返回给他们原先的GPU，并且通过反向排列恢复原线的序列顺序。
    • 如果配置了一个共享专家，它的输出也将在这一阶段加入。

有关共享专家
在 DeepSeekMoe 中我们有所提到，共享专家为了防止多个路由专家学习到相同知识造成的冗余，因此采用共享专家的模式来学习在多个上下文中存在的共性的知识，防止多个专家之间的参数冗余。实现则通过绕过门函数直接对token进行计算输出。因此可以与其他专家并行计算。
img

2.1.2 路由：token到专家的传递

  • 路由器将会把一个全局的token batch转换成针对不同专家的工作负载。通过以下两个操作：
    • 门函数：线性变换 \(\mathbf{W}_r\in\mathbb{R}^{h\times E}\) 将每个token的隐藏层 \(\mathbf{x}\in \mathbb{R}^h\) 映射成logits \(\mathbf{l}=\mathbf{W}_r^{\top}\mathbf{x}\in\mathbb{R}^{E}\).
    • TopK选择器：一个打分函数将logits映射成概率（归一化）：softmax(\(p_i=e^{l_i}/\sum_j e^{l_j}\)) 或者是 sigmoid(\(p_i=\sigma(l_i)/\sum_{j}\sigma(l_j)\))。后者用于dpsk-V3。逐token选择最高分数的K个专家。

这样就有可能导致专家利用率下降，影响训练性能和模型质量。在下图展示了Megatron Core路由器的流程以及负载均衡的方式。

img

2.1.3 Token分发器：通信抽象

dispatcher将会管理token在GPU之间的移动。主要包括以下的六阶段流水线：

分发预处理 -> 分发 -> 分发后处理（前向传播）-> 合并预处理 -> 合并 -> 合并后处理(反向传播)。

  • 通信后端：主要有以下三个类型的分发器：
    • Allgather：每个GPU收集所有的token和filters，提供给本地的专家。内存需求强但是简单，适合小EP。
    • All-to-All：标准的基于NCCL的点对点通信。每个GPU仅发送目的端需要的tokens。拓展性强但是同步开销大。
    • Flex：支持DeepEP的HT模式（通过交叠隐藏NVLink通信延迟），以及HybridEP（高带宽Moe通信kernel，在NvLink较多的拓扑上表现更好（NVL72））。

2.1.4 专家：计算模块

每个专家都是一个二层的MLP，以及有一个可选的门函数（SwiGLU/GeGLU激活）。Grouped GEMM允许了高效的批量多专家计算。

  • TEGroupedMLP：TransformerEngine优化了支持FP8/FP4量化的部署。
  • SequentialMLP：在一个循环内仅执行一个专家的计算。对debug有效但是更慢。

共享专家：Qwen/DpskV3 包含了共享专家，绕开了门函数对所有token进行处理。共享专家的计算可以与dispatch-计算-combine流水线共同执行，隐藏其计算开销。

2.2 系统集成：并行与优化器

接下来将介绍Moe层是如何继承到分布式训练系统中的：

  • 并行进程组
  • 优化器处理

2.2.1 并行组管理

Moe层要求与稠密层中不同的进程组，因为具有不同的通信方式。Megatron-Core通过ProcessGroupCollection来管理这些进程组。

  • ProcessGroupCollection
    • 注意力层进程组：TP/CP/DP/PP
    • 专家层进程组：EP/ETP/EDP/PP

每个Moe组成需要特殊的进程组，基于如下的通信要求：

组成成分 使用的进程组 原因
Router TP/CP/TP+CP 权重将跨 EP rank 复制
Token Dispatcher ep/tp+ep all-to-all 将跨越多个专家ranks
专家 ep/etp/edp 专家间分片(sharded)，梯度在EDP中规约
共享专家 TP 和稠密的MLP相同

这样的分离组实现了并行折叠技术：注意力层和Moe层可以使用不同的TP/DP配置。这样可以针对不同层的类型进行优化。

2.2.2 优化器和梯度处理

专家参数在并行优化中需要特殊的处理。Megatron-Core 通过使用 ChainedOptimizer 分别为稠密和专家参数封装独立的优化器。以下为三个关键的设计决策：

  • 参数识别。专家参数标记为allreduce=False，将它们从使用标准DP并行参数规约的稠密层参数重区分开来。
  • 分离规约组。稠密层跨 dp_cp_group（全DP）来规约梯度、专家跨expt_dp_group（专家数据并行）进行规约。这样保证了梯度是在正确数量的副本上平均的。
  • 梯度缩放。专家梯度可以通过 edp_size/dp_size 进行缩放，解决不同的有效batch大小在专家上（通过route后仅处理一部分）和在稠密层上（全部处理）的区别。

这样我们可以实现类似ZeRO的优化：在原先的稠密组中，每个组只负责 1/batch 的梯度存储，而专家组只负责分配到自己上的token的梯度存储，也就是 token/batch 的梯度存储。

3 缩放Moe：并行折叠和多维度并行

3.1 为什么并行？为什么Moe很特别？

3.1.1 为什么LLM需要并行

  • 内存墙的限制：每个GPU都具有内存限制，但是训练LLM需要很多的内存存储参数+优化器状态+梯度+激活值。
  • 计算吞吐率的效率考量。单个GPU具有计算能力上限。增加GPU以提升吞吐，降低开销。

3.1.2 需要考虑的权衡

  • 引入通信开销：GPU之间数据交换需要时间和带宽。
  • 同步。快的GPU需要等待慢的GPU。
  • 流水线气泡。在PP中引入了空闲时间。
  • 减少计算密度。分片减少了每个GPU的矩阵大小，降低GEMM的效率。

衡量指标：MFU（模型浮点运算使用率）将会受到并行策略的影响。有效的并行设计可以减少理想和实际MFU之间的差距。

  • 针对稠密模型：MFU随着模型增长保持稳定。
  • 针对稀疏（Moe）模型：参数量以E速度增长，但是计算量以K速度增长。需要更多的GPU内存，但是每个token的计算增长速度却不匹配导致通信开销暴露。这样的不对称增长打破了原有的并行策略的假设。

3.2 Moe并行的挑战

传统的并行策略针对稠密的Transformers设计。Moe模型具有完全不同的计算模式，需要新的并行策略——EP并行，并且EP并行与其他策略结合时带来了新的挑战。

3.2.1 Moe并行悖论

在Moe模型中，模型参数的增长与计算要求的增长不对称，造成了如下的影响：

  • 内存增长迅速：要求更多的GPU容纳所有专家E的模型参数/梯度/优化器状态。更高的topk也通过token副本增加了激活值内存。
  • 更多的通信：将专家通过EP进行训练，引入了all-to-all流量，以K为系数进行缩放。
  • 计算增长速率低。每个token的FLOPs随着K增长而不是随着E增长，造成了不充分的计算，难以与通信进行交叠。

结论：Moe并行从根本上是通信受限的，除非并行策略针对其不对称性进行设计。这跟稠密模型有本质上的区别。

3.2.2 传统并行策略

  • 张量并行：将权重的矩阵沿着隐藏层进行分片，并分布在不同的GPU上。通过 all_gather/reduce_scatter 等合并结果。当矩阵很大足以抵消通信开销时表现良好。
  • 流水线并行：将模型按层进行分割，分布在不同的GPU上。小数据批量流过流水线，每一层都有一个P2P通信。针对PP的优化主要关注的是流水线气泡减少和P2P通信开销，其扩展能力比TP更好。
  • 数据并行：需要将模型复制到不同的GPU上。每个GPU处理不同的批量数据，梯度通过AllReduce进行规约。DP简单但是需要每个GPU拥有全部的模型参数。

FSDP对于DP的改进
针对DP并行中全参数加载的问题，FSDP提出了一种方案，其中最为简单的就是Fully Sharded技术。其流程是这样的：1.对于N个GPU，每个GPU将加载模型每层参数的1/N；2.在数据到达对应层（对应GPU）后，通过all-gather的方式加载当前层的全部参数；进行计算激活值或者计算梯度更新参数；3.释放内存，仅保留每层参数的1/N准备下一步前向或后向传播。这样显著降低了内存开销，假设每层具有参数量 \(O(L_i)\), 那么最多的占用量为 \(O(\sum_i L_i / N + \max_i{L_i})\).

  • 上下文并行：沿着序列维度，将输入序列沿着GPU维度进行划分。每个GPU处理一个连续的序列块，通信仅在注意力机制中使用，因为此处需要token跨越序列块进行运算。CP对于长上下文训练非常重要，因为激活值内存与序列长度呈二次方增长。

为什么传统的并行策略在Moe上失效：

  • TP在Moe上：EP在隐藏层维度上很小，TP会导致更小的分片，降低计算效率，同时增加通信时间占比。TP维数增加提升注意力效率但是影响了Moe。
  • PP与其他许多专家：Moe模型具有大量的参数，单独使用PP需要许多流水线阶段。这样将会造成大量的流水线气泡，降低吞吐。
  • 单独DP：DP需要复制模型的全部参数，模型太大了。

3.2.3 专家并行：新的维度

  • 曾经的并行通过划分一下维度：

    • 层间划分（PP）、权重矩阵（TP）、序列（CP）、数据（DP）
    • EP具有独特的维度：专家是独立的子网络。

  • 专家并行将专家分布在GPU上。每个GPU拥有E/EP个专家。前向传播包括以下阶段：

    • 路由：路由器为每个token选择top-k个专家。
    • 分发：all-to-all通信将tokens发送给对应的专家。
    • 计算：每个GPU处理token，仅使用本地专家。
    • 合并：all-to-all 通信将返回结果给最初的GPU。

这样EP将会出现特殊的权衡如下：

  • 通信：all-to-all集合源语。通信量通过token数目来缩放，而不取决于专家的数量。
  • 计算：每个GPU运行很少量的专家，但是每个专家将使用全部的隐藏层参数。
  • 内存：更高的EP并行=在每个GPU中的专家数目降低=更低的内存压力。

img

EP并行提供了两个关键的优势：

  • 将tokens从不同的GPU聚合起来提供给单个专家，增加了计算强度，提升了GEMM的计算效率。
  • all-to-all通信与GPU数目而不是与专家数目有关。

但是单独仅有EP是不足的。EP仅提供给Moe层，Attention层等不存在专家的稠密层仍然需要传统并行策略。并且，双方都需要PP并行技术来更进一步分割模型参数，用于大规模Moe训练。这样的异构性导致了我们前面讲过的稠密和稀疏的不匹配问题。

3.2.4 将EP和传统并行策略合并带来的挑战

小小的transformer块带来了两种截然不同的计算范式：

范式 注意力（稠密层） Moe（稀疏层）
计算 每个token都参与所有层和其他token的计算 每个token只路由到E个专家中的K个
TP 大型QKV矩阵从更高维度的TP中受益 小型的专家导致高TP降低GPU利用率
CP 长上下文受益于CP维度 与CP无关
EP 与EP无关（没有专家） 对许多专家的分布式训练很关键
  • 稠密与稀疏的错配：传统的框架强制一种并行配置方式，提供给双方；但是他们的最优配置是相互冲突的。如上表所示。曾经将EP视为DP的一种子维度：

\[\text{World Size} = TP\times CP\times PP \times DP, \text{ where } EP\subseteq DP \]

假定的单一并行策略导致稠密层使用 (TP,CP,PP,DP), Moe稀疏层使用从DP中延伸出来的EP导致如下的挑战：

  • 挑战1: 增加的GPU使用。
    • 传统的框架需要最少 \(TP\times CP\times PP\times DP\) GPU。
    • 因为 EP 是 DP 的子集，因此EP=8的时候强制要求 \(DP\geq 8\)。
    • 这样与CP=8的序列并行一起使用的时候，最少需要64个GPU，即使Attn和Moe可以在理论上拥有8个相同的GPU。这导致了Moe训练的准入限制提高。
  • 挑战2: 强制的次优并行策略。
    • 因为注意力和Moe共享相同的TP并行配置，就会导致相互影响的配置。
    • 高TP导致Moe计算不饱和；低TP导致attention层计算并行度不足。
  • 挑战3: 跨节点的通信。
    • EP受限于DP内，高维度的EP通常强制all-to-all跨越节点通信，带宽相对于NvLink下降5-10x。同时，CP通信在稠密层上也会跨节点进行。
posted @ 2026-03-13 17:18  木木ちゃん  阅读(72)  评论(0)    收藏  举报