论文状态:已完成

WeiPipe: Weight Pipeline Parallelism for Communication-Effective Long-Context Large Model Training

发表:2025/02/28
原文链接
价格:0.100000
已有 3 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

长上下文大型模型的训练面临通信开销瓶颈。本文提出了WeiPipe,采用权重流水线并行方法,通过将模型权重划分为流水线阶段并重叠通信与计算,显著降低了通信成本并最大化了训练效率。实验证明,WeiPipe在可扩展性和吞吐量上优于现有方法。

摘要

Training large models with long context lengths requires significant communication overhead, which becomes a bottleneck in distributed training. We propose WeiPipe, a weight pipeline parallelism method designed to reduce communication costs effectively. By dividing the model weights into pipeline stages and overlapping communication with computation, WeiPipe minimizes idle times and achieves a communication-efficient training paradigm. Experimental results demonstrate that WeiPipe significantly improves scalability and throughput in training large models with extensive context lengths compared to existing methods.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

WeiPipe: Weight Pipeline Parallelism for Communication-Effective Long-Context Large Model Training (WeiPipe:用于通信高效长上下文大模型训练的权重流水线并行)

1.2. 作者

  • JUNFENG LIN (林俊峰), 清华大学 (Tsinghua University), 中国北京
  • ZIMING LIU (刘子明), 新加坡国立大学 (National University of Singapore), 新加坡
  • YANG YOU (游洋), 新加坡国立大学 (National University of Singapore), 新加坡
  • JUN WANG (王军), 赛特开创集团有限公司 (CETHIK Group Co. Ltd.), 中国杭州
  • WEIHAO ZHANG (张维豪), 凝思科技有限公司 (Lynxi Technologies), 中国北京
  • RONG ZHAO (赵荣), 清华大学 (Tsinghua University), 中国北京

1.3. 发表期刊/会议

PPoPP '25: The 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming (第 30 届 ACM SIGPLAN 并行编程原理与实践年度研讨会)

1.4. 发表年份

2025年2月28日 (UTC)

1.5. 摘要

训练长上下文 (long context) 的大型模型 (large models) 需要大量的通信开销 (communication overhead),这在分布式训练 (distributed training) 中成为了一个瓶颈。本文提出了 WeiPipe,一种权重流水线并行 (weight pipeline parallelism) 方法,旨在有效降低通信成本。WeiPipe 通过将模型权重 (model weights) 划分为流水线阶段 (pipeline stages) 并重叠通信 (communication) 与计算 (computation),最大限度地减少了空闲时间 (idle times),实现了通信高效的训练范式。实验结果表明,与现有方法相比,WeiPipe 在训练具有长上下文的大型模型时,显著提高了可扩展性 (scalability) 和吞吐量 (throughput)。

1.6. 原文链接

/files/papers/694664ea769f2826079b7079/paper.pdf 该论文于 2025 年 2 月 28 日 (UTC) 发布在 PPoPP '25 会议上。

2. 整体概括

2.1. 研究背景与动机

随着 Transformer 架构的大型语言模型 (Large Language Models, LLMs) 取得革命性进展,其模型规模和上下文长度 (context length) 都在迅速扩展。这种扩展带来了巨大的计算资源需求,尤其是在分布式训练 (distributed training) 环境中。现有的分布式训练技术,如数据并行 (Data Parallelism, DP)、张量并行 (Tensor Parallelism, TP) 和流水线并行 (Pipeline Parallelism, PP),各有优劣。其中,流水线并行 (PP) 因其能够降低每个工作节点 (worker) 的内存需求且主要依赖点对点通信 (Peer-to-Peer, P2P communication) 而备受青睐,P2P 通信对带宽的要求相对较低,适用于连接不那么稳健的设备。

然而,随着长上下文 LLM 的普及以及混合精度训练 (mixed-precision training)、重计算 (recomputation,也称梯度检查点 gradient checkpointing) 和 Flash Attention 等内存优化技术成为标准配置,传统 PP 方法面临新的挑战:

  1. 激活值和激活值梯度的通信量激增: 长序列长度导致层间传递的激活值 (activations) 大幅增加。

  2. 大微批次 (micro-batches) 的通信负担: 内存优化技术使得 PP 可以使用更大的微批次,这进一步加剧了激活值和梯度通信的负担。

    论文指出,传统 PP 主要关注减少流水线气泡 (pipeline bubble ratio),但面对长上下文场景,激活值和梯度的通信开销可能超过单个流水线阶段的计算负载。作者观察到,一个 Transformer 阶段的输出激活值大小约为 GSH (微批次大小 GG,序列长度 SS,隐藏维度 HH),而单层权重大小约为 12H212H^2。当激活值与权重的比率 GS/(12H)G S / (12H) 超过 1 时,权重传递 (weight-passing) 方式可能比传统的激活值传递 (activation-passing) 方式更具通信效率。

基于此,本文旨在解决长上下文 LLM 训练中因激活值和梯度通信量过大导致的瓶颈问题,通过改变流水线中传递的数据类型来提高通信效率和整体吞吐量。

2.2. 核心贡献/主要发现

本文的核心贡献在于提出并验证了一种创新的分布式训练范式——权重流水线并行 (WeiPipe),主要发现和贡献如下:

  • 提出了 WeiPipe 框架: 引入了从传统的激活值传递流水线 (activation-passing pipeline)权重传递流水线 (weight-passing pipeline) 的转变。在这种新范式中,工作节点之间仅传递模型权重 (weights) 及其梯度 (gradients),而非激活值和激活值梯度,从而显著降低了通信开销。
  • WeiPipe 不依赖集体通信: 与数据并行 (DP) 或全分片数据并行 (FSDP) 不同,WeiPipe 仅依赖点对点通信 (P2P communication),这确保了其在高可扩展性 (scalability) 场景下的优势,尤其是在网络连接不那么稳健的集群中。
  • 设计了 WeiPipe 的多种变体:
    • WeiPipe-Naive: 阐述了权重传递流水线的核心概念和机制。
    • WeiPipe-Interleave: 通过交错前向 (forward) 和反向 (backward) 传播,有效减少了流水线气泡 (bubble ratio),并将通信需求减半,显著提高了通信效率。
    • WeiPipe-zero-bubble: 探索了将 WeiPipe 与零气泡流水线并行 (zero-bubble pipeline parallelism) 结合的潜力,旨在实现接近零气泡的训练配置。
  • 显著提升了训练效率和可扩展性: 在多达 32 个 GPU 的实验中,针对长上下文 LLM 训练,WeiPipe-Interleave 相较于最先进的流水线并行方法 (如 1F1B、zero-bubble) 和全分片数据并行 (FSDP),在吞吐量 (throughput) 上实现了 30% 至 80% 的显著提升。特别是在通信受限的场景(如 PCIe 和 Ethernet 连接),WeiPipe 展现出更强的可扩展性 (scalability)。
  • 降低了对高带宽通信基础设施的依赖: WeiPipe 的设计使得 LLM 训练对昂贵的高带宽互连技术(如 NVLink)的依赖降低,使其在更广泛的硬件环境下具有实用性。

3. 预备知识与相关工作

