CVPR Oral | 南京大学李武军教授课题组推出分布式训练算法UniAP，大模型训练最高加速3.8倍

李武军教授为通讯作者，硕士生林昊（已毕业，现工作于阿里巴巴）、吴轲、李杰为共同第一作者，博士生李俊为参与作者。

训练成本高昂已经成为大模型和人工智能可持续发展的主要障碍之一。

大模型的训练往往采用多机多卡的分布式训练，大模型的分布式训练挑战巨大，即使硬件足够，不熟悉分布式训练的人大概率

（实验中验证有 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 能够通过自动搜索找到高效能的分布式训练方案，既解决了效率和成本问题

（实验中，比已有的最好方法最高快 3.8 倍，比不采用并行策略优化的算法最高快 9 倍），也解决了很多人在大模型分布式训练时因为超参数设置

（模型怎么切分和排布、数据怎么切分和排布等）不合理而无法成功运行训练过程的问题，即易用性问题。

此外，还实现了 UniAP 跟国产 AI 计算卡的适配。相关工作为大模型训练的降本增效提供了核心技术、（国产）平台和框架。

论文被 CVPR 2025 录用为 Oral（所有投稿论文的 0.7%，所有录用论文的 3.3%）。

方法简介

并行策略的设置是影响分布式训练算法效能的重要因素。当前主流的并行策略包括流水线并行、数据并行、张量并行、全分片数据并行等四种并行策略。

这些并行策略可以被分成以下两类：

• 层内并行策略：仅切分模型的层内张量，包括以数据并行、张量并行、全分片数据并行等为代表的并行策略；

• 层间并行策略：仅切分模型的层为多个互斥子集，包括流水线并行等并行策略。

基于已有的并行策略，大量的研究工作集中于并行方法的设计。这些并行方法可以按照是否需要用户手动指定并行策略划分为两类：手动并行方法和自动并行方法。

传统的手动并行方法不仅耗时耗力，而且难以适应复杂的硬件环境。

而现有的自动并行方法存在的问题是它们要么只考虑层内或层间两类并行策略中的一类并行策略，

要么把两类并行策略做分阶段优化而不是联合优化，求解得到的并行策略的训练效率存在提升空间。

UniAP 使用混合整数二次规划进行建模，实现对层内与层间并行策略的联合优化。这种联合优化使得 UniAP 有更大的策略探索空间。

UniAP 的架构图如下：

UniAP 首先对硬件和模型进行性能评估。然后，UniAP 会根据性能评估的结果和模型的计算图构建代价模型。

根据代价模型和模型的计算图，UniAP 将优化问题建模为一个混合整数二次规划问题并进行优化。

最后，UniAP 会将优化结果由向量转化成以计算图形式表达的并行计划，交由已有深度学习平台（如 PyTorch）进行训练。

性能评估和代价模型

因为自动并行框架要求在执行分布式训练前优化并行策略，所以框架需要对分布式训练的性能和开销进行模拟，再在模拟的结果上进行优化。

因此，对环境和任务进行性能评估是自动并行框架的重要组成部分。

具体地，在性能评估部分，UniAP 将收集硬件和模型的性能信息，如 P2P 通信效率、All-Reduce 集合通信效率、模型每一层的前向计算的时间开销和显存开销等。

出于时间效率考虑，一个自动并行框架只能完成有限的性能评估任务，然后根据性能评估的结果构建代价模型。

UniAP 的代价模型分为时间代价模型和显存代价模型。

前者用于估计模型采用不同并行策略的时间开销，包括计算时间开销和通信时间开销；后者用于估计模型采用不同并行策略的显存占用开销。

混合整数二次规划形式

显存约束：因为分布式机器学习系统要求在训练深度学习模型时不能发生显存溢出（Out-of-Memory，简称 OOM）异常，

所以 UniAP 为混合整数二次规划表达式引入显存约束。对于同构集群，因为所有工作设备均同构，所以它们的显存容量相同，记作 𝑚。

UniAP 方法的显存约束如下（原论文公式 5）：

除显存约束外，还需要建立流水线保序约束、放置策略约束、策略选择约束。详情可见原论文。

统一优化过程

根据混合整数二次规划的表达式，现有的优化器可以直接解得给定流水线的度 pp_size 和微批量数量 𝑐 的情况下最优的并行策略组合。

但因为 pp_size 和 c 是流水线并行的超参数，所以 UniAP 也需要统一优化这两个变量才能求得最优的并行策略组合。

UniAP 通过枚举这两个变量来解决这个问题，算法伪代码如下（原文算法 1）：

实验结果

论文在 4 种典型的硬件环境（NVIDIA GPU V100、TITAN XP、A100，和国产卡海光 DCU），5 种不同的基于 Transformer 的模型

（BERT、T5、ViT、Swin-Transformer、Llama）上进行实验。

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

其中 ENVA 是一台 8 卡 V100，ENVB-8 是两台 4 卡 TITAN XP，ENVC 是一台 8 卡 A100。

可以发现，在这 3 种硬件环境下，UniAP 的吞吐量均不差于两个 baseline，最大的提升达到 3.8 倍；

而 UniAP 的策略优化时间更是远远小于两个 baseline，最大缩短 107 倍。

然后是在国产 AI 计算卡上和手动并行方法的对比。选取的 baseline 是国际主流的大模型训练框架 Megatron 和 DeepSpeed。

两个框架中均有分布式训练的相关参数需要设置，实验中，枚举所有可能的设置，每个设置实际跑一定的轮次记录吞吐量，

选取性能最好的做为吞吐量结果，选取整个过程的时间为策略优化时间。结果如下（原论文表 2）：

其中硬件设置是 8 个 4 卡 DCU 节点。从表中可见，UniAP 找到了所有可行策略中的最优解，同时相较于手动暴力搜索，大大节约了策略优化时间。

在可拓展性方面，论文在最大 64 卡的集群上进行实验，验证了近线性拓展性（原论文图 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 计算卡：已适配海光 DCU 等国产 AI 计算卡，为提升国产卡的效能和易用性提供了潜在的解决方案，有望加速国产 AI 基础设施的普及进程。

文章来自于微信公众号 “机器之心”，作者 ：机器之心