EP并行与稀疏模型

EP并行与稀疏模型#

Author by: 于嘉阔

近年来，大语言模型的发展似乎遵循着一种“暴力美学”——模型参数量越大，能力越强。从GPT-3的1750亿参数到动辄数千亿乃至万亿参数的模型，参数规模的竞赛将人工智能推向了新的高度。然而，这种扩张也带来了巨大的挑战：训练一个高质量稠密模型（Dense Model）不仅需要数万块GPU日夜不停地运行数月，其高昂的计算和能源成本也让许多研究机构和企业望而却步。

这种“越大越好”的路径是否可持续？我们是否必须让模型中的每一个参数都参与每一次计算？

答案或许是否定的。正如人类专家在解决问题时只会调动相关领域的知识一样，模型或许也可以变得更加“专注”和“高效”。这催生了一种全新的模型设计范式——稀疏模型（Sparse Model），其中最具代表性的便是专家混合模型（Mixture-of-Experts,MoE）。MoE架构通过“按需激活”一小部分专家参数来处理输入，巧妙地实现了模型容量与计算成本的解耦。

然而，拥有了设计精巧的稀疏模型，我们又该如何指挥这个由成百上千位“专家”组成的庞大团队在分布式集群上协同工作呢？这便是专家并行（Expert Parallelism,EP）技术大显身手的舞台。EP是一种专为稀疏模型量身定制的并行计算策略，它解决了如何高效训练和运行这些巨型模型的关键工程难题。

本文旨在深入探讨稀疏模型与专家并行这两项相辅相成的技术。我们将从MoE的“专家团队”理念讲起，详细剖析EP并行背后的“数据大交换”机制，并最终阐明它们是如何携手，引领大语言模型从“暴力美学”迈向“精妙分工”的新时代。

稀疏模型介绍#

在引言中，我们将传统的大语言模型比作一位“全才专家”，其知识渊博但计算开销巨大。这种模型，我们称之为稠密模型。如果说稠密模型是一位单打独斗的专家，那么稀疏模型就是一支组织精良的“专家团队”。这支团队由多位各领域的专家组成，并由一位极其聪明的“接待员”负责调度。当问题出现时，接待员会迅速判断问题领域，并精准地引荐给最相关的一两位专家。这种“专家按需工作”的模式，正是当前最主流的稀疏模型架构——专家混合模型的核心思想。


图1：专家混合模型（MoE）层工作原理示意图

MoE模型的精妙之处在于其内部协同工作的两大关键组件：专家网络 (Experts) 与门控网络 (Gating Network)（如图 1 所示）。专家是真正执行计算的主体，在Transformer架构中，它们通常是一个个独立的前馈神经网络 (FFN)。一个MoE层可以包含从8个到超过64个不等的专家，每个专家通过独立学习形成各自的“专业领域”。

而门控网络，或称为“路由器”，则扮演着“接待员”的角色，是整个架构的智能调度核心。对于输入的每一个词元（token），路由器会通过一个可学习的权重矩阵 \(W_g\) 计算出每个专家的得分，并通过Softmax函数将其转换为一组概率权重 \(G(x)\)。路由器并不会激活所有专家，而是选择得分最高的Top-K个（\(K\) 通常为 1 或 2）。最终的输出 \(y\) 是这 \(K\) 个专家输出 \(E_i(x)\) 根据其权重 \(G(x)\) 的加权求和，其过程可表示为：

\[ y = \sum_{i \in \text{Top}-K} G(x)_i \cdot E_i(x) \]

这种基于输入动态决定计算路径的机制，被称为条件计算（Conditional Computation），是MoE实现高效率的关键。

当然，一个设计良好的系统必须考虑潜在的失效模式。如果路由器倾向于将大部分任务分配给少数几个“明星专家”，就会导致严重的负载不均衡，降低整个系统的效率。为了解决这一问题，研究者们在训练MoE模型时引入了负载均衡损失函数 (Load Balancing Loss)。

在主流MoE模型中，这个额外的辅助损失(\(L_{aux}\)) 通常被定义为：

\[ L_{aux} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i \]

