南京大学提出UniAP算法：加速大模型分布式训练，降低训练成本

南京大学提出高效能分布式训练算法 UniAP：加速大模型训练，降低训练成本

图片

该研究由南京大学计算机学院李武军教授课题组完成，李武军教授为通讯作者，硕士生林昊（已毕业，现就职于阿里巴巴）、吴轲、李杰为共同第一作者，博士生李俊为参与作者。

大模型训练成本居高不下，已成为制约人工智能可持续发展的一大瓶颈。究其原因，大模型训练通常需要借助多机多卡的分布式训练，而其中的挑战却十分巨大。即使拥有足够的硬件资源，对于不熟悉分布式训练的研究者来说，也很可能因为超参数设置不当（例如模型和数据的切分、排布方式等），导致训练过程无法顺利进行。实验数据显示，这种失败的概率高达 64%-87%。

此外，很多研究者在遇到大模型训练速度慢的问题时，往往只想到增加 GPU 数量等横向扩展（scale-out）的方法，而忽略了分布式训练算法本身在纵向扩展（scale-up）方面的潜力。

事实上，分布式训练算法对硬件算力的利用率有着显著影响。高效的分布式训练算法能够充分发挥硬件的性能。在同等硬件算力条件下，使用高效能算法训练同一个模型，速度可能比低效能算法快数倍甚至数十倍，从而大幅降低算力成本。然而，目前已有的许多分布式训练算法效率低下，甚至可能出现机器和 GPU 越多，训练速度反而越慢的怪象。

针对上述问题，南京大学计算机学院李武军教授课题组研发了一种高效能分布式训练算法 UniAP，并基于 UniAP 研发了相应的大模型分布式训练平台和框架。

图片

  • 论文标题： UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming
  • 论文地址： https://arxiv.org/abs/2307.16375

UniAP 的核心创新在于，它首次实现了层内并行策略（如张量并行）和层间并行策略（如流水线并行）的联合优化。对于给定的模型和硬件平台，UniAP 能够通过自动搜索，找到高效的分布式训练方案。这不仅解决了训练效率和成本问题（实验表明，UniAP 比已有的最佳方法最高快 3.8 倍，比未采用并行策略优化的算法最高快 9 倍），还解决了大模型分布式训练中因超参数设置不合理而导致训练失败的难题，显著提升了易用性。

更重要的是，UniAP 实现了与国产 AI 计算卡的适配，为大模型训练的降本增效提供了关键技术、平台和框架支持。

凭借其卓越的性能和创新性，该论文已被 CVPR 2025 接收为 Oral 报告（仅占所有投稿论文的 0.7%，所有录用论文的 3.3%）。

方法简介

并行策略的选择是影响分布式训练算法性能的关键因素。目前主流的并行策略包括：

  • 流水线并行（Pipeline Parallelism）
  • 数据并行（Data Parallelism）
  • 张量并行（Tensor Parallelism）
  • 全分片数据并行（Fully Sharded Data Parallelism, FSDP）

这些策略可以归纳为两大类：

  • 层内并行策略：仅对模型层内的张量进行切分，如数据并行、张量并行、全分片数据并行等。
  • 层间并行策略：将模型层划分为多个互斥的子集，如流水线并行。

基于已有的并行策略，大量研究工作集中于并行方法的设计。这些方法可以分为两类：

  • 手动并行方法：需要用户手动指定并行策略，费时费力，且难以适应复杂的硬件环境。
  • 自动并行方法：旨在自动寻找最优并行策略。

然而，现有的自动并行方法要么只考虑层内或层间两类并行策略中的一种，要么将两类策略分阶段优化，而非进行联合优化，导致训练效率仍有提升空间。

UniAP 采用混合整数二次规划（Mixed Integer Quadratic Programming, MIQP）进行建模，实现对层内与层间并行策略的联合优化，从而获得更大的策略探索空间。

UniAP 的整体架构图如下所示：

图片

UniAP 首先对硬件和模型进行性能评估，然后根据评估结果和模型的计算图构建代价模型。基于代价模型和计算图，UniAP 将优化问题建模为一个混合整数二次规划问题并求解。最后，UniAP 将优化结果转化为以计算图形式表达的并行计划，并交由现有的深度学习平台（如 PyTorch）进行训练执行。

性能评估和代价模型

由于自动并行框架需要在执行分布式训练之前优化并行策略，因此需要对分布式训练的性能和开销进行模拟，并在模拟结果的基础上进行优化。

性能评估是自动并行框架的重要组成部分。UniAP 通过性能评估收集硬件和模型的性能信息，例如 P2P 通信效率、All-Reduce 集合通信效率以及模型每一层的前向计算时间开销和显存开销等。

考虑到时间效率，自动并行框架只能完成有限的性能评估任务，然后根据评估结果构建代价模型。UniAP 的代价模型分为时间代价模型和显存代价模型。时间代价模型用于估计模型采用不同并行策略的时间开销，包括计算时间和通信时间开销；显存代价模型用于估计模型采用不同并行策略的显存占用开销。

混合整数二次规划形式

UniAP 的目标是通过设定并行策略，最小化训练中每次迭代所消耗的时间（Time-Per-Iteration, TPI）。

假设模型的计算图为 图片。层间并行策略可以由流水线的级数 pp_size、流水线的微批量数量 𝑐，以及计算图的层 𝑢 是否放置在第 𝑖 级流水线上的放置策略图片表示，图片；层内并行策略可以由层 𝑢 的层内并行策略集合图片和层 𝑢 是否选择第 𝑘 个层内并行策略的图片表示，图片。

