FarSkip-Collective: Unhobbling Blocking Communication in Mixture of Experts Models

文章标题:FarSkip-Collective: Unhobbling Blocking Communication in Mixture of Experts Models
作者/机构:Yonatan Dukler, Guihong Li, Deval Shah, Vikram Appia, Emad Barsoum Advanced Micro Devices Inc. (AMD)

A1 主要贡献

混合专家(MoE)模型已成为近年来主流大语言模型(LLM)的事实架构。由于其稀疏激活和计算成本降低的特性,MoE模型得以扩展到更大的参数规模,但这也要求在训练和推理时采用更广泛的分布式设置。然而,分布式环境中的阻塞式通信是一个主要障碍,它会导致计算加速器在等待通信完成时处于空闲状态,尤其是在专家、序列和张量并行等技术中。随着MoE架构规模的扩大和硬件计算速度的提升,通信瓶颈问题愈发突出。

本文提出了一种名为 FarSkip-Collective 的方法,通过修改模型架构来解决这一问题。其核心思想是在通信调用开始时,利用当前可用的(可能是过时或部分物化的)激活值来启动下一步的计算,从而将通信操作与计算过程重叠。具体而言,方法将本应等待通信结果的计算提前进行,并将通信结果“远跳”(far-skip)到下一层的残差连接上,供更未来的层使用。只要后续计算的持续时间长于通信时间,就可以避免计算资源的空闲。

这种方法在数学上相当于“丢弃”了网络中的一些连接,这可能会损害模型的能力。因此,本文重点评估了这种修改后的架构是否能保持与常规MoE模型相当的性能,特别是在大规模、最先进的模型上。通过开发一种名为 FarSkip-Collective Self-Distill (FCSD) 的自蒸馏方法,研究证明了这一点是可行的。FCSD是一种简单而有效的知识蒸馏方案,适用于任何缺乏强大教师模型的场景。

主要贡献如下:
* 提出FarSkip-Collectives方法:通过转换模型层的执行依赖关系,消除了MoE模型中的阻塞式通信模式,从而提升了推理和训练速度。
* 验证大规模模型能力:在100B+参数规模上证明了使用FarSkip-Collective架构的MoE模型能够保留现代MoE模型的能力,同时加速其在分布式环境中的执行。特别是,成功转换了Llama 4 Scout MoE (109B)模型,在保持模型平均准确率与指令调优的开源版本相比差距在1.0%以内。
* 开发FCSD蒸馏方案:开发了一种高效、通用的知识自蒸馏方案(FCSD),使用少于10B的训练token即可将现有LLM转换为FarSkip模型。已成功转换DeepSeek-V2 Lite (16B)、Qwen-3-30B MoE (30B)和Llama 4 Scout (109B),性能均在原始模型的2.5%范围内。
* 大规模训练实现:将该方法集成到Megatron-LM中,在MoE专家并行的前向和后向传播过程中,实现了先前阻塞的all-to-all通信集体操作88.4%的计算-通信重叠。
* 模型推理优化实现:在vLLM & SGLang中提供了FarSkip的优化实现,重叠了分布式推理中的通信。例如,对于修改后的Llama-4 Scout模型,在首次令牌生成时间(Time To First Token)上实现了18.5%的加速。


图 1: FarSkip-Collective修改了子模块之间的连接,以避免集合通信中的阻塞。计算继续使用可用的激活(部分,例如块N的输出;或过时,例如激活)。

A3 背景知识

2.1 MoE在训练和推理中的并行化

张量并行(TP)和专家并行(EP)是MoE训练和推理中的两种关键并行技术。在张量并行中,一个MLP或自注意力子块通过沿列或行均匀切分其密集权重矩阵来进行分割。对于一个现代MLP层,其形式为:

$$\text{MLP}(A) = \sigma(AW_1^\top \cdot g(AW_2^\top))W_3^\top,$$

其中$g$和$\sigma$是逐元素的非线性函数,$W_1, W_2 \in R^{c \times d}, W_3 \in R^{d \times c}$。大小为k的TP会将这些矩阵切分为:

$$\begin{aligned} \begin{aligned} W_{\{1,2\}}^{i} & =\left(W_{\{1,2\}}\right)_{[i * c / k:(i+1) * c / k,:]} \in \mathbb{R}^{(c / k) \times d}, \\ W_{3}^{i} & =\left(W_{3}\right)_{[:, i * c / k:(i+1) * c / k]} \in \mathbb{R}^{d \times(c / k)}, \end{aligned} \end{aligned}$$

