使用 Data-Centric Parallel 训练可变长序列

# 使用 Data-Centric Parallel 训练变长序列

Source: https://arxiv.org/html/2608.07524
Geng Zhang, Xuanlei Zhao∗, Kai Wang†, Yang You† 新加坡国立大学 \{zhangg,xuanlei,kai\.wang,youy\}@comp\.nus\.edu\.sg

###### 摘要

在变长序列上训练深度学习模型会带来巨大的计算挑战。现有方法被迫在效率和易用性之间进行艰难的权衡。简单的方法使用静态配置，会导致工作负载不均衡、效率低下；而复杂的方法则会给新模型带来显著的复杂性和代码改动。为了打破这种权衡，我们引入了 Data-Centric Parallel（DCP）。其核心原则是让数据本身驱动运行时。它通过根据每个批次的序列长度动态调整直接运行时设置（例如并行度、梯度累积、重计算）来实现这一点。实验结果表明，我们的方法在 32 块 H200 GPU 上最高可取得 2.88× 的加速。该方法专为泛化而设计，只需 10 行代码即可集成到任何模型中。我们预计这种简单而有效的方法将作为一个稳健的基线，促进变长序列分布式训练未来的发展。

## 1 引言

参见图1：变长序列训练的并行方法比较，包括 bucket parallel、packed parallel 和 data-centric parallel（本文方法）。Di指第i个设备。

