李武军教授为通讯作者，硕士生林昊（已毕业，现工作于阿里巴巴）、吴轲、李杰为共同第一作者，博士生李俊为参与作者。
训练成本高昂已经成为大模型和人工智能可持续发展的主要障碍之一。
大模型的训练往往采用多机多卡的分布式训练，大模型的分布式训练挑战巨大，即使硬件足够，不熟悉分布式训练的人大概率（实验中验证有 64%-87% 的概率）会因为超参数设置（模型怎么切分和排布、数据怎么切分和排布等）不合理而无法成功运行训练过程。
此外，不熟悉分布式训练的人在碰到大模型训练慢时容易只想到增加 GPU 硬件等横向拓展（scale-out）方法，而忽略了分布式训练算法的纵向拓展（scale-up）作用。
实际上，分布式训练算法会极大地影响硬件的算力利用率。高效能分布式训练算法具有高算力利用率。用同样的硬件算力训练同一个模型，高效能分布式训练算法会比低效能分布式训练算法速度快，最高可能会快数倍甚至数十倍以上。
也就是说，训练同一个模型，高效能分布式训练算法会比低效能分布式训练算法成本低，最高可能会节省数倍甚至数十倍以上的算力成本。很多已有的分布式训练算法的效能较低，甚至可能导致机器和 GPU 卡越多、训练速度越慢的结果。
南京大学计算机学院李武军教授课题组研发了高效能分布式训练算法 UniAP，并基于 UniAP 研发了相应的大模型分布式训练平台和框架。
论文标题：UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming
论文地址：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 的代价模型分为时间代价模型和显存代价模型。前者用于估计模型采用不同并行策略的时间开销，包括计算时间开销和通信时间开销；后者用于估计模型采用不同并行策略的显存占用开销。
混合整数二次规划形式
UniAP 的混合整数二次规划的目标是设定并行策略，使得训练中每次迭代所消耗的时间（Time-Per-Iteration，简称 TPI）最小化。
设模型的计算图为
表示，
和层
