|
| 1 | +<!--Copyright © ZOMI 适用于[License](https://github.com/Infrasys-AI/AIInfra)版权许可--> |
| 2 | + |
| 3 | +# EP并行与稀疏模型 |
| 4 | + |
| 5 | +> Author by: 于嘉阔 |
| 6 | +
|
| 7 | +近年来,大语言模型的发展似乎遵循着一种“暴力美学”——模型参数量越大,能力越强。从GPT-3的1750亿参数到动辄数千亿乃至万亿参数的模型,参数规模的竞赛将人工智能推向了新的高度。然而,这种扩张也带来了巨大的挑战:训练一个高质量稠密模型(Dense Model)不仅需要数万块GPU日夜不停地运行数月,其高昂的计算和能源成本也让许多研究机构和企业望而却步。 |
| 8 | + |
| 9 | +这种“越大越好”的路径是否可持续?我们是否必须让模型中的每一个参数都参与每一次计算? |
| 10 | + |
| 11 | +答案或许是否定的。正如人类专家在解决问题时只会调动相关领域的知识一样,模型或许也可以变得更加“专注”和“高效”。这催生了一种全新的模型设计范式——稀疏模型(Sparse Model),其中最具代表性的便是专家混合模型(Mixture-of-Experts,MoE)。MoE架构通过“按需激活”一小部分专家参数来处理输入,巧妙地实现了模型容量与计算成本的解耦。 |
| 12 | + |
| 13 | +然而,拥有了设计精巧的稀疏模型,我们又该如何指挥这个由成百上千位“专家”组成的庞大团队在分布式集群上协同工作呢?这便是专家并行(Expert Parallelism,EP)技术大显身手的舞台。EP是一种专为稀疏模型量身定制的并行计算策略,它解决了如何高效训练和运行这些巨型模型的关键工程难题。 |
| 14 | + |
| 15 | +本文旨在深入探讨稀疏模型与专家并行这两项相辅相成的技术。我们将从MoE的“专家团队”理念讲起,详细剖析EP并行背后的“数据大交换”机制,并最终阐明它们是如何携手,引领大语言模型从“暴力美学”迈向“精妙分工”的新时代。 |
| 16 | + |
| 17 | +## 稀疏模型介绍 |
| 18 | + |
| 19 | +在引言中,我们将传统的大语言模型比作一位“全才专家”,其知识渊博但计算开销巨大。这种模型,我们称之为稠密模型。如果说稠密模型是一位单打独斗的专家,那么稀疏模型就是一支组织精良的“专家团队”。这支团队由多位各领域的专家组成,并由一位极其聪明的“接待员”负责调度。当问题出现时,接待员会迅速判断问题领域,并精准地引荐给最相关的一两位专家。这种“专家按需工作”的模式,正是当前最主流的稀疏模型架构——专家混合模型的核心思想。 |
| 20 | + |
| 21 | +<div align="center"> |
| 22 | + <img src="images/04.EP01.png" > |
| 23 | + <br> |
| 24 | + <em>图1:专家混合模型(MoE)层工作原理示意图</em> |
| 25 | +</div> |
| 26 | + |
| 27 | +MoE模型的精妙之处在于其内部协同工作的两大关键组件:专家网络 (Experts) 与门控网络 (Gating Network)(如图 1 所示)。专家是真正执行计算的主体,在Transformer架构中,它们通常是一个个独立的前馈神经网络 (FFN)。一个MoE层可以包含从8个到超过64个不等的专家,每个专家通过独立学习形成各自的“专业领域”。 |
| 28 | + |
| 29 | +而门控网络,或称为“路由器”,则扮演着“接待员”的角色,是整个架构的智能调度核心。对于输入的每一个词元(token),路由器会通过一个可学习的权重矩阵 $W_g$ 计算出每个专家的得分,并通过Softmax函数将其转换为一组概率权重 $G(x)$。路由器并不会激活所有专家,而是选择得分最高的Top-K个($K$ 通常为 1 或 2)。最终的输出 $y$ 是这 $K$ 个专家输出 $E_i(x)$ 根据其权重 $G(x)$ 的加权求和,其过程可表示为: |
| 30 | + |
| 31 | +$$ |
| 32 | +y = \sum_{i \in \text{Top}-K} G(x)_i \cdot E_i(x) |
| 33 | +$$ |
| 34 | + |
| 35 | +这种基于输入动态决定计算路径的机制,被称为条件计算(Conditional Computation),是MoE实现高效率的关键。 |
| 36 | + |
| 37 | +当然,一个设计良好的系统必须考虑潜在的失效模式。如果路由器倾向于将大部分任务分配给少数几个“明星专家”,就会导致严重的负载不均衡,降低整个系统的效率。为了解决这一问题,研究者们在训练MoE模型时引入了负载均衡损失函数 (Load Balancing Loss)。 |
| 38 | + |
| 39 | +在主流MoE模型中,这个额外的辅助损失($L_{aux}$) 通常被定义为: |
| 40 | + |
| 41 | +$$ |
| 42 | +L_{aux} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i |
| 43 | +$$ |
| 44 | + |
| 45 | +此公式中,$f_i$代表该批次中被分配给第$i$个专家的token实际占比,而$P_i$则是门控网络赋予该专家的平均路由权重。通过最小化这两项的乘积,该损失函数会惩罚那些分配不均的路由行为,激励路由器将任务尽可能均匀地分配给所有专家。通过这种机制,可以确保整个“专家团队”协同工作,避免因个别专家过载而产生的瓶颈,从而最大化模型的整体性能和吞吐量。 |
| 46 | + |
| 47 | +MoE的思想并非一蹴而就,它的雏形最早可追溯至1991年,由Jacobs等研究者提出,但受限于当时的计算能力和训练技术,这一超越时代的设计在近三十年里未能成为主流<sup>[1]</sup>。直到2020年,谷歌发布的里程碑式工作GShard<sup>[2]</sup>,成功地将MoE架构与大规模Transformer模型结合,并解决了其在分布式环境下的训练难题,MoE自此“王者归来”。此后,谷歌通过Switch Transformer<sup>[3]</sup>将模型参数推向了1.6万亿的惊人规模,充分展示了MoE在模型扩展上的巨大潜力。而真正引爆开源社区的,是Mistral AI于2023年发布的Mixtral 8x7B模型<sup>[4]</sup>,它催生了一大批高质量的开源MoE模型,标志着稀疏模型进入了新的时代。 |
| 48 | + |
| 49 | +近年来,各大研究团队不再满足于简单的专家堆叠,而是围绕架构效率、训练稳定性展开了一系列探索,将MoE模型推向了新的高度。 |
| 50 | + |
| 51 | +<div align="center"> |
| 52 | + <img src="images/04.EP02.png" > |
| 53 | + <br> |
| 54 | + <em>图2:从MHA到MLA的注意力机制演进示意图</em> |
| 55 | +</div> |
| 56 | + |
| 57 | +DeepSeek V2系列便是架构效率探索的代表。其拥有高达6710亿的总参数,但激活参数却仅有370亿,稀疏程度令人惊叹。为了在有限的内存中支持更长的上下文,它采用了多头潜在注意力(MLA)技术。从图2中可以清晰地看到,传统的MHA为每个查询头(Query)都保留了独立的键(Key)和值(Value),导致推理时缓存量巨大。而MLA则另辟蹊径,它通过一个投影(projection)操作,将所有头的键和值信息压缩成一个极小的、高信息密度的潜在向量(Compressed Latent KV)。这一设计,极大地缩减了KV Cache的体积,使其能够在有限的内存中支持更长的上下文窗口。同时,DeepSeek V2摒弃了传统的负载均衡损失函数,转而采用一种无损的均衡策略,避免了辅助任务对主学习目标的干扰,实现了更优的综合性能。在其强大的“基座”之上,团队还利用强化学习进一步激发模型的复杂推理能力,展现了架构创新带来的巨大潜力。 |
| 58 | + |
| 59 | +然而,当模型规模迈向万亿级别时,一个更严峻的挑战浮出水面——训练稳定性。仅仅是数值上的“蝴蝶效应”,就足以让耗资巨大的训练任务瞬间崩溃。月之暗面的Kimi(K2)针对其进行的探索。其模型拥有惊人的1万亿总参数,每个MoE层配置了多达384个专家。为了驯服这头“性能巨兽”,团队研发了名为MuonClip的优化器,其稳定训练的秘诀在于一个精巧的两阶段过程,如算法1所示。 |
| 60 | + |
| 61 | +<div align="center"> |
| 62 | + <img src="images/04.EP03.png" > |
| 63 | + <br> |
| 64 | + <em>算法1:MuonClip优化器</em> |
| 65 | +</div> |
| 66 | + |
| 67 | +在每个训练步骤中,首先执行标准的优化器权重更新(行 3-7)。而真正的核心技术——qk-clip,则在更新之后立刻生效(行 9-17)。该技术会检查在前向传播中计算出的最大注意力得分($S_{max}$)是否超过了预设的稳定阈值$T$。一旦超过,它便会计算出一个缩放因子$\tau$,并用它来同时重新缩放注意力机制中的查询($W_Q$)和键($W_K$)权重矩阵。这种在每次迭代后主动进行“削峰”干预的做法,如同为注意力分数设置了一个“安全上限”,从源头上钳制了其尺度,有效防止了数值爆炸。这正是MuonClip能够成为稳定训练万亿参数模型“定海神针”的关键所在。 |
| 68 | + |
| 69 | +从Mixtral的经典范式,到DeepSeek的极致效率,再到Kimi的万亿稳定性,MoE架构的演进日新月异。这些愈发精妙和复杂的稀疏模型,也对底层的分布式系统提出了前所未有的要求。现在,让我们将目光从模型架构转向幕后英雄——专家并行,看看它是如何支撑起这场波澜壮阔的技术革新的。 |
| 70 | + |
| 71 | +## EP并行介绍 |
| 72 | + |
| 73 | +本节内容承接前文稀疏模型的讨论,聚焦于把 “专家团队” 真正跑在多机多卡上的方法论——EP(Expert Parallelism,专家并行)。 |
| 74 | + |
| 75 | +EP并行提出了一种与传统并行策略截然不同的解决方案。面对大型MoE模型,要求每个GPU都拥有完整模型副本的数据并行策略会立刻失效。因此,EP并行采用了“专家分家,各管一摊”的核心思想(如图3所示)。随着模型规模的增长,所有的专家已无法被容纳在单一设备上进行训练。EP并行将专家们分散到不同的工作节点(workers),每个节点只负责自己“领地”内的一个或几个专家。而模型的非专家部分则在每个GPU上依然保留副本。 |
| 76 | + |
| 77 | +<div align="center"> |
| 78 | + <img src="images/04.EP04.png" > |
| 79 | + <br> |
| 80 | + <em>图3:专家并行(EP)核心流程图</em> |
| 81 | +</div> |
| 82 | + |
| 83 | +EP并行虽然设计巧妙,但并非没有代价。All-to-All通信极易造成网络拥堵,而门控网络的任何“偏心”都会导致负载失衡。因此,在这一基础范式之上,一场围绕通信、负载与效率的优化就此展开。研究者们首先从路由策略与专家布局本身入手。例如,随后的Switch Transformer<sup>[3]</sup>为了最大化计算效率,将GShard的Top-2路由简化为更激进的Top-1路由。而DeepSpeed-MoE<sup>[5]</sup>则探索了更灵活的专家布局,例如在每个节点上设置“共享专家”,并在模型的更深层放置更多专家,以在通信成本与模型精度之间取得平衡。 |
| 84 | + |
| 85 | +EP并行并非孤立存在,它可以与数据、张量、流水线并行有效结合,被视为一种正交的并行维度。DeepSpeed-TED<sup>[6]</sup>便是一个典型案例,它提出了一套混合并行算法,并设计了如重复Token丢弃(DTD)等通信优化技术,来消除All-to-All过程中的冗余数据。更进一步地,由于MoE的路由具有动态性,为其选择一个静态的最优并行方案极具挑战性。为此,像Tutel<sup>[7]</sup>这样的研究设计了自适应的并行切换算法,能够在训练的每一次迭代中动态切换并行策略,且不带来额外开销,极大地提升了灵活性和效率。 |
| 86 | + |
| 87 | +一个更深入的挑战在于底层计算效率。由于负载不均,分配到每个专家的token数量往往不同。但标准的矩阵乘法(GeMMs)要求输入尺寸一致,这迫使框架进行不必要的token丢弃和填充,造成计算浪费。为了从根本上解决这个问题,MegaBlocks<sup>[8]</sup>和ScatterMoE<sup>[9]</sup>等框架通过实现定制化的计算核心,支持为不同专家处理不同批次的输入,从而根除了硬件层面的计算浪费。 |
| 88 | + |
| 89 | +从GShard确立范式,到各类研究在路由、布局、混合并行、动态调度乃至底层计算上不断深耕,EP并行技术已愈加成熟化、精密化。正是有了这个强大的“幕后英雄”,我们才能真正驾驭那些由成百上千位专家组成的稀疏模型,引领大语言模型迈向一个更加高效、更加可扩展的未来。 |
| 90 | + |
| 91 | +## 参考与引用 |
| 92 | + |
| 93 | +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. |
| 94 | +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. |
| 95 | +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. |
| 96 | +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*. |
| 97 | +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. |
| 98 | +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. |
| 99 | +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. |
| 100 | +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. |
| 101 | +9. Tan S, Shen Y, Panda R, et al. Scattered mixture-of-experts implementation[J]. *arXiv preprint arXiv:2403.08245*, 2024. |
0 commit comments