如何搞定一個可以支持多芯混合訓(xùn)練的AI集群
由于外部環(huán)境的變化,適用于大模型訓(xùn)練任務(wù)的 GPU 整體規(guī)模無法繼續(xù)增長。這些存量 GPU 組成的集群,仍然是當(dāng)前加速大模型訓(xùn)練的主要 AI 算力來源。同時,各類國產(chǎn) AI 芯片開始大規(guī)模投入實際生產(chǎn)任務(wù)。在未來一段時間內(nèi),數(shù)據(jù)中心的 AI 算力將保持多種芯片并存的現(xiàn)象。
但是,當(dāng)前基礎(chǔ)大模型訓(xùn)練所需要的最大 AI 算力集群規(guī)模,已經(jīng)從單一集群千卡逐步提升至萬卡量級。同時,很多智算中心已經(jīng)部署的 GPU 集群,通常是十幾臺至數(shù)百臺服務(wù)器不等,難以滿足未來行業(yè)大模型訓(xùn)練的需求。
所以,在已有 AI 算力集群的基礎(chǔ)上,構(gòu)建由 GPU、昆侖芯、昇騰等不同芯片混合組成的單一集群,為大模型訓(xùn)練提供更大 AI 算力,成為了一個自然的選擇。
大家都知道,成功建設(shè)一個全部由 GPU 芯片組成的集群,讓他們作為一個整體可以高效率地跑起來,就已經(jīng)足夠復(fù)雜。(詳情可參考:AI 大底座,大模型時代的答卷)
如果還要在GPU 集群中再加上其他類型的 AI 芯片,他們各自講著不同的語言,擁有完全不同的能力,讓他們像同一種芯片組成的集群一樣,實現(xiàn)多芯混合訓(xùn)練,加速單一大模型訓(xùn)練任務(wù),那真是復(fù)雜到頭了。
這里的挑戰(zhàn),和一個經(jīng)常出現(xiàn)在團建中的競技項目「多人多足」類似:大家肩并肩站成一排,人和人之間依次綁著小腿,大家齊頭向前跑沖向終點。如果這個時候參賽隊中有些人只會外語,有些人步子邁得很長……
1.如何建立和加速一個 GPU 集群
為了讓大家對「如何搞定一個支持多芯混合訓(xùn)練的 AI 集群」有更清晰的理解,我們首先以 GPU 集群為例,簡單介紹建立和加速一個AI 集群的三個關(guān)鍵方面。
1.1實現(xiàn) GPU 互聯(lián)互通
為了建設(shè)一個多卡集群,首先要完成 GPU 卡在物理層面進行連接。在單臺服務(wù)器內(nèi) 8 塊 GPU 卡通過NVLink連接。不同服務(wù)器之間的 GPU 卡通過 RDMA 網(wǎng)絡(luò)連接。
在完成網(wǎng)絡(luò)的搭建后,借助 NVIDA 開發(fā)的集合通信庫 NCCL,GPU 就能通過網(wǎng)絡(luò)實現(xiàn)相互通信完成數(shù)據(jù)同步,使得訓(xùn)練任務(wù)可以一輪一輪地往下推進,直到完成大模型的訓(xùn)練。
1.2制定分布式并行策略
為了加速大模型訓(xùn)練任務(wù),我們需要將這個任務(wù)拆分到集群的所有 GPU 中,使得這些 GPU 能夠共同完成任務(wù)。這就是常說的分布式并行策略。分布式并行策略有很多種,比如從訓(xùn)練數(shù)據(jù)維度進行切分的數(shù)據(jù)并行,按照模型的不同層面進行切分的流水線并行等。
我們需要依據(jù)集群的物理拓撲和大模型的參數(shù),找到最優(yōu)的分布式并行策略,充分發(fā)揮集群效能。在任務(wù)運行過程中,集群中所有 GPU 步驟保持一致,同時開始計算,同時開始通信,不存在一些 GPU 空轉(zhuǎn)等待另外一部分 GPU 的情況。
1.3部署 AI加速套件
按照分布式并行策略被拆分開的模型和數(shù)據(jù),將會以一個個算子的形式部署在 GPU 進行計算。為了加速 GPU 對算子的計算過程,我們還需要一個 AI 加速套件。
這個AI 加速套件需要包含數(shù)據(jù)加載、算子計算(各種 CUDA 庫)、多卡通信( NCCL 集合通信庫)等各個方面的優(yōu)化。比如采用數(shù)據(jù)預(yù)取策略,使得 I/O 的過程和 GPU 上的計算充分并行起來;使用 NVIDIA 優(yōu)化后的 GPU 算子或者全新的算子加速 GPU 卡的計算效率;更新 NCCL 能力提升 GPU 卡相互通信的效率。
我們這里總結(jié)一下,為了建設(shè)一個能夠高效訓(xùn)練大模型的集群,需要在卡間和機間建立高效的互聯(lián)互通,將大模型訓(xùn)練任務(wù)按照合適的并行策略拆分到 GPU 卡中,最后通過各種優(yōu)化方法,加速 GPU 對算子的計算效率,完成大模型訓(xùn)練。
2.建立不同芯片集群的差異
當(dāng)前,在數(shù)據(jù)中心的多芯算力的運用方式上,主流仍然是采用一種芯片對應(yīng)一個集群的思路,這需要根據(jù)每一種芯片的特點進行量身定制。
參照上文提到的三個方面,一起來看看基于昆侖芯和昇騰 910B,建設(shè)和加速這些集群的差異。
在互聯(lián)互通上,昆侖芯服務(wù)器內(nèi)部通過 XPU Link 進行連接,服務(wù)器之間通過標準的 RDMA 網(wǎng)卡進行連接,卡和卡之間使用 XCCL 通信庫進行相互通信。昇騰 910B 服務(wù)器內(nèi)部通過 HCCS 進行連接,服務(wù)器之間通過華為自研的內(nèi)置 RDMA 進行連接,卡和卡之間使用 HCCL 通信庫進行相互通信。
在并行策略上,NVIDIA GPU 和昆侖芯采用單機 8 卡的部署方式,昇騰910B 則是機內(nèi) 16 卡分為 2 個 8 卡通信組 。這意味著在 AI 框架下形成不同的集群拓撲,需要有針對性地制定分布式并行策略。
在 AI 加速套件上,由于 GPU、昆侖芯、昇騰等芯片在計算能力,顯存大小,I/O 吞吐,通信庫等均存在差異,故需要面向具體芯片進行特定優(yōu)化。最后的結(jié)果,就是每一種芯片,有一個各自對應(yīng)的算子庫,以及相應(yīng)的加速策略。
3建立和加速多芯混合集群的挑戰(zhàn)和方案
現(xiàn)在,我們回到今天本文討論的重點:建設(shè)一個支持多芯混合訓(xùn)練的 AI 集群,并加速運行一個大模型訓(xùn)練任務(wù)。(這里有個背景知識:要完成一個大模型的訓(xùn)練任務(wù),只能是在單一集群中完成,而不能拆分到不同集群中進行。)
為了實現(xiàn)這個目標,在前文提到的三個維度我們都將遇到挑戰(zhàn):
不同類型的卡的物理連接方式和集合通信庫是完全不同的,這導(dǎo)致他們之間無法直接互聯(lián)互通。我們需要為跨芯片的互聯(lián)互通設(shè)計一套新的物理網(wǎng)絡(luò)架構(gòu),并配套相應(yīng)的集合通信庫方案,使得通信效果達到最優(yōu)。
在集群建設(shè)完成后,由于不同種類卡的性能不一致,這就導(dǎo)致在單一芯片集群中,基于均勻切分計算量制定分布式并行策略的方法,將無法確保不同芯片可以按照相同節(jié)奏計算和通信,導(dǎo)致算力浪費。我們需要為此找到一套最優(yōu)的分布式并行策略,使得這些不同種類的芯片能夠在計算和通信的節(jié)奏上保持一致,使得集群算力效能最大化。
由于不同卡的加速方法存在差異,生態(tài)能力和調(diào)優(yōu)策略發(fā)展程度不一,這將導(dǎo)致集群中的各類算力效能無法充分發(fā)揮,即 MFU 未達到理想值( MFU 名詞解釋詳見文末)。我們需要為所有芯片提供統(tǒng)一的加速抽象以便屏蔽這些差異性,同時又能提供最優(yōu)的 AI 加速套件,確保各種芯片均可以達到 MFU 理想值,實現(xiàn)高效能運行。
3.1跨芯片的互聯(lián)互通,構(gòu)建多芯混合集群
傳統(tǒng)的觀點認為,不同芯片是很難互聯(lián)互通,無法支撐大模型訓(xùn)練。從上文可知,NVIDIA GPU、昆侖芯、昇騰 910B 的物理連接方式,以及使用的集合通信庫都不一樣。
百度百舸為了實現(xiàn)跨芯的互聯(lián)互通,使用了 CPU 轉(zhuǎn)發(fā)來實現(xiàn)跨昇騰 910B 子集群和 GPU 子集群的連接。借助百度自研的集合通信庫 BCCL,可以實現(xiàn) GPU、昆侖芯等標準 RDMA 設(shè)備的互聯(lián)互通,使得通信效果達到最優(yōu)。
3.2自適應(yīng)并行策略搜索,提升多芯混合訓(xùn)練任務(wù)的整體效能
傳統(tǒng)的分布式并行策略,都是按照等分的方式將大模型和訓(xùn)練數(shù)據(jù)拆開。其中確保這種等分方式有效的前提,在于集群中的芯片是同型號。這樣在算量相同的情況下,所有卡都可以同節(jié)奏地運行。
假設(shè)集群中存在兩種以上的卡,按照等分的方式制定分布式并行策略,則存在高性能卡等待其他卡的過程,產(chǎn)生算力的浪費。所以在同一個多芯集群中,我們需要按照不同芯片子集群的算力對比,分配合適的算量,將過去均勻切分的分布式并行策略改成按芯片算力大小適配的非均勻分布式并行策略方式。
在確定了非均勻并行策略的大方針后,還需要解決的就是具體怎么分,分多少的問題,確定分布式并行策略的最優(yōu)解。這個最優(yōu)解包括:采用什么樣組合的分布式并行策略;在不同芯片的子集群中分配多少算量,比如分配多少訓(xùn)練數(shù)據(jù),多少模型層數(shù)等。
百度百舸的 AI 加速套件 AIAK–LLM 實現(xiàn)了針對單一任務(wù)多種芯片的自適應(yīng)并行策略搜索功能,通過計算各種并行策略所需要的計算量、存儲量、通信量以及不同芯片的計算和 I/O 效率等(這里的所有數(shù)值均源自百度基于數(shù)十年 AI 技術(shù)積累產(chǎn)生的手冊:AI 芯片效能矩陣圖譜),從而快速計算出最優(yōu)的任務(wù)切分策略,保證在各種芯片配比下的單一集群,在運行大模型多芯混合訓(xùn)練任務(wù)時整體效能最大化。
3.3Accelerator抽象,屏蔽硬件差異充分發(fā)揮不同芯片的算力效能
傳統(tǒng)的加速方案,如上文所說,均是面向特定芯片進行優(yōu)化。如果同一個大模型任務(wù),運行在多芯混合集群上,我們需要為不同芯片配置相應(yīng)的加速策略。
國產(chǎn)化 AI 芯片由于生態(tài)還在不斷完善過程中,調(diào)優(yōu)策略仍然需要繼續(xù)優(yōu)化,所以導(dǎo)致算力沒有充分發(fā)揮出來。
百度百舸基于在 GPU 上的長期投入與沉淀(比如,百度百舸平臺的 A800 MFU 值達到 80。該值等于 A800 MFU 的理想值,即百度百舸完全發(fā)揮了 A800 的能力),在 AI 加速套件 AIAK-LLM 中構(gòu)建了「Accelerator 抽象層」,使得國產(chǎn)化 AI 芯片充分發(fā)揮各自算力。(關(guān)于 MFU 的含義,請參考文末名詞解釋)
其中,「Accelerator 抽象層」面向應(yīng)用,屏蔽底層芯片在硬件層面的差異,將芯片算子與上層策略解耦開來,芯片廠商僅需要進行各自芯片的算子調(diào)優(yōu)。百度百舸在 GPU 上沉淀的上層各項策略(比如通信 overlap、顯存 overlap 、張量并行、數(shù)據(jù)并行等)都會平滑地遷移到各種芯片上,確保了各種國產(chǎn)芯片在百度百舸上能達到一個非常高的運行效率。
3.4多芯混合訓(xùn)練的技術(shù)指標
在解決以上三個方面的挑戰(zhàn)后,我們就可以開始在這個集群上進行單一大模型下的多芯混合訓(xùn)練任務(wù)了。目前,百度百舸的百卡和千卡規(guī)模混合訓(xùn)練效能最大分別達到了 97% 和 95%。計算公式如下:
其中,芯片 A 集群吞吐量,為基于芯片 A 構(gòu)建的單一集群,在訓(xùn)練大模型時候的能力。MFU 的取值,使用的是經(jīng)過「AI 加速套件 AIAK-LLM」加速后在該集群獲得的 MFU 數(shù)值。
「Accelerator 抽象層」使得上層的各類策略都能平滑地遷移到不同卡層面,能夠有效提升各類芯片的 MFU 值。自適應(yīng)并行策略和百度集合通信庫 BCCL 等能夠有效提升訓(xùn)練任務(wù)的模型吞吐量。
拋開上文提到的具體參數(shù),關(guān)于混合訓(xùn)練效能指標的直觀理解:假設(shè)分別有 100 單位算力規(guī)模的 NVIDIA A800 集群,80 單位算力規(guī)模的由芯片 B 組成的 AI 集群,將這兩種芯片融合成為一個百卡規(guī)模的多芯混合集群后,新集群的訓(xùn)練效能相當(dāng)于 180*0.97=174.6 單位算力。
隨著百度百舸能力的不斷升級,這個多芯混合訓(xùn)練效能的數(shù)值將繼續(xù)提升。
4新舊算力統(tǒng)一融合,滿足未來業(yè)務(wù)增長
百度百舸的多芯混合訓(xùn)練方案,屏蔽了底層復(fù)雜的異構(gòu)環(huán)境,將各類芯片融合成為了一個大集群,可以實現(xiàn)存量不同算力的統(tǒng)一,整合發(fā)揮這些算力的最大效能,支持更大模型訓(xùn)練任務(wù)。同時,支持新增資源的快速融入,滿足未來業(yè)務(wù)增長的需要。
該方案不僅通過百度智能云的公有云提供服務(wù),同時還可以通過 ABC Stack 專有云進行交付。如果您的智算中心已經(jīng)或者計劃多種算力部署的打算,歡迎使用百度百舸的這項全新能力,打破單一算力的局限,實現(xiàn)算力的統(tǒng)一融合。
5名詞解釋
MFU,Model FLOPs Utilization,MFU =(實際觀測到的模型吞吐量)/ (假設(shè)峰值 FLOPs 下的理論最大吞吐量)
一個較高的 MFU 值意味著 GPU 的浮點運算能力被高效利用,模型訓(xùn)練的速度更快;而較低的 MFU 值則表明存在資源浪費,導(dǎo)致 GPU 的實際計算能力未得到充分施展。