对于$0 \le i \le k-1$。每个TP rank的计算可以独立进行,直到子块末尾,此时应用一个all-reduce通信集合来构建最终输出。

$$\text{MLP}(A)_i = \sigma(AW_1^{i\top} \cdot g(AW_2^{i\top}))W_3^{i\top},$$ $$\operatorname{MLP}(A) = \text{all-reduce}(\operatorname{MLP}(A)_i).$$

通过先对$W_{\{1,2\}}^i$进行列并行切分,再对$W_3^i$进行行并行切分,使得通信只需在MLP子块的末尾进行一次。对于多头自注意力,TP将计算分解为跨不同注意力头的独立计算分区,采用了类似的列并行(Q, K, V)后接行并行(O)和单个all-reduce的实现。

专家并行(EP)是MoE层的关键并行组件。一个拥有E个专家的MoE层将MLP泛化为:

$$\text{MoE}(A) = \sum_{j=1}^{E} G(A)_{j} \cdot \text{MLP}^{j}(A),$$

其中$G(A) = s(AW_R^\top)$是一个线性分类层后接一个稀疏路由器选择函数$s()$,它只激活MLP$_j$的一个子集。这里的MLP$_j(A)$指的是第$j$个独立的“专家”。在rank为$k$的EP中,大约$E/k$个专家的子集会分布在$k$个并行rank上。与TP不同,训练期间不同的输入token激活会根据路由器选择$G(A)$被映射到不同的专家。具体来说,特定的tokens $A[l,:] \in R^d$会被分组并映射到专家子集$P_l \subset \{1, \dots, E\}$,这需要对$A$进行置换,然后通过一个all-to-all集合操作,根据路由器定义的数据分区图在ranks之间发送和接收数据。

其中$R_i \in \{0, 1\}^{B_i \times d}$是$G(A)$的指示符;这被称为“Dispatch”,其带宽需求取决于专家并行大小$k$和$s$的稀疏度等因素。在每个专家接收到其专用token并计算完其rank上相关专家的MLP$_j(A)$后,会应用一个对应的all-to-all集合操作,称为“Combine”,来聚合和求和被路由的专家激活,从而得到MoE层的输出激活。此外,现代MoE设计通常还包括一个“共享专家”MLP层,它会处理所有token。

综合来看,一个典型的MoE Transformer层的训练执行流程如下:1) 注意力子块计算(包括层归一化);2) 如果启用了TP或CP,则可能进行注意力后的通信集合操作;3) 层归一化、门控和路由器分数计算(在每个rank上重复);4) Dispatch通信;5) 被路由的专家计算以及可能的共享专家计算;6) Combine集合通信,以聚合来自不同rank的被路由专家。这个流程会导致三个潜在的阻塞通信气泡:(a) 使用TP/CP时的注意力后阻塞通信,(b) Dispatch期间,以及(c) Combine期间,尽管如果存在共享专家,(c)可以被部分重叠。

对于推理,MoE的执行有不同的方法,其中只有一些涉及all-to-all的“Dispatch + Combine”【索引8,Deepseek-v3 technical report,2024,arXiv】。本文关注vLLM和SGLang【索引45,SGLang: Efficient Execution of Structured Language Model Programs,2024,NeurIPS;索引19,Efficient Memory Management for Large Language Model Serving with PagedAttention,2023,SOSP】实现的优化all-reduce方法,它取代了all-to-all集合,并且被路由专家的输入激活在专家并行rank之间被复制和索引。

2.2 模型蒸馏

FarSkip-Collective通过修改模型架构并进行自蒸馏来恢复原始模型的能力。一种基础方法是直接使用高质量数据通过监督微调(SFT)来微调模型,其损失函数如下:

其中$\theta$表示模型参数,$(x, y) \sim D$是输入-输出监督对。当转换模型时有一个目标模型(例如,本文旨在恢复原始模型),可以使用知识蒸馏【索引14,Distilling the Knowledge in a Neural Network,2015,arXiv;索引17,Sequence-Level Knowledge Distillation,2016,EMNLP】,通过Kullback-Leibler (KL)散度来对齐一个固定的教师模型$q$。基于logit的知识蒸馏通过优化$p_\theta$来匹配教师模型的预测分布,损失函数如下:

$$ \mathcal{L}_{\mathrm{KD}}(\theta)=\mathbb{E}_{x \sim \mathcal{D}}\left[\sum_{t=1}^{|y|} \mathrm{KL}\left(q\left(\cdot \mid x, y_{<t}\right) \| p_{\theta}\left(\cdot \mid x, y_{<t}\right)\right)\right] . $$ <p>除了对模型输出进行约束,还可以通过对齐模型的中间表示来增强其与教师模型q的一致性【索引39,Zebrallama: Towards extremely efficient hybrid models,2025,NeurIPS】,损失函数如下:

