Junfeng Lin∗ (清华大学), Ziming Liu∗ (新加坡国立大学), Yang You (新加坡国立大学), Jun Wang (中电海康集团有限公司), Weihao Zhang† (Lynxi Technologies Co. Ltd), Rong Zhao† (清华大学)

A1 主要贡献

本文旨在解决传统流水线并行(Pipeline Parallelism, PP)在训练长上下文大语言模型(LLMs)时遇到的通信挑战。随着模型上下文长度的增加和内存优化技术(如混合精度训练、重计算)的普及,传统PP中激活值及其梯度的通信开销成为新的瓶瓶颈。为了应对这一问题,本文提出了一种名为WeiPipe(Weight-Pipeline Parallelism,权重流水线并行)的新方法,将传统流水线从传递激活值转变为传递模型权重,从而降低通信成本并实现更均衡的资源利用。

核心问题与研究目标
* 问题:在长上下文和大微批次(micro-batch)场景下,传统基于“激活值传递”的流水线并行会因巨大的激活值(activations)和激活值梯度(gradients of activations)而产生严重的通信瓶颈。当激活值与权重的体积之比($bsh/12h^2$)超过1时,传递权重的通信效率更高。
* 目标:设计一种新的“权重传递”流水线并行策略(WeiPipe),该策略仅在工作节点(worker)之间以流水线方式传递权重(weights, $W$s)及其梯度(gradients of weights, $gW$s),从而降低通信开销,并保证高可扩展性(不依赖集体通信原语)。

主要创新与贡献
* 提出基础WeiPipe策略:阐述了权重传递流水线的基本概念和机制,即worker持有单个微批次的全部激活值,并在worker之间循环传递各层的权重进行计算。
* 提出WeiPipe-Interleave策略:对基础WeiPipe进行了改进,通过交错执行前向和后向传播,减少了流水线气泡(bubble),并将通信需求减半。
* 探索WeiPipe与零气泡并行:探讨了将WeiPipe与零气泡(zero-bubble)流水线相结合的潜力,展示了权重传递流水线可以实现接近零的气泡率。
* 实现与验证:从零开始在PyTorch上实现了WeiPipe-Interleave,并应用于主流的Llama风格模型。实验证明,在长上下文场景下,WeiPipe相比最先进的PP和FSDP(完全分片数据并行)可将训练效率提升约30%-80%,并表现出良好的弱扩展性和强扩展性。

为了更全面地分析现有方法和WeiPipe,本文标准化了符号表示,如下表所示。


表1. 本文中使用的符号含义。

A3 背景知识与相关工作

大模型与长上下文训练

Transformer架构的普及与模型规模的增长。基于Transformer【【38,Attention Is All You Need,2023,arXiv】】的模型已成为自然语言处理(NLP)和计算机视觉(CV)等领域的基础。模型规模在两个维度上显著增长:一是参数量急剧增加,例如Llama-3.1模型【【11,The Llama 3 Herd of Models,2024,arXiv】】拥有4050亿参数,对GPU内存和分布式训练提出了巨大挑战;二是训练过程中的激活值规模因序列长度增加而扩大。

长上下文的应用场景与挑战。在NLP中,长上下文对于多轮对话、长文档摘要【【2,Longformer: The Long-Document Transformer,2020,arXiv】】和大型代码库处理【【7,PaLM: Scaling Language Modeling with Pathways,2022,arXiv】】至关重要。在CV中,生成高分辨率和长时视频【【3,Video generation models as world simulators,2024,openai.com】】同样会产生长序列。在科学AI领域,如AlphaFold【【18,Highly accurate protein structure prediction with AlphaFold,2021,Nature】】,也需要处理长氨基酸序列。这些长上下文场景给并行方案的设计带来了新的挑战,特别是在实现高效和均衡的内存管理方面。

训练中的并行技术

