FlexPipe: Maximizing Training Efficiency for Transformer-based Models with Variable-Length Inputs

文章标题：FlexPipe：利用可变长度输入最大化基于Transformer模型的训练效率
作者/机构：Hairui Zhao (吉林大学 & 加州大学河滨分校), Qi Tian (吉林大学), Hongliang Li (吉林大学 & 教育部符号计算与知识工程重点实验室), Zizhong Chen (加州大学河滨分校)

A1 主要贡献

本文针对Transformer模型在处理可变长度输入时训练效率低下的问题，提出了一种从分布式系统角度出发的解决方案FlexPipe。现有方法主要集中在单次迭代内的优化，但忽略了由可变长度输入导致的跨迭代计算和内存需求的巨大波动，这在现有分布式框架的静态分区策略下会导致资源利用率低下。

核心问题：
1. 资源利用率低：可变长度输入导致不同迭代间的计算和内存需求剧烈波动。现有分布式框架采用静态分区策略，通常基于最大序列长度来配置资源，这使得处理较短序列的迭代中大量资源被闲置。如下图1所示，与固定长度数据集相比，在可变长度数据集上训练的吞吐量显著下降。
2. 静态并行策略的局限性：现有流水线并行（PP）框架在整个训练过程中使用固定的并行策略（如并行度组合），无法适应可变序列长度带来的动态资源需求，导致训练效率低下。

研究目标：
设计一个能够动态调整并行策略的灵活流水线并行框架，以适应可变长度输入带来的资源需求变化，从而最大化训练吞吐量和资源利用率。

创新点/主要贡献：
- 提出FlexPipe框架：首次引入一个灵活的流水线并行（PP）框架，通过“实时灵活性机制”（Live Flexibility Mechanism, LFM），能够在不中断训练、不损失精度的情况下动态调整PP配置，从而提高可变长度输入的训练效率。
- 定义新的优化问题：提出了一个新颖的“灵活内存优化问题”，旨在通过动态重新配置并行策略并利用“冗余”GPU来最大化训练效率。
- 设计高效启发式算法：设计了“启发式边界搜索算法”（Heuristic Bound Search Algorithm, HBSA）来解决上述优化问题。该算法通过综合考虑LFM的开销和其他内存优化方法（如重计算、内存虚拟化），在计算效率和内存使用之间取得了良好的平衡。
- 实验验证：在PyTorch上实现了FlexPipe，并通过大量实验证明其有效性。结果显示，与当前最先进（SOTA）的方法相比，FlexPipe的平均训练吞吞吐量提升了1.25倍。

图1：固定长度和可变长度训练的吞吐量对比。
图1：固定长度和可变长度训练的吞吐量对比。

A3 背景知识与动机

2.1 背景知识

Transformer模型与可变长度输入：Transformer模型【43, Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NIPS 2017】已成为各种深度学习应用的基础，其架构通常基于编码器和解码器，由多个同质层组成，每层资源消耗相似。为了执行多任务学习，Transformer模型通常在混合数据集上进行训练，这引入了可变长度的输入。例如，图2(a)展示了混合数据集FLANv2【30, Shayne Longpre, Le Hou, Tu Vu, Albert Webson, Hyung Won Chung, Yi Tay, Denny Zhou, Quoc V. Le, Barret Zoph, Jason Wei, and Adam Roberts. The flan collection: Designing data and methods for effective instruction tuning. ICML 2023】中序列长度的高度变化，随机采样时不同迭代的样本长度差异可达60000。

