HybridEP: Scaling Expert Parallelism to Cross-Datacenter Scenario via Hybrid Expert/Data Transmission
HybridEP: Scaling Expert Parallelism to Cross-Datacenter Scenario via Hybrid Expert/Data Transmission
作者/机构: Weihao Yang†, Hao Huang†, Donglei Wu‡, Ningke Li†, Yanqi Pan†, Qiyang Zheng†, Wen Xia†, Shiyi Li† and Qiang Wang† (Harbin institute of Technology, Shenzhen†, Guangzhou University‡)
A1 主要贡献
本文探讨了在跨数据中心(cross-DC)训练场景下,扩展专家并行(Expert Parallelism, EP)所面临的挑战。随着大型语言模型(LLM)的规模持续增长,单一数据中心的训练能力已达到瓶颈,跨数据中心训练成为必然趋势。然而,EP作为混合专家(MoE)模型的核心组件,在跨数据中心有限的带宽条件下,其通信开销成为训练效率的主要瓶颈。现有EP优化方法通常侧重于将计算与通信进行重叠,但在低带宽场景下,通信时间远超计算时间,导致这些优化方法收效甚微。
为了解决这一紧迫挑战,本文提出了一个名为HybridEP的建模指导框架,旨在优化带宽受限下的EP。
- 核心问题:在跨数据中心训练MoE模型时,由于数据中心间带宽有限,EP中频繁的All-to-All通信成为性能瓶颈,严重影响训练效率。
- 研究目标:提出一种新方法,从结构上消除EP的通信开销,而不仅仅是重叠或隐藏它,从而在带宽受限的环境下高效扩展EP。
- 核心创新点:本文的核心思想是通过动态改变专家的空间位置来重塑通信拓扑,从而减少数据通信的流量和频率。这种方法引入了数据和专家混合传输的复杂模式。为了找到最优解,本文构建了一个基于流的性能模型(Stream-Based Modeling),用于确定专家和数据之间的最佳传输比例。
- 关键技术实现:
- 基于域的划分(Domain-Based Partition):该技术用于在GPU级别构建混合通信模式与特定通信拓扑之间的映射,以确保与现有的分层硬件架构兼容。
- 参数高效的迁移(Parameter-Efficient Migration):该技术通过减少专家的传输开销和扩大域(domain)的规模来进一步优化通信拓扑。
实验结果表明,在带宽受限的情况下,HybridEP的性能比现有的最先进MoE训练系统(如Tutel、FasterMoE和SmartMoE)高出5.6倍。在大规模仿真中,当连接1000个数据中心时,HybridEP相较于传统EP实现了高达1.45倍的加速。
本文的主要贡献如下:
- 揭示了跨数据中心MoE训练是扩展LLM容量的必然趋势,并指出EP是阻碍训练效率的主要瓶颈。
- 提出通过空间重塑专家布局来有效扩展EP,这引入了数据和专家的混合通信模式。
-
实现了HybridEP框架,包含三大核心技术:
- 基于流的建模:确定混合模式的最佳比例。
- 基于域的划分:构建具体的通信拓扑。
- 参数高效的迁移:优化拓扑以提高效率。
-
实验证明,在真实场景和仿真中,HybridEP的性能均优于当前最先进的工作。
图 1. 真实训练配置下(2个节点 × 2个GPU)的Transformer和MoE块结构。每个MoE块将FFN复制为多个专家,并仅激活其中一部分(通过门控网络)来处理数据,从而实现恒定的计算成本。然而,All-to-All的频繁通信使其成为训练过程中的主要瓶颈。
A3 背景知识与动机
A. 并行策略、MoE及其通信瓶颈
常见的并行策略。分布式训练中存在五种常见的并行策略,各自具有不同的特性:
图 2. 常见设置下并行策略的分析及EP的可扩展性。(a)展示了更多数据中心和更大规模EP的趋势,这需要更多的数据中心间互连和跨数据中心扩展EP。(b)展示了在不同带宽下EP的开销比例。这一趋势揭示了在低带宽下扩展EP所带来的沉重开销。
- 张量并行(Tensor Parallelism, TP):将每层的权重分割到多个GPU上,在正向和反向传播过程中都需要大量的通信和高流量。
- 序列并行(Sequence Parallelism, SP):将序列在多个GPU上进行分割【索引23,Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models,2023,CoRR】【索引34,Ring attention with blockwise transformers for near-infinite context,2023,CoRR】,通常在预训练的后期用于支持更大的上下文窗口【索引19,The llama 3 herd of models,2024,CoRR】。
- 流水线并行(Pipeline Parallelism, PP):将模型层划分为多个阶段,每个阶段分配给不同的GPU。它在阶段之间使用点对点通信来发送/接收激活值和梯度。
- 数据并行(Data Parallelism, DP):在每个GPU上复制完整的模型,不同批次的数据被独立处理,梯度在反向传播期间跨副本进行同步。
- 专家并行(Expert Parallelism, EP):将MoE模型中的专家分布在多个GPU上。具体来说,专家是从Transformer结构【索引51,Attention is all you need,2017,NIPS】的基本模块FFN扩展而来的,如图1所示。在此基础上,MoE模型使用一个门控网络作为路由器,然后使用两次All-to-All(A2A)通信将数据交换到被激活的专家部分进行计算。然而,EP的大流量和高频率使其目前几乎只能部署在高性能计算(HPC)环境中。
跨数据中心扩展EP的必要性。首先,从基础设施角度看,LLM正在将现有数据中心推向极限。然而,扩展单个数据中心面临着电力限制和易受断电影响等挑战【索引3】【索引4】【索引5】。因此,近期的报告【索引8,Crosspipe: Towards optimal pipeline schedules for cross-datacenter training,2025,USENIX ATC】建议一个更实用和有弹性的解决方案是部署多个规模较小的数据中心。其次,从模型角度看,随着MoE的广泛应用,一些近期的代表性MoE模型【索引12,DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models,2024,CoRR】【索引13,Deepseek-v3 technical report,2024,CoRR】【索引17,Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,2022,JMLR】【索引26,Mixtral of experts,2024,CoRR】【索引39,Minimax-01: Scaling foundation models with lightning attention,2025,CoRR】【索引55,Qwen3 technical report,2025,CoRR】已显示出将专家数量扩展至数百甚至数千的迅速发展趋势,这超出了单个数据中心的容量。因此,结合容量扩展的挑战和单个数据中心的容量限制,EP将不可避免地扩展到多个数据中心,如图2(a)所示。
EP的通信瓶颈。如图2(b)所示,在相对较低的带宽下,EP会导致极其严重的瓶颈。现有研究也表明,在带宽受限的情况下,EP几乎占据了整个迭代时间的50%-90%【索引20,FasterMoE: modeling and optimizing training of large-scale dynamic pretrained models,2022,PPoPP】【索引38,Hetu: a highly efficient automatic parallel distributed deep learning system,2023,Sci. China Inf. Sci.】【索引47,Schemoe: An extensible mixture-of-experts distributed training system with tasks scheduling,2024,EuroSys】,成为提高训练效率的主要瓶颈。与DP和PP等对带宽友好的并行策略不同,EP的通信密集型设计使其对网络条件特别敏感。它依赖于数据到专家的路由模式(即通过A2A将每个数据发送到对应的专家),需要在设备间频繁通信以进行同步,并且通信流量与激活专家的数量成正比。因此,EP并非为跨数据中心扩展而设计,迫切需要高效的EP扩展方案。
B. 动机与挑战
在图3中,我们用一个简单的例子展示了数据和专家的布局如何影响EP的通信开销,并以此引出HybridEP的动机。
图 3. 数据和专家的空间布局如何影响EP效率的示例。(a)显示了原始EP的布局,存在7次低带宽通信,训练速度最慢。(b)显示了通过重塑数据布局而没有任何低带宽通信的理想情况,实现了最佳效率。(c)显示了一种实践中可能的优化,仅通过重塑专家布局,只有4次低带宽通信。这在理想情况和原始EP之间取得了良好的折衷。
原始EP的情况。在图3(a)中,原始EP的基本逻辑是将数据发送到由路由结果决定的相应专家。当专家被放置在由低带宽连接的不同节点上时,会引入多次慢速通信(例如,示例中的7次),这成为主要瓶颈并降低了EP的效率。
理想的完全消除EP的情况。理想情况下,如果我们预先知道数据将被路由到哪个专家,我们可以在训练前重新排列数据,将发送到同一专家的数据放在具有高带宽互连的位置,从而避免巨大的低带宽通信延迟,如图3(b)所示。这实质上是将EP的通信移至训练开始之前,减少了其频率和流量,从而实现了最快的速度。然而,问题在于数据路由结果是动态并实时决定的,因此几乎不可能预先知道路由结果并在训练前重新排列数据。同时,由于数据到专家的分配在不同层之间是变化的,假设在训练前存在完美的布局是不现实的。因此,这几乎是不可能实现的。
可能的部份消除EP的情况。考虑到EP的基本逻辑是将数据与相应的专家进行匹配,我们可以通过改变专家的空间布局来优化EP效率,如图3(c)所示。在专家处理完数据后,我们重新排列专家的布局,然后继续计算。虽然这无法达到理想情况,但它仍然减少了低带宽通信的次数(例如,示例中减少到5次),从而提高了效率。值得注意的是,我们区分了数据和专家的传输,因为我们发现,与数据相比,专家在轻量级传输方面具有两个相当好的优势,而这些优势尚未被先前的工作充分利用来优化EP的效率。具体来说:
1. 更好的可压缩性。专家权重通常表现出紧凑和集中的分布,与数据相比离群值更少,如图4所示,这一点也得到了先前工作【索引21,Kvquant: Towards 10 million context length LLM inference with KV cache quantization,2024,NeurIPS】【索引33,AWQ: activation-aware weight quantization for on-device LLM compression and acceleration,2024,GetMobile】【索引36,KIVI: A tuning-free asymmetric 2bit quantization for KV cache,2024,ICML】【索引54,SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models,2023,ICML】的证实。这使得专家可以被更积极地压缩以减少流量,而不会降低准确性。
2. 异步通信潜力。与每个计算操作符都紧密耦合的数据不同,专家只是间歇性地参与计算。这允许在EP之前独立地预传输专家,从而减少同步开销。
总而言之,这些特点表明专家天生适合进行轻量级、早期和稳定的通信,这构成了HybridEP的基础。
图 4. 数据和专家的可压缩性分析。数据分布幅度大且有许多离群值(红色部分),而专家权重的分布相对较小且平坦,导致更高的可压缩性。
为什么这种方法有效? 关键原因是专家是稀疏激活的,这也是EP与其他并行策略最大的区别。只有在EP中,数据由少数专家(部分模型参数)处理,而在其他并行模式中,所有数据都需要由所有模型参数处理。类比于图3(c),如果EP不具备稀疏激活的特性,即假设每个数据都需要与所有专家进行计算,那么无论如何改变专家的布局,低带宽EP通信的数量都不会改变。因此,只有像EP这样具有稀疏激活特性的并行策略才能通过改变专家布局来提高效率。
挑战。上述分析为高效的EP扩展提供了一个愿景,引入了专家和数据的混合传输。然而,在真实的训练场景中,要最大化其潜力仍然面临三个主要挑战:
- ① 如何确定数据-专家混合通信模式的最佳比例? 本质上,专家以All-Gather(AG)模式进行通信,因为每个设备都从其他设备收集专家(详见§III)。我们需要在AG模式与数据的A2A模式之间取得平衡。
- ② 如何构建具体的通信拓扑并确保与现有架构的兼容性? 即使确定了两种通信的比例,最基本的通信粒度是GPU,而不是DC。因此,需要构建一个特定的拓扑来对齐通信粒度。
- ③ 如何有效地利用高可压缩性和异步性来实现轻量级专家迁移? 尽管专家传输可以做到轻量级,我们仍然需要设计相应的压缩和通信机制来实现理论上的效率。
为此,我们提出了HybridEP来应对上述挑战。首先,我们设计了一个基于流的建模来解决第一个挑战,并为整体设计提供高层指导。接着,我们引入了两项技术来释放其潜力。我们通过划分专家域来确定混合通信拓扑并与现有分层架构兼容,从而解决第二个挑战。然后,我们设计了一个专家压缩算法和一个异步通信器来实现轻量级专家迁移,从而解决第三个挑战。
A2 方法细节
III. 基础:基于流的建模
建模过程如图5(a)所示。我们首先将MoE训练分解为计算流和通信流(§III-A, §III-B)。然后,我们对两个流之间的重叠进行建模(§III-C),并综合考虑计算、通信和重叠,以最小化训练延迟(§III-D, §III-E)。我们总结了建模中常用的符号于表I。我们的模型假设每个GPU拥有相同数量的专家,每个专家大小相同,并且门控网络均匀地激活专家。为简化起见,我们首先假设每个DC中只有一个GPU(用GPU代表DC)来统一通信粒度,这不影响模型的准确性。
TABLE I 建模中频繁使用的符号。
A. 计算建模
GeMM建模。计算延迟主要来源于通用矩阵乘法(GeMM)操作。我们遵循先前的工作【索引20,FasterMoE: modeling and optimizing training of large-scale dynamic pretrained models,2022,PPoPP】【索引32,Path forward beyond simulators: Fast and accurate gpu execution time prediction for dnn workloads,2023,MICRO】,使用一个线性模型来估计延迟。给定两个要相乘的矩阵,大小分别为(L, H)和(H, M),单个GeMM操作的延迟可以表示为:
$$Lat_{comp}^{GeMM} = \frac{LMH}{C},$$其中$C$代表GPU的平均计算吞吐量。注意,如果GeMM操作太小以至于无法充分利用GPU的计算能力,$C$会减小。然而,由于其开销很小,这不会影响整体建模的有效性,这一点已由【索引20,FasterMoE: modeling and optimizing training of large-scale dynamic pretrained models,2022,PPoPP】证实。
图 5. 训练过程的建模以及A2A和AG的通信分解。(a) 展示了使用分治法进行的建模过程。训练过程首先被分解为对计算流和通信流的独立建模,然后考虑它们的重叠关系进行合并。(b) 表明A2A的流量保持不变(即O(1)),而AG的流量则乘以GPU的数量(即O(n))。
计算流建模。假设一个MoE块之前有$m$个Transformer块,计算延迟可以表示为:
$$\begin{aligned} \begin{aligned} Lat_{comp} &= mLat_{comp}^{TF} + Lat_{comp}^{MoE} \\ &= (m+1)Lat_{comp}^{Att} + mLat_{comp}^{FFN} + nLat_{comp}^{Ep} \end{aligned} \end{aligned}$$其中,$T F, MoE$分别代表Transformer和MoE块。$Att, F F N, Ep$分别代表注意力、FFN和专家的计算过程,这些过程由多个GeMM操作组成。这里我们可以将它们的延迟视为常数。$n$是单个GPU上的专家数量,因此专家计算延迟需要重复$n$次。为简洁起见,我们将$(m+1)Lat_{Att_{comp}}+mLat_{F F N_{comp}}$视为Pre-Expert,记为$Lat_{P E_{comp}}$。
B. 通信建模
All-to-All通信建模。图5(a)展示了A2A通信的细节【索引50,Automatically tuned collective communications,2000,ICS】。具体来说,给定$G$个GPU,每个GPU上的数据$D$将被分割成$G$个大小为$D_G$的数据块。然后,$G-1$个数据块将通过A2A发送给其他GPU,而1个数据块保留在本地GPU。因此,对于参与A2A通信的GPU集合$G_{A2A}$,总通信量表示为:
$$V^{A2A} = \frac{D}{|G^{A2A}|} * (|G^{A2A}| - 1), Lat_{comm}^{A2A} = \frac{V^{A2A}}{B}$$其中$B$是带宽。由于$D$和$B$是常数,随着$|G_{A2A}|$的增加,$Lat_{A2A_{comm}}$几乎保持不变。
All-Gather通信建模。图5(b)展示了AG通信的细节。具体来说,每个GPU上的专家参数$P_E$将被发送给其他$G-1$个GPU。因此,对于参与AG的GPU集合$G_{AG}$,通信流量表示为:
因此,$Lat_{AG_{comm}}$随$|G_{AG}|$线性增加。
两种通信之间的关系。A2A可以无缝地转换为AG。如果一个专家已经通过AG获取,那么相应的数据块就不需要通过A2A传输。当第$i$个GPU $G_i$使用AG从$G_j$收集专家$P_E$时,A2A的流量从$D * \frac{G-1}{G}$变为$D * \frac{G-1}{G} - D_G$,而AG的流量从0变为$P_E$。因此,当A2A的流量减少$D_G$时,AG的流量增加$P_E$。
通信流建模。其延迟来自于A2A和AG两部分,可以表示为:
$$Lat_{comm} = Lat_{comm}^{AG} + 2 Lat_{comm}^{A2A}.$$其中A2A在专家计算前后各执行一次,而AG只执行一次,因为专家不需要被送回其原始GPU。
C. 重叠建模
两流的重叠建模。计算和通信之间的重叠时间来自三个部分,如图5(a)中红色虚线所分割。具体来说,① pre-expert计算(即$Lat_{P E_{comp}}$)与AG;② 专家计算($Lat_{Ep_{comp}}$)与AG;③ 专家计算($Lat_{Ep_{comp}}$)与A2A。需要注意的是,pre-expert计算不能与A2A重叠,因为数据依赖于pre-expert的计算结果。因此,重叠时间可以写为:
$$Lat_{ovlp} = Lat_{ovlp}^{PE,AG} + Lat_{ovlp}^{Ep,AG} + Lat_{ovlp}^{Ep,A2A} ,$$注意,情况②和③已被先前的工作【索引35,Janus: A unified distributed training framework for sparse mixture-of-experts models,2023,SIGCOMM】【索引46,Pipemoe: Accelerating mixture-ofexperts through adaptive pipelining,2023,INFOCOM】优化,因此专家计算与AG和A2A完全重叠($Lat_{ovlp}^{Ep, AG} + Lat_{ovlp}^{Ep, A2A} = nLat_{Ep_{comp}}$)。对于情况①,重叠时间为$min(Lat_{P E_{comp}}, Lat_{AG_{comm}})$。
$$Lat_{ovlp} = min(Lat_{comp}^{PE}, Lat_{comm}^{AG}) + nLat_{comp}^{EP}$$D. 问题形式化
为了最小化延迟,我们有以下定义:
定义1。给定一个有$G$个GPU($G > 1$)的集群,$p_i$是通过A2A传输的数据块(离开$G_i$)的比例,而$1-p_i$是转化为专家并通过AG传输的数据块的比例,其中$p_i \in \{ \frac{0}{G-1}, \dots, \frac{G-1}{G-1} \}$。
当$p_i = \frac{G-1}{G-1}$时,只有A2A通信;当$p_i = \frac{0}{G-1}$时,只有AG通信。训练延迟可以表示为:
图 6. 方程式10解的可视化。两个红点表示在两种情况下延迟最小的最优p。
需要注意的是,每个GPU都有其自己的方程式8,并且它们应该同步。因此,系统延迟是所有GPU中最大的延迟,可以表示为:
$$Lat_{all} = \max_{0 \le i < G} \{ \min Lat_{final}(p_i) \}$$最终,方程式9仅依赖于参数$p_i$,我们的目标是通过选择最优的$p_i$来最小化$Lat_{all}$。
E. 问题求解
为简单起见,我们假设所有的$p_i$都是相同的。因此,方程式9可以简化为一个易于求解的形式:
$$\begin{aligned} \begin{aligned} & Lat_{all} = \min Lat_{final}(p) \\ & = \begin{cases} \min(Lat_{comp}^{PE} + 2Lat_{comm}^{A2A}), \text{if } Lat_{comp}^{PE} \ge Lat_{comm}^{AG} \\ \min(Lat_{comm}^{AG} + 2Lat_{comm}^{A2A}), \text{if } Lat_{comp}^{PE} < Lat_{comm}^{AG}. \end{cases} \end{aligned} \end{aligned}$$最终解可以分为两种情况。
情况1: 当$Lat_{P E_{comp}} \geq Lat_{AG_{comm}}$时,方程式10简化为:
注意$Lat_{P E_{comp}}$, D, B都是正常数。因此,为了最小化$Lat_{final}$,我们需要配置最小的$p$。
情况2: 当$Lat_{P E_{comp}} < Lat_{AG_{comm}}$时,方程式10简化为:
注意,$(G-1)(2D-G P_E) / BG$的符号会影响$Lat_{final}$的最小值。
① 当$2D - G P_E < 0$时,我们配置最大的$p$,记为情况2.1。
② 当$2D - G P_E \geq 0$时,我们配置最小的$p$,记为情况2.2。
总结。我们的模型找到了最小延迟的最佳比例(即$p$),可以总结为两种情况,如图6所示。具体来说,当$2D - G P_E < 0$时,总延迟的变化由情况1和情况2.1组成。因此,最优的$p$被配置为$1 - \frac{B Lat_{P E_{comp}}}{P_E(G-1)}$,此时我们同时使用AG和A2A。当$2D - G P_E \geq 0$时,总延迟的变化由情况1和情况2.2组成。因此,最优的$p$被配置为0,此时我们只使用AG。注意,当$p=1$时,HybridEP退化为标准的EP,这表明EP是我们的框架的一个特例。
IV. 设计与实现
HybridEP的概览如图7所示。在训练开始前,HybridEP首先将环境配置作为输入,利用模型找到传输数据和专家的最佳比例。以此为导向,HybridEP接着引入基于域的划分,为A2A和AG通信划分GPU(§IV-A),从而在GPU层面构建通信拓扑。此外,HybridEP设计了参数高效迁移,通过更好的划分来优化已确定的通信拓拓扑(§IV-B)。
图 7. HybridEP概览。在模型确定了数据和专家的传输比例后,HybridEP使用基于域的划分来构建特定的GPU通信拓扑。此外,参数高效迁移通过减少开销来实现更好的划分。
A. 基于域的划分
专家域与基于域的通信。专家域(Expert Domain)是一组仅在内部使用AG通信的数据中心(DC)集合。专家域的大小定义为其包含的DC数量,记为$S_{ED}$。HybridEP假设每个域的大小相同。图8(a)右侧展示了一个例子,我们设置$S_{ED} = 2$,并依次将每2个DC分组到同一个域中。借助专家域,我们制定了以下基于域的通信规则:AG仅发生在域内通信中,而A2A仅发生在域间通信中。这样简单的规则可以有效地将两种通信模式分开,以便更好地管理。
图 8. 基于域的通信和多级拓扑构建。(a)展示了专家域如何影响通信,将通信分为域内AG和跨域A2A。(b)展示了通过三个关键步骤(多级描述、位置重编号和拓扑构建)将拓扑映射到多级划分。
扩展到多级的必要性。在实际场景中,训练环境通常由分层架构组成,基本的通信粒度是GPU。因此,尽管DC之间的通信方式已经确定,但每个GPU的具体行为仍不明确。与GPU粒度对齐是真实训练场景中的关键一步。为了弥合这一差距,HybridEP首先将分层结构抽象为多级描述(Multilevel Description),以处理现实中复杂多变的环境。然后,它通过位置重编号(Location Renumbering)来重新编号全局GPU,以适应多级结构。最后,它执行拓扑构建(Topology Construction)算法来确定GPU级别的具体拓扑。其工作流程如图8(b)所示。
多级描述。我们首先定义工作者(Worker)为一个物理实体(例如,DC、节点或GPU)。通常,我们将GPU视为工作者的最小粒度。级别(Level)是一组通过同构带宽连接的工作者。因此,我们将级别$l$的专家域大小扩展为该域中的工作者数量,记为$S_{ED}^l$。为了描述不同级别之间的关系,我们使用扩展因子(scaling factor)$SF_i$来表示级别$i-1$的一个工作者可以扩展为级别$i$的$SF_i$个子工作者。注意,我们将$SF_0$设为级别0的总工作者数。以图8(b)为例,给定一个包含4个DC,每个DC有4个GPU的环境,它被分为两个级别,其中$SF_0 = 4, SF_1 = 4$,每个级别的域大小分别为$S_{ED}^0 = 2, S_{ED}^1 = 4$。
位置重编号。为了明确详细的通信规则,我们首先为多级架构中的每个GPU重新编号其位置。具体来说,我们遵循Pytorch【索引41,PyTorch: An Imperative Style, HighPerformance Deep Learning Library,2019,NeurIPS】为每个GPU分配一个全局索引$m$。然后,给定一个$L-1$级的划分,我们将全局索引$m$重新编号为多级位置$(x_0, x_1, \dots, x_{L-1})$。利用扩展因子列表$[SF_0, \dots, SF_{L-1}]$,重编号函数$f : m \to (x_0, x_1, \dots, x_{L-1})$可以表示为:
因此,GPU $m$的级别$i$工作者编号可以通过$f(m)[i]$获得。此外,利用专家域大小$S_{ED}^i$,GPU $m$的级别$i$域可以通过$\lfloor \frac{f(m)[i]}{S_{ED}^i} \rfloor$获得。
拓扑构建。相关伪代码如算法1所示。具体来说,给定两个全局索引为$m$和$n$的GPU,我们决定在不同级别需要哪种类型的通信。我们首先通过$f(m)$和$f(n)$获得它们的多级位置。为了限制多次通信导致的效率低下,我们限制了可以相互通信的GPU范围。具体来说,只有当$f(m)[l] \neq f(n)[l]$且后续层级的索引相同时,两个GPU才能相互通信。在每个级别上,GPU之间的通信遵循基于域的通信规则。
算法1:拓扑构建
1: 输入: GPU m, GPU n, 当前级别 l, 扩展因子 SFl, 专家域大小 SlED
2: 输出: 通信类型 (None 或 AG 或 A2A)
3: Loc_m ← f(m)
4: Loc_n ← f(n)
5: W_m, W_n ← Loc_m[l], Loc_n[l]
6: ED_m, off_m ← ⌊W_m / SlED⌋, W_m mod SlED
7: ED_n, off_n ← ⌊W_n / SlED⌋, W_n mod SlED
8: if Loc_m[l + 1:] == Loc_n[l + 1:] then
9: if ED_n == ED_m and off_n ≠ off_m then
10: return AG
11: end if
12: if ED_m ≠ ED_n and off_m == off_n then
13: return A2A
14: end if
15: end if
16: return None
B. 参数高效的迁移
轻量级迁移如何优化通信拓扑。本质上,轻量级迁移减小了$P_E$的大小,从而导致可以形成更大的专家域,进而实现更高的效率。具体来说,如图9所示,一个更小的$P_E$可以导致一个更小的$p$,这主要有两个方面的原因:1. 当$2D < G P_E$时,最优点(即红点)对应的$p$会减小。2. 它使得更多的训练配置可以从$2D < G P_E$转换为$2D \ge G P_E$。一个更小的$p$意味着一个更大的专家域,这改变了所构建的通信拓扑。此外,域越大,能达到的效率就越高。这是因为方程式11和方程式12从理论上表明,在$p$减小后,整体延迟会降低。因此,我们将参数高效迁移视为一个通过扩展专家域来优化通信拓扑的过程,旨在进一步提高效率。
图 9. 专家间的冗余性及基于SR的专家压缩工作流程。在SRDecode中,我们实际上融合了恢复和加法操作以提高效率。
专家间的冗余性。除了图4中展示的专家更好的可压缩性外,我们进一步探索了专家间的冗余性以提高压缩比。我们发现,不同专家之间的主要差异集中在少数参数上。这表明不同的专家可能从数据中学到了相似的知识,这一点在其他相关工作中也有报道【索引12,DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models,2024,CoRR】。如图9(a)所示,在对专家权重求平均并从原始权重中减去后,结果(带有“res”后缀)的分布比原始分布更集中。这表明残差是稀疏的,专家之间的关键差异集中在少数参数上。
基于SR的专家压缩。我们受此启发,将专家分为共享部分和残差部分,分别学习冗余知识和特定知识。具体来说,共享专家由所有GPU共享,并通过对所有专家求平均来初始化。在每个训练迭代中,它将在反向传播阶段通过异步All-Reduce进行同步。我们的专家压缩有两个阶段,如图9(b)所示。在编码阶段,压缩器首先通过目标专家减去共享专家来获得专家残差。然后,它通过Top-k压缩专家残差。压缩后的专家残差以值-索引格式保存,以传输给其他GPU。在解码阶段,压缩器首先恢复压缩后的专家残差。然后,它通过将共享专家和残差专家相加来恢复目标专家。注意,在实践中,我们融合了解码阶段的上述两个步骤以减少开销。
图 10. 异步通信器的两个阶段。(a)显示了初始化阶段,该阶段与上一个优化器步骤融合。每个MoE层依次将其经过SREncode处理的专家发送到发送队列。(b)显示了异步通信阶段,该阶段与pre-expert计算重叠。每个MoE层的通信结果存储在接收队列中,并由SRDecode处理以供后续计算使用。
异步通信器的机制。我们使用一个异步通信器来尽可能地实现我们建模的理论效果。为了在不过多增加额外开销的情况下,将异步特性与SR压缩充分结合,我们将异步通信器的行为分为两个阶段,类似于SR压缩。如图10所示,通信器将模型视为一个由(pre-expert, expert)对组成的堆栈,并设有一个发送队列和一个接收队列。在初始化阶段,模型中的所有专家依次由SREncode处理,压缩结果将被送入发送队列。接收队列被设置为空。注意,这个过程在每次迭代开始前发生,我们将其与上一次迭代的更新过程(优化器步骤)融合以减少开销。在异步通信阶段,发送队列依次弹出专家残差进行AG通信。接收队列接收相应的结果并发送给SRDecode,用于后续的专家计算。注意,这个通信过程与模型的pre-expert计算过程并行,因此我们可以将它们重叠。此外,我们还将SRDecode与专家计算融合以提高效率。
A4 实验环境
-
硬件配置:
- Cluster-S: 单个数据中心(DC)内的8块 NVIDIA A800 GPU。
- Cluster-M: 2个DC上的16块 NVIDIA A800 GPU。
- Cluster-L: 4个DC上的32块 NVIDIA A800 GPU。
- DC(节点)内部通过PCIe3.0 x16(128 Gbps)连接,DC之间通过以太网(10 Gbps)连接。
-
软件配置:
- 实现基础: Tutel【索引22,Tutel: Adaptive Mixture-of-Experts at Scale,MLSys】 和 Pytorch v.1.12.1。
- 操作系统: Ubuntu-18.04。
- 依赖库: CUDA-11.3, cuDNN-7.6, NCCL-2.10。
- 优化器: Adam,学习率为1e-4。
- 梯度同步: Pytorch DDP (All-Reduce)。
-
模型与数据集:
- Llama-Tiny: 在PennTreebank【索引2,Penn Treebank,2020,https://paperswithcode.com/dataset/penn-treebank】数据集上进行测试 。
- Mistral-Small: 在wikitext2【索引37,Pointer sentinel mixture models,2017,ICLR】数据集上进行测试。
- GPT-Medium: 在OpenWebText-10k【索引1,openwebtext-10k,2019,https://huggingface.co/datasets/stas/openwebtext10k】数据集上进行测试 。
- GPT-Large: 在WikiText103【索引37,Pointer sentinel mixture models,2017,ICLR】数据集上进行测试。
-
对比方法:
- Tutel【索引22,Tutel: Adaptive Mixture-of-Experts at Scale,MLSys】
- FasterMoE【索引20,FasterMoE: modeling and optimizing training of large-scale dynamic pretrained models,2022,PPoPP】
- SmartMoE【索引58,SmartMoE: Efficiently training Sparsely-Activated models through combining offline and online parallelization,2023,USENIX ATC】
A4 实验结果
建模验证
- 计算与通信延迟估计:实验通过调整数据流量和专家大小来验证模型的准确性。图11显示,模型估计的计算、A2A和AG延迟(红色标记)与真实延迟(黑线)非常接近,表明基于流的建模能够有效地模拟系统延迟。
- 最优比例p的验证:实验通过四种不同的配置(Mix-1, Mix-2, AG-only-1, AG-only-2)验证模型能否找到最优的A2A与AG通信比例($p$)。如图12所示,模型预测的最优$p$值(红色方块)在所有候选$p$值(黑色三角)中实现了最低的平均迭代延迟,证实了模型的有效性。例如,对于Mix-1和Mix-2,模型找到了混合A2A和AG通信的最优比例;对于AG-only-1和AG-only-2,模型正确地选择了纯AG通信($p=0$)。
图 11. 计算和通信延迟验证。由于估计的计算、A2A、AG延迟(红色标记)接近真实延迟(黑线),我们的基于流的建模可以有效地模拟系统延迟。
图 12. 建模验证。结果表明,我们的模型可以找到最优的p(红色标记),在候选配置(黑色标记)中具有最短的迭代时间。
端到端加速效果
- 不同数据流量下的性能:在Cluster-M和Cluster-L上,通过改变数据流量(从6MB到192MB),测试了HybridEP的性能。表V结果显示,随着数据流量增大、带宽降低和DC数量增多,EP的瓶颈愈发明显,而HybridEP通过优化A2A和AG的比例,实现了高达5.60倍的平均加速。
- 不同专家大小下的性能:固定数据流量为16MB,改变专家大小(从32MB到2MB)。图13显示,HybridEP实现了1.18倍到2.57倍的加速。随着专家大小减小,HybridEP能够以更小的流量传输更多专家,从而扩大专家域,减少EP开销。
图 13. 不同专家大小下的平均迭代时间。结果表明,随着专家大小减小,计算成本降低,HybridEP的迭代延迟也随之降低。然而,尽管计算开销减少,对比方法的迭代延迟几乎没有变化。
TABLE V 不同数据流量大小下的平均迭代时间(秒)和平均加速比(×)。
消融研究
- 组件贡献分析:通过对比仅使用“基于域的划分”(基线)和同时使用“参数高效迁移”(即完整的HybridEP),评估了各组件的贡献。表VI显示,在不同配置下,完整的HybridEP相比基线实现了1.25倍至2.82倍的加速。这表明参数高效迁移对性能提升至关重要,特别是在数据流量大、专家尺寸小的场景下。
TABLE VI 消融研究。
迁移分析
- 对准确性的影响:实验评估了SR专家压缩算法对模型准确性的影响。图14显示,使用共享专家的HybridEP(HybridEP w/ S)在高达50倍的压缩率下,其损失曲线与Tutel、FasterMoE等基线方法基本重合,表明该压缩算法在保持高压缩比的同时,几乎不影响模型精度。相比之下,不使用共享专家的直接压缩方法(HybridEP w/o S)导致了明显的精度下降。
- 时间开销分析:图15展示了SR压缩的编码(SREncode)和解码(SRDecode)阶段的时间开销。通过将这两个阶段与其他计算操作(如参数更新、专家计算)融合,其开销可以分别降低高达30%和45%,使得这部分额外计算开销在可接受范围内。
图 14. 损失分析。HybridEP的专家压缩比为50倍。此外,w/o S表示HybridEP直接压缩专家,而w/ S表示HybridEP使用共享专家压缩专家(§IV-B)。
图 15. 参数高效迁移两个阶段的时间分解。在不同专家大小下,(a)显示了SREncode与上一迭代的参数更新融合的效果,可减少30%的开销。(b)显示了SRDecode与多个专家计算融合的效果,可减少45%的开销。
特性比较:EP vs. HybridEP
- 通信流量分析:图16显示,原始EP的通信流量随token数量线性增长。相比之下,HybridEP的通信流量是固定的,与输入无关,并且有明确的上限。这使得HybridEP在低带宽或对突发流量敏感的环境中更具可预测性和稳定性。
- 通信频率分析:表VII展示了不同专家域大小下的通信频率。随着专家域的扩大,A2A的通信频率二次方下降,而AG的频率相应增加。尽管通信模式发生了转变,但由于AG更适合异步执行且可通过压缩显著减少流量,HybridEP实现了更高的效率。
图 16. 流量可扩展性分析。在带宽受限的情况下,HybridEP的通信流量更少,从而具有更好的可扩展性。配置是一个三元组,表示EP的大小以及专家权重的两个维度(H和M)。
TABLE VII 不同EP大小下的通信频率。
大规模仿真
-
可扩展性验证:由于实验环境限制,使用SimAI仿真器在更大规模上验证了HybridEP的有效性。图17显示了两种情况下的仿真结果。
- 固定专家域大小,扩展DC数量:在这种更贴近实际的场景中,随着DC数量增加,加速效果有所减弱。在1000个DC的设置下,HybridEP仍能实现1.05倍至1.45倍的加速。
- 固定通信比例,扩展DC数量:在这种理想情况下,随着DC数量和专家域规模的扩大,加速效果增强。在1000个DC时,HybridEP可实现1.31倍至3.76倍的加速。
-
结论:仿真结果表明,在各种带宽条件下,HybridEP都比原始EP具有更好的可扩展性,尤其是在带宽较低时,加速效果更显著。
图 17. HybridEP在大规模仿真中的加速效果。
A7 补充细节
基于SR的专家压缩的存储开销。我们提出的专家压缩算法所引入的额外存储开销是可控的。具体来说,它包含专家残差和共享专家两部分。① 专家残差由于其高可压缩性,仅消耗少量GPU内存。② 共享专家会与本地专家竞争GPU内存,这个问题可以通过将本地专家卸载到CPU内存,同时将共享专家保留在GPU内存中来解决。将本地专家卸载到CPU内存是一种有效的策略,已有深入研究(例如,Zero-Offload【索引44,ZeRO-Offload: Democratizing Billion-Scale model training,2021,USENIX ATC】)并且可以直接集成到HybridEP中。
MoE训练中的反向传播。反向传播阶段有独特的All-Reduce通信来同步模型参数,这会与其他类型的通信竞争,从而影响我们的建模。尽管如此,我们的建模对于反向传播仍然有效,因为All-Reduce通信的流量与模型大小相关,当模型配置确定时,其延迟可以视为一个常数。因此,我们的建模可以通过简单地增加一个常数来处理反向传播。
A5 结论
本文提出了HybridEP,一个由建模指导的框架,旨在有效扩展带宽受限环境下的专家并行(EP)。通过建模分析,我们发现最优的通信策略完全取决于数据和专家传输流量的比例。基于该模型,我们引入了两项关键技术:基于域的划分和参数高效的迁移,使我们的建模方法变得实用、高效且可扩展。实验结果表明,HybridEP的性能比当前最先进的MoE系统高出最多5.6倍。在超过1000个数据中心和不同带宽下的大规模仿真中,与原始EP相比,HybridEP实现了高达1.45倍的加速。
💬 评论讨论
欢迎在这里分享您的想法和见解!