现有并行技术的概述。为应对上述挑战,多种并行技术被开发出来。数据并行(Data Parallelism, DP)是最直接的方法,它在设备间划分数据,并通过all-reduce操作同步梯度。完全分片数据并行(Fully Sharded Data Parallelism, FSDP)【【41,PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,2023,arXiv】】和零冗余优化器(Zero Redundancy Optimizer, ZeRO)【【33,ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,2020,arXiv】】进一步优化了DP,通过分片模型权重和优化器状态来节省内存。张量并行(Tensor Parallelism, TP)在层内部分割张量,但需要频繁的集体通信。流水线并行(Pipeline Parallelism, PP)在层级上划分模型,并将批次数据切分为微批次,通过点对点(P2P)通信传递激活值。序列并行(Sequence Parallelism, SP)【【14,LoongTrain: Efficient Training of Long-Sequence LLMs with Head-Context Parallelism,2024,arXiv】;【16,DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models,2023,arXiv】;【24,Sequence Parallelism: Long Sequence Training from System Perspective,2022,arXiv】;【25,Ring Attention with Blockwise Transformers for Near-Infinite Context,2023,arXiv】;【27,WallFacer: Guiding Transformer Model Training Out of the Long-Context Dark Forest with N-body Problem,2024,arXiv】】则沿着序列维度分割激活值。在这些策略中,PP因其仅依赖P2P通信而适用于网络连接较弱的设备。

流水线并行(Pipeline Parallelism)。本文主要关注同步流水线。其中最著名的是GPipe【【15,GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism,2018,arXiv】】,它在前向传播完成后才开始后向传播。Dapple(也称为1F1B)【【12,DAPPLE: A Pipelined Data Parallel Approach for Training Large Models,22020,arXiv】】通过立即启动已完成前向传播的微批次的后向传播来改进GPipe,减少了峰值内存。Chimera【【22,Chimera: efficiently training large-scale neural networks with bidirectional pipelines,2021,SC】】结合不同方向的多个流水线来减少气泡率。Hanayo【【26,Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency,2023,SC ’23】】在此基础上使用波浪形流水线,将Chimera与模型复制解耦。零气泡PP(Zero-bubble PP)【【32,Zero Bubble Pipeline Parallelism,2023,arXiv】】将权重梯度和激活值梯度的计算分开,允许更灵活的流水线安排。尽管现有PP专注于减少气泡率,但在长上下文场景下,激活值和梯度的通信开销可能超过计算负载,成为主要瓶颈。

完全分片数据并行(Fully Sharded Data Parallelism)。FSDP的功能与ZeRO-3类似,它将数据、模型权重和优化器状态分片到不同GPU上。权重和梯度仅在计算需要时通过all-gather操作收集,计算完成后立即丢弃,并通过reduce-scatter操作更新。FSDP利用异步通信来增强重叠效果。尽管易于使用且内存高效,FSDP在前向和后向传播中需要大量的集体通信,这在规模化场景中可能成为挑战。

总结。现有技术在减少长上下文场景下的通信开销方面准备不足,尤其是在网络连接较弱的集群中。

A2 方法细节

WeiPipe-Naive

基础思想与前向传播。为了介绍WeiPipe的基本思想,本文以常用的环形拓扑PP为例。首先,将模型层数 $L$ 平均分配给所有 $N$ 个worker,假设 $L=N$,每个worker持有一层的权重。训练过程包括前向、后向和更新三个阶段。在前向传播中,worker 0持有第一层的权重 $W_0$ 并开始计算,同时将 $W_0$ 传递给worker 1,并从worker $N-1$ 接收 $W_1$。随后,worker 0丢弃 $W_0$,但在内存中保留第一层的激活值 $A_0$,并使用 $W_1$ 开始第二层的前向计算。此过程持续进行,直到worker 0完成所有层的前向计算。图1($t=0$ 到 $t=3$)展示了 $N=4$ 的情况。图中,圆圈表示数据的排列和传输,每个worker持有当前位置和对角线位置的数据。初始时,worker 0到3分别持有 $W_0$ 到 $W_3$。圆圈逆时针旋转,每个worker处理对应层的前向计算,并在各自内存中累积激活值 $A_0$ 到 $A_3$。