$$\mathcal{L}_{L2}(\theta) = \sum_{i=1}^{L} \| o_i(\theta) - t_i \|_2^2,$$

其中$o_i$和$t_i$分别表示学生模型和教师模型在$L$个层中匹配的隐藏激活。


图 2: FarSkip-Collective MoE层主要操作符执行流程。g、d和c分别指门控+路由、Dispatch开始和Combine开始。对于通信操作,我们仅标示操作的起点。FarSkip-Collective启用的重叠窗口由操作符上方的阴影区域表示。

A2 方法细节

3 FarSkip-Collective 框架

现代网络利用残差连接,将自注意力、MLP或MoE子块的输出增量地加到残差激活上。若网络经过$k$层后的输出激活为$o_k$,网络的第$i$个子块(层)为$f_i$,则输出$o_k$的计算方式为:

$$o_k = f_1(o_0) + f_2(o_1) + \cdots + f_k(o_{k-1}).$$

由于$f_k$的计算可能涉及阻塞式通信,导致$o_k$在准备好之前无法作为$f_{k+1}$和$o_{k+1}$的输入,这在子块通信期间会造成计算资源的闲置。

我们提出使用一个可用的激活,记为$o_k^*$,作为$f_{k+1}$的输入,并在生成$o_k$的通信集体操作运行时计算下一层。$o_{k+1}$将在通信完成的$o_k$准备好后用其进行更新;但现在,$o_k$的通信可以被“远跳”(far-skipped)并与以$o_k^*$为输入的$f_{k+1}(o_k^*)$的计算过程重叠。

$$o_k=o_0+f_1(o_0)+f_2(o_1^*)+\cdots+f_k(o_{k-1}^*)$$

我们考虑了两种修改后的$o_k^*$输入选项

$$\begin{aligned} \begin{cases} o_{k}^{*}=o_{k-1} \quad \text { ("outdated") } \\ \text { or } \\ o_{k}^{*}=o_{k-1}+f_{k}^{*}\left(o_{k-1}^{*}\right) \quad \text { ("partial") }\end{cases} \end{aligned}$$

这里,$f^*$表示块$f_k$中在集合通信前已准备好的独立计算部分,例如公式2中的$MLP_i(A)$。在这两种情况下,输出激活$o_k$将包含与之前相同数量的块,但区别在于每个块的输入激活。$f_{k+1}$仍然可以访问之前所有的块表示,除了块$f_k$的完整表示;然而,所有未来的层$f_j$($j \ge k+2$)都将能够访问完整的$f_k$。

跳过更多计算可以实现更大的重叠,但会省略更多输入,可能降低模型能力。在用FarSkip-Collective转换预训练的MoE模型时,我们结合使用“部分”(partial)和“过时”(outdated)方法,以最大化重叠机会和准确性。具体来说,对于第$k$层的注意力子块,输入$o_k^*$将对应于(8b)中的部分激活,即MLP计算的共享专家部分,但不包括需要通过Combine集合通信的路由专家。对于MoE子块的输入,我们遵循(8a),传递上一层的最终输出,该输出包含除最近的第$k$个注意力子块输出之外的所有先前输出。令$o_k^*(attn)$为第$k$个注意力子块的输入,我们建议使用部分激活:

$$\text{attn-in}_k := o_k^*(\text{attn}) = o_{k-2} + \text{attn-out}_{k-1} + \text{shared-exp-out}_{k-1}$$

对于MLP块的输入,$o_k^*(mlp)$是过时激活:

attn-in_k相比,mlp-in_k多一个输入,可以分解为mlp-in_k = attn-in_k + routed-exp-out_{k-1}。通过这种连接,我们确保每个输入最多只相差一个子块,以获得更好的准确性。同时,共享专家和注意力的计算独立于前面的路由专家,这意味着Combine操作可以与子块的操作重叠;MLP的输入现在也独立于注意力的计算,这意味着Dispatch也可以与注意力子块重叠。我们将在下一节提供FarSkip-Collective执行和重叠的详细信息。

我们将更激进的多块“远跳”变体作为未来工作,这在通信时长超过整个子块计算时间的情况下可能有用,例如在极度稀疏和大规模的MoE模型中,或在考虑网络拓扑的MoE模型中。

