CVPR Oral | 南京大学李武军教授课题组推出分布式训练算法UniAP，大模型训练最高加速3.8倍
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
李武军教授为通讯作者，硕士生林昊（已毕业，现工作于阿里巴巴）、吴轲、李杰为共同第一作者，博士生李俊为参与作者。
训练成本高昂已经成为大模型和人工智能可持续发展的主要障碍之一。
大模型的训练往往采用多机多卡的分布式训练，大模型的分布式训练挑战巨大，即使硬件足够，不熟悉分布式训练的人大概率（实验中验证有64%-87%的概率）会因为超参数设置（模型怎么切分和排布、数据怎么切分和排布等）不合理而无法成功运行训练过程。
此外，不熟悉分布式训练的人在碰到大模型训练慢时容易只想到增加GPU硬件等横向拓展（scale-out）方法，而忽略了分布式训练算法的纵向拓展（scale-up）作用。
实际上，分布式训练算法会极大地影响硬件的算力利用率。高效能分布式训练算法具有高算力利用率。用同样的硬件算力训练同一个模型，高效能分布式训练算法会比低效能分布式训练算法速度快，最高可能会快数倍甚至数十倍以上。
也就是说，训练同一个模型，高效能分布式训练算法会比低效能分布式训练算法成本低，最高可能会节省数倍甚至数十倍以上的算力成本。很多已有的分布式训练算法的效能较低，甚至可能导致机器和GPU卡越多、训练速度越慢的结果。
南京大学计算机学院李武军教授课题组研发了高效能分布式训练算法UniAP，并基于UniAP研发了相应的大模型分布式训练平台和框架。
论文标题：UniAP:UnifyingInter-andIntra-LayerAutomaticParallelismbyMixedIntegerQuadraticProgramming
论文地址：https ://arxiv.org/abs/2307.16375
UniAP是首个能实现层内并行策略（张量并行等）和层间并行策略（流水线并行等）联合优化的工作。给定模型和硬件平台，UniAP能够通过自动搜索找到高效能的分布式训练方案，既解决了效率和成本问题（实验中，比已有的最好方法最高快3.8倍，比不采用并行策略优化的算法最高快9倍），也解决了很多人在大模型分布式训练时因为超参数设置（模型怎么切分和排布、数据怎么切分和排布等）不合理而无法成功运行训练过程的问题，即易用性问题。
此外，还实现了UniAP跟国产AI计算卡的适配。相关工作为大模型训练的降本增效提供了核心技术、（国产）平台和框架。
论文被CVPR2025录用为Oral（所有投稿论文的0.7%，所有录用论文的3.3%）。
方法简介
并行策略的设置是影响分布式训练算法效能的重要因素。当前主流的并行策略包括流水线并行、数据并行、张量并行、全分片数据并行等四种并行策略。这些并行策略可以被分成以下两类：
层内并行策略：仅切分模型的层内张量，包括以数据并行、张量并行、全分片数据并行等为代表的并行策略；
层间并行策略：仅切分模型的层为多个互斥子集，包括流水线并行等并行策略。
基于已有的并行策略，大量的研究工作集中于并行方法的设计。这些并行方法可以按照是否需要用户手动指定并行策略划分为两类：手动并行方法和自动并行方法。传统的手动并行方法不仅耗时耗力，而且难以适应复杂的硬件环境。
而现有的自动并行方法存在的问题是它们要么只考虑层内或层间两类并行策略中的一类并行策略，要么把两类并行策略做分阶段优化而不是联合优化，求解得到的并行策略的训练效率存在提升空间。
UniAP使用混合整数二次规划进行建模，实现对层内与层间并行策略的联合优化。这种联合优化使得UniAP有更大的策略探索空间。
UniAP的架构图如下：
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）：
实验结果
首先是在NVIDIAGPU上和其他自动并行方法（Galvatron和Alpa）的对比（原论文表1）：
其中ENVA是一台8卡V100，ENVB-8是两台4卡TITANXP，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基础设施的普及进程。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