为了方便描述，引入以下额外记号：对于一个给定的层图片，图片表示它的第 𝑘 个层内并行策略的前向传播和反向传播时间代价之和，图片代表了该层的第 𝑘 个层内并行策略在训练中占用的峰值显存量。对于一个给定的边图片，如果该边的源点和终点位于同一个流水线的计算阶段中，那么它的通信代价为图片。否则，如果该边的源点和终点位于不同流水线的计算阶段中，那么它的通信代价为图片。此处，图片均由 UniAP 的代价模型给定，在 UniAP 算法的优化过程中始终为常量。

流水线阶段内的时间开销：流水线某个阶段内部的时间开销由该流水线阶段每一层内部的时间开销和层间通信时间组成。记流水线的第 i 个阶段在一个微批量上的时间代价为图片，图片。可建立如下约束（原论文公式 3）：

图片

跨流水线阶段的时间开销：跨流水线阶段时间开销由通信开销组成。记跨第 i 个流水线阶段和第 i+1 个流水线阶段的时间开销为图片，图片。可建立如下约束（原论文公式 4）：

图片

式中图片指示计算图中相邻的层 𝑢 和层 𝑣 是否放置在相邻的计算阶段上，图片则指示了层 𝑢 和层 𝑣 之间的通信代价。

目标函数：有了图片和图片，就可以得到图片，形式根据不同的流水线调度而有所不同，以 GPipe 调度为例，则目标函数为（原论文公式 2）：

图片

详细解释可见原文 3.3.1。

显存约束：由于分布式机器学习系统要求在训练深度学习模型时避免显存溢出（Out-of-Memory, OOM）异常，UniAP 为混合整数二次规划表达式引入了显存约束。对于同构集群，所有工作设备的显存容量相同，记作 𝑚。UniAP 方法的显存约束如下（原论文公式 5）：

图片

除了显存约束外，还需要建立流水线保序约束、放置策略约束和策略选择约束。详细内容请参考原论文。

统一优化过程

根据混合整数二次规划的表达式，现有的优化器可以直接解得给定流水线的级数 pp_size 和微批量数量 𝑐 的情况下最优的并行策略组合。但由于 pp_size 和 c 是流水线并行的超参数，UniAP 也需要统一优化这两个变量才能求得最优的并行策略组合。UniAP 通过枚举这两个变量来解决这个问题，算法伪代码如下所示（原文算法 1）：

图片

实验结果

该论文在四种典型的硬件环境（NVIDIA GPU V100、TITAN XP、A100，以及国产卡海光 DCU）和五种不同的基于 Transformer 的模型（BERT、T5、ViT、Swin-Transformer、Llama）上进行了实验验证。

首先，在 NVIDIA GPU 上与其他的自动并行方法（Galvatron 和 Alpa）进行了对比（原论文表 1）：

图片

其中 ENVA 是一台 8 卡 V100，ENVB-8 是两台 4 卡 TITAN XP，ENVC 是一台 8 卡 A100。实验结果表明，在这三种硬件环境下，UniAP 的吞吐量均优于两个 baseline，最大提升达到 3.8 倍；同时，UniAP 的策略优化时间也远小于两个 baseline，最多缩短 107 倍。

其次，在国产 AI 计算卡上与手动并行方法进行了对比。选取的 baseline 是国际主流的大模型训练框架 Megatron 和 DeepSpeed。这两个框架中都有分布式训练的相关参数需要设置。实验中，枚举所有可能的设置，对每个设置实际运行一定的轮次并记录吞吐量，选取性能最佳的设置作为吞吐量结果，并将整个过程的时间作为策略优化时间。结果如下（原论文表 2）：

图片

硬件设置为 8 个 4 卡 DCU 节点。从表中可以看出，UniAP 找到了所有可行策略中的最优解，并且相比于手动暴力搜索，大大节省了策略优化时间。

在可扩展性方面，该论文在最大 64 卡的集群上进行了实验，验证了 UniAP 的近线性扩展性（原论文图 5 和表 4）：

图片

图片

此外，该论文还通过对实验中 Megatron 的策略空间的分析，深入探讨了自动并行的重要性（原论文表 3）：

图片

从表中可以看出，对于没有分布式训练经验的研究者来说，从所有支持的并行策略中随机选择一个，有 64.1%（41/64）到 87.5%（42/48）的概率会因为策略选择不合理而导致模型训练无法成功运行（例如出现显存溢出）。即使选择了能够成功运行训练过程的策略，最快策略和最慢策略之间的速度也相差高达 2.01/0.22≈9 倍。更多细节讨论请参考原文。

全文总结

UniAP 具有以下显著优点：

  • 高效性：在相同硬件条件下，UniAP 的训练速度比已有的最佳方法最高快 3.8 倍，意味着成本降低 3.8 倍。
  • 易用性：用户无需理解复杂的概念，例如流水线划分和张量分片。UniAP 平台能够自动生成最优方案，避免 64%-87% 的无效策略，从而让没有分布式训练经验的用户也能像使用单卡训练模型一样轻松地利用集群的大规模算力。
  • 适配国产 AI 计算卡：UniAP 已经适配了海光 DCU 等国产 AI 计算卡，为提升国产卡的性能和易用性提供了潜在的解决方案，有望加速国产 AI 基础设施的普及。