3.1 使用FarSkip-Collective框架蒸馏现有模型

FarSkip-Collective方法修改了架构连接,但未改变模型的参数布局,这使得可以使用现有的检查点,通过相对较少的模型定义修改来应用FarSkip-Collective连接,并使用相同的主要核函数。如图3所示,我们将原始的Qwen-3-30B MoE模型检查点加载到激活了不同数量FarSkip-Collective层的模型中,并评估其在不同基准上的性能。我们观察到,随着转换层数的增加,模型性能显著下降,当所有层都被转换时,模型在MMLU上达到随机基线准确率,在HumanEval+上为0%。这是意料之中的,因为我们传递了与模型训练时不同的输入激活,导致了分布外的输出。


图 3: 未经训练的Qwen-3-30B MoE模型在修改了N个FarSkip-Collective层后的准确率。修改的层应用于最后N层(蓝色)和前N层(红色)。

然而,我们证明了通过使用典型的指令调优数据,以KL知识蒸馏的方式继续训练原始检查点,能够恢复原始模型的性能,而所需计算量远小于用FarSkip-Collective架构从头重新训练(便宜约100-1000倍)。我们系统地研究了不同的蒸馏训练方法(见表2),发现使用原始模型作为教师的KL知识蒸馏(自蒸馏)表现最好或与其他测试方法相当。我们还研究了批量大小和学习率等因素的影响,发现它们对最终模型性能和训练稳定性起着重要作用。基于我们的实证评估,我们提出了一个简单而稳健的“FarSkip-Collective Self-Distillation” (FCSD)方案,用于将任何MoE模型转换为FarSkip-Collective连接。与直接用指令调优数据进行微调相比,自蒸馏到原始模型的概率分布为恢复模型的现有表示提供了更精细的信号。使用教师模型的概率分布作为训练信号也减少了对精心策划和高质量SFT数据的依赖。由于修改后的架构只改变了原始模型的连接性,KL知识蒸馏能有效快速地与原始模型粗略对齐,并随着持续训练而改善。尽管如此,在训练后期,当学生和教师模型已经大致对齐时,KL知识蒸馏可能导致训练不稳定,因为教师和学生模型之间的微小差异偶尔会产生大的梯度。我们测试了不同的方法来克服这个问题,但发现使用早停法能避免该问题。对于早停验证,我们使用MBPP+【索引23,Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation,2023,NeurIPS】数据集作为快速代理来检测不稳定性,每1000个训练步骤评估一次,耐心值为20次评估,性能下降阈值为2%。MBPP+评估速度快,且作为一个代码生成数据集,对训练不稳定性引起的破坏性分布偏移很敏感。在评估其他方法时,我们也应用了相同的早停程序以进行公平比较。

4 FarSkip模型的显式重叠

现代GPU配备了数百个独立的计算单元(流式多处理器),能够通过在不同的处理单元集上同时调度工作来独立处理多个队列(流)的核函数【索引44,PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,2023,VLDB;索引30,Distributed w/ TorchTitan: Introducing Async Tensor Parallelism in PyTorch,2024,PyTorch Discussion Forum;索引8,Deepseek-v3 technical report,2024,arXiv】。计算和通信操作都利用计算单元来运行,通信操作只利用总可用单元的一小部分,从而允许与计算重叠。然而,计算-通信重叠需要专门的实现,除了标准的预定义模式外,像PyTorch和JAX这样的现代框架不会自动实现这一点。因此,即使修改后的FarSkip-Collective模型在逻辑上促进了计算图的并行和非阻塞流,但如果没有精心的显式实现,模型也不会自动重叠通信和计算。下面,我们描述了我们对FarSkip-Collective框架的显式实现,它实现了通信调用与计算的高度重叠。作为设计选择,我们旨在使我们的实现具有通用性并尽可能地与硬件无关,通过坚持在框架层面而不是底层核函数或Triton层面进行实现。为了实现非阻塞通信调度的调度,我们依赖于torch.distasync_op=True参数或使用torch.cuda.Stream()上下文,后者为在非主队列上调度核函数提供了更精细的控制。请注意,通过重叠操作,会占用一些处理单元,这可能导致计算速度相比于它们单独在硬件上执行时出现不可避免的减慢。

4.1 训练