3.1. 基础概念

  • 大型语言模型 (Large Language Models, LLMs): 基于 Transformer [38] 架构的巨型神经网络模型,拥有数十亿甚至数千亿参数,能够处理和生成人类语言。它们在自然语言处理 (NLP) 任务中表现出色,例如文本生成、问答和翻译。随着模型规模的增长,其计算资源需求也呈指数级上升。
  • Transformer 架构: LLM 的核心,由多层编码器 (encoder) 和解码器 (decoder) 组成,每层包含自注意力机制 (self-attention mechanism) 和前馈网络 (feed-forward network)。自注意力机制允许模型在处理序列中的每个元素时,考虑序列中所有其他元素的重要性。其核心计算公式如下: Attention(Q,K,V)=softmax(QKTdk)V \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V 其中:
    • QQ (Query) 是查询矩阵。
    • KK (Key) 是键矩阵。
    • VV (Value) 是值矩阵。
    • dkd_k 是键向量的维度,用于缩放点积,防止梯度过小。
    • softmax\mathrm{softmax} 函数将结果归一化为概率分布。
    • QKTQK^T 计算查询和键之间的相似度。
  • 分布式训练 (Distributed Training): 为了训练参数量巨大或上下文长度很长的模型,需要将训练任务分配到多个计算设备(如 GPU)上并行执行。
    • 数据并行 (Data Parallelism, DP): 最常见的分布式训练方法。它将训练数据分成多个批次,每个设备复制一份完整的模型,独立处理一个数据批次。在每个训练迭代结束时,所有设备的梯度 (gradients) 会通过集体通信 (collective communication) 操作(如 All-reduce)进行聚合,然后更新所有设备的模型参数。
    • 张量并行 (Tensor Parallelism, TP): 也称为模型内并行 (intra-layer parallelism)。它将模型中的大型张量(如权重矩阵)沿着特定维度切分成多个小块,由不同的设备负责计算。这要求设备之间在层内进行频繁的集体通信来同步中间结果。
    • 流水线并行 (Pipeline Parallelism, PP): 也称为层间并行 (inter-layer parallelism)。它将神经网络的模型层 (layers) 划分为多个阶段 (stages),每个阶段由一个或一组设备负责。数据批次被进一步细分为微批次 (micro-batches),这些微批次像流水线一样依次流过各个阶段。每个设备只存储部分模型层,从而减少了单设备的内存占用。PP 主要依赖点对点通信 (P2P communication) 来传递激活值 (activations) 和激活值梯度 (gradients) 在不同阶段之间。
    • 序列并行 (Sequence Parallelism, SP): 将输入序列 (sequence) 沿着其长度维度进行分片,并将不同的序列片段分配给不同的设备进行处理。这种方法在处理超长序列时可以显著减少激活值的内存占用。
  • 激活值 (Activations): 神经网络前向传播 (forward pass) 过程中,每一层输出的中间结果。它们是下一层的输入。
  • 梯度 (Gradients): 反向传播 (backward pass) 过程中计算出的损失函数 (loss function) 对模型参数(权重)的偏导数。梯度用于通过优化器 (optimizer) 更新模型参数。
    • 激活值梯度 (Gradients of Activations): 损失函数对激活值的偏导数,用于计算前一层的权重梯度。
    • 权重梯度 (Gradients of Weights): 损失函数对模型权重的偏导数,用于更新权重。
  • 点对点通信 (Peer-to-Peer, P2P Communication): 两个计算设备之间直接进行数据传输,而不涉及所有设备之间的协调。PP 主要依赖这种通信方式。
  • 集体通信 (Collective Communication): 一组计算设备之间进行的复杂数据传输和聚合操作,如 All-reduce (所有设备共享总和结果)、Reduce-scatter (所有设备将自己的数据部分发送到其他设备并接收部分总和) 和 All-gather (所有设备将各自的数据部分发送给其他所有设备)。FSDP 和 DP 广泛使用集体通信。
  • 重计算 (Recomputation) / 梯度检查点 (Gradient Checkpointing): 一种内存优化技术。在前向传播时,故意不存储所有激活值,而是在反向传播时根据需要重新计算它们。这减少了内存占用,但增加了计算量。
  • 混合精度训练 (Mixed Precision Training): 在训练神经网络时,同时使用单精度 (FP32) 和半精度 (FP16 或 BF16) 浮点数。这可以减少内存占用、加速计算,并降低通信带宽需求,同时通常能保持与全精度训练相似的模型精度。
  • Flash Attention: 一种针对 Transformer 注意力机制的优化算法 [8],通过重新组织计算和内存访问模式,显著减少了注意力计算的内存占用和运行时间,尤其在处理长序列时效果显著。

3.2. 前人工作

本文主要关注流水线并行 (PP)全分片数据并行 (FSDP) 两种分布式训练技术,并与它们进行比较。

3.2.1. 流水线并行 (Pipeline Parallelism, PP)

传统 PP 的主要目标是减少流水线气泡 (pipeline bubbles),即设备在等待数据时处于空闲状态的时间。本文主要关注同步流水线 (synchronous pipelines),以保持模型收敛性。

  • GPipe [15]: 最早的 PP 方法之一。它要求所有微批次 (micro-batches) 的前向传播 (forward pass) 全部完成后,才开始所有微批次的反向传播 (backward pass)。这导致了较大的气泡。

  • Dapple [12] (也称 1F1B): 改进了 GPipe,在一个微批次的前向传播完成后立即启动其反向传播。这通过在时间上重叠前向和反向计算,减少了流水线气泡和某些设备上的内存峰值。

  • Chimera [22]: 通过结合多个不同方向的流水线进一步优化,以降低气泡率。

  • Hanayo [26]: 在 Chimera 的基础上,利用波浪形流水线 (wave-shaped pipeline) 解耦了模型复制,提高了效率。

  • Zero-bubble PP [32]: 最先进的 PP 策略之一。它将反向传播 (backward computation) 分解为两个独立的部分:B pass (计算激活值梯度) 和 W pass (计算权重梯度)。B pass 以反向传播的方式进行,而 W pass 可以延迟执行以填充流水线中的空闲时间,从而实现接近零气泡的配置。该方法提供了两种方案:一种优化内存峰值,另一种优化气泡率。

    传统 PP 面临的挑战: 尽管这些方法在减少气泡方面取得了进展,但在长上下文 LLM 训练中,由于激活值和激活值梯度的通信开销巨大,可能超出单个流水线阶段的计算负载,这成为其瓶颈。

3.2.2. 全分片数据并行 (Fully Sharded Data Parallelism, FSDP)

FSDP [41] 与 ZeRO Stage 3 (ZeRO-3) [33] 类似,它将模型权重 (model weights)、梯度 (gradients) 和优化器状态 (optimizer states) 分片 (sharding) 到不同的 GPU 上。

  • 工作机制: 权重和梯度只在需要计算时才通过集体通信操作 (如 All-gather 和 Reduce-scatter) 聚合到本地 GPU,计算完成后立即被丢弃或重新分片。
  • 异步通信: FSDP 融入了异步通信 (asynchronous communication) 以隐藏通信开销。
  • 特点: 易于使用且内存效率高。然而,FSDP 在前向和反向传播过程中都需要大量的集体通信,这在扩展到大规模集群时可能面临挑战,尤其是在网络连接不佳的情况下。

3.3. 差异化分析

本文提出的 WeiPipe 与上述现有方法的核心区别在于:

  • 数据传递类型: 传统 PP(如 1F1B、Zero-bubble PP)是激活值传递流水线 (activation-passing pipeline),即在流水线阶段之间传递激活值和激活值梯度。而 WeiPipe 是一种权重传递流水线 (weight-passing pipeline),它在流水线阶段之间传递模型权重和权重梯度。

  • 通信瓶颈焦点: 传统 PP 侧重于减少流水线气泡,但在长上下文场景下,其激活值通信量成为主要瓶颈。FSDP 虽然内存效率高,但其大量集体通信可能在高扩展性或通信受限场景下成为瓶颈。WeiPipe 则直接瞄准了长上下文场景下激活值通信量过大的问题,通过传递相对较小的权重来降低通信开销。

  • 通信原语: 传统 PP 和 WeiPipe 都主要依赖点对点通信 (P2P communication),这使其在高扩展性场景下具有优势。FSDP 则大量依赖集体通信 (collective communication)

  • 对内存优化技术的适应性: 混合精度训练、重计算和 Flash Attention 等技术虽然减少了激活值的内存占用,允许使用更大的微批次,但同时也间接加大了激活值/梯度通信的总量,这正是 WeiPipe 希望通过改变传递数据类型来解决的问题。

    总之,WeiPipe 在分布式训练技术演进的脉络中,通过创新性地改变流水线中传递的数据类型,为解决长上下文 LLM 训练中的通信瓶瓶颈提供了一个新的视角和解决方案,尤其是在通信受限的环境中。

4. 方法论

本文引入了权重流水线并行 (WeiPipe) 框架,旨在解决传统激活值传递流水线 (activation-passing pipeline) 在长上下文大型模型训练中遇到的通信瓶颈。核心思想是,在流水线阶段之间传递模型权重 (weights) 及其梯度 (gradients),而非激活值 (activations) 及其梯度。

4.1. 方法原理