后向传播的实现。后向传播更为复杂,因为它需要以相反的顺序使用权重。为维持权重流水线,我们将 $W_3$ 到 $W_0$ 填充在圆圈的另一侧。当worker 0完成前向传播后,它会依次接收 $W_3$ 到 $W_0$ 并执行从第3层到第0层的后向计算(图1中 $t=4$ 到 $t=8$ 的蓝色块)。在WeiPipe-Naive中,索引较小的worker会比索引较大的worker更早完成前向传播并开始后向传播,因此需要同时传输用于前向和后向的权重。为实现这一点,每个worker都持有用于前向和后向的权重。

更新阶段与梯度累积。在更新阶段,每个worker使用权重梯度($gW$s)来更新权重。由于权重在worker之间循环,且每个worker只为其微批次生成一部分 $gW$,因此更新某一层权重需要收集所有对应的 $gW$。与DP中使用all-reduce不同,WeiPipe仅使用P2P通信,因此 $gW$ 在生成后也随权重一起在环中循环。例如,在图1的 $t=4$ 时,worker 0对第3层进行后向计算,生成 $gW_3$,然后 $gW_3$ 与 $W_3$ 一起在环中传递。每个处理第3层后向计算的worker都会将新生成的 $gW_3$ 与已有的 $gW_3$ 进行平均、求和或归一化,而通信量保持不变。此过程重复进行,直到所有微批次处理完毕,此时每个worker使用指定的优化器更新其持有的 $W$ 和 $gW$。由于每个worker只更新一层权重,它也存储相应层的优化器状态,无需在worker间传输。


图1. WeiPipe-Naive策略。worker下方的文本表示内存中保留的数据。圆圈内的文本表示在worker之间逆时针循环的数据。在 $t$ 时刻,每个worker拥有其位置和橙色线另一侧的数据。例如,在 $t=5$ 时,worker 1拥有 $A_0, W_3, gW_3$,并从worker 0接收 $W_1, W_2, gW_2$。在 $t=6$ 时,worker 1持有 $A_1, W_2, gW_2$。绿色块表示worker正在执行前向传播,蓝色块表示后向传播。

WeiPipe-Interleave:更低的气泡率

WeiPipe-Naive的缺陷与改进思路。WeiPipe-Naive存在两个主要缺陷:一是同时循环两股权重流,但只有一个用于当前计算,导致冗余的传输成本;二是后向传播的计算时间约是前向传播的两倍,当一个worker执行后向传播时,会给其他仍在执行前向传播的worker带来显著的流水线气泡。为了解决这些问题,本文通过交错执行前向和后向传播,提出了WeiPipe-Interleave策略。

交错执行的机制。其核心思想是利用圆圈对角线位置的权重进行另一个微批次的前向或后向计算。初始的前向过程与WeiPipe-Naive相同。如图2所示,在 $t=4$ 时,worker 0开始执行后向传播,同时它利用 $W_0$ 对一个新的微批次输入执行前向传播。因此,在这个位置,worker 0会在圆圈转动一次前执行一次后向和一次前向计算。在此期间,其他worker只能执行一次前向操作并等待worker 0的计算完成。在随后的轮次中,当计算工作负载在worker之间达到平衡时,所有worker都进入前向-后向交错阶段。

WeiPipe-Interleave的优势。在前向-后向交错阶段,无论后向和前向的工作负载差异多大,都没有流水线气泡。每个worker还可以根据其数据传输情况决定前向和后向的执行顺序。此外,前向和后向传播期间的空闲内存被用来存储新微批次产生的激活值,从而实现了跨不同worker和时间的均衡存储利用。WeiPipe-Interleave进一步降低了通信开销。对于Llama结构(每层有 $12h^2$ 个参数),在前向-后向交错阶段,每个worker在每一轮接收两层权重和一层 $gW$,通信量为 $36h^2$,与WeiPipe-Naive相比,计算/通信比提高了一倍。