对于训练,我们考虑使用MegatronLM【索引27,Efficient large-scale language model training on GPU clusters using Megatron-LM,2021,SC】进行带有MoE层的GPT训练。我们主要考虑包含共享专家、多头潜在注意力(MLA)并运行EP(专家并行)而不对注意力使用TP(张量并行)的设置,遵循DeepSeek的V3模型训练方案。通常在这种设置下,all-to-all集合通信会在每层的前向和后向传播中,作为Dispatch和Combine的一部分导致两个通信气泡。通过我们修改的架构消除了阻塞依赖后,我们修改了第2节讨论的执行顺序。具体来说,我们将注意力子块的计算分为两部分:a) MLA的(q, k, v)准备,和 b) 核心注意力+输出投影。这使我们能轻松地在这两部分之间异步启动一个通信核函数,然后立即继续进行注意力计算。对于DeepSeek模型设置,我们执行FarSkip-Collective MoE层的前向传播,步骤如下:1) 注意力部分(a)计算;2) 如果上一层是FarSkip-Collective MoE层,则同步Combine通信;3) MoE门控和路由器分数计算;4) 启动Dispatch(async-op模式会立即将操作排队并返回);5) 注意力部分(b)计算;6) 同步Dispatch通信,然后进行路由专家的计算;7) 启动Combine(async-op模式);8) 运行共享专家计算。我们在图2中提供了此执行流程与常规操作流程的视觉对比,这使我们能够最大化前向调用期间的通信重叠窗口,如下所示:

使用FarSkip-Collective,层中唯一不能被重叠的计算是路由专家和门控操作。这是因为在修改后的架构中,路由专家需要前一个Dispatch的输出,并作为下一个Combine调用的输入;而门控操作需要前一个Combine调用的输出,并作为下一个Dispatch调用的输入。

对于反向传播,我们希望重叠Combine和Dispatch的梯度调用,这些调用也会触发阻塞的all-to-all通信集合。如果天真地运行,将需要在非async-op模式下运行反向传播通信,因为其输出需要为计算图中的下一个梯度同步,这会再次使通信变为阻塞。显式控制操作符顺序并确保正确同步的标准方法是为整个MoE Transformer块层的反向传播计算使用自定义的torch.autograd.Function(仅为Dispatch和Combine或其子部分实现自定义反向传播无法定义其外部的同步点,而这正是重叠所需要的)。然而,手动实现如此大一个层的反向传播计算图是繁琐且容易出错的,因为每个操作和模型权重都需要正确地连接到它们的下一个输入。

相反,我们提出了两种创新技术,它们继续依赖自动autograd进行反向传播,同时干净地实现了重叠。首先,我们实现了一个异步all-to-all自定义autograd函数,它带有一个包含前向和后向通信句柄的状态字典。在前向传播期间,前向all-to-all通信句柄由async-op模式下的集合操作生成;然而,后向all-to-all通信句柄此时尚不存在,但在层的状态字典中有一个专用的键。当对该操作符调用backward时,它会在async-op模式下运行反向传播,并用相应的后向通信句柄填充字典。这使我们能够在不直接控制反向调用执行的情况下存储和访问后向句柄。然后,我们实现了一个backward-hook,将其挂钩到all-to-all通信操作符的输入张量节点上,它将在处理输入之前触发通信的同步。这使得在反向传播中异步运行all-to-all通信成为可能,同时确保梯度在被访问时是就绪的。

由于我们在前向传播中试图优化重叠,因此前向传播中的通信实现为:一旦导致通信输入的计算完成,通信就立即启动,这使我们能够最大化前向传播的重叠窗口。然而,在反向传播中,这会导致相反的效果,因为通信调用的输入现在将在反向通信调用后立即启动,句柄在启动后被迫立即同步并等待通信完成。为了解决这个问题,我们通过PyTorch autograd的内部实现“劫持”了torch.autograd的优先级排序。在torch autograd中,计算图将根据节点间依赖关系的拓扑排序算法进行处理;但是,当多个节点同时准备好处理时,autograd使用基于节点在前向传播中创建顺序的序列号来决定哪个节点将首先被处理。利用FarSkip-Collective带来的依赖关系下降,我们可以在到达通信调用输入的依赖障碍之前处理整个子块的反向传播autograd节点。利用这一点,我们通过重新分配自定义序列号,重新确定autograd优先级队列的优先级,以处理子块计算中的节点,并降低导致通信调用输入的计算的处理优先级。这样,这些节点将在子块反向传播计算发生后才启动,从而在无需手写大型反向传播函数的情况下实现了大的重叠机会。使用我们的优化实现,在单节点上训练一个带有EP8的DeepSeek-V2 Lite时,我们实现了88.4%的all-to-all通信时间重叠,如表3所示。请注意,反向传播中的第一个all-to-all和前向传播中的最后一个all-to-all无法重叠,因为没有额外的潜在计算可以与它们重叠。