WeiPipe 的核心直觉基于一个观察:在长上下文模型中,单个微批次 (micro-batch) 的激活值大小 GSH (微批次大小 GG、序列长度 SS、隐藏维度 HH) 可能会显著大于单层权重的大小 12H212H^2 (Llama2 架构中,注意力模块 4H24H^2,前馈网络 8H28H^2)。当 GSH/(12H2)=GS/(12H)GSH / (12H^2) = GS / (12H) 比率超过 1 时,即当 GS>12HGS > 12H 时,传递权重而非激活值将更为通信高效。WeiPipe 通过环形拓扑 (ring topology) 在工作节点 (workers) 之间流式传递权重和权重梯度,将通信开销从与数据量(批次大小、序列长度)相关的激活值转移到与模型结构(隐藏维度)相关的权重上。

WeiPipe 与传统的流水线并行 (PP) 策略不同,它不关注流水线气泡 (pipeline bubbles) 的优化,而是强调数据(权重及其梯度)的流畅流动,以最小化通信开销。

4.2. 核心方法详解

4.2.1. WeiPipe-Naive

WeiPipe-Naive 是 WeiPipe 最基础的策略,它在环形拓扑中展示了权重传递流水线的基本概念和机制。假设有 PP 个工作节点,模型有 LL 层,且 L=PL = P,每个工作节点最初持有一层权重。

4.2.1.1. 前向传播 (Forward Pass)

  1. 初始化: 每个工作节点 pp 最初持有模型层 pp 的权重 WpW_p
  2. 数据流动:
    • Worker 0 首先使用其持有的权重 W0W_0 对传入的微批次输入进行前向计算。
    • 同时,Worker 0 会将 W0W_0 传递给 Worker 1,并从 Worker P-1 接收 W1W_1
    • Worker 0 在完成 W0W_0 的计算和传输后,会丢弃 W0W_0(不再存储),但会保留其生成的激活值 A0A_0 在本地内存中。
    • 接着,Worker 0 使用新接收到的 W1W_1 对当前微批次继续进行第二层的计算。
    • 这个过程持续进行,直到 Worker 0 完成所有层的计算。
  3. 激活值存储: 每个工作节点在完成其负责的所有层的前向计算后,都会在本地内存中累积该微批次的激活值 A0,A1,,AL1A_0, A_1, \ldots, A_{L-1}

4.2.1.2. 反向传播 (Backward Pass)

反向传播的复杂性在于权重需要以反向顺序(从高层到低层)进行。

  1. 权重逆序传递: 为了维护权重流水线,当一个微批次的前向传播完成后,权重会以反向顺序(例如从 WL1W_{L-1}W0W_0)继续在环形拓扑中传递。
  2. 计算梯度: Worker 0 完成前向后,将依次接收 WL1W_{L-1}W0W_0,并执行从高层到低层的反向计算,生成对应的权重梯度 DjD_j 和激活值梯度 BjB_j
  3. 并发性: 在 WeiPipe-Naive 中,索引较低的工作节点会比索引较高的节点更早完成前向传播,从而更早开始反向传播。这意味着在某些时间步,某些工作节点可能正在进行反向传播,而另一些还在进行前向传播。为了支持这种并发性,每个工作节点需要同时持有用于前向和反向的权重。

4.2.1.3. 更新阶段 (Update Pass)

  1. 梯度聚合: 每个工作节点在反向传播过程中会生成其处理的层的权重梯度 DjD_j。由于权重是在工作节点之间循环的,要更新某一层权重,需要聚合所有工作节点为该层生成的梯度。
  2. P2P 梯度传递: 与传统的数据并行 (DP) 使用 All-reduce 聚合梯度不同,WeiPipe 旨在只使用 P2P 通信。因此,生成的权重梯度 DjD_j 也会随其对应的权重 WjW_j 一起在环中循环传递。
  3. 本地聚合: 当一个工作节点接收到一个 DjD_j(例如来自 Worker 0),它会将其与自己为同一层生成的 DjD_j 进行平均、求和或归一化,并继续传递更新后的 DjD_j
  4. 权重更新: 经过所有微批次的处理后,每个工作节点将利用其最终聚合的 DjD_j 来更新它所负责的层 jj 的权重 WjW_j,使用指定的优化器。每个工作节点也存储其对应层的优化器状态 (optimizer state),无需在工作节点之间传输。

4.2.1.4. 图示分析 (Figure 1)

下图(原文 Figure 1)展示了 WeiPipe-Naive 的工作流程。

该图像是示意图,展示了 WeiPipe 方法在不同时间步(t=0 到 t=7)下的权重流水线并行处理过程。每个工人在时间步内处理不同的权重,通过有效重叠通信与计算,减少了通信开销,并优化了训练过程。 该图像是示意图,展示了 WeiPipe 方法在不同时间步(t=0 到 t=7)下的权重流水线并行处理过程。每个工人在时间步内处理不同的权重,通过有效重叠通信与计算,减少了通信开销,并优化了训练过程。

  • 环形表示: 图中用一个圆环表示数据的排列和传输。对于 P=4P=4 个工作节点,圆环被分成 8 个位置,每个位置填充权重 WW 或权重梯度 DD
  • 数据持有: 每个工作节点持有其当前位置的数据和对角线位置的数据(由橙色线表示)。
  • 前向阶段 (t=0 到 t=3):
    • 初始时,Worker 0 到 Worker 3 分别持有 W0W_0W3W_3 (黑色文本)。
    • 圆环逆时针旋转,每个旋转步长表示 t+1t+1
    • 在每个 tt 中,绿色块表示工作节点正在执行前向传播。例如,t=0t=0 时,Worker 0 执行 W0W_0 的前向;t=1t=1 时,Worker 0 执行 W1W_1 的前向(因为 W1W_1 已从 Worker 3 传过来),同时 Worker 1 执行 W0W_0 的前向(因为 W0W_0 已从 Worker 0 传过去)。
    • 工作节点在本地累积激活值 A0A_0A3A_3
  • 反向阶段 (t=4 到 t=8):
    • 反向传播需要逆序权重。为了实现这一点,圆环的另一侧填充了 W3W_3W0W_0 (蓝色文本),以支持反向计算。
    • 蓝色块表示工作节点正在执行反向传播。例如,t=4t=4 时,Worker 0 使用 W3W_3 执行反向传播,并生成 D3D_3
    • D3D_3 随后与 W3W_3 一起在环中传递,供其他处理过 Layer 3 的工作节点聚合。
    • t=8t=8 之后,Worker 0 可以为新的微批次启动另一个循环。

4.2.2. WeiPipe-Interleave for Lower Bubble Ratio

WeiPipe-Naive 存在两个主要缺点:

  1. 冗余传输: 两个权重流同时循环,但只有一个用于当前计算,增加了传输成本。

  2. 气泡问题: 反向传播的计算时间大约是前向的两倍。当一个工作节点执行反向时,会显著增加其他正在进行前向传播的工作节点的流水线气泡。

    为了解决这些问题,WeiPipe-Interleave 策略通过交错 (interleaving) 前向和反向传播来优化。

4.2.2.1. 核心思想

魏管道交错 (WeiPipe-Interleave) 的核心思想是利用环形结构中对角线位置的权重进行另一个微批次的前向或反向计算。

4.2.2.2. 流程

  1. 初始前向: 最初的前向过程与 WeiPipe-Naive 相同。
  2. 交错阶段:
    • 当一个工作节点(例如 Worker 0)完成其所有层的前向传播后,在它开始反向传播的同时,它会利用对角线位置的权重(例如 W0W_0)对一个新的微批次输入进行前向传播。
    • 因此,在这个阶段,一个工作节点(例如 Worker 0)会在圆环转动一圈之前,执行一次反向计算和一次前向计算
    • 在此期间,其他工作节点可能只能执行一次前向操作,并等待 Worker 0 的计算完成。
    • 在随后的轮次中,所有工作节点(Worker 0 到 Worker 3)都将进入这种前向-反向交错阶段,从而实现计算负载的平衡。
  3. 无流水线气泡: 在前向-反向交错阶段,无论反向和前向的计算负载差异多大,都不会产生流水线气泡,因为工作节点可以根据数据传输情况自行决定前向和反向的执行顺序。
  4. 内存利用: 前向和反向传播过程中的空闲内存被用于存储新微批次生成的激活值,从而实现了不同工作节点和时间点上的存储利用率平衡。
  5. 通信效率提升: WeiPipe-Interleave 进一步降低了通信开销。对于 Llama 结构,单层有 12H212H^2 个参数。在前向-反向交错阶段,每个工作节点每轮接收两个层的权重 (2×12H22 \times 12H^2) 和一个层的权重梯度 (1×12H21 \times 12H^2),总通信量为 36H236H^2。这使得计算/通信比 (compute/commute ratio) 相较于 WeiPipe-Naive 翻倍。