分布式训练并行技术：由于计算和内存需求的指数级增长，分布式训练变得不可或缺。常用的并行技术包括数据并行（DP）、张量并行（TP）、流水线并行（PP）和序列并行（SP）。本文主要关注PP的弹性，因其通信开销相对较低【12, Chaoyang He, Shen Li, Mahdi Soltanolkotabi, and Salman Avestimehr. Pipetransformer: Automated elastic pipelining for distributed training of large-scale models. ICML 2021】【16, Chenyu Jiang, Zhen Jia, Shuai Zheng, Yida Wang, and Chuan Wu. Dynapipe: Optimizing multi-task training through dynamic pipelines. EuroSys 2024】【54, Yujia Zhai, Chengquan Jiang, Leyuan Wang, Xiaoying Jia, Shang Zhang, Zizhong Chen, Xin Liu, and Yibo Zhu. Bytetransformer: A high-performance transformer boosted for variable-length inputs. IPDPS 2023】，在商业和学术深度学习集群中被广泛采用。在PP中，模型被按层划分为多个阶段（stage）。为了提高并发性，PP将一个mini-batch拆分为多个micro-batch【9, Shiqing Fan, Yi Rong, Chen Meng, Zongyan Cao, Siyu Wang, Zhen Zheng, Chuan Wu, Guoping Long, Jun Yang, Lixue Xia, Lansong Diao, Xiaoyong Liu, and Wei Lin. DAPPLE: a pipelined data parallel approach for training large models. PPoPP 2021】【14, Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. NeurIPS 2019】【16, Chenyu Jiang, Zhen Jia, Shuai Zheng, Yida Wang, and Chuan Wu. Dynapipe: Optimizing multi-task training through dynamic pipelines. EuroSys 2024】进行流水线处理。

2.2 动机

观察1：可变长度训练导致跨迭代的计算和内存资源利用率低下。传统框架通过零填充（zero-padding）来处理可变长度输入，但这会浪费计算和内存。Packing方法【8, Hantian Ding, Zijian Wang, Giovanni Paolini, Varun Kumar, Anoop Deoras, Dan Roth, and Stefano Soatto. Fewer truncations improve language modeling. ICML 2024】【38, Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR 2020】将短样本拼接成一个长样本，但会因自注意力机制的二次方复杂度而带来计算和内存开销，并可能引入截断和交叉污染问题。Bucketing方法【34, Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. NAACL-HLT 2019】则破坏了样本批次的随机性。尽管现有工作在单次迭代内进行了优化（如核融合、非矩阵乘法FLOPs减少），但仍无法解决跨迭代的巨大波动问题。如图2(b)所示，在使用packing方法训练GPT(3.35B)时，由于样本长度变化大，大多数迭代都在使用为处理最长样本而设计的设备，导致平均计算吞吐量和内存利用率分别低于55%和39%。这揭示了从分布式系统角度进行跨迭代优化的机会。

观察2：PP的静态分区导致可变长度训练效率低下。现有的PP框架通常根据数据集中最长的序列长度来静态配置并行策略（如批大小、阶段数），以避免内存溢出（OOM）。这种静态策略对于处理较短序列的迭代是低效的。如图2(c)所示，训练GPT(3.35B)时，不同固定长度的数据集所需的高效训练设备数量是不同的。例如，对于长度为3k的样本，使用3个GPU的训练时间最短（4.7秒/迭代），而使用8个GPU则因为过度细粒度的分区导致通信开销增大和GPU资源利用率下降，训练时间反而增加到7.08秒/迭代。在可变长度训练中，大部分迭代处理的是短样本（如FLANv2中95%的迭代最大样本长度低于4k），这使得静态策略的性能下降问题更加严重。这启发我们探索跨迭代动态重构并行策略，并利用“冗余”GPU来加速训练。

图2：训练GPT(3.35B)模型时处理可变长度输入的观察。(a) FLANv2数据集中的序列长度分布。(b) 不同序列长度下的计算效率和内存使用情况。(c) 使用不同数量GPU训练时的训练时间。
图2：训练GPT(3.35B)模型时处理可变长度输入的观察。(a) FLANv2数据集中的序列长度分布。(b) 不同序列长度下的计算效率和内存使用情况。(c) 使用不同数量GPU训练时的训练时间。

2.3 挑战

  1. 动态调整PP的巨大开销：现有分布式框架通过“暂停-恢复”或“迭代间停顿”机制进行调整，这会引入初始化和通信的巨大开销，停顿时间甚至可能超过一次迭代的持续时间。
  2. 调整时机与策略选择的复杂性：频繁的重构或不当的并行策略选择会降低性能。因此，何时调整PP以及如何找到一个能最大化吞吐量的合适并行策略，是一个关键挑战。

A2 方法细节

3 FlexPipe系统设计

3.1 概览

FlexPipe的核心思想是在训练期间根据每次迭代中输入的变量长度动态调整并行策略。它引入了一个灵活的流水线并行（PP）框架，能够根据迭代间的计算和内存需求选择合适的设备数量并配置PP。同时，FlexPipe将闲置的GPU用于增加数据并行（DP）的程度，从而提高资源利用率和训练效率。图3以一个使用4个A100 GPU训练GPT(3.35B)的例子说明了FlexPipe的工作方式。