图2. WeiPipe-Interleave策略。$A_i^j$ 表示第 $j$ 个微批次在第 $i$ 层的激活值。

WeiPipe-zero-bubble

与零气泡流水线的结合。为了进一步探索WeiPipe的潜力,本文提出了将WeiPipe与零气泡流水线相结合的设计。零气泡策略是当前最先进的PP技术,它将后向计算分解为计算激活值梯度的B pass和计算权重梯度的W pass,以实现几乎零气泡。基于这种解耦方法,本文介绍了两种策略:WeiPipe-zero-bubble 1 (WZB1)和WeiPipe-zero-bubble 2 (WZB2)。

WeiPipe-zero-bubble 1 (WZB1)。WZB1相比WeiPipe-Interleave能略微降低气泡率,且存储和通信开销相对较低。如图3所示,在 $t=4$ 时,worker 0执行两个任务:使用 $W_0$ 对新微批次4进行前向计算,生成 $A_4^0$;使用 $W_3$ 对旧微批次0进行B pass,产生激活值梯度 $gA_3^0$。worker 0在 $t=4$ 后仍持有 $A_0^0$ 的一部分以支持未来的W pass。与WeiPipe-Interleave不同,用于后向传播的 $gW_1$ 与 $W_3$ 放置在一起。在 $t=5$ 时,worker 0执行W pass,消耗 $A_0^0$ 和 $gA_1^0$(图中未画出,但概念上存在)生成 $gW_3$,同时使用 $W_1$ 对微批次4进行前向计算。此过程以“一个前向-一个B”和“一个前向-一个W”交替进行,直到微批次4的前向传播完成。接着,从 $t=8$ 开始,worker 0交替执行两次B pass和两次W pass,直到微批次0的所有计算完成,并开始新微批次8的前向传播。


图3. WeiPipe-zero-bubble 1。绿色块代表前向传播,蓝色块代表B pass,紫色块代表W pass。$gA_i^j$ 表示第 $j$ 个微批次在第 $i$ 层的激活值梯度。层间的中间梯度被忽略。

WZB1的数据排列与通信。图3展示了环中的数据排列。用于前向的 $W_0$ 到 $W_3$ 与之前策略一致。用于B pass的权重和W pass生成的梯度成对放置,通常 $W_{N-1}$ 和 $gW_{N-1}$ 一起, $W_{N-2}$ 和 $gW_{N-2}$ 一起,以此类推。这种排列确保每个worker在每一轮执行2个数据块的操作,同时向下一个worker传输3个数据块。

WeiPipe-zero-bubble 2 (WZB2)。WZB2提供了一个更简单的流程,其 $W$s和 $gW$s的排列与WeiPipe-Interleave相同,如图4所示。在WZB2中,一个worker内顺序执行所有层的前向、B pass和W pass。特别是,W pass从第0层到第3层进行,与前向顺序一致。在B pass期间,旧版本的权重可以被丢弃以减少传输量。最后一个worker(worker 3)聚合所有的 $gW$s并更新权重。例如在 $t=11$ 时,worker 3持有聚合后的 $gW_0$ 并更新 $W_0$。然后worker 3发送更新后的 $W_0$,从而在 $t=12$ 时启动一个新的前向传播。这种无缝交接允许用更少的微批次实现频繁的权重更新,并达到几乎零气泡。WZB2的通信和存储成本更高,每一轮执行一个数据块的操作,同时向下一个worker传输两个数据块。


图4. WeiPipe-zero-bubble 2。worker 3在每轮结束时暂存更新后的权重,并逐一传递给worker 0用于下一轮。这种安排不会超过worker 3的峰值内存。

理论分析

与其他流水线策略的比较。图5对比了WeiPipe与其他PP策略。WeiPipe-Interleave的气泡率与1F1B相同。两种零气泡策略通过在气泡率公式中为 $M$ 乘以一个大于1的因子来减少气泡。ZB2和WZB2在迭代过程中几乎没有气泡。