此公式中，\(f_i\)代表该批次中被分配给第\(i\)个专家的token实际占比，而\(P_i\)则是门控网络赋予该专家的平均路由权重。通过最小化这两项的乘积，该损失函数会惩罚那些分配不均的路由行为，激励路由器将任务尽可能均匀地分配给所有专家。通过这种机制，可以确保整个“专家团队”协同工作，避免因个别专家过载而产生的瓶颈，从而最大化模型的整体性能和吞吐量。

MoE的思想并非一蹴而就，它的雏形最早可追溯至1991年，由Jacobs等研究者提出，但受限于当时的计算能力和训练技术，这一超越时代的设计在近三十年里未能成为主流[1]。直到2020年，谷歌发布的里程碑式工作GShard[2]，成功地将MoE架构与大规模Transformer模型结合，并解决了其在分布式环境下的训练难题，MoE自此“王者归来”。此后，谷歌通过Switch Transformer[3]将模型参数推向了1.6万亿的惊人规模，充分展示了MoE在模型扩展上的巨大潜力。而真正引爆开源社区的，是Mistral AI于2023年发布的Mixtral 8x7B模型[4]，它催生了一大批高质量的开源MoE模型，标志着稀疏模型进入了新的时代。

近年来，各大研究团队不再满足于简单的专家堆叠，而是围绕架构效率、训练稳定性展开了一系列探索，将MoE模型推向了新的高度。


图2：从MHA到MLA的注意力机制演进示意图

DeepSeek V2系列便是架构效率探索的代表。其拥有高达6710亿的总参数，但激活参数却仅有370亿，稀疏程度令人惊叹。为了在有限的内存中支持更长的上下文，它采用了多头潜在注意力（MLA）技术。从图2中可以清晰地看到，传统的MHA为每个查询头（Query）都保留了独立的键（Key）和值（Value），导致推理时缓存量巨大。而MLA则另辟蹊径，它通过一个投影（projection）操作，将所有头的键和值信息压缩成一个极小的、高信息密度的潜在向量（Compressed Latent KV）。这一设计，极大地缩减了KV Cache的体积，使其能够在有限的内存中支持更长的上下文窗口。同时，DeepSeek V2摒弃了传统的负载均衡损失函数，转而采用一种无损的均衡策略，避免了辅助任务对主学习目标的干扰，实现了更优的综合性能。在其强大的“基座”之上，团队还利用强化学习进一步激发模型的复杂推理能力，展现了架构创新带来的巨大潜力。

然而，当模型规模迈向万亿级别时，一个更严峻的挑战浮出水面——训练稳定性。仅仅是数值上的“蝴蝶效应”，就足以让耗资巨大的训练任务瞬间崩溃。月之暗面的Kimi（K2）针对其进行的探索。其模型拥有惊人的1万亿总参数，每个MoE层配置了多达384个专家。为了驯服这头“性能巨兽”，团队研发了名为MuonClip的优化器，其稳定训练的秘诀在于一个精巧的两阶段过程，如算法1所示。


算法1：MuonClip优化器

在每个训练步骤中，首先执行标准的优化器权重更新（行 3-7）。而真正的核心技术——qk-clip，则在更新之后立刻生效（行 9-17）。该技术会检查在前向传播中计算出的最大注意力得分（\(S_{max}\)）是否超过了预设的稳定阈值\(T\)。一旦超过，它便会计算出一个缩放因子\(\tau\)，并用它来同时重新缩放注意力机制中的查询（\(W_Q\)）和键（\(W_K\)）权重矩阵。这种在每次迭代后主动进行“削峰”干预的做法，如同为注意力分数设置了一个“安全上限”，从源头上钳制了其尺度，有效防止了数值爆炸。这正是MuonClip能够成为稳定训练万亿参数模型“定海神针”的关键所在。

从Mixtral的经典范式，到DeepSeek的极致效率，再到Kimi的万亿稳定性，MoE架构的演进日新月异。这些愈发精妙和复杂的稀疏模型，也对底层的分布式系统提出了前所未有的要求。现在，让我们将目光从模型架构转向幕后英雄——专家并行，看看它是如何支撑起这场波澜壮阔的技术革新的。

EP并行介绍#

本节内容承接前文稀疏模型的讨论，聚焦于把 “专家团队” 真正跑在多机多卡上的方法论——EP（Expert Parallelism，专家并行）。