图3：FlexPipe在扩增和缩减阶段时的示例。图下方的条形图比较了内存使用（M）和训练吞吐量（T）。E代表使用FlexPipe，S代表使用基于4K序列长度的静态分区（即图3(a)的情况）。
图3：FlexPipe在扩增和缩减阶段时的示例。图下方的条形图比较了内存使用（M）和训练吞吐量（T）。E代表使用FlexPipe，S代表使用基于4K序列长度的静态分区（即图3(a)的情况）。

3.2 设计

FlexPipe被设计为介于典型训练系统（如FlashAttention）和其底层执行引擎（如PyTorch）之间的中间件系统。它能在线动态调整PP阶段的数量和配置。其架构如图4所示，包含三个主要模块：

图4：FlexPipe系统架构概览。
图4：FlexPipe系统架构概览。

4 实时灵活性机制 (Live Flexibility Mechanism)

实现快速且实时的灵活性机制是FlexPipe的核心挑战。现有PP框架不支持运行时调整PP阶段数量和配置，依赖于“暂停-恢复”机制，这会带来包括检查点、引导和数据加载在内的沉重开销。即使是在迭代间隙进行调整，初始化和通信开销仍然会导致长时间的停顿。FlexPipe设计的实时灵活性机制（LFM）旨在通过减少初始化开销和重叠通信与计算时间，来透明地调整PP阶段数量和配置，而无需停顿。

LFM的核心思想是引入“孪生层”（TwinLayer）。每个节点的主机内存中存储了该节点内所有层的副本及其对应的优化器状态。这带来了两个优势：首先，孪生层支持细粒度的层级管理，避免了在调整时需要先将阶段（stage）分解为层、跨设备传输、再在目标设备上重组的巨大开销（例如，对GPT(3.35B)模型，此过程开销为1.2秒）。其次，通过在主机内存中存储副本，减少了频繁的内存分配和释放开销。

LFM通过孪生层管理器（TwinLayer Manager）高效地协调通信和计算操作。
- 减少PP阶段（Shrinking）：如图5所示，当需要减少一个PP阶段时，空闲出的GPU（如GPU2和GPU3）将被用于增加DP的程度。决策在前一个迭代开始时做出。数据重分布和PP/DP的初始化过程与数据传输过程重叠。在前向传播（BP）步骤之后，微批次间的依赖关系（执行顺序）被修改。上一个迭代更新完成后，GPU2和GPU3将其更新后的L4和L5层参数发送到主机内存的孪生层中，然后分别从孪生层中复制L5和L4的参数。这个复制过程是流水线化的，参数一部分到达后即可开始复制。之后，下一个迭代以新的配置开始。

图5：在具有4个GPU的节点内使用FlexPipe增加一个阶段的示例。
图5：在具有4个GPU的节点内使用FlexPipe增加一个阶段的示例。

图6：在具有4个GPU的节点内使用FlexPipe移除一个阶段的示例。
图6：在具有4个GPU的节点内使用FlexPipe移除一个阶段的示例。

FlexPipe还支持层在设备间的迁移以实现灵活的重分区。如图6所示，L3层从GPU1迁移到GPU2以满足内存需求。利用PP的特性（BP需要FP产生的激活值），只要内存允许，参数传输可以在前一个迭代中进行。决策做出后，参数立即开始传输，而已在GPU1上计算出的微批次的激活值也被传输到GPU2。剩余微批次的前向和后向传播将在GPU2上继续进行。评估表明，非实时的迁移方法平均执行停顿为7.16秒，而FlexPipe仅需0.79秒。跨服务器的LFM与服务器内类似，但由于节点间带宽较低，FlexPipe会通过优先进行节点内重映射和在必要时预传输关键数据来最小化跨节点传输。

5 高效的可变长度训练

5.1 问题建模

系统模型：一个由$L = \{l_1, l_2, ..., l_{|L|}\}$个Transformer层组成的模型，被划分为多个阶段$S = \{s_1, s_2, ..., s_{|S|}\}$。这些阶段被部署在多个设备组$G = \{g_1, g_2, ..., g_{|G|}\}$上。

训练时间估计：给定一个阶段到设备组的映射$S\Pi_G = \{s_1 \rightarrow g_1, ..., s_{|S|} \rightarrow g_{|G|}\}$，训练时间$T(S\Pi_G, n_{micro})$可由以下公式表示：

