李武軍教授為通訊作者，碩士生林昊（已畢業(yè)，現(xiàn)工作于阿里巴巴）、吳軻、李杰為共同第一作者，博士生李俊為參與作者。
訓練成本高昂已經(jīng)成為大模型和人工智能可持續(xù)發(fā)展的主要障礙之一。
大模型的訓練往往采用多機多卡的分布式訓練，大模型的分布式訓練挑戰(zhàn)巨大，即使硬件足夠，不熟悉分布式訓練的人大概率（實驗中驗證有 64%-87% 的概率）會因為超參數(shù)設(shè)置（模型怎么切分和排布、數(shù)據(jù)怎么切分和排布等）不合理而無法成功運行訓練過程。
此外，不熟悉分布式訓練的人在碰到大模型訓練慢時容易只想到增加 GPU 硬件等橫向拓展（scale-out）方法，而忽略了分布式訓練算法的縱向拓展（scale-up）作用。
實際上，分布式訓練算法會極大地影響硬件的算力利用率。高效能分布式訓練算法具有高算力利用率。用同樣的硬件算力訓練同一個模型，高效能分布式訓練算法會比低效能分布式訓練算法速度快，最高可能會快數(shù)倍甚至數(shù)十倍以上。
也就是說，訓練同一個模型，高效能分布式訓練算法會比低效能分布式訓練算法成本低，最高可能會節(jié)省數(shù)倍甚至數(shù)十倍以上的算力成本。很多已有的分布式訓練算法的效能較低，甚至可能導致機器和 GPU 卡越多、訓練速度越慢的結(jié)果。
南京大學計算機學院李武軍教授課題組研發(fā)了高效能分布式訓練算法UniAP，并基于 UniAP 研發(fā)了相應的大模型分布式訓練平臺和框架
- 論文標題：UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming
- 論文地址：https://arxiv.org/abs/2307.16375
UniAP 是首個能實現(xiàn)層內(nèi)并行策略（張量并行等）和層間并行策略（流水線并行等）聯(lián)合優(yōu)化的工作。給定模型和硬件平臺，UniAP 能夠通過自動搜索找到高效能的分布式訓練方案，既解決了效率和成本問題（實驗中，比已有的最好方法最高快 3.8 倍，比不采用并行策略優(yōu)化的算法最高快 9 倍），也解決了很多人在大模型分布式訓練時因為超參數(shù)設(shè)置（模型怎么切分和排布、數(shù)據(jù)怎么切分和排布等）不合理而無法成功運行訓練過程的問題，即易用性問題。
此外，還實現(xiàn)了 UniAP 跟國產(chǎn) AI 計算卡的適配。相關(guān)工作為大模型訓練的降本增效提供了核心技術(shù)、（國產(chǎn)）平臺和框架。
論文被 CVPR 2025 錄用為 Oral（所有投稿論文的 0.7%，所有錄用論文的 3.3%）。
方法簡介
并行策略的設(shè)置是影響分布式訓練算法效能的重要因素。當前主流的并行策略包括流水線并行、數(shù)據(jù)并行、張量并行、全分片數(shù)據(jù)并行等四種并行策略。這些并行策略可以被分成以下兩類：
- 層內(nèi)并行策略：僅切分模型的層內(nèi)張量，包括以數(shù)據(jù)并行、張量并行、全分片數(shù)據(jù)并行等為代表的并行策略；
- 層間并行策略：僅切分模型的層為多個互斥子集，包括流水線并行等并行策略。
基于已有的并行策略，大量的研究工作集中于并行方法的設(shè)計。這些并行方法可以按照是否需要用戶手動指定并行策略劃分為兩類：手動并行方法和自動并行方法。傳統(tǒng)的手動并行方法不僅耗時耗力，而且難以適應復雜的硬件環(huán)境。
而現(xiàn)有的自動并行方法存在的問題是它們要么只考慮層內(nèi)或?qū)娱g兩類并行策略中的一類并行策略，要么把兩類并行策略做分階段優(yōu)化而不是聯(lián)合優(yōu)化，求解得到的并行策略的訓練效率存在提升空間。
UniAP 使用混合整數(shù)二次規(guī)劃進行建模，實現(xiàn)對層內(nèi)與層間并行策略的聯(lián)合優(yōu)化。這種聯(lián)合優(yōu)化使得 UniAP 有更大的策略探索空間。
UniAP 的架構(gòu)圖如下：
UniAP 首先對硬件和模型進行性能評估。然后，UniAP 會根據(jù)性能評估的結(jié)果和模型的計算圖構(gòu)建代價模型。根據(jù)代價模型和模型的計算圖，UniAP 將優(yōu)化問題建模為一個混合整數(shù)二次規(guī)劃問題并進行優(yōu)化。最后，UniAP 會將優(yōu)化結(jié)果由向量轉(zhuǎn)化成以計算圖形式表達的并行計劃，交由已有深度學習平臺（如 PyTorch）進行訓練。
性能評估和代價模型
因為自動并行框架要求在執(zhí)行分布式訓練前優(yōu)化并行策略，所以框架需要對分布式訓練的性能和開銷進行模擬，再在模擬的結(jié)果上進行優(yōu)化。
因此，對環(huán)境和任務(wù)進行性能評估是自動并行框架的重要組成部分。具體地，在性能評估部分，UniAP 將收集硬件和模型的性能信息，如 P2P 通信效率、All-Reduce 集合通信效率、模型每一層的前向計算的時間開銷和顯存開銷等。
出于時間效率考慮，一個自動并行框架只能完成有限的性能評估任務(wù)，然后根據(jù)性能評估的結(jié)果構(gòu)建代價模型。UniAP 的代價模型分為時間代價模型和顯存代價模型。前者用于估計模型采用不同并行策略的時間開銷，包括計算時間開銷和通信時間開銷；后者用于估計模型采用不同并行策略的顯存占用開銷。
混合整數(shù)二次規(guī)劃形式
UniAP 的混合整數(shù)二次規(guī)劃的目標是設(shè)定并行策略，使得訓練中每次迭代所消耗的時間（Time-Per-Iteration，簡稱 TPI）最小化。
統(tǒng)一優(yōu)化過程
根據(jù)混合整數(shù)二次規(guī)劃的表達式，現(xiàn)有的優(yōu)化器可以直接解得給定流水線的度 pp_size 和微批量數(shù)量 的情況下最優(yōu)的并行策略組合。但因為 pp_size 和 c 是流水線并行的超參數(shù)，所以 UniAP 也需要統(tǒng)一優(yōu)化這兩個變量才能求得最優(yōu)的并行策略組合。UniAP 通過枚舉這兩個變量來解決這個問題，算法偽代碼如下（原文算法 1）：
實驗結(jié)果
論文在 4 種典型的硬件環(huán)境（NVIDIA GPU V100、TITAN XP、A100，和國產(chǎn)卡海光 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。可以發(fā)現(xiàn)，在這 3 種硬件環(huán)境下，UniAP 的吞吐量均不差于兩個 baseline，最大的提升達到 3.8 倍；而 UniAP 的策略優(yōu)化時間更是遠遠小于兩個 baseline，最大縮短 107 倍。
然后是在國產(chǎn) AI 計算卡上和手動并行方法的對比。選取的 baseline 是國際主流的大模型訓練框架 Megatron 和 DeepSpeed。兩個框架中均有分布式訓練的相關(guān)參數(shù)需要設(shè)置，實驗中，枚舉所有可能的設(shè)置，每個設(shè)置實際跑一定的輪次記錄吞吐量，選取性能最好的做為吞吐量結(jié)果，選取整個過程的時間為策略優(yōu)化時間。結(jié)果如下（原論文表 2）：
其中硬件設(shè)置是 8 個 4 卡 DCU 節(jié)點。從表中可見，UniAP 找到了所有可行策略中的最優(yōu)解，同時相較于手動暴力搜索，大大節(jié)約了策略優(yōu)化時間。
在可拓展性方面，論文在最大 64 卡的集群上進行實驗，驗證了近線性拓展性（原論文圖 5 和表 4）：
另外，論文還通過對實驗中 Megatron 的策略空間的分析深度探討了自動并行的必要性（原論文表 3）：
從表中可以看出，對于一個沒有分布式訓練經(jīng)驗的人來說，從所有支持的并行策略中隨機選擇一個，有 64.1%（41/64）到 87.5%（42/48）的概率會因為策略選擇不合理而導致模型無法成功運行訓練過程（出現(xiàn)顯存溢出等）；即使選擇到了能成功運行訓練過程的策略，最快的策略和最慢的策略間的速度最大相差了 2.01/0.22≈9 倍。更多討論可見原文。
全文小結(jié)
UniAP具有如下優(yōu)點：
高效性：在相同硬件條件下，UniAP 訓練速度比已有最好方法最高快 3.8 倍，即成本降低 3.8 倍；
易用性：用戶無需理解流水線劃分、張量分片等復雜概念，UniAP 平臺能自動生成最優(yōu)方案，自動規(guī)避 64%-87% 的無效策略，讓沒有分布式訓練經(jīng)驗的用戶，也能像使用單卡訓練模型般輕松使用集群的大規(guī)模算力；
適配國產(chǎn) AI 計算卡：已適配海光 DCU 等國產(chǎn) AI 計算卡，為提升國產(chǎn)卡的效能和易用性提供了潛在的解決方案，有望加速國產(chǎn) AI 基礎(chǔ)設(shè)施的普及進程。
熱門跟貼