點選下方卡片，關注“CVer”公眾號
AI/CV重磅乾貨，第一時間送達
五一節日福利來了！CVer學術知識星球優惠券！現在贈送50元新使用者優惠券(左圖領取)，7折+20元老使用者續費券(右圖領取)，大家快掃碼加入！每天分享最新最優質的AI工作，假期不停更！最強助力你的科研和工作！ideas拉滿！
AI/CV重磅乾貨，第一時間送達
五一節日福利來了！CVer學術知識星球優惠券！現在贈送50元新使用者優惠券(左圖領取)，7折+20元老使用者續費券(右圖領取)，大家快掃碼加入！每天分享最新最優質的AI工作，假期不停更！最強助力你的科研和工作！ideas拉滿！
李武軍教授為通訊作者，碩士生林昊（已畢業，現工作於阿里巴巴）、吳軻、李傑為共同第一作者，博士生李俊為參與作者。
訓練成本高昂已經成為大模型和人工智慧可持續發展的主要障礙之一。
大模型的訓練往往採用多機多卡的分散式訓練，大模型的分散式訓練挑戰巨大，即使硬體足夠，不熟悉分散式訓練的人大機率（實驗中驗證有 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、流水線的微批次數量 𝑐，計算圖的層 𝑢 是否放置在第 𝑖 級流水線上的放置策略
表示，
；層內並行策略可以由層 𝑢 的層內並行策略集合
和層 𝑢 是否選擇第 𝑘 個層內並行策略的
表示，
。
為方便描述，需要引入如下額外記號：對於一個給定的層
，
表示它的第 𝑘 個層內並行策略的前向傳播和反向傳播時間代價之和，
代表了該層的第 𝑘 個層內並行策略在訓練中佔用的峰值視訊記憶體量。對於一個給定的邊
，如果該邊的源點和終點位於同一個流水線的計算階段中，那麼它的通訊代價為
。否則，如果該邊的源點和終點位於不同流水線的計算階段中，那麼它的通訊代價為
。此處，
均由 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）：
實驗結果
論文在 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 基礎設施的普及程序。
五一節日福利來了！CVer學術知識星球優惠券！現在贈送50元新使用者優惠券(左圖領取)，7折+20元老使用者續費券(右圖領取)，大家快掃碼加入！每天分享最新最優質的AI工作，假期不停更！最強助力你的科研和工作！ideas拉滿！
CVPR 2025 論文和程式碼下載
CVPR 2025 論文和程式碼下載
ECCV 2024 論文和程式碼下載
ECCV 2024 論文和程式碼下載
CV垂直方向和論文投稿交流群成立
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群
▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）來了！想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料，歡迎掃描下方二維碼，加入CVer計算機視覺（知識星球），已彙集上萬人！
▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看
▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）來了！想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料，歡迎掃描下方二維碼，加入CVer計算機視覺（知識星球），已彙集上萬人！
▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看