图片描述
图片描述

其中，$n_{micro}$是微批次数，$t_i$是PP单个操作的时间，包括FP和BP计算以及设备间通信。

内存使用估计：为了有效利用“冗余”GPU，需要估计每次迭代的峰值内存消耗。令$B$, $\mu$, 和$M$分别表示全局批大小、最大序列长度和单个GPU的内存容量。峰值内存$M_{peak}$计算如下：

图片描述
图片描述

其中，$M_{p,g,o}$是参数、梯度和优化器状态的内存，$M_{act}$和$M_{buf}$是激活值和其他缓冲区的内存。

灵活内存优化问题 (FMOP)：为缓解内存压力，LFM并非总是最优选择，还需综合考虑其他内存优化方法，如重计算【5, Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/1604.06174, 2016】【14, Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. NeurIPS 2019】或内存虚拟化【39, Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W. Keckler. vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design. MICRO 2016】【57, Shixiong Zhao, Fanxin Li, Xusheng Chen, Xiuxian Guan, Jianyu Jiang, Dong Huang, Yuhao Qing, Sen Wang, Peng Wang, Gong Zhang, Cheng Li, Ping Luo, and Heming Cui. vpipe: A virtualized acceleration system for achieving efficient and scalable pipeline parallel DNN training. IEEE Trans. Parallel Distributed Syst., 33(3):489–506, 2022】。我们定义了一系列条件来决定应用哪种机制：
- (A) 新配置SΠG的峰值内存小于等于当前设备组G的总内存：

图片描述
图片描述

- (B) 从S0ΠG0切换到SΠG的开销加上新配置的计算时间大于等于原配置的计算时间：
图片描述
图片描述

- (C) 在原配置S0ΠG0下使用内存优化方案$O_{plan}$后的内存小于等于当前设备组G0的总内存：
图片描述
图片描述

- (D) 切换开销加上新配置的计算时间大于等于原配置计算时间与内存优化（重计算$T_{rc}$和卸载$T_{swap}$）开销之和：
图片描述
图片描述

迭代时间$T$表示为：

图片描述
图片描述

其中条件 $E = (\neg A \land \neg D) \lor (A \land \neg B)$。

定义1. 灵活内存优化问题 (FMOP)：给定模型$L$、GPU集群$D$、当前映射$S_0 \Pi_{G_0}$及下一迭代的最大序列长度$\mu$，FMOP旨在寻找一个合适的内存优化方案$O_{plan}$和阶段到设备组的映射$S\Pi_G$，使得在考虑可变长度数据输入时，总迭代时间$T$最小化。

图片描述
图片描述

图片描述
图片描述

5.2 启发式边界搜索算法 (HBSA)

由于FMOP是NP-hard问题，我们设计了启发式边界搜索算法（HBSA）来高效求解。该算法包含三部分：
1. 边界计算 (算法伪代码行3-11)：首先，通过穷举搜索确定满足峰值内存$M_{peak}(L,B,\mu)$所需的最小设备数，即阶段数$N_{stage}$。然后，基于$N_{stage}$与当前阶段数$\|S_0\|$的比较，计算每个阶段中层数的边界（$[n_l^{low}, n_l^{up}]$）和每个设备组中GPU数量的边界（$[n_g^{low}, n_g^{up}]$）。这些边界基于Transformer模型的同质性和分区均衡的启发式思想，有效减少了搜索空间。
2. 导航高效映射 (行14-26)：利用计算出的边界作为递归搜索的起止条件，来寻找高效的阶段到设备的映射。Search函数通过在边界内对模型层$L$和设备$D$进行分区，避免了因OOM或LFM开销过大而导致的无效和低效映射。对于边界内的候选映射，通过公式1进行成本建模以估计训练时间，从而找到最优映射。
3. 确定触发时机 (行28-37)：一旦找到有效映射，利用5.1节中定义的条件A, B, C, D来决定是否触发LFM。如图7所示，HBSA会选择四种方案之一：①保持原映射，不进行内存优化；②在原映射下应用内存优化技术；③增加阶段数（grow）；④减少阶段数（shrink）。例如，当条件A和B都满足时，不采取任何行动（①）。当条件A不满足时，必须采用内存优化（②），此时再根据条件C和D判断是否需要同时触发LFM（③）。经验表明，在训练GPT-3.35B时，系统平均每37次迭代触发一次伸缩（③或④），而内存优化（②）则作为主要优化手段，平均每4次迭代执行一次。