通信效率。通信效率通过总接收数据量除以计算时间来评估。对于激活值传递的PP,其通信量与微批次大小 $b$ 和序列长度 $s$ 呈线性关系。相比之下,WeiPipe的通信量由权重大小决定,与 $b$ 和 $s$ 无关。以Llama风格模型为例,每层权重为 $12h^2$。WeiPipe-Interleave在每个轮次需要传输两块 $W$ 和一块 $gW$,通信量为 $36h^2$,耗时为 $T_{fwd} + T_{bwd}$。

内存消耗。内存消耗受具体实现影响很大。本文基于 $W, A, gW, gA$ 的数量进行粗略估计。1F1B和WeiPipe-Interleave的主要内存消耗是激活值存储,约为 $Mbsh$。WeiPipe-Interleave的内存消耗与1F1B相似但更均衡。在分离B pass和W pass的策略中,B pass会消耗部分激活值并产生梯度。对于零气泡流水线,本文评估了最后一个worker的峰值存储。ZB1和WZB2都需要存储 $M(bsh+gAs)$ 的数据。ZB2的需求几乎翻倍。然而,WZB1可以实现约 $1.5Mbsh$ 的最大内存消耗,当 $T_{fwd}/T_{bwd} > 1.5 - \mu$ 时,这比其他零气泡策略要低。


图5. WeiPipe与其他流水线策略的比较。在流水线图中,块内的上层数字表示层索引,下层数字是微批次索引。方程表中的分布图显示了这些策略在设备和时间上的理论带宽和内存使用情况。具体数值基于Llama结构计算,其中 $h=4096, s=4096, M=12, N=32$,以及32个头的Transformer。内存分布是基于 $\mu=0.6$ 并且不考虑重计算和Flash Attention的统计数据估计的。实际分布可能因实现而异。

A7 补充细节:实现

实现框架。当前分布式训练框架多提供激活值传递的流水线策略,但很少支持实现权重流水线所需的高级修改。因此,本文从零开始在PyTorch【【29,PyTorch: An Imperative Style, High-Performance Deep Learning Library,2019,arXiv】】中实现了WeiPipe-Interleave。该实现专为训练LLMs定制,但方法不限于Transformer架构。WeiPipe-zero-bubble因其控制复杂精细,本文仅作理论探讨,实现留待未来工作。

混合精度。与许多仅关注全精度实现的研究不同,本文的实现集成了广泛采用的混合精度训练,以进一步减少存储和通信开销。具体来说,$W$、$A$ 和 $gA$ 使用fp16精度,$gW$ 使用bf16精度。优化器状态使用fp32精度,并分布在各个worker上。

通信重叠。WeiPipe通过精心调整计算和权重传递的安排来平衡通信和计算负载。为确保高性能,实现了通信隐藏优化。在前向-后向交错期间,使用PyTorch分布式库提供的batch_isend_irecv函数进行异步通信,预取 $W$s 和 $gW$s。

重计算和Flash Attention。重计算(梯度检查点)是一种通过在后向传播中重新计算部分前向传播的激活值来节省存储空间的技术。Flash Attention【【8,FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,2022,arXiv】】是一种优化注意力模块内存访问和节省内存的常用技术。通过采用这些技术,可以增加微批次大小,从而凸显WeiPipe的优势。为了公平比较,这些优化也应用于所有其他对比策略中。

A4 实验环境

模型配置
* 模型架构: 基于开源的Llama-2结构【【37,Llama 2: Open Foundation and Fine-Tuned Chat Models,2023,arXiv】】(也即GPT风格模型)。
* 关键参数: 固定注意力头数为32,层数 $L$ 为32。通过改变隐藏层维度 $h$、序列长度 $s$、微批次大小 $b$ 和迭代中的微批次数量 $M$ 来构建不同规模的模型。