4.2 推理

对于推理,我们在vLLM中实现了FarSkip,后来将其扩展到SGLang。SGLang和vLLM是现代LLM推理引擎,支持MoE模型(如DeepSeek)的TP、EP和PP。与其他使用一对all-to-all集合进行Dispatch和Combine的MoE EP实现不同,在vLLM和SGLang中,模型激活在各rank间被复制,但模型权重(包括专家权重)仍通过EP分布。这种方法消除了Dispatch和Combine的需求,并用优化后的all-reduce操作实现,该操作在MLP层计算完成后应用于激活。对于注意力子块,vLLM和SGLang采用常规的TP方法,带有一个all-reduce集合。这两个all-reduce调用通常是阻塞的,因为下一层需要这些激活。

为了实现MoE层的FarSkip-Collective,我们在async-op模式下运行all-reduce,并且仅在下一次MoE计算之前同步它,因为这些激活不再需要用于Dispatch。对于注意力层,我们专注于修改输出投影层,该层通常会运行一个包含all-reduce的RowParallelLinear层,我们将其修改为在async-op模式下运行,并且仅在下一个注意力层之前应用同步调用。对于像DeepSeek模型中的MLA这样的特殊注意力,预填充(prefill)和生成(generation)会运行不同的融合核函数,我们分别处理每种情况,但在每种场景中都应用一个async-op模式的all-reduce调用。为了将FarSkip与HIP/CUDA-graphs集成,我们使用了与图兼容的通信API调用,并使用了直接的Python绑定(PyNCCL)。我们使用通过FCSD微调的自蒸馏模型测试了我们的推理流水线,并观察到我们的蒸馏恢复了模型在基于聊天的生成任务中的性能(生成示例见附录)。

A4 实验环境

  • 模型

    • 能力评估:DeepSeek-V2-Lite (16B, 3B激活参数), Qwen3-30B MoE (30B, 3B激活参数), Llama-4 Scout (109B, 17B激活参数)。所有模型均为指令调优/聊天版本。
    • 性能评估:DeepSeek-V2 (235B), Llama 4 Scout (109B), DeepSeek-V3 (671B), 以及一个6层的短版DeepSeek-V3 (71B)。
  • 数据集

    • 蒸馏训练:使用最多10B tokens的SFT数据集组合【索引43,GenQA: Generating millions of instructions from a handful of prompts,2024,arXiv;索引22,Infinity Instruct: Scaling Instruction Selection and Synthesis to Enhance Language Models,2025,arXiv】。
    • 能力评估:在11个下游任务上进行基准测试,包括PIQA, ARC-E, ARC-C, HS, CSQa, WG, HEval+, MMLU, OBook, GSM-8K, MBPP+。
    • 预训练:从一个高质量公开数据混合语料库中随机抽样50B tokens。
  • 硬件配置

    • 单节点训练/推理:1台配备8x MI325X GPU的机器,或1台配备8x MI300X GPU的机器。
    • 多节点训练:4个节点,每个节点配备4x MI325X(总计32个GPU),节点间通信带宽为400Gbps。
    • 多节点推理:2个节点,通过8个400Gb/s的NIC连接进行节点间通信。
  • 软件配置

    • 训练框架:Megatron-LM【索引27,Efficient large-scale language model training on GPU clusters using Megatron-LM,2021,SC】。
    • 推理引擎:vLLM【索引19,Efficient Memory Management for Large Language Model Serving with PagedAttention,2023,SOSP】和SGLang【索引45,SGLang: Efficient Execution of Structured Language Model Programs,2024,NeurIPS】。
    • 核心库:PyTorch, PyNCCL(用于CUDA-graphs集成)。
    • 训练参数:使用AdamW优化器,余弦退火学习率调度器,1000步预热期。对批量大小和学习率进行了扫描。序列长度为4096。
    • 推理参数:采用FP8量化和融合的MoE前向核函数。

A4 实验结果

5.1 模型能力评估

FCSD蒸馏结果
在表1中,我们展示了对三种不同规模的开源MoE模型(DeepSeek-V2-Lite 16B, Qwen3-30B MoE, Llama-4 Scout 109B)进行FCSD蒸馏的主要结果。我们将FarSkip-Collective应用于所有模型层,并使用最多10B tokens的SFT数据进行训练。结果显示,FCSD能够成功转换这些模型,性能损失极小。例如,Llama-4 Scout的平均性能仅下降1.0%,而三个模型的最大平均性能下降为2.5%。相比之下,标准的SFT微调方法导致了灾难性的遗忘,尤其是在生成任务上,表现明显差于FCSD。这证明了即使模型最初不是为此连接方式设计的,修改后的架构仍具有很强的内在能力。