算法伪代码

输入: 模型 L, GPU集群 D, 时间估计函数, 内存估计函数, 最大序列长度 μ. 当前状态: S0ΠG0, μ0.
输出: 内存优化方案 Oplan, 阶段到组的映射 SΠG.
1: 初始化: Oplan = ∅, S = ∅, G = ∅.
2: // 穷举搜索.
3: for i = 1 to ||D|| do
4:     if Mpeak(B, μ) < i × M then
5:         Nstage = i
6:         break
7: // 计算边界.
...
13: // 导航高效映射.
14: function Search(L, D, St, Gt, ν)
15:     for i = nllow to nlup do
16:         for j = nglow to ngup do
17:             st = Partition(L, i), gt = Partition(D, j)
18:             Lt = L − st, Dt = D − gt
19:             if ||Lt|| > nlup ∧ ||Dt|| > ngup then
20:                 Search(Lt, Dt, St ∪ {st}, Gt ∪ {gt}, ν)
21:             if ||Lt|| ∈ [nllow, nlup] ∧ ||Dt|| ∈ [nglow, ngup] then
22:                 Sˆ = St ∪ {st} ∪ {Lt}
23:                 Gˆ = Gt ∪ {gt} ∪ {Dt}
24:                 if T(SˆΠGˆ) + T(S0ΠG0 → SˆΠGˆ) < ν then
25:                     S = Sˆ, G = Gˆ
26:                     ν = T(SˆΠGˆ) + T(S0ΠG0 → SˆΠGˆ)
27: // 确定触发时机.
28: if A ∧ B then
29:     return Oplan, S0ΠG0  // ①
30: else if A ∧ ¬B then
31:     return Oplan, SΠG  // ④
32: if ¬A ∧ ¬B then
33:     Oplan = init(L, S0ΠG0)
34:     if C ∧ D then
35:         return Oplan, S0ΠG0  // ②
36:     else if ¬D then
37:         return Oplan, SΠG  // ③

图7：结合重计算和内存虚拟化的灵活性触发条件。
图7：结合重计算和内存虚拟化的灵活性触发条件。

6 系统实现

FlexPipe使用8K行Python和2K行C++/CUDA代码实现。
- 数据预处理：实现数据预取策略，提前获取下一个全局批次并计算其最大序列长度，以避免数据输入停顿。
- 模型重构：设计的TwinLayer通过直接引用张量地址而非复制张量本身来最小化管理开销。TwinLayer Manager在每个服务器上作为独立进程运行，使用C++线程以避免Python的全局解释器锁（GIL）限制。
- 通信传输：使用PyTorch的DDP（基于NCCL后端）在GPU间和GPU-主机间传输张量。权重和梯度被分区并以流水线方式传输。通过使用cudaMemcpyAsync和isend等异步接口实现非阻塞数据传输，以消除实时灵活性对训练的潜在干扰。

A4 实验环境与结果

实验环境

表1: 用于评估的Transformer模型

表1：用于评估的Transformer模型
表1：用于评估的Transformer模型

实验结果

不同最大序列长度下的性能

如图8所示，在不同最大序列长度下，FlexPipe的平均训练吞吐量比Zero-Bubble高40.4%，比FlashAttention高22.7%，比DynaPipe高13.9%。随着最大序列长度增加，所有方法的吞吐量都下降，但FlexPipe的性能下降更为平缓，因为它能更有效地处理由更大序列长度范围带来的迭代间波动。

图8：不同最大序列长度下的平均训练吞吐量。
图8：不同最大序列长度下的平均训练吞吐量。

不同全局批大小下的性能

如图9所示，随着全局批大小的增加，所有方法的性能都得到提升。FlexPipe和DynaPipe的性能增长速度更快，因为它们都支持根据迭代间的可变长度动态调整分布式训练策略。对于更大的模型，FlexPipe的性能优势更明显，例如，在GPT(13B)上性能提升57%，而在BERT24上为9%。这是因为大模型在可变长度输入下波动更大，为FlexPipe提供了更多优化机会，且更高的计算通信比使得FlexPipe能更好地重叠计算与通信。综合来看，FlexPipe在不同配置下比SOTA方法平均吞吐量高1.25倍。

