Scale-Up 擴(kuò)展與故障容錯(cuò):NVIDIA 非均勻張量并行 精華
一、背景
最近華為推出了超節(jié)點(diǎn) CloudMatrix 384,進(jìn)一步引發(fā)業(yè)內(nèi)對(duì) Scale-Up 和 Scale-Out 的廣泛討論。不可避免地也會(huì)涉及與 NVIDIA 超節(jié)點(diǎn) NVL72 的對(duì)比。Scale-Up 和 Scale-Out 各自具有不同的優(yōu)劣勢(shì)和局限性。除了擴(kuò)展性和成本問題外,故障和容錯(cuò)也是一個(gè)不可忽略的挑戰(zhàn)。本文中,我們介紹一個(gè) NVIDIA 最近在這一領(lǐng)域的研究工作,著重探討隨著 Scale-Up 域的擴(kuò)展,如何應(yīng)對(duì)相應(yīng)的容錯(cuò)問題。
對(duì)應(yīng)的論文為:[2504.06095] Nonuniform-Tensor-Parallelism: Mitigating GPU failure impact for Scaled-up LLM Training [1]
二、摘要
LLM 預(yù)訓(xùn)練通常會(huì)采用 DP(Data Parallelism)、TP(Tensor Parallelism)、PP(Pipeline Parallelism)和 EP(Expert Parallelism)等混合并行策略,以便將訓(xùn)練任務(wù)擴(kuò)展到數(shù)千甚至數(shù)萬(wàn) GPU 集群中。實(shí)現(xiàn)高效訓(xùn)練的關(guān)鍵在于將 TP 執(zhí)行部署到緊密耦合的 GPU 子集(即 Scale-Up 域,比如一個(gè) Node 的 8 GPU),且 Scale-Up 域越大性能越優(yōu)。
當(dāng)前數(shù)據(jù)中心架構(gòu)的一個(gè)主要演進(jìn)趨勢(shì)就是 Scale-Up 域的擴(kuò)展,比如 NVIDIA 的 NVL72 方案,以及華為的 CloudMatrix 384 方案。然而,更大的 Scale-Up 域也會(huì)擴(kuò)大故障域的影響范圍:?jiǎn)蝹€(gè) GPU 的故障可能波及整個(gè) Scale-Up 域的 TP 執(zhí)行,從而顯著降低 LLM 整體訓(xùn)練吞吐量。當(dāng)僅有 0.1% 的 GPU 處于故障狀態(tài)時(shí),高 TP 維度的作業(yè)可能遭遇 19% 的訓(xùn)練吞吐下降。
為此,論文作者提出了非均勻 TP(Nonuniform-TP,NTP)機(jī)制以緩解 GPU 故障的級(jí)聯(lián)效應(yīng)。在 NTP 架構(gòu)下,遭遇 GPU 故障的 DP 副本會(huì)以降低后的 TP 維度繼續(xù)執(zhí)行,其貢獻(xiàn)的吞吐量與剩余正常 GPU 的比例相匹配。作者還設(shè)計(jì)了一種具備增強(qiáng)供電與散熱能力的機(jī)架方案,可對(duì)發(fā)生故障的 Scale-Up 域維持功率提升。結(jié)合 NTP 技術(shù),這種設(shè)計(jì)能確保 TP 維度降低的 DP 副本(即包含故障 GPU 的副本)與其他副本保持同步,最終實(shí)現(xiàn)大規(guī)模 LLM 訓(xùn)練中接近零損失的吞吐表現(xiàn)。
三、引言
3.1 常見 8 GPU 節(jié)點(diǎn)
如下圖所示為常見的 8 GPU 服務(wù)器拓?fù)?,通常包括?/p>
- 2 個(gè) CPU,CPU 之間通過 UPI 連接。
- 每個(gè) CPU 下有 2 個(gè) PCIe Switch(比如常見的 Broadcom PEX89104,PEX89144),總共 4 個(gè)。
- 每個(gè) PCIe Switch 下連接 2 個(gè) GPU,2 個(gè) NIC(如果一臺(tái)服務(wù)器 4 個(gè)后向 NIC,這里每個(gè) PCIe Switch 有 1 個(gè) NIC),一些 NVMe。
- 8 個(gè) GPU 通過 NVLink + NVSwitch 實(shí)現(xiàn)全互聯(lián)。
3.2 NVIDIA NVL72
常見的 NVLink + NVSwitch 全互聯(lián)通常在一臺(tái)單機(jī)內(nèi),比如常見的單機(jī) 8 GPU 服務(wù)器。而 Superpod 中,比如常見的 NVL72,將 NVLink + NVSwitch 全互聯(lián)的范圍進(jìn)一步擴(kuò)展到單個(gè)機(jī)柜甚至多個(gè)機(jī)柜內(nèi)。如下圖所示,DGX GB200 NVL72 將其擴(kuò)展到一個(gè)機(jī)柜的 72 個(gè) B200 GPU。
3.3 CloudMatrix 384
華為最也發(fā)布了 CloudMatrix 384 超節(jié)點(diǎn),可以將全互聯(lián)進(jìn)一步擴(kuò)展到 384 個(gè) Ascend 910C NPU。成為 NVIDIA NVL72 的有力競(jìng)爭(zhēng)者。
3.4 阿里 HPN
阿里在 Alibaba HPN: A Data Center Network for Large Language Model Training [3] 中介紹了其萬(wàn)卡集群的互聯(lián)方案。如下圖所示為其發(fā)表 Paper 之前介紹的拓?fù)浞绞剑▓D片來(lái)自 Revolutionizing Data Center Networks: Alibaba’s SONiC Journey [2]),是一個(gè)完全無(wú)收斂的方案。下圖的拓?fù)渲校?/p>
- 每個(gè) Segment 有 128 個(gè)節(jié)點(diǎn),共 1024 GPU(單層千卡)。
- 每個(gè) Pod 有 8 個(gè) Segment,也就是每個(gè) Pod 有 8192 GPU。
- 總共有 128 個(gè) Pod,也就是可以支持 1,048,576 個(gè) GPU(三層 100 萬(wàn))。
而在 HPN Paper 中的拓?fù)浞绞脚c上圖稍有不同(雙上聯(lián)、雙平面等思路都是完全一樣的),如下圖 Figure 7 所示:
- 下面的拓?fù)渲邪饲跋蚓W(wǎng)絡(luò)(Frontend Network)和后向網(wǎng)絡(luò)(Backend Network):
后向網(wǎng)絡(luò):有收斂,使用每個(gè)節(jié)點(diǎn) 9 個(gè) NIC 中的 NIC1-NIC9 這 8 個(gè)互聯(lián),主要用于大規(guī)模分布式訓(xùn)練,并且一個(gè) GPU 連接一個(gè) NIC。
前向網(wǎng)絡(luò):無(wú)收斂,使用每個(gè)節(jié)點(diǎn) 9 個(gè) NIC 中的 NIC0 互聯(lián)。為了支持更多的場(chǎng)景,比如訓(xùn)練/推理混部,模型傳輸,數(shù)據(jù)加載等場(chǎng)景。
- 后向網(wǎng)絡(luò)依然是 3 層:
- Segment:依然采用雙上聯(lián)方式,一個(gè) NIC 上有 2 個(gè) 200Gbps 的 Port(PS:沒有采用之前介紹的 2 個(gè) 200 Gbps NIC 的方式),會(huì)連接兩個(gè)不同的 ToR 交換機(jī)。
一個(gè) Segment 里面依然有 16 個(gè) ToR 交換機(jī),每個(gè)交換機(jī) 128 個(gè) 400Gbps Port,但是有 60 連接 Spine 交換機(jī),68 個(gè)連接節(jié)點(diǎn)的 NIC。
68 個(gè) 400Gbps Port 可以對(duì)應(yīng) 136 個(gè) 200Gbps NIC Port,也就是一個(gè) Segment 里面 136 個(gè)節(jié)點(diǎn),共 138*8=1104 個(gè) GPU。
實(shí)際上 136 個(gè)節(jié)點(diǎn)中有 8 個(gè)是備份,以便節(jié)點(diǎn)故障(比如 GPU、網(wǎng)卡、硬盤、CPU 等)時(shí)可以快速替換。實(shí)際使用 128 個(gè)節(jié)點(diǎn),共 1024 GPU,對(duì)應(yīng)的網(wǎng)絡(luò)收斂比為 (1024*400)/(60*400*16)=1.067:1。
- Pod:一個(gè) Pod 中的 Segment 從 8 個(gè)變成 15 個(gè),所以最多能支持 15*1024=15K GPU。
在 Spine(Agg)交換機(jī)上采用 15:1 的收斂比,因此可以有更多的下行 Port 連接 Leaf 交換機(jī)。
具體來(lái)說,每個(gè) Spine 交換機(jī)有 120 個(gè) Port 連接 Leaf 交換機(jī),也就可以連接 120/8=15 個(gè) Segment(每個(gè) Segment 里面同一平面的 8 個(gè) Leaf 交換機(jī)連接到同一個(gè) Spine 交換機(jī))。
Cluster:一個(gè) Cluster 可以包含多個(gè) Pod,通過 Core 交換機(jī)連接。
Spine(Agg) 交換機(jī)有 8 個(gè) Port 連接 Core 交換機(jī)。這個(gè)是為了支持更大規(guī)模的 GPU,比如 8 個(gè) Pod,則可以支持 120K GPU。
在大規(guī)模模型訓(xùn)練時(shí),可以將 PP(Pipeline Parallelism)中的不同切片放在不同的 Pod,這樣跨 Pod 的通信量比較小,也就不容易出現(xiàn)瓶頸。
3.5 張量并行 TP
3.5.1 Column Parallelism
如下圖所示為 Column Parallelism,其中的 Column 就是指權(quán)重參數(shù) W 按照 Column 維度切分。每個(gè) GPU 都包含一部分權(quán)重參數(shù),并使用整個(gè)輸入 X 計(jì)算,得到 Y 的一部分,最后通過 AllGather 操作可以獲得全量結(jié)果。
3.5.2 Row Parallelism
如下圖所示為 Row Parallelism,其中的 Row 就是指權(quán)重參數(shù) W 按照 Row 維度切分。每個(gè) GPU 都包含一部分權(quán)重參數(shù),并使用部分輸入 X 計(jì)算,結(jié)果和 Y 的 Shape 相同,但結(jié)果不完整,最后通過 AllReduce 操作可以獲得全量結(jié)果。因?yàn)?AllReduce 可以通過 ReduceScatter 和 AllGather 的方式實(shí)現(xiàn),而 Column Parallelism 中的 AllGather 和 Row Parallelism 中 AllGather 通信量是一樣的,因此,總體來(lái)說 Column Parallelism 的通信量更少:
3.5.3 Column Parallelism + Row Parallelism
在 Transformer 等模型中會(huì)存在連續(xù)兩個(gè)矩陣乘法(Linear Layer)的情況,此時(shí)通常都會(huì)采用先 Column Parallelism,之后 Row Parallelism 的方式切分,可以在兩個(gè) Linear 之間減少一次通信操作。如下圖所示,W 是第一個(gè) Linear 權(quán)重,V 是第二個(gè) Linear 權(quán)重。只用在最后進(jìn)行一次 AllReduce 操作即可:
3.5.4 TP 擴(kuò)展
無(wú)論采用何種具體的混合并行配置,LLM 訓(xùn)練任務(wù)在擴(kuò)展至多設(shè)備運(yùn)行時(shí),通信開銷會(huì)逐漸成為性能瓶頸。隨著 Scale-Up 域的擴(kuò)展,提供給各種分布式策略的機(jī)會(huì)也就更多,比如可以使用更大的 TP(AllReduce),EP(All2All),相應(yīng)通信均在 Scale-Up 域內(nèi)。
如下圖 Figure 2a 展示了在不同 NVL 域規(guī)模的集群上訓(xùn)練 480B 參數(shù) LLM 的結(jié)果(序列長(zhǎng)度 8K,每 Batch 16M token):在較小規(guī)模(8K GPU)下,增大 NVL 域?qū)π阅芴嵘邢蓿驗(yàn)榇藭r(shí)通信尚未成為主要瓶頸。但隨著規(guī)模擴(kuò)大,更大的 NVL 域能有效避免通信瓶頸——在 32K GPU規(guī)模下,NVL32(87%)與 NVL8(68%)的單 GPU 利用率差異接近 20%。需要說明的是,在更大的 Scale-Up 域中,將 TP 直接設(shè)置為域規(guī)模并非總是最優(yōu)配置。如下圖 Figure 2b 展示了固定 NVLink 域?yàn)?16 時(shí),相同訓(xùn)練任務(wù)和集群規(guī)模下不同混合并行配置的對(duì)比:作者搜索了配置空間,并繪制出 TP 限制為 8、16 及無(wú)限制時(shí)的最優(yōu)單 GPU 吞吐量曲線。隨著訓(xùn)練規(guī)模擴(kuò)大,必須采用更大 TP 以維持單 GPU 吞吐量——這是因?yàn)樵诟笠?guī)模下,繼續(xù)增加 DP 或 PP 會(huì)加大 Bubble Ratio,而過度增加 DP 則會(huì)增加 AllReduce 通信時(shí)間。
PS:對(duì)于現(xiàn)在常見的 MoE 模型而言,細(xì)粒度 Expert 變得越來(lái)越流行,而細(xì)粒度 Expert 其實(shí)不適合 TP 切分,會(huì)降低算術(shù)強(qiáng)度,不利于 GPU 的高效計(jì)算。為此,可以采用 EP,擴(kuò)展 Scale-Up 域?qū)?All2All 也是有幫助的。
3.6 異常 & 容錯(cuò)
對(duì)于大規(guī)模任務(wù),通常都會(huì)采用整機(jī)調(diào)度的方式,一個(gè)節(jié)點(diǎn)內(nèi)的 GPU 有 NVLink + NVSwitch 高速互聯(lián),也會(huì)將通信量較大的 TP 放在一個(gè)節(jié)點(diǎn)內(nèi),當(dāng)然,這也限制了 TP 的大小通常不會(huì)超過單節(jié)點(diǎn)的 8 個(gè) GPU。同時(shí),當(dāng)一個(gè) GPU 異常時(shí),為了盡可能保持分布式訓(xùn)練的高效性,會(huì)屏蔽整個(gè)節(jié)點(diǎn)(PS:如果只屏蔽單個(gè) GPU,可能導(dǎo)致 TP 跨機(jī),會(huì)極大影響訓(xùn)練性能)。
當(dāng)前 Scale-Up 域的 GPU 規(guī)模也在不斷增大,比如 NVL72 達(dá)到了 72 個(gè) GPU 的 NVLink + NVSwitch 高效互聯(lián),也就為 TP 等分布式策略提供了更大的空間。然而,這也進(jìn)一步擴(kuò)大了單點(diǎn)故障的影響范圍——當(dāng)某個(gè) GPU 異常時(shí),整個(gè) TP 域都會(huì)受到影響。以 TP64 為例,當(dāng) 0.1% GPU 處于異常狀態(tài)時(shí),可能導(dǎo)致近 10% 的 GPU 將無(wú)法充分發(fā)揮算力。對(duì)于一個(gè) 32K GPU 的訓(xùn)練任務(wù)而言,意味著約 3K GPU 將無(wú)法充分發(fā)揮算力。
為說明該問題,作者以大型 NVLink 域 GPU 系統(tǒng)為例開展研究。如下圖 Figure 3 展示了 32K GPU 集群在不同 TP 配置下,其可用性隨均勻分布故障 GPU 數(shù)量變化的關(guān)系。在相同故障 GPU 數(shù)量條件下,較大 TP(需更大 Scale-Up 域)會(huì)因故障域規(guī)模放大而顯著降低可用性。以 TP64 為例,僅需 0.1% 的 GPU 處于故障狀態(tài)即可使集群可用性降至 94%。
此類故障場(chǎng)景并不罕見。如下圖 Figure 4 所示,作者基于 Llama 3 訓(xùn)練報(bào)告(The Llama 3 Herd of Models | Research - AI at Meta [4])中 NVIDIA H100 GPU 集群的故障率數(shù)據(jù)進(jìn)行了故障模擬:設(shè)定 78% 為硬件故障(如報(bào)告所述),恢復(fù)時(shí)間 3/5天(對(duì)高性能硬件更換而言可能偏短),其余 22% 為軟件故障(恢復(fù)時(shí)間 3 小時(shí))。在 15 天的追蹤期內(nèi),集群 81%的時(shí)間存在 > 0.1% 的 GPU 故障,該故障量足以使 TP64 配置的可用性降至 94%。
隨著最新硬件(如 TPU-POD、GB200-NVL72)復(fù)雜度的提升——這類系統(tǒng)包含更多組件(例如更高容量的 HBM、更多用于擴(kuò)展帶寬的線纜),其故障率預(yù)計(jì)將顯著高于 LLaMA 3 訓(xùn)練報(bào)告中 NVL8 系統(tǒng)的記錄值。此外,故障率隨時(shí)間呈現(xiàn)顯著波動(dòng)并可能出現(xiàn)峰值突變:Meta 的 [2410.21680] Revisiting Reliability in Large-Scale Machine Learning Research Clusters [5] 中指出 16K 規(guī)模 A100 集群的故障率波動(dòng)幅度可達(dá) 7 倍。假設(shè)某場(chǎng)景故障率較 LLaMA 3 報(bào)告觀測(cè)值提升 3 倍時(shí),峰值并發(fā)故障數(shù)將增加 2 倍,足以將系統(tǒng)可用性降至 80%。由此可見,面對(duì)日益復(fù)雜的現(xiàn)代硬件架構(gòu),必須設(shè)計(jì)具備更強(qiáng)故障容忍能力的系統(tǒng)。
理想的故障解決方案應(yīng)滿足以下三個(gè)核心要求:
- 無(wú)需或僅需極少量備用資源即可維持固定 Mini-Batch 規(guī)模。
- 吞吐量下降幅度與 GPU 故障率嚴(yán)格成正比(即不會(huì)因 NVL域規(guī)模/TP 規(guī)模導(dǎo)致故障放大)。
- 通過 GPU 共享機(jī)制恢復(fù)因故障放大損失的 GPU 利用率。
要實(shí)現(xiàn)目標(biāo) 1 和 2,唯一途徑是通過降低 TP 來(lái)利用部分故障的 NVL 域,從而抑制故障放大效應(yīng)。但該方案面臨三大技術(shù)挑戰(zhàn):
- 單一 PP 階段的 TP 降低可能導(dǎo)致 DP 副本內(nèi)其他 PP Stage 出現(xiàn)計(jì)算瓶頸。
- 單個(gè) DP 副本 的 TP 降低會(huì)拖慢參數(shù)同步,進(jìn)而制約其他(正常)DP 副本的運(yùn)行效率。
- 不同 TP 配置的副本間參數(shù)同步機(jī)制尚屬未知領(lǐng)域,其可能引入的額外開銷也會(huì)拖慢訓(xùn)練進(jìn)程。
四、系統(tǒng)設(shè)計(jì)
4.1 概述
本文提出非均勻張量并行(Nonuniform Tensor Parallelism, NTP),作為一種彈性分布式訓(xùn)練方法來(lái)解決上述問題。在 NTP 框架下,遭遇一個(gè)或多個(gè) GPU 故障的 DP 副本將重新配置為以低于其他完整健康 DP 副本的 TP 規(guī)模運(yùn)行(例如,當(dāng)某 Sacle-Up 域內(nèi) 64 個(gè) GPU 中有 2 個(gè)異常時(shí),在故障修復(fù)期間,該副本將采用 TP62 運(yùn)行)。顯然,由于重構(gòu)后的 DP 副本使用較少 GPU 運(yùn)行相同模型,其吞吐量預(yù)期會(huì)下降。為避免這種持續(xù)性滯后拖累整體同步執(zhí)行效率,最簡(jiǎn)解決方案是為重構(gòu)副本減少輸入樣本數(shù)量——這在傳統(tǒng)機(jī)器架構(gòu)下可將訓(xùn)練吞吐量影響降至最低。
通過采用支持故障 GPU 所在 Scale-Up 域動(dòng)態(tài)超頻的機(jī)架設(shè)計(jì)方案,近零吞吐?lián)p失成為可能:僅加速單個(gè)重構(gòu) TP 實(shí)例即可使其與其他副本保持同步,且平均能耗增幅近乎為零。該系統(tǒng)顯著降低了對(duì)備用設(shè)備的依賴需求,但仍保留必要時(shí)啟用備用設(shè)備的容災(zāi)能力。
4.2 Nonuniform Tensor Parallelism
在 Transformer 層中,張量并行(Tensor Parallelism)被應(yīng)用于 MLP 模塊和 Attention 模塊。如下所示,其切分策略我們?cè)谇懊嬉呀?jīng)詳細(xì)介紹,這里不再贅述。
通常,矩陣 A 和 B 的列/行會(huì)被分成連續(xù)的片段,以構(gòu)成 TP 的分片,但并不是嚴(yán)格要求的。實(shí)際上,A/B 的任何行或列都可以放在在任意地分片中,只要 B/A 中對(duì)應(yīng)的行或列被放在同一個(gè)分片中即可。換句話說,不管每個(gè) Zi 在哪里計(jì)算,只要最后能做 AllReduce,就能得到最終的 Z。
例如,如果某個(gè) TP 副本必須在 TP=N2 下運(yùn)行(原始 TP=N1,N1 > N2),那么只需對(duì) A/B 的列或行做分片(可以連續(xù)的也可以非連續(xù)的),在 N1 個(gè)分片中,那么只需將異常的 N1 - N2 分片對(duì)應(yīng)的數(shù)據(jù)分給 N2 計(jì)算即可。
這樣做的挑戰(zhàn)在于不同 DP 副本之間的梯度同步:當(dāng) N1 = N2 且 A/B 在各副本間采用完全相同地分片策略時(shí),每個(gè)分片只需與存有完全相同 A/B 的 列/行(及對(duì)應(yīng)梯度)的其他副本中的單一對(duì)應(yīng)分片同步。若簡(jiǎn)單的將 N1 > N2 分片對(duì)應(yīng)的數(shù)據(jù)分給 N2,則會(huì)引入一些小的時(shí)延敏感的通信,進(jìn)而導(dǎo)致負(fù)載的不均衡,并且 N2 越接近 N1 越明顯。以隱藏層維度 12K 為例,假設(shè):
- DP1 副本對(duì)應(yīng) TP 為 N1=32,每個(gè)切片維度 375。
- DP2 副本對(duì)應(yīng) TP 為 N2=30,每個(gè)切片維度為 400。
此時(shí)通信變?yōu)椋?/p>
- DP1 中的 30 個(gè)切片要與 DP2 的 30 個(gè)切片進(jìn)行 1 對(duì) 1 通信,通信維度為 375。
- DP1 中的另 2 個(gè)切片要與 DP2 的 30 個(gè)切片進(jìn)行 1 對(duì)多通信,通信維度為 375/15=25。
分片映射算法:理想情況下,期望在梯度同步時(shí)(兩個(gè) DP 副本均使用 N2 個(gè) GPU)保持分片間的一一映射關(guān)系。同時(shí)需保持同步分片的連續(xù)性,以便融合為單一操作來(lái)最小化時(shí)延開銷。這要求將健康副本中的梯度從 N1 分片組重新切分為 N2 分片組。由于該重分區(qū)操作在 TP 組內(nèi)完成(通常屬于 Scale-Up 域),若采用重疊執(zhí)行則不會(huì)成為同步瓶頸,但仍需通過利用更高帶寬實(shí)現(xiàn)并行最大化來(lái)最小化重分區(qū)時(shí)間。
Attention 模塊:Attention 模塊也包含兩個(gè)矩陣乘操作,也可以應(yīng)用 TP。不過在 Attention 模塊通常是沿著 Head 維度進(jìn)行并行切分,假設(shè) Head 個(gè)數(shù)是 32,則通常會(huì)按照 TP8、TP4 或 TP2 切分。此時(shí)如果存在 GPU 異常,則不均衡問題會(huì)很明顯。
4.3 Dynamic power allocation
當(dāng)某個(gè) TP 分區(qū)中的部分 GPU 發(fā)生故障時(shí),NTP 技術(shù)會(huì)將計(jì)算任務(wù)重新分配給剩余的 GPU,同時(shí)保持 Local Batch 大小不變。這種機(jī)制實(shí)際上增加了單個(gè) GPU 的計(jì)算負(fù)載。例如,在一個(gè)包含 8 個(gè) GPU的 TP 組中,若一個(gè) GPU 故障,其余每個(gè) GPU 需額外承擔(dān) 14.3% 的計(jì)算操作。這種簡(jiǎn)單的任務(wù)重分配方式可能導(dǎo)致該 GPU 組成為大規(guī)模同步訓(xùn)練的性能瓶頸。
為緩解由此產(chǎn)生的性能下降,作者提出了一種創(chuàng)新的機(jī)架電源動(dòng)態(tài)分配方案。該設(shè)計(jì)允許將故障 GPU 的功率預(yù)算重新調(diào)配給同機(jī)架內(nèi)正常工作的 GPU,使其在不降低 Local Batch 規(guī)模的前提下維持全吞吐量運(yùn)行。這種機(jī)架設(shè)計(jì)可為剩余工作 GPU 提供最高 TDP 30% 的額外功率,通過提升運(yùn)行頻率實(shí)現(xiàn)單 GPU 性能增益。實(shí)驗(yàn)表明,該方案無(wú)需依賴機(jī)架備用 GPU 即可近乎完全消除因 GPU 故障導(dǎo)致的性能損失。
上述技術(shù)在現(xiàn)代數(shù)據(jù)中心中比較容易實(shí)現(xiàn)。需要說明的是,英偉達(dá) GH200 系列 GPU 已具備動(dòng)態(tài)功率平衡功能,允許 GPU 突破 700W 額定功耗限制(H200 型號(hào)可支持 1000W),實(shí)際運(yùn)行功率最高可達(dá) 900W。進(jìn)一步來(lái)看,從 Ampere 架構(gòu)(A100-SXM,400W)到 Hopper 架構(gòu)(H100-SXM,700W),再到 Blackwell 架構(gòu)(B200-SXM/GB200,1000W/1200W),NVIDIA GPU 的功耗預(yù)算已增長(zhǎng)超過 50%。這表明,盡管 GPU 散熱是一個(gè)難題,但風(fēng)冷和液冷技術(shù)的創(chuàng)新正推動(dòng)芯片制造商在每一代產(chǎn)品中將 TDP 推向更高極限。因此,作者認(rèn)為所增加的電氣與熱力學(xué)要求,并不會(huì)對(duì)提出的機(jī)架設(shè)計(jì)方案構(gòu)成不合理的挑戰(zhàn)。
通過將上述優(yōu)化的電氣與熱管理方案與 NTP 所提供的計(jì)算靈活性相結(jié)合,其設(shè)計(jì)在局部故障場(chǎng)景下仍能保持穩(wěn)定的訓(xùn)練吞吐量,且無(wú)需冗余硬件帶來(lái)的額外開銷。這一成果驗(yàn)證了動(dòng)態(tài)功率分配技術(shù)在數(shù)據(jù)中心系統(tǒng)中的實(shí)際可行性及其顯著優(yōu)勢(shì)。
4.4 Resource manager
若某訓(xùn)練任務(wù)采用 PP 策略,且某個(gè) DP 副本內(nèi)存在部分失效的 PP Stage,剩余的正常 PP Stage 將受這些部分失效 Stage 的瓶頸制約。緩解此問題的一種途徑是通過 PP Stage 的 Re-Balancing 技術(shù),但該技術(shù)復(fù)雜度極高,可能無(wú)法兼容更復(fù)雜的 PP 調(diào)度方案(如 1F1B),還會(huì)引發(fā)高度復(fù)雜的(即多對(duì)多)PP Stage 參數(shù)同步問題(因不同 DP 副本會(huì)采用不同的 PP Stage 劃分方案)。
為此,作者選擇將部分失效的 Scale-Up 域"打包"至盡可能少的 DP 副本中,并使包含任何部分失效 Scale-Up 域的 DP 副本以降級(jí)的域/TP 規(guī)模運(yùn)行。故障發(fā)生時(shí),任務(wù)必須重啟;重啟時(shí)進(jìn)程組 Rank 會(huì)被重新分配,通過將故障機(jī)架集中置于最低 Rank 實(shí)現(xiàn)故障域聚合。該策略最小化了受故障影響的 DP 副本數(shù)量,從而最大限度緩解(盡管無(wú)法徹底消除)PP Stage 瓶頸問題。剩余的正常但受瓶頸制約的 Scale-Up 域被迫以低于其潛在能力的域/TP 規(guī)模運(yùn)行,但這些正常域閑置的 GPU 資源可被重新分配執(zhí)行其他工作負(fù)載,避免資源空置。
PS:阿里在 FALCON: Pinpointing and Mitigating Stragglers for Large-Scale Hybrid-Parallel Training [6] 中也提到過類似的解決方案。
五、實(shí)現(xiàn)
5.1 NTP 實(shí)現(xiàn)
如下圖 Figure 5(上)展示了 NTP 重分片與梯度同步重疊的概覽。作者在 NVIDIA Megatron 框架上實(shí)現(xiàn) NTP。AllReduce 前的重分片作為 PyTorch Backward Hook 的一部分實(shí)現(xiàn)(即與最終 Backward 過程重疊);該 Hook 用于將梯度標(biāo)記為“準(zhǔn)備同步”(當(dāng)桶中所有梯度均被標(biāo)記為就緒時(shí),整個(gè)桶進(jìn)行同步),在標(biāo)記梯度就緒前先對(duì)其進(jìn)行重分片。同步后的重分片操作與最后一個(gè)桶的 AllReduce 同步執(zhí)行——這是因?yàn)?Megatron 為保障性能穩(wěn)定性/可預(yù)測(cè)性設(shè)置了 CUDA_DEVICE_MAX_CONNECTIONS=1,為確保同步后重分片前所有梯度已完成同步,需等待最終桶的 AllReduce 完成。在評(píng)估中,將實(shí)現(xiàn)的開銷分解為:
- 同步前重分片對(duì)最終 Backward 的開銷。
- AllReduce 數(shù)據(jù)量增加(通信設(shè)備數(shù)降低)導(dǎo)致的通信開銷。
- AllReduce 時(shí)間內(nèi)的同步后重分片開銷。
如下圖 Figure 5(下)為 PyTorch Trace 結(jié)果,同步前重分片操作(與 TP 通信算子同 Stream 執(zhí)行)與 Backward 重疊,同步后重分片操作則與最終 AllReduce 重疊。
流水行并行(PP):在 Transformer 層中,MLP/Attention 模塊的輸出會(huì)在 TP 分片間復(fù)制。而 PP Stage 邊界始終位于 Transformer 層之間,因此 PP Stage 的輸出也在 TP 分片間復(fù)制。而 PP Stage 間發(fā)送激活受限于跨 Stage 的 IB/以太網(wǎng)帶寬,實(shí)際上 PP 引發(fā)的 P2P 通信在端到端延遲中占比極小。不過,若某個(gè) PP Stage 的 TP 規(guī)模縮減,則會(huì)按比例降低跨 Stage 聚合帶寬。對(duì)于 NTP(無(wú)功耗重分配),所有 Stage 都在縮減的 TP 規(guī)模下運(yùn)行,因此只需以較低帶寬發(fā)送激活即可。而對(duì)于 NTP-PW,縮減的 TP Stage 可能需要與正常 TP Stage 互傳激活——此時(shí)激活以縮減 TP 規(guī)模的相應(yīng)比例帶寬進(jìn)行交換,隨后(如有需要)再?gòu)V播至更大 TP 組的額外 GPU 中;廣播發(fā)生在 Scale-Up 擴(kuò)展域內(nèi),可與接收激活操作重疊,實(shí)際上不產(chǎn)生額外開銷。在 NTP 和 NTP-PW 的模擬中,作者計(jì)入了以較低跨階段帶寬 Send/Recive 激活的開銷,并將所有必要的廣播視為完全重疊操作。
5.2 性能建模和預(yù)估
作者在現(xiàn)有系統(tǒng)上進(jìn)行了概念驗(yàn)證設(shè)計(jì)的評(píng)估,但 NTP 的主要應(yīng)用場(chǎng)景針對(duì) B200-NVL72 等具備更大 Scale-Up 能力的系統(tǒng)。鑒于此類系統(tǒng)尚未廣泛普及,作者采用高性能模擬器來(lái)評(píng)估 NTP 的優(yōu)勢(shì)——該模擬器能精準(zhǔn)建模大規(guī)模多節(jié)點(diǎn) GPU 系統(tǒng)的運(yùn)行表現(xiàn)。該模擬器具備高度復(fù)雜性,其特點(diǎn)包括:
- 對(duì)底層 GPU 微架構(gòu)的精細(xì)化建模。
- LLM 應(yīng)用并行映射的精確模擬。
- 通信/網(wǎng)絡(luò)行為的真實(shí)再現(xiàn)。
考慮到 LLM 中每個(gè) GPU 的工作負(fù)載具有高度一致性,模擬器會(huì)根據(jù)所采用的并行配置策略,以單 GPU 為單位進(jìn)行任務(wù)劃分。在后續(xù)的結(jié)果分析中,通過對(duì)比模擬器預(yù)測(cè)性能與實(shí)際系統(tǒng)測(cè)量數(shù)據(jù)的相關(guān)性研究,證實(shí)了模擬器的保真度,從而為預(yù)測(cè)數(shù)值提供可靠的理論支撐。
具體實(shí)現(xiàn)層面:
- 將 LLM 建模為可分割的計(jì)算圖結(jié)構(gòu),基于并行策略劃分并嵌入相應(yīng)的通信操作。
- 綜合考量 GPU 微架構(gòu)特性和系統(tǒng)規(guī)模,對(duì)單 GPU 上的計(jì)算與通信操作進(jìn)行性能建模。
- 支持不同計(jì)算-通信重疊策略的仿真模擬。
該模擬器具備雙重?cái)U(kuò)展能力:既可模擬 NVL72(及以上規(guī)格)的大規(guī)模 Scale-Up 系統(tǒng),也能仿真 Scale-Out 的大型計(jì)算集群。除應(yīng)用程序的大規(guī)模性能預(yù)估外,其功能還包括:
- 系統(tǒng)功耗估算。
- 通過類電源管理機(jī)制提升設(shè)備性能表現(xiàn)。
六、實(shí)驗(yàn) & 結(jié)果
6.1 原型評(píng)估
作者構(gòu)建了一個(gè)系統(tǒng)原型作為概念驗(yàn)證,并針對(duì) NTP 開銷進(jìn)行了一系列敏感性研究。該原型在 2×DGX-A100 計(jì)算節(jié)點(diǎn)上進(jìn)行評(píng)估,每個(gè)節(jié)點(diǎn)配備 8 個(gè) 80GB 顯存的 A100 GPU(NVLink 帶寬600GBps)、8 個(gè) 200Gbps IB NIC。實(shí)驗(yàn)對(duì)象為隱層維度 12288 與 6144 的 LLM 訓(xùn)練過程,其 Attention Head 維度為 128,F(xiàn)FN 維度為隱層維度的 4 倍,序列長(zhǎng)度介于 4K 至 16K 之間。
為精準(zhǔn)測(cè)量本方法中 DP 與 TP 的協(xié)同開銷(PP 開銷可忽略不計(jì)),采用1 個(gè) PP Stage 配置,2 個(gè) DP 副本(每個(gè)節(jié)點(diǎn) 1 個(gè))。每個(gè) DP 副本在參數(shù)同步前,基于 1-2 個(gè)樣本的 Local Batch 對(duì) 1-3 個(gè)網(wǎng)絡(luò)層進(jìn)行訓(xùn)練。通過 NTP 技術(shù)動(dòng)態(tài)調(diào)整每個(gè)副本的 TP 規(guī)模,以量化分析并行策略的開銷特性。
作者在大規(guī)模不同 TDP 條件下驗(yàn)證了性能模型的準(zhǔn)確性。針對(duì)這些實(shí)驗(yàn),作者在 DGX-H100 平臺(tái)上進(jìn)行了訓(xùn)練過程的性能剖析與模擬。
6.2 大規(guī)模模擬 & 敏感性研究
為了在大規(guī)模、大 NVL 域環(huán)境下評(píng)估 NTP 與 NTP-PW 性能,作者采用自主研發(fā)的仿真系統(tǒng)進(jìn)行實(shí)驗(yàn)。實(shí)驗(yàn)?zāi)M了基于 32,000 個(gè) B200 GPU(單卡顯存 189 GB)的訓(xùn)練集群,其中 NVL 域規(guī)模為 32 個(gè) GPU(單卡帶寬 1.8TB/s),每個(gè) GPU 配備 1 個(gè) 800Gbps IB NIC。訓(xùn)練模型參數(shù)量達(dá) 480B,具體架構(gòu)為:隱含層維度 20,480,Attention Head 個(gè)數(shù)為 128,F(xiàn)FN 維度為隱含層的 4 倍,共 100 個(gè) Layer。訓(xùn)練設(shè)置包括 16K 的序列長(zhǎng)度,每個(gè) Batch 1600 萬(wàn) Token。
如前所述,包含部分失效設(shè)備的 DP 副本組必須通過以下兩種方式避免對(duì)健康 DP 副本組形成性能瓶頸:降低 Local Batch 規(guī)模(DP-DROP)或?qū)Σ糠质У?NVL 域?qū)嵤┕β侍嵘?。本?shí)驗(yàn)在 TP32 配置下運(yùn)行,同時(shí)支持降級(jí)至 TP30 和 TP28 模式。針對(duì)降級(jí)后的 TP 配置,通過仿真計(jì)算得出兩種解決方案的關(guān)鍵參數(shù):非功率提升模式下的最大 Local Batch 規(guī)模,以及功率提升模式下的最低運(yùn)行功耗,確保降級(jí) TP 副本的迭代時(shí)間不超過健康副本的基準(zhǔn)值。如下圖 Table 1 完整列出了所有配置方案及其對(duì)應(yīng)的性能指標(biāo)。
6.3 主要結(jié)果
如下圖 Figure 6 所示,根據(jù) Figure 4 觀測(cè)到的故障比例,沿 x 軸調(diào)整 GPU 故障比例參數(shù)。針對(duì)每個(gè)故障比例值,分別計(jì)算各容錯(cuò)方法的吞吐量損失值。由于 GPU 故障分布模式會(huì)影響最終吞吐量,作者對(duì)每個(gè)比例參數(shù)進(jìn)行大量故障場(chǎng)景采樣,并繪制各場(chǎng)景的均值曲線。實(shí)驗(yàn)數(shù)據(jù)顯示:
- DP-DROP 方案的最大吞吐量降幅達(dá) 12%;
- NTP 方案將降幅控制在 3% 以內(nèi);
- NTP-PW 方案在故障比例 ≤4×10?3 時(shí)仍能保持 <1% 的吞吐量損失。
與之對(duì)應(yīng),如下圖 Figure 7 將 Mini-Batch 規(guī)模設(shè)為固定參數(shù)。當(dāng)故障導(dǎo)致實(shí)際 Mini-Batch 規(guī)模低于目標(biāo)值時(shí),訓(xùn)練進(jìn)程將暫停直至足夠多故障 GPU 恢復(fù)運(yùn)行。本實(shí)驗(yàn)采用與 LLaMA 3.1 觀測(cè)值相同的故障率參數(shù),硬件故障恢復(fù)時(shí)間設(shè)為 5 天。通過增加備用 NVL 域數(shù)量,繪制單 GPU 吞吐量變化曲線。結(jié)果表明:
- NTP 方案僅需 16 個(gè)備用 NVL 域即可實(shí)現(xiàn)連續(xù)訓(xùn)練(無(wú)暫停),這是因?yàn)樵撚?xùn)練任務(wù)每個(gè) DP 副本使用 8 個(gè) NVL 域,且基礎(chǔ) NTP 方案(無(wú)備用域)的吞吐量降幅始終不超過等效丟失 2 個(gè) DP 副本的水平。
- DP-DROP 方案需要 90 個(gè)備用 NVL 域才能實(shí)現(xiàn)不間斷訓(xùn)練,這反而會(huì)降低單 GPU 吞吐量——因?yàn)橐_(dá)到與 16 個(gè)備用 NVL 域的 NTP 方案相同吞吐量,需要額外配置 90 個(gè)NVL 域。
- NTP-PW 方案在此故障場(chǎng)景下無(wú)需備用域即可實(shí)現(xiàn):
- 訓(xùn)練不中斷;
- 相較無(wú)故障基準(zhǔn)的吞吐量損失 <1%。
6.4 原型系統(tǒng)評(píng)估
作者進(jìn)一步對(duì)原型系統(tǒng)進(jìn)行性能評(píng)估,以量化 NTP 機(jī)制的開銷。由于 pre-sync 重分片操作與最終 Backward 過程存在重疊,重點(diǎn)測(cè)量其對(duì) Backward 階段的性能影響。如下圖 Figure 8 所示,作者測(cè)試了不同隱藏層維度和序列長(zhǎng)度的多種工作負(fù)載。實(shí)驗(yàn)采用一個(gè) TP8 并行度、Local Batch 為 2 的 DP 副本,與另一個(gè)降低 TP 并行度(Local Batch 為 1)的副本進(jìn)行對(duì)比訓(xùn)練,并在必須執(zhí)行重分片的 TP8 副本上測(cè)量 Backward 時(shí)延(基準(zhǔn)為:2 副本均采用 TP8,Local Batch 2 的訓(xùn)練配置)。
作者提出假設(shè):性能下降主要受兩個(gè)參數(shù)影響:
- 總 Backward 計(jì)算量——計(jì)算量越大,重分片操作獲得重疊優(yōu)化的機(jī)會(huì)越多;
- 單 GPU 在重分片過程中的最大數(shù)據(jù)傳輸量——該值直接決定重分片耗時(shí)。
其中 1)受模型規(guī)模和序列長(zhǎng)度影響,2)則與模型規(guī)模及降低后的 TP 并行度相關(guān)(TP 并行度降幅越大,重分片引發(fā)的通信開銷越高)。通過將 2)除以 1)計(jì)算通信計(jì)算比,并將其作為 Figure 8 橫坐標(biāo),縱坐標(biāo)表示 Backward 時(shí)延增幅。
實(shí)驗(yàn)證實(shí):
- 通信計(jì)算比與性能降幅呈顯著線性相關(guān),這意味著模型規(guī)模越大或序列越長(zhǎng),性能降幅越小。
- TP 并行度降幅越大(即初始 TP 并行度失效比例越高),性能降幅越顯著。這是因?yàn)楸M管實(shí)現(xiàn)方案將大部分重分片操作與 Backward 重疊,仍存在部分無(wú)法隱藏的重分片操作;當(dāng) TP 降幅增大(需要更多重分片數(shù)據(jù)傳輸)時(shí),這些未隱藏操作的出現(xiàn)概率也隨之增加。
- 在所有工作負(fù)載和設(shè)置中,最終 Backward 的減速幅度最多僅為4%。
6.5 仿真驗(yàn)證
在如下圖 Figure 11b 中,在 DGX-H100 節(jié)點(diǎn)上采用 FP8 與 BF16 兩種精度格式,針對(duì)多種模型規(guī)模(8B 至 175B)、不同序列長(zhǎng)度(2K 至 8K)以及在多個(gè)計(jì)算規(guī)模(8 至 512 個(gè) GPU)條件下展開訓(xùn)練實(shí)驗(yàn)。通過系統(tǒng)性地搜索各工作負(fù)載的最佳并行配置方案,并將實(shí)際硬件平臺(tái)觀測(cè)到的吞吐量與模擬器預(yù)測(cè)值進(jìn)行對(duì)比分析??梢钥闯觯M器的預(yù)測(cè)結(jié)果與實(shí)際性能表現(xiàn)高度吻合。
如下圖 Figure 11a 展示了采用 DGX-H100 GPU 集群訓(xùn)練兩種規(guī)模模型(15B 和 340B)時(shí),在不同單 GPU 功耗限制條件下的實(shí)驗(yàn)結(jié)果。將實(shí)測(cè)訓(xùn)練性能數(shù)據(jù)與模擬器預(yù)測(cè)值進(jìn)行對(duì)比繪制,結(jié)果表明模擬器的預(yù)測(cè)輸出與實(shí)測(cè)數(shù)據(jù)間具有極高的相關(guān)性。
6.5 敏感性研究
功率提升機(jī)制允許部分失效的 NVL 域以更高吞吐量運(yùn)行(以避免對(duì)健康域形成瓶頸)。理論上也可通過提升健康域功率來(lái)進(jìn)一步提高其吞吐量。但功率提升會(huì)帶來(lái)性能/watt 比下降及數(shù)據(jù)中心能耗需求增加的代價(jià)。如 Table 1 中 TP30 配置所示,1.1 倍基準(zhǔn)功率時(shí)性能/watt 降低 2.8%,1.2 倍時(shí)降幅達(dá) 6.5%。鑒于 NTP-PW 方案僅對(duì)異常機(jī)架實(shí)施功率提升,最壞情況下僅 12% 的 NVL 域會(huì)承受這種效率損失。且該機(jī)制通過重新分配故障 GPU 的功率資源,不會(huì)額外增加供電需求。
實(shí)驗(yàn)設(shè)計(jì)中通常假設(shè)單 GPU 故障僅影響所在 NVL 域的一個(gè) GPU。[2503.11901] Characterizing GPU Resilience and Impact on AI/HPC Systems [7] 中表明:91% 的故障為不可控的內(nèi)存錯(cuò)誤或 MMU 錯(cuò)誤(僅影響單個(gè) GPU),5% 為可能擴(kuò)散的 NVLink 錯(cuò)誤。但在 GB200-NVL72 等架構(gòu)中,整個(gè)節(jié)點(diǎn)棄置(含 36 CPU + 72 GPU)比部分 GPU 隔離更易實(shí)施。如下圖 Figure 10 展示了故障影響范圍(單 GPU 故障導(dǎo)致的連帶失效 GPU 數(shù)量)對(duì) NTP 方案的影響:
- DP-DROP 因需丟棄整個(gè) DP 副本而始終維持高有效影響范圍;
- NTP 與 NTP-PW 雖隨影響范圍擴(kuò)大出現(xiàn)吞吐?lián)p失加劇,但性能仍顯著優(yōu)于 DP-DROP 方案。
七、參考鏈接:
- [1] https://arxiv.org/abs/2504.06095
- [2] https://sonicfoundation.dev/revolutionizing-data-center-networks-alibabas-sonic-journey/
- [3] https://ennanzhai.github.io/pub/sigcomm24-hpn.pdf
- [4] https://ai.meta.com/research/publications/the-llama-3-herd-of-models/
- [5] https://arxiv.org/abs/2410.21680
- [6] https://arxiv.org/pdf/2410.12588
- [7] https://arxiv.org/abs/2503.11901?
本文轉(zhuǎn)載自?????????AI閑談?????????,作者:AI閑談
