李武軍教授為通訊作者，碩士生林昊（已畢業，現工作于阿里巴巴）、吳軻、李杰為共同第一作者，博士生李俊為參與作者。

訓練成本高昂已經成為大模型和人工智能可持續發展的主要障礙之一。

大模型的訓練往往采用多機多卡的分布式訓練，大模型的分布式訓練挑戰巨大，即使硬件足夠，不熟悉分布式訓練的人大概率（實驗中驗證有 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 的代價模型分為時間代價模型和顯存代價模型。前者用于估計模型采用不同并行策略的時間開銷，包括計算時間開銷和通信時間開銷；后者用于估計模型采用不同并行策略的顯存占用開銷。

混合整數二次規劃形式

UniAP 的混合整數二次規劃的目標是設定并行策略，使得訓練中每次迭代所消耗的時間（Time-Per-Iteration，簡稱 TPI）最小化。

統一優化過程

根據混合整數二次規劃的表達式，現有的優化器可以直接解得給定流水線的度 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 基礎設施的普及進程。