硬件配置
* 平台: Colossal Cloud。
* GPU: A800 GPU,配备80GB HBM显存和312 TFlops的fp16/bf16张量核心性能。A800的NVLink带宽被限制为400GB/s(相较于A100的600GB/s)。
* 网络基础设施: 实验在两种不同的通信基础设施上进行:
1. NVLink + 以太网 (16/32卡): 16个A800 GPU(分属2个集群)或32个A800 GPU(分属4个集群)。集群内部使用NVLink【【13,Ultra-Performance Pascal GPU and NVLink Interconnect,2017,IEEE Micro】】连接,集群之间使用10Gb以太网连接。
2. PCIe + 以太网 (16卡): 16个A800 GPU(分属2个集群),集群内部通过PCIe连接。
3. 纯NVLink (8卡): 8个A800 GPU通过NVLink连接,用于评估通信约束较小的场景。
* 拓扑: 所有并行策略均保持环形拓扑。

软件配置
* 代码实现:
* WeiPipe-Interleave:使用PyTorch从零开始实现。
* 1F1B 和 zero-bubble 策略:通过Megatron-LM实现。
* FSDP:通过DeepSpeed中的ZeRO-3优化实现。
* 通信库: 使用NCCL作为底层通信库。
* 通用优化: 所有策略均应用了相同的模型配置、微批次大小、混合精度设置、Flash Attention。除zero-bubble策略外,所有策略均应用了重计算(recomputation),因为重计算对zero-bubble策略无存储收益反而增加计算开销。

A4 实验结果

吞吐量与内存消耗

NVLink + 以太网环境(表2)。在16个GPU上,针对不同隐藏层维度(4096-16384)和序列长度(4096-16384)的模型进行测试。在所有配置中,WeiPipe-Interleave均表现出比其他策略更高的性能。特别是在长上下文场景下,例如当 $h=4096$ 且 $s=16384$ 时,WeiPipe的吞吐量相比1F1B和FSDP分别提升了22.3%和78.4%。内存方面,1F1B是所有流水线策略中内存占用最小的。Zero-bubble策略由于其机制,在启用Flash Attention和重计算后,峰值内存显著高于1F1B,这与原始论文的结论有所出入,因为优化技术改变了激活值的主要来源(从Attention转向FFN),导致峰值内存出现在最后一个Rank的第一个W Pass之前,约为第一个Rank峰值的两倍。WeiPipe的内存使用略高于FSDP,主要因为需要为通信分配额外的缓冲区。

表2. 在16个GPU(NVLink和以太网连接)上训练Llama风格模型的吞吐量和峰值内存。由于内存限制,对于ZB策略,当S=4096时G设为4,当S=8192或16384时G设为1。

PCIe + 以太网环境(表3)。在通信压力更大的PCIe和以太网环境中,WeiPipe的优势更加明显。当 $h=16384$ 且 $s$ 分别为2048和4096时,WeiPipe-Interleave的吞吐量相比其他策略的最佳性能分别提升了31.7%和22.2%。这些结果表明WeiPipe使得LLM训练对高带宽网络的依赖性降低。

表3. 在16个GPU(PCIe和以太网连接)上训练Llama风格模型的吞吐量和峰值内存。由于内存限制,对于ZB策略,当S=4096时G设为4,当S=8192或16384时G设为1。

纯NVLink环境(表4)。在通信约束较小的8 GPU NVLink环境中,传统方法可能具有优势。这表明WeiPipe的优势主要体现在通信受限的场景。

表4. 在8个GPU(NVLink连接)上训练层数为16的Llama风格模型的吞吐量。对于ZB策略,当S=4096时G设为4,当S=8192或16384时G设为1。

弱扩展性

弱扩展性实验通过同时增加worker数量和任务规模来评估策略效率。实验分为小规模(4-16 GPU)和大规模(8-32 GPU)两组,均在NVLink+以太网环境下进行,序列长度固定为16384。结果(图6和图7)显示,WeiPipe在小规模和大规模测试中均表现出优于其他策略的弱扩展性,无论是总吞吐量还是单GPU吞吐量(tokens/sec/gpu)都保持在较高水平。