Table 1: 原始模型和蒸馏后的FarSkip-Collective模型在下游评估任务上的性能。

消融研究
在表2中,我们研究了不同蒸馏技术和部分转换对Qwen-3-30B MoE模型的影响。实验表明,仅使用KL损失目标的FCSD方法最为有效。将FarSkip-Collective应用于部分层(例如,模型的最后75%层)比完全转换更容易,尤其是在生成任务(如HEval+)上能更快恢复性能。

Table 2: Qwen-3-30B MoE模型不同FarSkip-Collective蒸馏训练设置的下游性能。我们评估了训练5亿token的不同训练和转换设置。(* 对于0.25×BS和4×BS,我们保持相同的训练步数)

未经训练的转换影响
在图3中,我们评估了未经训练直接修改Qwen-3-30B MoE模型层数的影响。结果显示,随着转换层数的增加,模型性能急剧下降。修改模型的初始层比修改末尾层对性能的损害更大,这可能是因为早期层的错误会逐层传递,并且深层网络对依赖关系的丢失更不敏感。

5.2 显式重叠性能评估

训练性能
在表3中,我们测量了在Megatron-LM中单节点训练的性能。对于DeepSeek-V2 Lite和短版DeepSeek-V3模型,FarSkip-Collective在前向传播(87.6%, 92.9%)和后向传播(89.0%, 84.1%)中均实现了高度的all-to-all通信重叠,分别带来了11%和4%的端到端训练加速。在图5的多节点强扩展性测试中,随着专家并行(EP)规模从8增加到32,FarSkip-Collective带来的端到端训练加速比从1.1倍提升至1.22倍,显示了良好的可扩展性。

Table 3: 使用EP=8的重叠式FarSkip-Collective MegatronLM训练中all-to-all集合的计算-通信重叠。我们评估了DeepSeek-V2 Lite和一个包含6层的缩短版DeepSeek-V3模型的训练。


图 5: DeepSeek-V3 (L=6) FarSkipCollective训练在不同专家并行大小和批量大小配置下的加速效果。

推理性能

  • vLLM Prefill (TTFT): 如图4所示,在vLLM推理引擎中,对于DeepSeek-V2 (235B) 和 Llama 4 Scout (109B) 模型的prefill阶段,FarSkip-Collective分别带来了8.2%-16.8%和12.2%-18.5%的TTFT加速。all-reduce通信的重叠率分别达到了97.6%和95.3%。Llama-4 Scout的专家数量较少,计算速度更快,因此暴露的通信瓶颈更严重,加速效果更明显。


图 4: 使用vLLM推理引擎在不同提示长度下的首次令牌生成时间(prefill阶段)。每个模型使用EP=8进行MLP子块服务,TP=8进行注意力服务,处理16个并发请求。

  • SGLang Prefill & Decode: 如图6所示,在SGLang中对DeepSeek-V3 (671B)进行推理,prefill阶段实现了高达1.34倍的加速。在多节点解码阶段(图7),由于分布式设置(EP=16, TP=16)增加了通信开销并允许更大的批量(BS=1024),FarSkip-Collective带来了显著且一致的解码速度提升(高达1.25倍),解决了大规模MoE解码中计算-通信平衡变化带来的新瓶颈。


图 6: 使用SGLang在不同批量大小和提示长度下的DeepSeek-V3首次令牌生成时间(prefill阶段)。每个模型使用EP=8进行MLP子块服务,TP=8进行注意力服务。


图 7: 使用SGLang对DeepSeek-V3进行2节点推理时,在不同提示长度下的令牌间时间(解码阶段)。每个模型使用EP=16进行MLP子块服务,TP=16进行注意力服务。

A5 结论

本文提出了一种针对MoE模型的修改后连接架构——FarSkip-Collective,它通过移除网络块之间的模型依赖关系,从而有潜力解决训练和执行中的阻塞式通信问题。我们证明了这种修改后的架构能够在大规模模型上表现良好,并通过知识蒸馏技术,我们能够高效地转换一系列参数规模高达109B的最先进MoE模型,同时保持了99%的模型准确率。在证明了修改后架构的可行性后,我们通过开发优化的训练和推理实现,发挥了FarSkip-Collective的优势:在MoE训练中实现了88.9%的all-to-all通信重叠,在推理方面则使首次令牌生成时间最多加速了18.5%。展望未来,通过使模型执行更易于重叠,我们希望可以重新审视现有的MoE架构配置和并行设置,从而为探索更广阔的并行化和架构搜索空间提供可能。