EP并行提出了一种与传统并行策略截然不同的解决方案。面对大型MoE模型，要求每个GPU都拥有完整模型副本的数据并行策略会立刻失效。因此，EP并行采用了“专家分家，各管一摊”的核心思想（如图3所示）。随着模型规模的增长，所有的专家已无法被容纳在单一设备上进行训练。EP并行将专家们分散到不同的工作节点（workers），每个节点只负责自己“领地”内的一个或几个专家。而模型的非专家部分则在每个GPU上依然保留副本。


图3：专家并行（EP）核心流程图

EP并行虽然设计巧妙，但并非没有代价。All-to-All通信极易造成网络拥堵，而门控网络的任何“偏心”都会导致负载失衡。因此，在这一基础范式之上，一场围绕通信、负载与效率的优化就此展开。研究者们首先从路由策略与专家布局本身入手。例如，随后的Switch Transformer[3]为了最大化计算效率，将GShard的Top-2路由简化为更激进的Top-1路由。而DeepSpeed-MoE[5]则探索了更灵活的专家布局，例如在每个节点上设置“共享专家”，并在模型的更深层放置更多专家，以在通信成本与模型精度之间取得平衡。

EP并行并非孤立存在，它可以与数据、张量、流水线并行有效结合，被视为一种正交的并行维度。DeepSpeed-TED[6]便是一个典型案例，它提出了一套混合并行算法，并设计了如重复Token丢弃（DTD）等通信优化技术，来消除All-to-All过程中的冗余数据。更进一步地，由于MoE的路由具有动态性，为其选择一个静态的最优并行方案极具挑战性。为此，像Tutel[7]这样的研究设计了自适应的并行切换算法，能够在训练的每一次迭代中动态切换并行策略，且不带来额外开销，极大地提升了灵活性和效率。

一个更深入的挑战在于底层计算效率。由于负载不均，分配到每个专家的token数量往往不同。但标准的矩阵乘法（GeMMs）要求输入尺寸一致，这迫使框架进行不必要的token丢弃和填充，造成计算浪费。为了从根本上解决这个问题，MegaBlocks[8]和ScatterMoE[9]等框架通过实现定制化的计算核心，支持为不同专家处理不同批次的输入，从而根除了硬件层面的计算浪费。

从GShard确立范式，到各类研究在路由、布局、混合并行、动态调度乃至底层计算上不断深耕，EP并行技术已愈加成熟化、精密化。正是有了这个强大的“幕后英雄”，我们才能真正驾驭那些由成百上千位专家组成的稀疏模型，引领大语言模型迈向一个更加高效、更加可扩展的未来。

参考与引用#

  1. Jacobs R A, Jordan M I, Nowlan S J, et al. Adaptive mixtures of local experts[J]. Neural computation, 1991, 3(1): 79-87.

  2. Lepikhin D, Lee H J, Xu Y, et al. Gshard: Scaling giant models with conditional computation and automatic sharding[J]. arXiv preprint arXiv:2006.16668, 2020.

  3. Fedus W, Zoph B, Shazeer N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity[J]. Journal of Machine Learning Research, 2022, 23(120): 1-39.

  4. Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., ... & Team, G. E. (2024). Mixtral of experts. arXiv preprint arXiv:2401.04088.

  5. Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., ... & He, Y. (2022, June). Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale. In International Conference on Machine Learning (pp. 18332-18346). PMLR.

  6. Singh S, Ruwase O, Awan A A, et al. A hybrid tensor-expert-data parallelism approach to optimize mixture-of-experts training[C]//Proceedings of the 37th International Conference on Supercomputing. 2023: 203-214.

  7. Hwang C, Cui W, Xiong Y, et al. Tutel: Adaptive mixture-of-experts at scale[J]. Proceedings of Machine Learning and Systems, 2023, 5: 269-287.

  8. Gale T, Narayanan D, Young C, et al. Megablocks: Efficient sparse training with mixture-of-experts, 2022[J]. URL https://arxiv. org/abs/2211.15841.

  9. Tan S, Shen Y, Panda R, et al. Scattered mixture-of-experts implementation[J]. arXiv preprint arXiv:2403.08245, 2024.

本节目录 Contents