图6. 小规模弱扩展性。GPU数量从4扩展到16(每台服务器4个GPU),批次大小从64相应增加到256。左坐标对应条形图表示的总吞吐量,右坐标表示每个GPU每秒处理的token数。


图7. 大规模弱扩展性。GPU数量从8扩展到32(每台服务器8个GPU),批次大小从128相应增加到512。

强扩展性

强扩展性实验通过保持任务规模不变而增加worker数量来衡量可扩展性。模型配置与弱扩展性实验相同,小规模组批次大小固定为128,大规模组固定为256。结果(图8和图9)表明,WeiPipe展现出比其他策略更优的强扩展性。这揭示了WeiPipe有潜力利用更多GPU来实现更大的加速比,尤其是在长token长度和低带宽通信基础设施给传统PP方法带来严重通信压力的情况下。


图8. 小规模强扩展性。GPU数量从4扩展到16,批次大小保持为128。


图9. 大规模强扩展性。GPU数量从8扩展到32,批次大小保持为256。

A5 结论

本文引入了一种新颖的流水线并行策略WeiPipe,旨在通过从基于激活值的通信转向基于权重的通信来减少LLM训练中的通信开销。我们开发并实现了WeiPipe-Interleave,以最小化权重传递的开销、减少气泡率并平衡内存使用。为了进一步探索权重传递流水线的潜力,我们讨论了WeiPipe-zero-bubble方法以实现几乎无空闲的流水线。实验证明,WeiPipe-Interleave的性能优于当前的PP方法和FSDP。WeiPipe为使用PP训练长上下文模型开辟了新的前景,并降低了对昂贵的高带宽通信基础设施的依赖。

方法细节相关引用汇总

本文中,主要的方法论章节(第3、4、5节)并未直接引用外部文献来支撑其核心机制设计。相关的背景和技术对比引用主要集中在第2节“背景与相关工作”中,这些引用为理解WeiPipe的动机和优势提供了上下文。以下是第2节中引用的关键文献及其在文中的作用:

  • [38] Attention Is All You Need, 2023, arXiv: 在2.1节中被引用,作为Transformer架构的基础性工作,说明了当前大模型的技术基石。
  • [11] The Llama 3 Herd of Models, 2024, arXiv: 在2.1节中被引用,用以举例说明现代LLM参数规模的巨大,突显分布式训练的必要性。
  • [15] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2018, arXiv: 在2.2.1节中被引用,作为最知名的传统流水线并行技术之一,是本文所要改进和对比的基础方法。
  • [12] DAPPLE: A Pipelined Data Parallel Approach for Training Large Models, 2020, arXiv: 在2.2.1节中被引用,作为GPipe的一种改进(即1F1B),通过交错前向和后向传播来减少峰值内存。
  • [22] Chimera: efficiently training large-scale neural networks with bidirectional pipelines, 2021, SC: 在2.2.1节中被引用,作为一种通过组合多向流水线来减少气泡率的优化PP方法。
  • [26] Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency, 2023, SC ’23: 在2.2.1节中被引用,作为Chimera的进一步发展,提高了效率。
  • [32] Zero Bubble Pipeline Parallelism, 2023, arXiv: 在2.2.1节中被引用,作为一种先进的PP技术,其将后向传播解耦为B pass和W pass的核心思想被WeiPipe-zero-bubble借鉴和探讨。
  • [41] PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, 2023, arXiv[33] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020, arXiv: 在2.2.2节中被引用,作为与WeiPipe进行性能比较的另一种主流并行策略(完全分片数据并行),突出了FSDP/ZeRO-3对集体通信的依赖性。
  • [8] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022, arXiv: 在第5节“实现”中被引用,说明了通过采用FlashAttention等内存优化技术,可以增大微批次大小,从而放大WeiPipe在通信效率上的优势。