4.2.2.3. 图示分析 (Figure 2)

下图(原文 Figure 2)展示了 WeiPipe-Interleave 的策略。 原始论文描述: Figure 2. WeiPipe-Interleave strategy. AjiA _ { j } ^ { i } denotes the activation values of jj th layer in ith micro-batch.

Figure 2. WeiPipe-Interleave strategy. \(A _ { j } ^ { i }\) denotes the activation values of \(j\) th layer in ith micro-batch. 该图像是示意图,展示了 WeiPipe-Interleave 策略的工作机制。图中显示了不同时间步 tt(从 4 到 11)中,多个工作节点(Worker 0 到 Worker 3)之间的激活值 AjiA _ { j } ^ { i } 的传递与计算重叠过程。

  • 激活值表示: AjiA_j^i 表示第 ii 个微批次的第 jj 层激活值。
  • 交错执行:t=4t=4 开始,Worker 0 在执行反向传播的同时,也开始了一个新微批次的前向传播。例如,t=4t=4 时,Worker 0 执行 A04A_0^4 的前向计算,并执行 B30B_3^0 的反向计算。
  • 并发计算与通信: 图中绿色块表示前向计算,蓝色块表示反向计算。在交错阶段,可以看到工作节点同时进行这两种计算,并且权重和梯度在不同工作节点之间传递,确保了计算的连续性。

4.2.3. WeiPipe-zero-bubble

为了进一步探索 WeiPipe 的潜力,本文提出了一种结合零气泡流水线 (zero-bubble pipeline) 思想的设计。零气泡策略将反向传播 (backward computation) 分解为 B pass (计算激活值梯度) 和 W pass (计算权重梯度),以实现几乎零气泡的流水线。

4.2.3.1. WeiPipe-zero-bubble 1 (WZB1)

WZB1 旨在略微减少气泡率,同时保持相对较低的存储和通信开销。

  1. 分解反向: 反向传播被分解为 B pass(蓝色块,计算激活值梯度)和 W pass(紫色块,计算权重梯度)。
  2. 并发任务: 如下图(原文 Figure 3)所示,从 t=4t=4 开始,Worker 0 执行两个任务:
    • 对新微批次 (micro-batch 4) 的输入进行第 0 层的前向传播 (forward),使用 W0W_0,生成 A04A_0^4
    • 对旧微批次 (micro-batch 0) 的第 3 层执行 B pass,使用 W3W_3,生成激活值梯度 B30B_3^0
  3. 激活值保留: Worker 0 在 t=4t=4 之后仍然保留部分 A30A_3^0,以支持未来对 W3W_3 的 W pass。
  4. 权重和梯度传输: 与 WeiPipe-Interleave 不同,用于反向传播的 W1W_1W3W_3 一起放置在环中。红色文本表示当前工作节点使用的数据。
  5. 交替执行:t=5t=5 时,Worker 0 执行 W pass,消耗 A30A_3^0B30B_3^0 来生成 D3D_3,然后将其发送给 Worker 1。同时,Worker 0 还使用 W1W_1 进行微批次 4 的前向传播。这个过程持续交替进行“一前向-一 B”和“一前向-一 W”,直到微批次 4 的前向传播完成。
  6. 后续 B/W Pass:t=8t=8 时,Worker 0 执行两个 B pass:一个用于微批次 0 的 W1W_1,另一个用于微批次 4 的 W3W_3。从 t=8t=8t=11t=11,Worker 0 交替执行两个 B pass 和两个 W pass,直到 t=12t=12 时微批次 0 的所有传播都完成,并开始一个新微批次 (micro-batch 8) 的前向传播。
  7. 数据排列: 如下图(原文 Figure 4)所示的环形数据排列,用于前向传播的 W0W_0W3W_3 与之前的策略一致。用于 B pass 的权重和 W pass 生成的梯度是成对放置的。通常,对于一个有 LL 层的网络,WL1W_{L-1}WL/21W_{L/2-1} 配对,WL2W_{L-2}WL/22W_{L/2-2} 配对,依此类推。梯度 DsD_s 也采用相同的配对方式。这种安排确保每个工作节点在每次轮转中执行 2 个块 (chunk) 操作,同时向下一个工作节点传输 3 个块的数据。

4.2.3.2. 图示分析 (Figure 3 & Figure 4)

下图(原文 Figure 3)展示了 WZB1 的流程。 VLM 描述: 该图像是示意图,展示了魏管道(WeiPipe)在不同时间步(t=4t=4t=11t=11)下,各个工作节点(Worker)如何进行权重传递与计算重叠。该方法通过分层次的方式,显著降低了通信开销,有助于提升大规模模型的训练效率。

该图像是示意图,展示了魏管道(WeiPipe)在不同时间步(\(t=4\) 到 \(t=11\))下,各个工作节点(Worker)如何进行权重传递与计算重叠。该方法通过分层次的方式,显著降低了通信开销,有助于提升大规模模型的训练效率。 该图像是示意图,展示了魏管道(WeiPipe)在不同时间步(t=4t=4t=11t=11)下,各个工作节点(Worker)如何进行权重传递与计算重叠。该方法通过分层次的方式,显著降低了通信开销,有助于提升大规模模型的训练效率。

  • 阶段分解: 蓝色块表示 B pass (计算激活值梯度),紫色块表示 W pass (计算权重梯度)。

  • 并发: Worker 0 在 t=4t=4 同时执行前向 (绿色) 和 B pass (蓝色)。在 t=5t=5 同时执行前向 (绿色) 和 W pass (紫色)。

  • 数据流: 红色文本指示当前工作节点使用的数据。可以看到 W0W_0W1W_1 以及 D3D_3 等在不同时间步和工作节点之间流转。

    下图(原文 Figure 4)展示了 WZB2 的数据排列。 VLM 描述: 该图像是示意图,展示了WeiPipe方法在不同时间t下,各个工作节点(Worker)在权重管道并行过程中所处理的模型权重。图中的圆环标示出不同Worker所负责的权重分配,并随时间变化而更新。

该图像是示意图,展示了WeiPipe方法在不同时间t下,各个工作节点(Worker)在权重管道并行过程中所处理的模型权重。图中的圆环标示出不同Worker所负责的权重分配,并随时间变化而更新。 该图像是示意图,展示了WeiPipe方法在不同时间t下,各个工作节点(Worker)在权重管道并行过程中所处理的模型权重。图中的圆环标示出不同Worker所负责的权重分配,并随时间变化而更新。

  • 权重排列: W0W_0W3W_3 用于前向传播。用于 B pass 的权重和 W pass 产生的梯度是成对放置的。例如,W3W_3W1W_1 (原文中是 WL1W_{L-1}WL/21W_{L/2-1}) 是配对的。
  • 数据循环: 权重和梯度在环中循环,确保每个工作节点都能在正确的时间获得所需数据。

4.2.3.3. WeiPipe-zero-bubble 2 (WZB2)

WZB2 提供了一种更简单的零气泡流水线方案,与 WeiPipe-Interleave 具有相同的权重 (Ws) 和权重梯度 (Ds) 排列方式。

  1. 顺序执行: 在 WZB2 中,所有层的前向 (forward)、B pass 和 W pass 在一个工作节点内按顺序执行。
  2. W pass 顺序: W pass 从第 0 层到第 3 层依次进行,与前向顺序匹配。
  3. 旧权重丢弃: 在 B pass 期间,旧版本的权重可以被丢弃以减少传输,这在 Figure 4 中由空白表示(t=8t=8t=11t=11)。
  4. 最终聚合与更新: 最后一个工作节点(Worker 3)会聚合所有的权重梯度 DsD_s 并更新权重。
    • 例如,在 t=11t=11 时,Worker 3 持有聚合后的 D0D_0,并使用指定的优化器更新 W0W_0
    • 然后,Worker 3 将更新后的 W0W_0 发送出去,在 t=12t=12 时启动新的前向传播。
  5. 优势: 这种无缝交接允许频繁的权重更新,且微批次数量较少,从而实现几乎零气泡。
  6. 劣势: WZB2 会产生更高的通信和存储成本,因为它在每次轮转中执行一个块操作,同时向下一个工作节点传输两个块的数据。

