Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 101 additions & 0 deletions 04Train/01ParallelBegin/04EP.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
<!--Copyright © ZOMI 适用于[License](https://github.com/Infrasys-AI/AIInfra)版权许可-->

# EP并行与稀疏模型

> Author by: 于嘉阔

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

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

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

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

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

## 稀疏模型介绍

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

<div align="center">
<img src="images/04.EP01.png" >
<br>
<em>图1:专家混合模型(MoE)层工作原理示意图</em>
</div>

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等研究者提出,但受限于当时的计算能力和训练技术,这一超越时代的设计在近三十年里未能成为主流<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模型,标志着稀疏模型进入了新的时代。

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

<div align="center">
<img src="images/04.EP02.png" >
<br>
<em>图2:从MHA到MLA的注意力机制演进示意图</em>
</div>

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所示。

<div align="center">
<img src="images/04.EP03.png" >
<br>
<em>算法1:MuonClip优化器</em>
</div>

在每个训练步骤中,首先执行标准的优化器权重更新(行 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上依然保留副本。

<div align="center">
<img src="images/04.EP04.png" >
<br>
<em>图3:专家并行(EP)核心流程图</em>
</div>

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

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

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

从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.
Binary file added 04Train/01ParallelBegin/images/04.EP01.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added 04Train/01ParallelBegin/images/04.EP02.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added 04Train/01ParallelBegin/images/04.EP03.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added 04Train/01ParallelBegin/images/04.EP04.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.