图9：不同全局批大小下的平均训练吞吐量。
图9：不同全局批大小下的平均训练吞吐量。

灵活性开销

如图10所示，与传统的调整方法（Suspend-Resume、Iteration Stalling）和移除TwinLayer的降级版FlexPipe（Flex w/o TL）相比，FlexPipe分别减少了35%、23.1%和18.8%的训练时间。这是因为TwinLayer机制和精细的计算通信重叠优化了动态调整流水线的初始化时间。图10(c)显示，在跨节点环境中，FlexPipe的开销增长最小。单次重构的通信开销约为0.95秒（GPT-3.35B），这个成本可以被后续多次迭代的吞吐量增益所摊销。

图10：与两种常用方法和降级版FlexPipe的灵活性开销对比。(c) 不同环境下的详细开销：intra-node（单节点4GPU），inter-node（4节点各1GPU），mixed-node（2节点各2GPU）。
图10：与两种常用方法和降级版FlexPipe的灵活性开销对比。(c) 不同环境下的详细开销：intra-node（单节点4GPU），inter-node（4节点各1GPU），mixed-node（2节点各2GPU）。

A7 补充细节

7.3.1 FlexPipe的并行策略

动态调整策略：表2展示了使用FlexPipe训练BERT96和GPT(13B)时的几个代表性迭代的详细信息。与传统框架使用固定数量的GPU（如16个GPU训练GPT-3.35B）不同，FlexPipe根据内存需求动态调整其并行策略。例如，对于序列长度为8192的输入，它使用8个流水线阶段；而对于长度为512的输入，则使用4个阶段，并有效利用“冗余”GPU来提升训练吞吐量。

策略选择的复杂性：即使峰值内存需求相同，流水线阶段数也可能不同。例如，在训练GPT(3.35B)、序列长度为2560的两个不同迭代中，FlexPipe采用了不同的并行策略。这是因为它会综合考虑前一迭代的状态以及其他内存优化技术（如重计算和交换）的开销与收益，以做出最优决策。

收敛性：FlexPipe旨在加速每个训练迭代同时保持严格的训练语义，因此重构不会影响训练的稳定性和模型收敛所需的迭代次数。更高的训练吞吐量反而能让训练更快完成。

表2：使用FlexPipe训练BERT96和GPT(3.35B)的配置和训练详情。μ和Mpeak分别表示序列长度和峰值内存使用。并行策略表示PP每个阶段的DP度数。Oplan(r, s)表示训练期间重计算r GB的激活值并卸载s GB的张量。
表2：使用FlexPipe训练BERT96和GPT(3.35B)的配置和训练详情。μ和Mpeak分别表示序列长度和峰值内存使用。并行策略表示PP每个阶段的DP度数。Oplan(r, s)表示训练期间重计算r GB的激活值并卸载s GB的张量。

7.3.2 算法有效性

HBSA算法的消融研究：如图11所示，通过与修改版FlexPipe进行比较，评估了HBSA的有效性。
- 与Flex w/o M的比较：图11(a)显示，FlexPipe比仅实现灵活性机制而不包括重计算和内存虚拟化的Flex w/o M性能高出32.3%。这表明，在内存波动适中时，结合轻量级的内存优化技术可以避免频繁且开销大的灵活性调整，从而提高效率。
- 与HBSA w/o BD的比较：图11(b)显示，HBSA（平均开销15ms）相比于基于全局空间进行暴力搜索的HBSA w/o BD（平均开销745ms），显著降低了决策开销。HBSA的计算开销随系统规模线性增长，是可接受的，并且由于序列长度预取，其执行可以与其他计算重叠，进一步减少了重构决策的延迟。

图11：算法的消融研究。
图11：算法的消融研究。

A5 结论

本文提出了FlexPipe，一个灵活的流水线并行（PP）框架，通过其“实时灵活性机制”（LFM）显著提升了可变长度输入的训练效率。我们定义了一个新颖的“灵活内存优化问题”（FMOP），并通过一个高效的“启发式边界搜索算法”（HBSA）来求解。该算法能够动态地重新配置并行策略，利用“冗余”GPU，并综合考虑LFM与其他内存优化方法的开销。实验结果表明，与现有的SOTA框架相比，FlexPipe能够实现平均1.25倍的训练吞吐量提升。