4.3. 理论分析

本节对 WeiPipe 与其他流水线并行 (PP) 策略在气泡率 (bubble ratio)、通信效率 (communication efficiency) 和内存消耗 (memory consumption) 方面进行了理论比较。

4.3.1. 气泡率

  • WeiPipe-Interleave1F1B 具有相同的气泡率。
  • Zero-bubble 策略 (ZB1, ZB2, WZB1, WZB2) 通过在气泡率方程中将微批次数量 NN 乘以一个大于 1 的因子来减少气泡。
  • ZB2 和 WZB2 在迭代过程中几乎没有气泡。

4.3.2. 通信效率

通信效率通过将接收到的数据总量除以计算时间来衡量(即带宽使用)。

  • 激活值传递 PP: 如图 5 所示,其通信效率方程(Zone 1)是基于中间工作节点传输激活值 As 和激活值梯度 Bs 的总和 (因此方程中有一个因子“2”)。激活值传递策略的通信量随微批次大小 GG 和序列长度 SS 呈线性增长。
  • WeiPipe: 其通信量由权重大小决定,与 GGSS 无关。以 Llama 风格模型为例,每层权重数量为 12H212H^2
    • WeiPipe-Interleave 需要传输两个 WW 块和一个 DD 块,每轮通信量为 36H236H^2
    • 其时间持续为 TF+TBP\frac{T_F + T_B}{P},其中 TFT_FTBT_B 分别是完成一次前向和反向传播的总时间。

4.3.3. 内存消耗

内存消耗受具体实现影响较大。这里提供一个粗略的估计,仅基于值计数,忽略通信缓冲区、嵌入层和混合精度训练等因素。

  • 1F1B 和 WeiPipe-Interleave: 主要的内存消耗是激活值的存储,估计为 GMAG M_A。WeiPipe-Interleave 具有与 1F1B 相似的内存消耗,但更平衡。
  • 分离 B pass 和 W pass 的策略: B pass 的一个阶段(完成整个 B pass 的 1P\frac{1}{P} 部分)会消耗前向传播存储的部分激活值 MA/PM_A/P,并生成梯度 MB/PM_B/P。假设在 B pass 的一个阶段后,剩余 αMA/P\alpha M_A/P 的激活值存储。
  • 零气泡流水线: 评估最后一个工作节点所需的峰值存储。
    • ZB1 和 WZB2 都需要存储 G(αMA+MB)G (\alpha M_A + M_B) 数据。

    • ZB2 几乎使这一需求翻倍。

    • WZB1 可以实现大约 1.5GMA1.5 G M_A 的最大内存消耗,当 MBMA>1.5α\frac{M_B}{M_A} > 1.5 - \alpha 时,这低于其他零气泡策略。

    • 注意: 零气泡 1 的峰值内存计算可能较为复杂。在较小的 α\alpha 和特定的 GPU 内存管理策略下,Worker 0 可能具有最高的内存消耗。

      下图(原文 Figure 5)对 WeiPipe 和其他 PP 策略进行了理论分析比较。 VLM 描述: 该图像是示意图,展示了WeiPipe及其变体在不同阶段(如前向、反向)和内存分布下的通信效率和资源利用情况。图中包括了多种泡沫方案的通信分布及峰值内存的计算公式,如 GMB+MW+MdG_{MB} + M_W + M_dMW+MdM_W + M_d

该图像是示意图,展示了WeiPipe及其变体在不同阶段(如前向、反向)和内存分布下的通信效率和资源利用情况。图中包括了多种泡沫方案的通信分布及峰值内存的计算公式,如 \(G_{MB} + M_W + M_d\) 和 \(M_W + M_d\)。 该图像是示意图,展示了WeiPipe及其变体在不同阶段(如前向、反向)和内存分布下的通信效率和资源利用情况。图中包括了多种泡沫方案的通信分布及峰值内存的计算公式,如 GMB+MW+MdG_{MB} + M_W + M_dMW+MdM_W + M_d

  • 气泡 (Bubble): 图中左侧部分展示了不同策略的气泡率公式,其中 ZB1WZB2 的气泡率通过乘因子减小。
  • 带宽使用 (Bandwidth Usage): 中间部分展示了 Zone 1(完全交错阶段)的带宽使用公式。激活值传递 PP (如 1F1B, ZB1, ZB2) 的带宽使用与 GSHG \cdot S \cdot H 相关,而 WeiPipe 的带宽使用与 H2H^2 相关,独立于 GGSS
  • 峰值内存 (Peak Memory): 右侧部分展示了不同策略的峰值内存估计公式。1F1BWeiPipe-Interleave 具有相对较低的内存需求 GMAG M_A。零气泡策略的内存需求更高。

4.4. 实施细节

由于当前的分布式训练框架通常不支持高层次的权重流水线修改,本文从头开始在 PyTorch [29] 中实现了 WeiPipe-Interleave。虽然实现针对 LLM 训练,但该方法不限于 Transformer 架构。WeiPipe-zero-bubble 仅作为潜力探讨和理论设计,暂未实现,留待未来探索。

  • 混合精度 (Mixed Precision): 实现中采用了广泛使用的混合精度训练,以进一步减少存储和通信开销。具体来说,激活值 AA、权重 WW 和权重梯度 DD 采用 fp16 精度,而激活值梯度 BB 采用 bf16 精度。优化器状态 (optimizer state) 采用 fp32 精度,并分布在各个工作节点之间。
  • 通信重叠 (Communication Overlap): WeiPipe 精心调整了计算和权重传递的安排,以平衡通信和计算负载。为了确保 WeiPipe 的高性能潜力,实现了通信隐藏 (communication hiding) 优化。在前向-反向交错过程中,权重 Ws 和权重梯度 Ds 通过 PyTorch 分布式库提供的 batch_isend_irecv 函数进行异步通信 (asynchronous communication) 预取 (prefetched)。
  • 重计算 (Recomputation) 和 Flash Attention:
    • 重计算 (Recomputation) / 梯度检查点 (Gradient Checkpointing): 是一种通过在前向传播时放弃部分激活值,而在反向传播时重新计算它们来节省存储的策略,从而显著减少内存消耗,但会增加冗余计算。
    • Flash Attention [8]: 一种广泛使用的技术,用于优化注意力模块中的内存访问并节省内存。
    • 通过采用这些技术,可以增加微批次大小 (micro-batch size),从而进一步发挥 WeiPipe 的优势。为了公平比较,这些优化也应用于其他并行策略。值得注意的是,对于零气泡流水线,重计算不会带来存储节省,反而会增加计算开销,因此在零气泡策略中未启用。

5. 实验设置

5.1. 数据集

论文的实验是针对大型语言模型 (LLMs) 训练场景,使用了基于开源 LLama-2 结构 [37] (也称为 GPT 风格模型) 的模型。虽然没有明确提及具体的数据集名称,但实验关注的是模型结构和训练配置对性能的影响,而非特定数据集的特点。

5.2. 评估指标

论文使用了以下指标来评估不同策略的性能:

5.2.1. 吞吐量 (Throughput)

  • 概念定义: 衡量模型训练效率的核心指标,表示在单位时间内,每个 GPU 能够处理的词元 (tokens) 总量。吞吐量越高,表示训练效率越高。
  • 数学公式: Throughput (Tokens/second/GPU)=Total Tokens ProcessedTotal Time×Number of GPUs \text{Throughput (Tokens/second/GPU)} = \frac{\text{Total Tokens Processed}}{\text{Total Time} \times \text{Number of GPUs}}
  • 符号解释:
    • Total Tokens Processed\text{Total Tokens Processed}:在整个测量时间内,模型处理的词元总数。
    • Total Time\text{Total Time}:完成处理这些词元所花费的总时间。
    • Number of GPUs\text{Number of GPUs}:参与训练的 GPU 数量。

5.2.2. 内存消耗 (Memory Consumption)

  • 概念定义: 衡量模型训练过程中,每个 GPU 上占用的显存 (GPU memory) 大小。内存消耗越低,表示内存效率越高,允许在相同的硬件条件下训练更大的模型或使用更大的批次大小。
  • 数学公式: 论文未提供具体的计算公式,通常指训练过程中的峰值内存 (peak memory) 占用。
  • 符号解释: 通常以千兆字节 (GB) 为单位。