A6 附录

A 预训练结果

我们通过从头开始预训练一个MoE模型来测试FarSkip-Collective架构,该模型使用了MegatronLM中的DeepSeek-V2 Lite模型架构和配置(64个专家)。我们在一个由高质量公开数据混合而成的大型语料库中随机抽样了50B的训练tokens,序列长度为4096 tokens,批量大小为4194304 tokens。由于该模型有2.4B的活跃参数,这对应于一个简单的Chinchilla扩展估计【索引15,Training compute-optimal large language models,2022,arXiv】。我们称该模型为“代理”模型,因为它使用了DeepSeek-V2-Lite的模型架构但非其训练方案和数据。我们使用完全相同的设置,但关闭了FarSkip-Collective(即原始的MegatronLM实现)作为基线进行了重复训练。我们观察到两个模型的损失曲线非常接近,FarSkip-Collective模型在最后50步的平均最终训练损失为2.205,而常规模型为2.187。在评估方面,我们使用了Llama-1论文【索引36,LLaMA: Open and Efficient Foundation Language Models,2023,arXiv】中的评估数据集来从预训练中获取早期信号,并观察到FarSkip-Collective的性能与常规模型架构相当,平均得分为54.7,而常规模型为54.4。这一结果进一步加强了FarSkip-Collective MoE模型修改的可行性,并补充了正文中更大规模自蒸馏的结果。


图 8: 使用常规和FarSkip-Collective架构的MoE预训练损失。我们使用DeepSeek-V2-Lite模型架构对一个16B参数的MoE进行了50B tokens的预训练。我们观察到FarSkip-Collective的损失与常规模型紧密匹配。

Table 4: 常规架构和FarSkip-Collective架构在预训练任务上的评估结果。

B 定性模型生成样本

我们通过在基于聊天的生成任务中对FarSkip-Collective模型进行定性测试,进一步验证了我们转换后的模型和推理实现,我们使用了与原始模型相同的生成参数。例如,在图9中,我们观察到当提示FarSkip-Collective DeepSeek-V2 Lite时,它给出了合理的回答。

DeepSeek-V2 Lite Instruct (原始)
提示 线性代数的主要主题是什么?
生成 线性代数是数学的一个分支,处理向量、向量空间(也称为线性空间)、线性变换和线性方程组。它在物理学、工程学、计算机科学、经济学和社会科学中有广泛的应用。线性代数的主要主题包括:

DeepSeek-V2 Lite Instruct (FarSkip-Collective)
提示 线性代数的主要主题是什么?
生成 线性代数是数学的一个分支,它涉及对向量、向量空间、线性变换和线性方程组的研究。线性代数的主要主题包括:1. 向量:理解向量的概念及其属性,如加法、减法和标量乘法。

图 9: 完全转换的FarSkip-Collective DeepSeek-V2-Lite模型(下)和原始模型检查点(上)的生成内容。两者都使用默认的生成超参数(τ = 0.3, p = 0.95)。

C FarSkip-Collective 层执行轨迹

我们展示了在训练和推理期间显式重叠的FarSkip模型的部分层执行轨迹。在每个轨迹中,我们展示了一层的持续时间。由于FarSkip减少了一层的总持续时间,与常规层的轨迹相比,FarSkip的轨迹被放大了。


图 10: DeepSeek-V2 vLLM prefill推理层执行。(上) 常规连接。(下) FarSkip-Collective。在下图中,all-reduce集合通过在第二个硬件队列上异步运行,与注意力和MoE子块重叠。


图 11: DeepSeek-V2-Lite预训练前向传播层执行。(上) 常规连接。(下) FarSkip-Collective。在下图中,all-to-all通信与计算重叠,第一个调用对应于Dispatch,与核心注意力计算重叠。在第二个调用中,all-to-all对应于Combine,与共享专家和下一层的注意力q, k, v计算重叠。


图 12: DeepSeek-V2-Lite预训练后向传播层执行。(上) 常规连接。(下) FarSkip-Collective。后向传播操作符的执行顺序被从默认的torch.autograd序列号排序中“劫持”,以重新排序操作以实现重叠。具体来说,路由专家的后向计算在Combine all-to-all后向梯度的同步点完成后立即启动,而Dispatch梯度在注意力第一部分(q, k, v计算)的梯度之前启动,以便在同步前与之重叠。