Skip to content

04EP.md

Latest commit

 

History

History
101 lines (65 loc) · 13.3 KB

04EP.md

File metadata and controls

101 lines (65 loc) · 13.3 KB

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.