5.2.3. 弱扩展性 (Weak Scaling)

  • 概念定义: 评估当计算资源(GPU 数量)和每个 GPU 负责的工作量(例如,每个 GPU 的微批次大小,或者总批次大小与 GPU 数量同比例增加)同时按比例增加时,系统的效率(例如,每个 GPU 的吞吐量)是否能够保持相对稳定。理想的弱扩展性意味着增加资源和工作量,每个 GPU 的效率不会下降。
  • 数学公式: 未提供具体公式,通过观察每个 GPU 的吞吐量随 GPU 数量和任务规模增加的变化来评估。

5.2.4. 强扩展性 (Strong Scaling)

  • 概念定义: 评估当总任务规模(例如,总批次大小)固定不变时,增加计算资源(GPU 数量)能否按比例缩短完成任务所需的时间,或者提高总的吞吐量。理想的强扩展性意味着 GPU 数量增加一倍,完成时间减半,或总吞吐量翻倍。
  • 数学公式: 未提供具体公式,通过观察总吞吐量或完成时间随 GPU 数量增加的变化来评估。

5.3. 对比基线

论文将 WeiPipe-Interleave 与以下主流和最先进的分布式训练策略进行了比较:

  • 1F1B (One-Forward-One-Backward): 一种经典的流水线并行 (Pipeline Parallelism, PP) 策略,通过 Megatron-LM [35] 实现。
  • Zero-bubble 1 (ZB1) & Zero-bubble 2 (ZB2): 最先进的零气泡流水线并行策略,通过 Megatron-LM 实现。
  • FSDP (Fully Sharded Data Parallelism): 全分片数据并行,通过 DeepSpeed [33] 的 ZeRO-3 优化实现。

5.4. 实验环境

  • 硬件: Colossal Cloud 提供的 A800 GPU
    • 每个 A800 GPU 具有 80GB HBM (高带宽内存)。
    • FP16/BF16 张量核心的计算能力为 312 TFlops。
    • A800 的 NVLink 带宽为 400 GB/s (低于 A100 的 600 GB/s)。
  • 通信基础设施: 论文使用了两种不同的通信基础设施配置来评估 WeiPipe 的通信效率:
    1. NVLink + Ethernet 混合连接 (16 GPUs): 16 个 A800 GPU 分布在两个集群中,每个集群内部 GPU 之间通过 NVLink 连接,而集群之间通过 10Gb Ethernet 连接。
    2. NVLink + Ethernet 混合连接 (32 GPUs): 32 个 A800 GPU 分布在四个集群中,每个集群内部 GPU 之间通过 NVLink 连接,而集群之间通过 10Gb Ethernet 连接。
    3. 纯 NVLink 连接 (8 GPUs): 8 个 A800 GPU 仅通过 NVLink 连接(通信约束较小)。
  • 通信库: 采用了 NCCL (NVIDIA Collective Communications Library) 作为底层通信库。NCCL 默认使用基于环形拓扑 (ring-based topology) 的实现来进行集体通信原语 (collective primitives),如 reduce-scatterall-gather。因此,所有实验中的并行策略都维持了环形拓扑。

5.5. 模型配置与训练设置

  • 基础模型: 基于 LLama-2 结构(GPT 风格)的模型。
  • 固定参数: 头部数量 (head number) 32,层数 (layer number) LL 32。
  • 可变参数:
    • 隐藏维度大小 HH:1024, 2048, 4096。
    • 词元序列长度 SS:4096, 8192, 16384。
    • 微批次大小 GG:根据不同实验和模型配置设置,如表 2 中所示,通常为 1 到 16。总批次大小通常为 64 或更高。
    • 模型大小范围:从 384M 到 6.1B 参数。
  • 统一设置: 所有策略都采用相同的模型配置、微批次大小、混合精度设置和 Flash Attention。
  • 重计算: 除了零气泡流水线策略(因为它们使用时不会节省存储反而增加计算开销),所有其他策略都应用了重计算 (recomputation)。

6. 实验结果与分析

6.1. 核心结果分析

下表(原文 Table 2)展示了在 16 个 GPU 使用 NVLink 连接环境下,不同模型配置下,WeiPipe 与其他策略的吞吐量和峰值内存。 原始论文描述:WetestonLlamastructurewith\text{原始论文描述}: We test on Llama-structure with H = 4 0 9 6 S = 4 0 9 6,N = 1 2 G = 3 2, and 32 heads Transformer. Memory distribution is estimated with\alpha = 0 . 6baaaneTe u cul e on implementation.

Model Config Throughput(Tokens/second/GPU) Memory(GB)
H S G 1F1B ZB1 ZB2 FSDP WeiPipe 1F1B ZB1 ZB2 FSDP WeiPipe
1024 4096 16 8581.7 7547.0 7638.5 11525.9 15138.8 13.0 20.4 39.3 8.6 9.4
8192 8 7403.8 6739.6 6768.1 9424.4 12122.3 9.9 10.7 20.5 8.6 9.4
16384 4 5641.2 5651.6 5651.9 6973.6 8188.3 9.1 21.6 42.2 8.6 9.4
2048 4096 16 4163.2 3823.3 OOM 4104.8 6499.7 18.7 44.3 OOM 17.9 19.9
8192 8 3791.3 3517.8 OOM 3706.8 6033.2 19.6 22.3 OOM 17.9 19.9
16384 4 3146.3 3050.1 OOM 3087.2 4607.8 22.9 42.9 OOM 17.9 19.9
4096 4096 16 1662.7 OOM OOM 1110.5 2023.1 40.5 OOM OOM 39 44.5
8192 8 1556.2 OOM OOM 1063.2 2059.4 41.6 OOM OOM 39 44.5
16384 4 1331.6 OOM OOM 944.2 1684.9 45.1 OOM OOM 39 44.5
  • 吞吐量优势: 在所有模型配置下,WeiPipe (WeiPipe-Interleave) 都展现出比其他策略更高的性能。特别是在 H=4096,S=16384H=4096, S=16384 这种极端长上下文场景下,WeiPipe 比 1F1B 吞吐量提高了 22.3% (1684.9 vs 1331.6),比 FSDP 提高了 78.4% (1684.9 vs 944.2)。这验证了 WeiPipe 在通信密集型长上下文训练中的有效性。
  • 内存消耗分析:
    • 1F1B 在所有流水线并行策略中内存占用最小。重计算和 Flash Attention 显著降低了 1F1B 和零气泡策略早期工作节点的内存占用。
    • 然而,零气泡 (Zero-bubble) 方法(ZB1, ZB2)的内存消耗显著高于 1F1B,导致在长上下文场景下容易出现内存不足 (Out-Of-Memory, OOM) 错误(如表格中所示,对于 H=2048H=2048H=4096H=4096 的配置,ZB1 和 ZB2 频繁 OOM)。这与 [32] 中 ZB1 内存不高于 1F1B,ZB2 内存是 1F1B 两倍的结论有所矛盾。论文解释称,原始零气泡方法没有使用 Flash Attention 或重计算,注意力计算产生的激活值在 B Pass 期间会释放,内存开销较小。但当启用 Flash Attention 后,注意力产生的激活值大大减少,此时 FFN 的计算以及 B Pass 产生的梯度成为主要内存消耗,峰值内存出现在最后一个 Rank 的第一个 W Pass 之前,大约是第一个 Rank 峰值内存的两倍。这导致零气泡方法虽然能降低气泡率,但高内存消耗限制了其微批次大小,影响了计算效率。
    • WeiPipe 的内存使用略高于 FSDP,主要原因是 WeiPipe-Interleave 需要额外的通信缓冲区(每个工作节点需要三个缓冲区用于接收 WsDs,三个缓冲区用于发送),而 FSDP 在操作级别创建缓冲区,可能更碎片化且更小。
    • Flash Attention 和重计算的采用,使得传统上使用小微批次的激活值传递 PP 可以使用更大的微批次,从而突出了 WeiPipe 的优势。

6.1.2. PCIe 和 Ethernet 环境下的吞吐量

下表(原文 Table 3)展示了在 16 个 A800 GPU (来自 2 个集群,集群内 NVLink,集群间 Ethernet) 的 PCIe 和 Ethernet 环境下,不同策略的吞吐量。 原始论文描述:TaTuyUVLh.Duetothelimitationofmemory,forZBstrategies,Gissetto4if\text{原始论文描述}: Ta Tu y U VLh. Due to the limitation of memory, for ZB strategies, G is set to 4 if S { = } 4 0 9 6and\mathrm { G } { = } 1ifS { = } 8 1 9 2or 16384.