处理长序列的能力是越来越多的深度学习应用的关键驱动力。这一趋势在多个领域都很明显，包括视频生成 (Brookset al.,2024 (https://arxiv.org/html/2608.07524#bib.bib8); Zhenget al.,2024 (https://arxiv.org/html/2608.07524#bib.bib1); Polyaket al.,2024 (https://arxiv.org/html/2608.07524#bib.bib45); Konget al.,2024 (https://arxiv.org/html/2608.07524#bib.bib56))、图像生成 (Esseret al.,2024 (https://arxiv.org/html/2608.07524#bib.bib18))、多模态感知 (Chenet al.,2024b (https://arxiv.org/html/2608.07524#bib.bib3); Wanget al.,2024 (https://arxiv.org/html/2608.07524#bib.bib2))、文本生成 (Touvronet al.,2023 (https://arxiv.org/html/2608.07524#bib.bib5); Baiet al.,2023 (https://arxiv.org/html/2608.07524#bib.bib4)) 以及科学计算 (Jumperet al.,2021 (https://arxiv.org/html/2608.07524#bib.bib6))。

然而，在这类数据上训练会带来巨大的计算挑战：1) 长序列：序列的大量长度消耗大量 GPU 内存，需要序列并行将一条序列分割到多个设备以降低内存成本。2) 变长：如图4 (https://arxiv.org/html/2608.07524#S2.F4) 所示，序列长度存在固有的较大差异，在分布式训练中会导致严重的工作负载不均匀，尤其是与序列并行结合时。

针对变长序列训练的并行方法可分为三类，如图1 (https://arxiv.org/html/2608.07524#S1.F1) 所示。Bucket parallel (Esseret al.,2024 (https://arxiv.org/html/2608.07524#bib.bib18)) 对所有序列设置固定的并行度，并通过降低慢批次的批大小来减少不均匀性，如图2 (https://arxiv.org/html/2608.07524#S1.F2) 所示。然而，这种简单的解决方案引入了两个问题。首先，对于很长的序列，批大小通常本来就受到限制，几乎没有调整空间来实现期望的负载均衡。其次，它在降低批大小的时候没有考虑短序列的计算效率，导致明显的速度损失。

参见图2：不同序列长度下 bucket parallel 的工作负载均衡情况。虚线表示批大小。

Packed parallel (Dehghaniet al.,2023 (https://arxiv.org/html/2608.07524#bib.bib17)) 通过将多个序列打包进一个批次来改善负载均衡，而不是减少批次，但这会引入通信开销，并且需要对序列级操作进行额外修改。基于编译器的并行方法 (Geet al.,2024 (https://arxiv.org/html/2608.07524#bib.bib53)) 应用编译器自动优化高效计划，但这需要大量的代码更改，并且难以适配新模型。

现有方法在效率和易用性之间存在艰难的权衡。我们认为，阻碍简单方法发挥效果的核心限制在于它们依赖预定义的运行时设置（例如并行度），这显著限制了通信成本和工作负载均衡的优化空间。通过直接正视这一点，无需更重量级的系统即可实现高效率。这引出了我们工作的核心问题：我们如何让数据本身以简单而有效的方式驱动运行时？

为了应对这一挑战，我们提出了 Data-Centric Parallel（DCP），这是第一种同时实现高效和易用的变长序列训练方法。DCP 不使用固定设置，而是根据每个数据的序列长度动态调整运行时设置，如并行策略、梯度累积和重计算。通过最小化每个批次的通信成本并平衡工作负载，DCP 显著提高了训练吞吐量。

DCP 由两个策略组成：DCP-inter 和 DCP-intra。DCP-inter 利用梯度累积来平衡工作负载，而不是减少批大小。DCP-intra 通过最小化重计算进一步加速。额外的内存成本则通过仔细调整序列并行度和批大小以可忽略的代价来解决。

实验结果表明，DCP 在 32 块 H200 GPU 上、2 个模型、3 个数据集上有效提高了变长序列训练的吞吐量，最高达 2.88×。DCP 专为泛化而设计，只需最多 10 行代码更改即可适配任何模型，如附录D (https://arxiv.org/html/2608.07524#A4) 所示。我们预计这个简单而有效的 DCP 将作为稳健的基线，促进变长序列分布式训练未来的发展。

## 2 变长序列训练中的挑战

参见图3：Panda-80M 的序列长度分布可视化。
参见图4：基于 bucket parallel 的 140 个数据集的不均衡比率分析。

### 2.1 数据长度的巨大差异

主要挑战是真实数据集中序列长度具有很大的方差。例如，图4 (https://arxiv.org/html/2608.07524#S2.F4) 用包含 7000 万个视频的 Panda-70M (Chenet al.,2024a (https://arxiv.org/html/2608.07524#bib.bib14)) 数据集说明了这一点。因此，训练系统必须处理工作负载的极端多样性。

一种简单的解决方案是将长度相近的数据分组到一个训练步骤中。然而，这种方法会损害模型质量，因为它破坏了稳定收敛所必需的独立同分布采样 (Waltz,1984 (https://arxiv.org/html/2608.07524#bib.bib16); Brookset al.,2024 (https://arxiv.org/html/2608.07524#bib.bib8))。因此，一个有效的训练系统必须在不牺牲统计效率的情况下处理这种极端的长度差异。

参见图5：Transformer-1D 在不同序列长度下，不同序列并行度的弱扩展通信开销。

### 2.2 序列并行的通信成本

对于数据集中最长的序列，需要较大的 SP（序列并行）规模以避免内存不足错误。这个最大所需并行度往往决定了整个训练过程的全局 SP 配置。

#### 长序列必须有大的 SP

对于最长序列，较大的 SP 规模不可避免以避免内存不足错误，并决定了全局序列并行度。

#### 大的 SP 对较短序列有害

对于通常占数据集主导地位的较短序列，这个大的 SP 成为主要瓶颈，如图5 (https://arxiv.org/html/2608.07524#S2.F5) 所示。对于这些序列，每个设备上的计算较少，但在设备间交换信息所需的通信固定成本仍然很高。这种通信开销很容易主导总处理时间，大幅降低训练效率。

这就形成了一个困境：大的固定 SP 规模对常见的短序列是一种浪费，而小规模又无法容纳必要的长序列。

### 2.3 分布式训练的工作负载不均匀

当使用数据并行（DP）时，序列长度的高方差会导致严重的工作负载不均匀，如图4 (https://arxiv.org/html/2608.07524#S2.F4) 所示。在 DP 系统中，每个工作进程处理不同的数据批次，但所有工作进程必须在开始下一步之前进行同步。

由于长度差异，被分配了长序列批次的工作进程完成计算所需的时间会显著更长。相反，处理短序列的工作进程很快完成，被迫闲置等待最慢的工作进程完成其任务。这导致计算资源的严重利用不足，因为昂贵的加速器在空转等待。整体训练吞吐量受限于每一步中最慢的工作进程，并且随着集群中工作进程数量的增加，这个问题会加剧。

### 2.4 均衡与通信之间的权衡

一种直观的缓解工作负载不均匀的方法是静态调整批大小：对长序列使用较小的批，对短序列使用较大的批，以均衡每一步的处理时间。然而，这种策略是无效的，并且引入了代价高昂的新权衡。

#### 长序列的小批大小

对于长序列，内存约束已经迫使批大小达到最小值，因此无法进一步减小。即使长序列只有少数几个批次，一旦缩小批大小，对于所有其他序列，也需要使用比 GPU 可容纳的更小的批大小，这导致昂贵硬件的利用率不足。

#### 减小批大小带来的通信开销

更关键的是，这种方法会产生新的瓶颈。通过使用较小的批次来均衡每个批次的工作负载，处理整个数据集所需的训练步数大幅增加。由于每次迭代都需要执行一次代价高昂的全局同步步骤来收集所有参数的梯度，这种策略将每迭代的工作负载不均匀换成了总通信成本的巨大增加。其结果往往是没有任何净性能提升，甚至出现回退。

要点：工作负载不均匀、通信成本和计算效率是有效训练变长序列的三个关键因素。

## 3 Data-Centric Parallel

### 3.1 概述

参见图6：DCP 的概述。它首先通过快速双层剖析收集速度和内存成本。然后根据序列长度动态调整并行度、批大小和梯度检查点以最大化效率。

如图6 (https://arxiv.org/html/2608.07524#S3.F6) 所示，Data-Centric Parallel（DCP）动态调整并行、批大小和梯度检查点以最大化效率。具体来说，我们首先将不同序列分组为若干个大小相近的组。对每个组，我们使用快速剖析来获取不同批大小和序列并行度下的速度和内存成本。然后，我们应用以下两种策略之一，根据数据动态调整设置：

#### DCP-inter

通过梯度累积优化负载均衡。在每次迭代之前，它根据剖析为传入批次的序列长度选择最优批大小和序列并行度。这种方法不是减少批大小，而是以零成本实现更好的均衡。

#### DCP-intra

进一步利用动态重计算来获得更快的速度。这是基于我们的分析：对于短序列，梯度检查点会引入相当多的不必要计算开销。基于这一洞察，DCP-intra 实施了一种战略性权衡：它部分停用梯度检查点。为了控制内存消耗的增加，它会以可忽略的代价动态调整序列并行度和批大小。

### 3.2 问题建模

给定nn个批次\{d1,d2,...,dn\}\\\{d\_\{1\},d\_\{2\},\.\.\.,d\_\{n\}\\\}，其中每个训练迭代的序列长度为\{s1,s2,...,sn\}\\\{s\_\{1\},s\_\{2\},\.\.\.,s\_\{n\}\\\}，DCP 根据各批次的序列长度动态调整其训练配置以提高训练吞吐量。关键是联合提高吞吐量并均衡不同序列长度在每次训练迭代中的执行时间。形式上，这转化为两个优化目标：

max（bi∗sipi∗T（di）），s.t.∑i=1npi=W,\displaystyle\text{max}(\frac{b_{i}*s_{i}}{p_{i}*T(d_{i})}),\text{s.t.}\sum_{i=1}^{n}{p_{i}}=W, (1)

min∑i=1n（maxj=1nT（dj）−T（di））/W.\displaystyle\text{min}\sum_{i=1}^{n}{(\text{max}_{j=1}^{n}T(d_{j})-T(d_{i}))/W}. (2)

其中bib\_\{i\}、pip\_\{i\}和T（ni）T(n\_\{i\})分别是nin\_\{i\}的批大小、序列并行度和执行时间，WW是 GPU 总数。

### 3.3 DCP-INTER

DCP-inter 首先根据剖析为每种序列长度确定最佳并行设置。然后，它利用梯度累积来均衡每次迭代中各数据批次的执行时间，而不是减少批大小。这种简单而有效的策略保持了慢批次的高吞吐量，并通过运行多个批次来填充快批次的空闲时间。

给定nn个批次，DCP-inter 旨在根据T（di）T(d\_\{i\})为每个批次寻找累积步数gig\_\{i\}，以实现均衡和吞吐量。T（di）T(d\_\{i\})是通过快速剖析获得的，使得每个bib\_\{i\}的吞吐量最优。为了实现最优工作负载均衡，方程2 (https://arxiv.org/html/2608.07524#S3.E2) 可改写为：

min∑i=1n（maxj=1n（gj∗T（dj））−gi∗T（di））/W.\displaystyle\text{min}\sum_{i=1}^{n}{(\text{max}_{j=1}^{n}(g_{j}*T(d_{j}))-g_{i}*T(d_{i}))/W}. (3)

为了搜索满足这一目标的gig\_\{i\}，DCP-inter 遍历每个批次所有可能的累积步数，并使用方程3 (https://arxiv.org/html/2608.07524#S3.E3) 作为性能模型评估nn个批次的gig\_\{i\}每种组合，以在每次迭代中为这nn个批次找到最优配置，详见附录C (https://arxiv.org/html/2608.07524#A3)。

### 3.4 DCP-INTRA

参见图7：DCP-intra 的概述。它对较短的序列使用更少的重计算，并增加序列并行度以应对额外的内存成本。

对于长序列训练，通常会完全应用激活检查点策略以减少 GPU 内存成本 (Yuanet al.,2024 (https://arxiv.org/html/2608.07524#bib.bib47))。然而，在变长序列情况下，我们有空闲 GPU 内存，可以对较短的序列启用更少的重计算层以减少执行时间T（bi）T(b\_\{i\})。此外，序列并行度越大，每个 GPU 的内存成本就越低，从而为减少重计算留出更多空间，以提升方程1 (https://arxiv.org/html/2608.07524#S3.E1) 中的吞吐量。基于这两个观察，我们提出了 DCP-intra。

对于一个LL层模型，bib\_\{i\}的执行时间T（bi）T(b\_\{i\})和内存开销M（bi）M(b\_\{i\})为：

T（di）=Tf（di）+Tb（di）+Tf′（di）+TC,\displaystyle T(d_{i})