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