Model Config Throughput(Tokens/second/GPU) Memory(GB)
H S G 1F1B ZB1 ZB2 FSDP WeiPipe 1F1B ZB1 ZB2 FSDP WeiPipe
1024 4096 16 8581.7 7547.0 7638.5 11525.9 15138.8 13.0 20.4 39.3 8.6 9.4
8192 8 7403.8 6739.6 6768.1 9424.4 12122.3 9.9 10.7 20.5 8.6 9.4
16384 4 5641.2 5651.6 5651.9 6973.6 8188.3 9.1 21.6 42.2 8.6 9.4
2048 4096 16 4163.2 3823.3 OOM 4104.8 6499.7 18.7 44.3 OOM 17.9 19.9
8192 8 3791.3 3517.8 OOM 3706.8 6033.2 19.6 22.3 OOM 17.9 19.9
16384 4 3146.3 3050.1 OOM 4607.8 22.9 42.9 OOM 17.9 19.9
4096 4096 16 1662.7 OOM OOM 1110.5 2023.1 40.5 OOM OOM 39 44.5
8192 8 1556.2 OOM OOM 1063.2 2059.4 41.6 OOM OOM 39 44.5
16384 4 1331.6 OOM OOM 944.2 1684.9 45.1 OOM OOM 39 44.5
  • 通信压力: 这种混合连接环境(集群间 Ethernet)进一步加剧了长上下文训练的通信压力。
  • WeiPipe 优势: 在这种通信受限的环境下,WeiPipe-Interleave 的优势更加明显。当 S=16384S=16384H=2048H=2048 时,WeiPipe 吞吐量为 4607.8,比 FSDP (3087.2) 提升了 49.2%,比 1F1B (3146.3) 提升了 46.5%。当 S=16384S=16384H=4096H=4096 时,WeiPipe 吞吐量为 1684.9,比 FSDP (944.2) 提升了 78.4%,比 1F1B (1331.6) 提升了 26.5%。
  • 低带宽适应性: 结果表明,WeiPipe 使得 LLM 训练对高带宽技术(如 NVLink)的依赖降低。WeiPipe 擅长将通信与计算重叠,即使在激活值传递方法中,通信也受数据依赖限制,而 WeiPipe 可以并发传输权重和计算。

下表(原文 Table 4)展示了在 8 个 A800 GPU 仅通过 NVLink 连接的环境下,不同策略的吞吐量。 原始论文描述:Table4.ThroughputoftrainingtheLLamastylemodelon8GPUswithNVLinkconnectionswithlayernumberas16.ForZBstrategies,Gissetto4if\text{原始论文描述}: Table 4. Throughput of training the LLama-style model on 8 GPUs with NVLink connections with layer number as 16. For ZB strategies, G is set to 4 if S { = } 4 0 9 6and\mathrm { G } { = } 1ifS { = } 8 1 9 2or 16384.

Model Config 1F1B Throughput(Kilo Tokens/second/GPU)
H S 4k 1k G ZB1 ZB2 FSDP WeiPipe
16 32.0 45.8 46.5 37.9 31.3
2k 16k 4 15.9 22.0 22.1 17.8 16.9
4k 16k 16 4 15.0 22.4 9.4 12.8 OOM 17.0 14.2
4k 4k 16 5.2 OOM OOM OOM 10.1 6.0 9.7 4.9
16k 4 3.7 OOM OOM 3.8 3.6
  • 通信不构成瓶颈时: 在纯 NVLink 且 GPU 数量较少(8 个 A800)的环境下,通信约束不那么显著。在此场景下,传统方法可能具有优势。
  • WeiPipe 仍有优势: 尽管通信约束较小,WeiPipe 仍然可以超越当前的 PP 方法,这归因于其计算与通信的良好重叠。

6.1.4. 弱扩展性 (Weak Scaling)

弱扩展性 (Weak Scaling) 评估当 GPU 数量和任务规模(例如,总批次大小随 GPU 数量同比例增加,保持每个 GPU 的工作量恒定)同时增加时,策略如何保持效率。

6.1.4.1. 小规模弱扩展性

下图(原文 Figure 6)展示了小规模弱扩展性实验结果,GPU 数量从 4 扩展到 16 (每服务器 4 个 GPU),总批次大小从 64 增加到 256。 原始论文描述: Figure 6. Small scale weak scaling. The number of GPUs scales from 4 to 16 (4 GPUs in 1 server), and the batch size increases proportionally from 64 to 256. The left coordinates correspond to the overall throughput as represented by the bar chart, while the coordinates on the right indicate tokens per second per GPU.

Figure 6. Small scale weak scaling. The number of GPUs scales from 4 to 16 (4 GPUs in 1 server), and the batch size increases proportionally from 64 to 256. The left coordinates correspond to the overall throughput as represented by the bar chart, while the coordinates on the right indicate tokens per second per GPU. 该图像是图表,展示了在不同 GPU 数量下的弱扩展性性能。X 轴表示 GPU 数量(4、8、16),Y 轴显示每秒处理的千个标记量(Kilo Tokens/s)和每个 GPU 的处理能力(Kilo Tokens/s/GPU)。不同颜色代表不同的方法,WeiPipe 在所有 GPU 配置中表现出更高的吞吐量。

  • WeiPipe 表现: WeiPipe 在 4 到 16 个 GPU 的扩展过程中,每 GPU 的吞吐量 (Tokens/second/GPU) 保持较高水平,显示出良好的弱扩展性。
  • 总吞吐量: WeiPipe 的总吞吐量随着 GPU 数量的增加而显著提高,且在所有 GPU 配置下都高于其他方法。

6.1.4.2. 大规模弱扩展性

下图(原文 Figure 7)展示了大规模弱扩展性实验结果,GPU 数量从 8 扩展到 32 (每服务器 8 个 GPU),总批次大小从 128 增加到 512。 原始论文描述: Figure 7. Large scale weak scaling. The number of GPUs scales from 8 to 32 (8 GPUs in 1 server) with the batch size increasing proportionally from 128 to 512.

Figure 7. Large scale weak scaling. The number of GPUs scales from 8 to 32 (8 GPUs in 1 server) with the batch size increasing proportionally from 128 to 512. 该图像是图表,展示了不同GPU数目(8、16和32)下的训练性能,包括每秒处理的Kilo Tokens数(Kilo Tokens/s)和每个GPU的处理能力(Kilo Tokens/s/GPU)。通过比较1F1B、FSDP和WeiPipe方法,可以看到WeiPipe在GPU数增多时性能提升明显。

  • WeiPipe 表现: 在大规模场景下,WeiPipe 同样展现出卓越的弱扩展性。其每 GPU 吞吐量在 GPU 数量增加时依然保持稳定,且远超 1F1B 和 FSDP。
  • 通信约束的影响: 在长上下文和低带宽通信基础设施下,传统 PP 方法面临严峻挑战,而 WeiPipe 因其通信效率而能更好地应对这些挑战。

6.1.5. 强扩展性 (Strong Scaling)

强扩展性 (Strong Scaling) 评估当总任务规模(例如,总批次大小)固定时,增加 GPU 数量能否按比例缩短完成时间或提高总吞吐量。

6.1.5.1. 小规模强扩展性

下图(原文 Figure 8)展示了小规模强扩展性实验结果,GPU 数量从 4 扩展到 16,总批次大小固定为 128。 原始论文描述: Figure 8. Small scale strong scaling. The number of GPUs scales from 4 to 16, with the batch size remains 128.

Figure 8. Small scale strong scaling. The number of GPUs scales from 4 to 16, with the batch size remains 128. 该图像是图表,展示了在小规模强扩展下,从4到16个GPU的数量变化,批量大小保持为128。不同颜色的条形代表不同的性能指标,表现出在扩展过程中各指标的波动情况。

  • WeiPipe 优势: WeiPipe 在 4 到 16 个 GPU 的扩展过程中,总吞吐量提升最为显著,表明其能更好地利用更多 GPU 来加速固定大小的任务。

6.1.5.2. 大规模强扩展性

下图(原文 Figure 9)展示了大规模强扩展性实验结果,GPU 数量从 8 扩展到 32,总批次大小固定为 256。 原始论文描述: Figure 9. Large scale strong scaling. The number of GPUs scales from 8 to 32, with the batch size remains 256.

Figure 9. Large scale strong scaling. The number of GPUs scales from 8 to 32, with the batch size remains 256. 该图像是一个图表,展示了使用不同数量的GPU(8、16、32)时,每秒处理的Kilo Tokens数,其中WeiPipe方法在32个GPU时表现最佳,达到最高的吞吐量,对比1F1B和FSDP方法的数据。

  • WeiPipe 卓越表现: 在大规模强扩展性测试中,WeiPipe 再次显示出优于其他策略的可扩展性。随着 GPU 数量从 8 增加到 32,WeiPipe 的总吞吐量增长更为明显,表明它有潜力通过更多的 GPU 实现更大的加速比。这进一步证明了 WeiPipe 在处理长上下文模型时,能有效利用更多的计算资源来克服通信瓶颈。

6.2. 数据呈现 (表格)

本节已将原文 Table 2, Table 3 和 Table 4 完整转录并穿插在上述分析中。

6.3. 消融实验/参数分析

论文中没有明确进行传统的消融实验来逐一验证模型各组件(如通信重叠、特定调度等)的贡献。然而,通过改变隐藏维度 HH、序列长度 SS 和微批次大小 GG 等模型配置,并评估在不同通信基础设施下的性能,论文实际上展示了 WeiPipe 在不同参数和通信约束条件下的鲁棒性和性能优势。特别强调了长上下文(大 SS 值)和通信受限环境(Ethernet)对传统方法的影响以及对 WeiPipe 优势的凸显。

7. 总结与思考

7.1. 结论总结

本文提出了 WeiPipe (Weight Pipeline Parallelism),一种新颖的流水线并行策略,旨在通过将模型层间传递的数据从激活值 (activations) 转向权重 (weights) 及其梯度 (gradients),从而有效降低大型语言模型 (LLMs) 训练中的通信开销,尤其是在长上下文 (long context) 场景下。

主要贡献和结论包括:

  • 创新性的权重传递范式: WeiPipe 改变了传统流水线并行中传递激活值的模式,转而传递模型权重及其梯度,从而绕开了长上下文和大数据量微批次带来的激活值通信瓶颈。

  • 通信高效且可扩展: WeiPipe 不依赖于集体通信 (collective communication),而是主要采用点对点通信 (P2P communication),这使其在通信受限的环境下具有出色的可扩展性。

  • WeiPipe-Interleave 的优越性能: 论文详细设计并实现了 WeiPipe-Interleave 策略,通过交错前向和反向传播,有效减少了流水线气泡 (pipeline bubble ratio),平衡了内存使用,并降低了通信需求。实验结果显示,WeiPipe-Interleave 在长上下文 LLM 训练中,相较于最先进的流水线并行方法 (如 1F1B、zero-bubble) 和全分片数据并行 (FSDP),在吞吐量方面取得了 30% 到 80% 的显著提升。

  • 降低对高带宽基础设施的依赖: 实验证明,WeiPipe 在通信带宽较低的环境(如跨集群的 Ethernet 连接)中表现出更强的优势,降低了对昂贵的高带宽互连(如 NVLink)的依赖。

  • WeiPipe-zero-bubble 的潜力: 论文还探讨了结合零气泡流水线思想的 WeiPipe-zero-bubble 策略,展示了实现近乎零气泡流水线的潜力,为未来的研究方向奠定了基础。

    总而言之,WeiPipe 扩展了利用流水线并行训练长上下文模型的可能性,为解决大规模分布式训练中的通信瓶颈提供了新的有效途径。

7.2. 局限性与未来工作

论文作者指出了当前的局限性以及未来可能的研究方向:

  • WeiPipe-zero-bubble 尚未实现: 论文中提出的 WeiPipe-zero-bubble 策略,尽管在理论上展示了实现几乎零气泡的潜力,但由于需要更精细和复杂的控制,尚未进行实际实现和实验验证。这被作者留作未来的探索工作。
  • 未深入探讨其他模型架构: 尽管论文提到其方法不限于 Transformer,但所有实验都在 Llama 风格模型上进行。在其他不同架构的模型(例如 Vision Transformer、生成模型等)上的性能和适应性有待进一步验证。
  • 对通信拓扑的假设: 论文假设并维护了环形拓扑 (ring topology) 用于所有并行策略。在更复杂的集群拓扑(如树状、网格)下,WeiPipe 的性能和通信模式可能需要进一步优化和评估。

7.3. 个人启发与批判

7.3.1. 个人启发

  1. 通信瓶颈的重新审视: 这篇论文给我最大的启发是,在分布式训练中,通信瓶颈并非一成不变地由模型参数同步引起。随着模型结构和训练技术的演进(例如长上下文、Flash Attention 和重计算),层间激活值和激活值梯度的通信反而可能成为新的、更显著的瓶颈。这种对核心通信对象的转变,提供了一个全新的优化思路。
  2. 数据流范式的创新: 从“激活值传递”到“权重传递”的范式转变是极具创意的。它提醒我们,解决问题的关键有时在于挑战基本假设。通过识别哪个数据流对当前任务的通信带宽影响最大,并改变其传输方式,可以带来颠覆性的性能提升。这种思维方式可以推广到其他分布式系统设计中,即不断审视数据流的组成和特点,以优化整体性能。
  3. 对硬件异构性的适应: WeiPipe 在 NVLink 和 Ethernet 混合连接这种通信受限环境下的卓越表现,表明其降低了对昂贵高带宽互连的依赖。这对于在更广泛、更具成本效益的硬件环境中部署和训练大规模模型具有重要意义。它提示我们,设计分布式系统时应充分考虑不同通信层级的特性和限制,并设计能够自适应或规避最弱环节的策略。
  4. 通信与计算的深度重叠: WeiPipe-Interleave 通过巧妙地交错前向和反向传播,并在其中嵌入异步通信,实现了通信与计算的深度重叠。这种细粒度的调度和资源利用是高性能分布式系统设计的核心。

7.3.2. 批判

  1. 零气泡策略的实现与验证: 论文提出了 WeiPipe-zero-bubble 的概念和理论框架,但并未进行实现和实验验证。这使得该部分内容停留在理论层面,其在实际系统中的复杂性、性能表现以及可能引入的新问题(如额外的内存开销或调度难度)仍是未知数。对于一项创新策略,完整的实现和验证是不可或缺的。
  2. 内存消耗的权衡: 论文指出零气泡策略(ZB1、ZB2)在启用 Flash Attention 和重计算后,内存消耗显著增加,甚至导致 OOM。虽然 WeiPipe-Interleave 在内存上优于这些零气泡策略,但其内存使用仍略高于 FSDP。这意味着在追求通信效率的同时,内存仍然是一个需要仔细权衡的关键资源。未来的工作可能需要探索 WeiPipe 与更激进的内存优化技术(如更细粒度的激活值卸载)的结合。
  3. 收敛性和训练稳定性: 论文主要关注吞吐量和内存效率等系统性能指标,但对于模型训练的收敛速度、最终精度、以及在权重传递过程中可能引入的梯度陈旧 (stale gradient) 或其他训练稳定性问题没有深入探讨。尽管论文提到关注同步流水线以保持收敛,但对于这种新颖的权重传递机制,更详细的收敛性分析和经验验证会更有说服力。
  4. 超大规模扩展的验证: 论文在引言中提及 Llama-3.1 需要 16,000 块 H100 GPU,暗示了 WeiPipe 在超大规模场景下的潜力。然而,当前的实验仅在最多 32 块 A800 GPU 上进行。尽管在 32 块 GPU 上展现了良好的扩展性,但与数千甚至上万块 GPU 的真实超大规模训练场景仍有巨大差距。WeiPipe 在极致规模下(例如,网络拓扑更复杂、延迟更高、容错需求更强)的表现,以及其 P2P 通信模式是否能持续保持优势,仍需更全面的验证。
  5. 特定模型结构的假设: 论文的分析和实验主要基于 Llama 风格的 Transformer 模型,并假设了每层 12H212H^2 的权重结构。对于具有不同层结构、不同权重/激活值比率的其他大型模型(例如 MoE 模型、CNN 模型或更复杂的生成式模型),WeiPipe 的相对优势可能需要重新评估。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。