一、背景

最近華為推出了超節(jié)點(diǎn) CloudMatrix 384，進(jìn)一步引發(fā)業(yè)內(nèi)對 Scale-Up 和 Scale-Out 的廣泛討論。不可避免地也會涉及與 NVIDIA 超節(jié)點(diǎn) NVL72 的對比。Scale-Up 和 Scale-Out 各自具有不同的優(yōu)劣勢和局限性。除了擴(kuò)展性和成本問題外，故障和容錯也是一個不可忽略的挑戰(zhàn)。本文中，我們介紹一個 NVIDIA 最近在這一領(lǐng)域的研究工作，著重探討隨著 Scale-Up 域的擴(kuò)展，如何應(yīng)對相應(yīng)的容錯問題。

對應(yīng)的論文為：[2504.06095] Nonuniform-Tensor-Parallelism: Mitigating GPU failure impact for Scaled-up LLM Training [1]

二、摘要

LLM 預(yù)訓(xùn)練通常會采用 DP（Data Parallelism）、TP（Tensor Parallelism）、PP（Pipeline Parallelism）和 EP（Expert Parallelism）等混合并行策略，以便將訓(xùn)練任務(wù)擴(kuò)展到數(shù)千甚至數(shù)萬 GPU 集群中。實現(xiàn)高效訓(xùn)練的關(guān)鍵在于將 TP 執(zhí)行部署到緊密耦合的 GPU 子集（即 Scale-Up 域，比如一個 Node 的 8 GPU），且 Scale-Up 域越大性能越優(yōu)。

當(dāng)前數(shù)據(jù)中心架構(gòu)的一個主要演進(jìn)趨勢就是 Scale-Up 域的擴(kuò)展，比如 NVIDIA 的 NVL72 方案，以及華為的 CloudMatrix 384 方案。然而，更大的 Scale-Up 域也會擴(kuò)大故障域的影響范圍：單個 GPU 的故障可能波及整個 Scale-Up 域的 TP 執(zhí)行，從而顯著降低 LLM 整體訓(xùn)練吞吐量。當(dāng)僅有 0.1% 的 GPU 處于故障狀態(tài)時，高 TP 維度的作業(yè)可能遭遇 19% 的訓(xùn)練吞吐下降。

為此，論文作者提出了非均勻 TP（Nonuniform-TP，NTP）機(jī)制以緩解 GPU 故障的級聯(lián)效應(yīng)。在 NTP 架構(gòu)下，遭遇 GPU 故障的 DP 副本會以降低后的 TP 維度繼續(xù)執(zhí)行，其貢獻(xiàn)的吞吐量與剩余正常 GPU 的比例相匹配。作者還設(shè)計了一種具備增強(qiáng)供電與散熱能力的機(jī)架方案，可對發(fā)生故障的 Scale-Up 域維持功率提升。結(jié)合 NTP 技術(shù)，這種設(shè)計能確保 TP 維度降低的 DP 副本（即包含故障 GPU 的副本）與其他副本保持同步，最終實現(xiàn)大規(guī)模 LLM 訓(xùn)練中接近零損失的吞吐表現(xiàn)。

三、引言

3.1 常見 8 GPU 節(jié)點(diǎn)

如下圖所示為常見的 8 GPU 服務(wù)器拓?fù)洌ǔ０ǎ?/p>

• 2 個 CPU，CPU 之間通過 UPI 連接。
• 每個 CPU 下有 2 個 PCIe Switch（比如常見的 Broadcom PEX89104，PEX89144），總共 4 個。
• 每個 PCIe Switch 下連接 2 個 GPU，2 個 NIC（如果一臺服務(wù)器 4 個后向 NIC，這里每個 PCIe Switch 有 1 個 NIC），一些 NVMe。
• 8 個 GPU 通過 NVLink + NVSwitch 實現(xiàn)全互聯(lián)。

3.2 NVIDIA NVL72

常見的 NVLink + NVSwitch 全互聯(lián)通常在一臺單機(jī)內(nèi)，比如常見的單機(jī) 8 GPU 服務(wù)器。而 Superpod 中，比如常見的 NVL72，將 NVLink + NVSwitch 全互聯(lián)的范圍進(jìn)一步擴(kuò)展到單個機(jī)柜甚至多個機(jī)柜內(nèi)。如下圖所示，DGX GB200 NVL72 將其擴(kuò)展到一個機(jī)柜的 72 個 B200 GPU。

3.3 CloudMatrix 384

華為最也發(fā)布了 CloudMatrix 384 超節(jié)點(diǎn)，可以將全互聯(lián)進(jìn)一步擴(kuò)展到 384 個 Ascend 910C NPU。成為 NVIDIA NVL72 的有力競爭者。

3.4 阿里 HPN

阿里在 Alibaba HPN: A Data Center Network for Large Language Model Training [3] 中介紹了其萬卡集群的互聯(lián)方案。如下圖所示為其發(fā)表 Paper 之前介紹的拓?fù)浞绞剑▓D片來自 Revolutionizing Data Center Networks: Alibaba’s SONiC Journey [2]），是一個完全無收斂的方案。下圖的拓?fù)渲校?/p>

• 每個 Segment 有 128 個節(jié)點(diǎn)，共 1024 GPU（單層千卡）。
• 每個 Pod 有 8 個 Segment，也就是每個 Pod 有 8192 GPU。
• 總共有 128 個 Pod，也就是可以支持 1,048,576 個 GPU（三層 100 萬）。

而在 HPN Paper 中的拓?fù)浞绞脚c上圖稍有不同（雙上聯(lián)、雙平面等思路都是完全一樣的），如下圖 Figure 7 所示：

• 下面的拓?fù)渲邪饲跋蚓W(wǎng)絡(luò)（Frontend Network）和后向網(wǎng)絡(luò)（Backend Network）：

后向網(wǎng)絡(luò)：有收斂，使用每個節(jié)點(diǎn) 9 個 NIC 中的 NIC1-NIC9 這 8 個互聯(lián)，主要用于大規(guī)模分布式訓(xùn)練，并且一個 GPU 連接一個 NIC。

前向網(wǎng)絡(luò)：無收斂，使用每個節(jié)點(diǎn) 9 個 NIC 中的 NIC0 互聯(lián)。為了支持更多的場景，比如訓(xùn)練/推理混部，模型傳輸，數(shù)據(jù)加載等場景。

• 后向網(wǎng)絡(luò)依然是 3 層：
• Segment：依然采用雙上聯(lián)方式，一個 NIC 上有 2 個 200Gbps 的 Port（PS：沒有采用之前介紹的 2 個 200 Gbps NIC 的方式），會連接兩個不同的 ToR 交換機(jī)。

一個 Segment 里面依然有 16 個 ToR 交換機(jī)，每個交換機(jī) 128 個 400Gbps Port，但是有 60 連接 Spine 交換機(jī)，68 個連接節(jié)點(diǎn)的 NIC。

68 個 400Gbps Port 可以對應(yīng) 136 個 200Gbps NIC Port，也就是一個 Segment 里面 136 個節(jié)點(diǎn)，共 138*8=1104 個 GPU。

實際上 136 個節(jié)點(diǎn)中有 8 個是備份，以便節(jié)點(diǎn)故障（比如 GPU、網(wǎng)卡、硬盤、CPU 等）時可以快速替換。實際使用 128 個節(jié)點(diǎn)，共 1024 GPU，對應(yīng)的網(wǎng)絡(luò)收斂比為 (1024*400)/(60*400*16)=1.067:1。

• Pod：一個 Pod 中的 Segment 從 8 個變成 15 個，所以最多能支持 15*1024=15K GPU。

• 在 Spine（Agg）交換機(jī)上采用 15:1 的收斂比，因此可以有更多的下行 Port 連接 Leaf 交換機(jī)。

• 具體來說，每個 Spine 交換機(jī)有 120 個 Port 連接 Leaf 交換機(jī)，也就可以連接 120/8=15 個 Segment（每個 Segment 里面同一平面的 8 個 Leaf 交換機(jī)連接到同一個 Spine 交換機(jī)）。

• Cluster：一個 Cluster 可以包含多個 Pod，通過 Core 交換機(jī)連接。

• Spine（Agg） 交換機(jī)有 8 個 Port 連接 Core 交換機(jī)。這個是為了支持更大規(guī)模的 GPU，比如 8 個 Pod，則可以支持 120K GPU。

• 在大規(guī)模模型訓(xùn)練時，可以將 PP（Pipeline Parallelism）中的不同切片放在不同的 Pod，這樣跨 Pod 的通信量比較小，也就不容易出現(xiàn)瓶頸。

3.5 張量并行 TP

3.5.1 Column Parallelism

如下圖所示為 Column Parallelism，其中的 Column 就是指權(quán)重參數(shù) W 按照 Column 維度切分。每個 GPU 都包含一部分權(quán)重參數(shù)，并使用整個輸入 X 計算，得到 Y 的一部分，最后通過 AllGather 操作可以獲得全量結(jié)果。

3.5.2 Row Parallelism

如下圖所示為 Row Parallelism，其中的 Row 就是指權(quán)重參數(shù) W 按照 Row 維度切分。每個 GPU 都包含一部分權(quán)重參數(shù)，并使用部分輸入 X 計算，結(jié)果和 Y 的 Shape 相同，但結(jié)果不完整，最后通過 AllReduce 操作可以獲得全量結(jié)果。因為 AllReduce 可以通過 ReduceScatter 和 AllGather 的方式實現(xiàn)，而 Column Parallelism 中的 AllGather 和 Row Parallelism 中 AllGather 通信量是一樣的，因此，總體來說 Column Parallelism 的通信量更少：

3.5.3 Column Parallelism + Row Parallelism

在 Transformer 等模型中會存在連續(xù)兩個矩陣乘法（Linear Layer）的情況，此時通常都會采用先 Column Parallelism，之后 Row Parallelism 的方式切分，可以在兩個 Linear 之間減少一次通信操作。如下圖所示，W 是第一個 Linear 權(quán)重，V 是第二個 Linear 權(quán)重。只用在最后進(jìn)行一次 AllReduce 操作即可：

3.5.4 TP 擴(kuò)展

無論采用何種具體的混合并行配置，LLM 訓(xùn)練任務(wù)在擴(kuò)展至多設(shè)備運(yùn)行時，通信開銷會逐漸成為性能瓶頸。隨著 Scale-Up 域的擴(kuò)展，提供給各種分布式策略的機(jī)會也就更多，比如可以使用更大的 TP（AllReduce），EP（All2All），相應(yīng)通信均在 Scale-Up 域內(nèi)。

如下圖 Figure 2a 展示了在不同 NVL 域規(guī)模的集群上訓(xùn)練 480B 參數(shù) LLM 的結(jié)果（序列長度 8K，每 Batch 16M token）：在較小規(guī)模（8K GPU）下，增大 NVL 域?qū)π阅芴嵘邢蓿驗榇藭r通信尚未成為主要瓶頸。但隨著規(guī)模擴(kuò)大，更大的 NVL 域能有效避免通信瓶頸——在 32K GPU規(guī)模下，NVL32（87%）與 NVL8（68%）的單 GPU 利用率差異接近 20%。需要說明的是，在更大的 Scale-Up 域中，將 TP 直接設(shè)置為域規(guī)模并非總是最優(yōu)配置。如下圖 Figure 2b 展示了固定 NVLink 域為 16 時，相同訓(xùn)練任務(wù)和集群規(guī)模下不同混合并行配置的對比：作者搜索了配置空間，并繪制出 TP 限制為 8、16 及無限制時的最優(yōu)單 GPU 吞吐量曲線。隨著訓(xùn)練規(guī)模擴(kuò)大，必須采用更大 TP 以維持單 GPU 吞吐量——這是因為在更大規(guī)模下，繼續(xù)增加 DP 或 PP 會加大 Bubble Ratio，而過度增加 DP 則會增加 AllReduce 通信時間。

PS：對于現(xiàn)在常見的 MoE 模型而言，細(xì)粒度 Expert 變得越來越流行，而細(xì)粒度 Expert 其實不適合 TP 切分，會降低算術(shù)強(qiáng)度，不利于 GPU 的高效計算。為此，可以采用 EP，擴(kuò)展 Scale-Up 域?qū)?All2All 也是有幫助的。

3.6 異常 & 容錯

對于大規(guī)模任務(wù)，通常都會采用整機(jī)調(diào)度的方式，一個節(jié)點(diǎn)內(nèi)的 GPU 有 NVLink + NVSwitch 高速互聯(lián)，也會將通信量較大的 TP 放在一個節(jié)點(diǎn)內(nèi)，當(dāng)然，這也限制了 TP 的大小通常不會超過單節(jié)點(diǎn)的 8 個 GPU。同時，當(dāng)一個 GPU 異常時，為了盡可能保持分布式訓(xùn)練的高效性，會屏蔽整個節(jié)點(diǎn)（PS：如果只屏蔽單個 GPU，可能導(dǎo)致 TP 跨機(jī)，會極大影響訓(xùn)練性能）。

當(dāng)前 Scale-Up 域的 GPU 規(guī)模也在不斷增大，比如 NVL72 達(dá)到了 72 個 GPU 的 NVLink + NVSwitch 高效互聯(lián)，也就為 TP 等分布式策略提供了更大的空間。然而，這也進(jìn)一步擴(kuò)大了單點(diǎn)故障的影響范圍——當(dāng)某個 GPU 異常時，整個 TP 域都會受到影響。以 TP64 為例，當(dāng) 0.1% GPU 處于異常狀態(tài)時，可能導(dǎo)致近 10% 的 GPU 將無法充分發(fā)揮算力。對于一個 32K GPU 的訓(xùn)練任務(wù)而言，意味著約 3K GPU 將無法充分發(fā)揮算力。

為說明該問題，作者以大型 NVLink 域 GPU 系統(tǒng)為例開展研究。如下圖 Figure 3 展示了 32K  GPU 集群在不同 TP 配置下，其可用性隨均勻分布故障 GPU 數(shù)量變化的關(guān)系。在相同故障 GPU 數(shù)量條件下，較大 TP（需更大 Scale-Up 域）會因故障域規(guī)模放大而顯著降低可用性。以 TP64 為例，僅需 0.1% 的 GPU 處于故障狀態(tài)即可使集群可用性降至 94%。

此類故障場景并不罕見。如下圖 Figure 4 所示，作者基于 Llama 3 訓(xùn)練報告（The Llama 3 Herd of Models | Research - AI at Meta [4]）中 NVIDIA H100 GPU 集群的故障率數(shù)據(jù)進(jìn)行了故障模擬：設(shè)定 78% 為硬件故障（如報告所述），恢復(fù)時間 3/5天（對高性能硬件更換而言可能偏短），其余 22% 為軟件故障（恢復(fù)時間 3 小時）。在 15 天的追蹤期內(nèi)，集群 81%的時間存在 > 0.1% 的 GPU 故障，該故障量足以使 TP64 配置的可用性降至 94%。

隨著最新硬件（如 TPU-POD、GB200-NVL72）復(fù)雜度的提升——這類系統(tǒng)包含更多組件（例如更高容量的 HBM、更多用于擴(kuò)展帶寬的線纜），其故障率預(yù)計將顯著高于 LLaMA 3 訓(xùn)練報告中 NVL8 系統(tǒng)的記錄值。此外，故障率隨時間呈現(xiàn)顯著波動并可能出現(xiàn)峰值突變：Meta 的 [2410.21680] Revisiting Reliability in Large-Scale Machine Learning Research Clusters [5] 中指出 16K 規(guī)模 A100 集群的故障率波動幅度可達(dá) 7 倍。假設(shè)某場景故障率較 LLaMA 3 報告觀測值提升 3 倍時，峰值并發(fā)故障數(shù)將增加 2 倍，足以將系統(tǒng)可用性降至 80%。由此可見，面對日益復(fù)雜的現(xiàn)代硬件架構(gòu)，必須設(shè)計具備更強(qiáng)故障容忍能力的系統(tǒng)。

理想的故障解決方案應(yīng)滿足以下三個核心要求：

1. 無需或僅需極少量備用資源即可維持固定 Mini-Batch 規(guī)模。
2. 吞吐量下降幅度與 GPU 故障率嚴(yán)格成正比（即不會因 NVL域規(guī)模/TP 規(guī)模導(dǎo)致故障放大）。
3. 通過 GPU 共享機(jī)制恢復(fù)因故障放大損失的 GPU 利用率。

要實現(xiàn)目標(biāo) 1 和 2，唯一途徑是通過降低 TP 來利用部分故障的 NVL 域，從而抑制故障放大效應(yīng)。但該方案面臨三大技術(shù)挑戰(zhàn)：

1. 單一 PP 階段的 TP 降低可能導(dǎo)致 DP 副本內(nèi)其他 PP Stage 出現(xiàn)計算瓶頸。
2. 單個 DP 副本 的 TP 降低會拖慢參數(shù)同步，進(jìn)而制約其他（正常）DP 副本的運(yùn)行效率。
3. 不同 TP 配置的副本間參數(shù)同步機(jī)制尚屬未知領(lǐng)域，其可能引入的額外開銷也會拖慢訓(xùn)練進(jìn)程。

四、系統(tǒng)設(shè)計

4.1 概述

本文提出非均勻張量并行（Nonuniform Tensor Parallelism, NTP），作為一種彈性分布式訓(xùn)練方法來解決上述問題。在 NTP 框架下，遭遇一個或多個 GPU 故障的 DP 副本將重新配置為以低于其他完整健康 DP 副本的 TP 規(guī)模運(yùn)行（例如，當(dāng)某 Sacle-Up 域內(nèi) 64 個 GPU 中有 2 個異常時，在故障修復(fù)期間，該副本將采用 TP62 運(yùn)行）。顯然，由于重構(gòu)后的 DP 副本使用較少 GPU 運(yùn)行相同模型，其吞吐量預(yù)期會下降。為避免這種持續(xù)性滯后拖累整體同步執(zhí)行效率，最簡解決方案是為重構(gòu)副本減少輸入樣本數(shù)量——這在傳統(tǒng)機(jī)器架構(gòu)下可將訓(xùn)練吞吐量影響降至最低。

通過采用支持故障 GPU 所在 Scale-Up 域動態(tài)超頻的機(jī)架設(shè)計方案，近零吞吐?lián)p失成為可能：僅加速單個重構(gòu) TP 實例即可使其與其他副本保持同步，且平均能耗增幅近乎為零。該系統(tǒng)顯著降低了對備用設(shè)備的依賴需求，但仍保留必要時啟用備用設(shè)備的容災(zāi)能力。

4.2 Nonuniform Tensor Parallelism

在 Transformer 層中，張量并行（Tensor Parallelism）被應(yīng)用于 MLP 模塊和 Attention 模塊。如下所示，其切分策略我們在前面已經(jīng)詳細(xì)介紹，這里不再贅述。

通常，矩陣 A 和 B 的列/行會被分成連續(xù)的片段，以構(gòu)成 TP 的分片，但并不是嚴(yán)格要求的。實際上，A/B 的任何行或列都可以放在在任意地分片中，只要 B/A 中對應(yīng)的行或列被放在同一個分片中即可。換句話說，不管每個 Zi 在哪里計算，只要最后能做 AllReduce，就能得到最終的 Z。

例如，如果某個 TP 副本必須在 TP=N2 下運(yùn)行（原始 TP=N1，N1 > N2），那么只需對 A/B 的列或行做分片（可以連續(xù)的也可以非連續(xù)的），在 N1 個分片中，那么只需將異常的 N1 - N2 分片對應(yīng)的數(shù)據(jù)分給 N2 計算即可。

這樣做的挑戰(zhàn)在于不同 DP 副本之間的梯度同步：當(dāng) N1 = N2 且 A/B 在各副本間采用完全相同地分片策略時，每個分片只需與存有完全相同 A/B 的 列/行（及對應(yīng)梯度）的其他副本中的單一對應(yīng)分片同步。若簡單的將 N1 > N2 分片對應(yīng)的數(shù)據(jù)分給 N2，則會引入一些小的時延敏感的通信，進(jìn)而導(dǎo)致負(fù)載的不均衡，并且 N2 越接近 N1 越明顯。以隱藏層維度 12K 為例，假設(shè)：

• DP1 副本對應(yīng) TP 為 N1=32，每個切片維度 375。
• DP2 副本對應(yīng) TP 為 N2=30，每個切片維度為 400。

此時通信變?yōu)椋?/p>

• DP1 中的 30 個切片要與 DP2 的 30 個切片進(jìn)行 1 對 1 通信，通信維度為 375。
• DP1 中的另 2 個切片要與 DP2 的 30 個切片進(jìn)行 1 對多通信，通信維度為 375/15=25。

分片映射算法：理想情況下，期望在梯度同步時（兩個 DP 副本均使用 N2 個 GPU）保持分片間的一一映射關(guān)系。同時需保持同步分片的連續(xù)性，以便融合為單一操作來最小化時延開銷。這要求將健康副本中的梯度從 N1 分片組重新切分為 N2 分片組。由于該重分區(qū)操作在 TP 組內(nèi)完成（通常屬于 Scale-Up 域），若采用重疊執(zhí)行則不會成為同步瓶頸，但仍需通過利用更高帶寬實現(xiàn)并行最大化來最小化重分區(qū)時間。

Attention 模塊：Attention 模塊也包含兩個矩陣乘操作，也可以應(yīng)用 TP。不過在 Attention 模塊通常是沿著 Head 維度進(jìn)行并行切分，假設(shè) Head 個數(shù)是 32，則通常會按照 TP8、TP4 或 TP2 切分。此時如果存在 GPU 異常，則不均衡問題會很明顯。

4.3 Dynamic power allocation

當(dāng)某個 TP 分區(qū)中的部分 GPU 發(fā)生故障時，NTP 技術(shù)會將計算任務(wù)重新分配給剩余的 GPU，同時保持 Local Batch 大小不變。這種機(jī)制實際上增加了單個 GPU 的計算負(fù)載。例如，在一個包含 8 個 GPU的 TP 組中，若一個 GPU 故障，其余每個 GPU 需額外承擔(dān) 14.3% 的計算操作。這種簡單的任務(wù)重分配方式可能導(dǎo)致該 GPU 組成為大規(guī)模同步訓(xùn)練的性能瓶頸。

為緩解由此產(chǎn)生的性能下降，作者提出了一種創(chuàng)新的機(jī)架電源動態(tài)分配方案。該設(shè)計允許將故障 GPU 的功率預(yù)算重新調(diào)配給同機(jī)架內(nèi)正常工作的 GPU，使其在不降低 Local Batch 規(guī)模的前提下維持全吞吐量運(yùn)行。這種機(jī)架設(shè)計可為剩余工作 GPU 提供最高 TDP 30% 的額外功率，通過提升運(yùn)行頻率實現(xiàn)單 GPU 性能增益。實驗表明，該方案無需依賴機(jī)架備用 GPU 即可近乎完全消除因 GPU 故障導(dǎo)致的性能損失。

上述技術(shù)在現(xiàn)代數(shù)據(jù)中心中比較容易實現(xiàn)。需要說明的是，英偉達(dá) GH200 系列 GPU 已具備動態(tài)功率平衡功能，允許 GPU 突破 700W 額定功耗限制（H200 型號可支持 1000W），實際運(yùn)行功率最高可達(dá) 900W。進(jìn)一步來看，從 Ampere 架構(gòu)（A100-SXM，400W）到 Hopper 架構(gòu)（H100-SXM，700W），再到 Blackwell 架構(gòu)（B200-SXM/GB200，1000W/1200W），NVIDIA GPU 的功耗預(yù)算已增長超過 50%。這表明，盡管 GPU 散熱是一個難題，但風(fēng)冷和液冷技術(shù)的創(chuàng)新正推動芯片制造商在每一代產(chǎn)品中將 TDP 推向更高極限。因此，作者認(rèn)為所增加的電氣與熱力學(xué)要求，并不會對提出的機(jī)架設(shè)計方案構(gòu)成不合理的挑戰(zhàn)。

通過將上述優(yōu)化的電氣與熱管理方案與 NTP 所提供的計算靈活性相結(jié)合，其設(shè)計在局部故障場景下仍能保持穩(wěn)定的訓(xùn)練吞吐量，且無需冗余硬件帶來的額外開銷。這一成果驗證了動態(tài)功率分配技術(shù)在數(shù)據(jù)中心系統(tǒng)中的實際可行性及其顯著優(yōu)勢。

4.4 Resource manager

若某訓(xùn)練任務(wù)采用 PP 策略，且某個 DP 副本內(nèi)存在部分失效的 PP Stage，剩余的正常 PP Stage 將受這些部分失效 Stage 的瓶頸制約。緩解此問題的一種途徑是通過 PP Stage 的 Re-Balancing 技術(shù)，但該技術(shù)復(fù)雜度極高，可能無法兼容更復(fù)雜的 PP 調(diào)度方案（如 1F1B），還會引發(fā)高度復(fù)雜的（即多對多）PP Stage 參數(shù)同步問題（因不同 DP 副本會采用不同的 PP Stage 劃分方案）。

為此，作者選擇將部分失效的 Scale-Up 域"打包"至盡可能少的 DP 副本中，并使包含任何部分失效 Scale-Up 域的 DP 副本以降級的域/TP 規(guī)模運(yùn)行。故障發(fā)生時，任務(wù)必須重啟；重啟時進(jìn)程組 Rank 會被重新分配，通過將故障機(jī)架集中置于最低 Rank 實現(xiàn)故障域聚合。該策略最小化了受故障影響的 DP 副本數(shù)量，從而最大限度緩解（盡管無法徹底消除）PP Stage 瓶頸問題。剩余的正常但受瓶頸制約的 Scale-Up 域被迫以低于其潛在能力的域/TP 規(guī)模運(yùn)行，但這些正常域閑置的 GPU 資源可被重新分配執(zhí)行其他工作負(fù)載，避免資源空置。

PS：阿里在 FALCON: Pinpointing and Mitigating Stragglers for Large-Scale Hybrid-Parallel Training [6] 中也提到過類似的解決方案。

五、實現(xiàn)

5.1 NTP 實現(xiàn)

如下圖 Figure 5（上）展示了 NTP 重分片與梯度同步重疊的概覽。作者在 NVIDIA Megatron 框架上實現(xiàn) NTP。AllReduce 前的重分片作為 PyTorch Backward Hook 的一部分實現(xiàn)（即與最終 Backward 過程重疊）；該 Hook 用于將梯度標(biāo)記為“準(zhǔn)備同步”（當(dāng)桶中所有梯度均被標(biāo)記為就緒時，整個桶進(jìn)行同步），在標(biāo)記梯度就緒前先對其進(jìn)行重分片。同步后的重分片操作與最后一個桶的 AllReduce 同步執(zhí)行——這是因為 Megatron 為保障性能穩(wěn)定性/可預(yù)測性設(shè)置了 CUDA_DEVICE_MAX_CONNECTIONS=1，為確保同步后重分片前所有梯度已完成同步，需等待最終桶的 AllReduce 完成。在評估中，將實現(xiàn)的開銷分解為：

• 同步前重分片對最終 Backward 的開銷。
• AllReduce 數(shù)據(jù)量增加（通信設(shè)備數(shù)降低）導(dǎo)致的通信開銷。
• AllReduce 時間內(nèi)的同步后重分片開銷。

如下圖 Figure 5（下）為 PyTorch Trace 結(jié)果，同步前重分片操作（與 TP 通信算子同 Stream 執(zhí)行）與 Backward 重疊，同步后重分片操作則與最終 AllReduce 重疊。

流水行并行（PP）：在 Transformer 層中，MLP/Attention 模塊的輸出會在 TP 分片間復(fù)制。而 PP Stage 邊界始終位于 Transformer 層之間，因此 PP Stage 的輸出也在 TP 分片間復(fù)制。而 PP Stage 間發(fā)送激活受限于跨 Stage 的 IB/以太網(wǎng)帶寬，實際上 PP 引發(fā)的 P2P 通信在端到端延遲中占比極小。不過，若某個 PP Stage 的 TP 規(guī)模縮減，則會按比例降低跨 Stage 聚合帶寬。對于 NTP（無功耗重分配），所有 Stage 都在縮減的 TP 規(guī)模下運(yùn)行，因此只需以較低帶寬發(fā)送激活即可。而對于 NTP-PW，縮減的 TP Stage 可能需要與正常 TP Stage 互傳激活——此時激活以縮減 TP 規(guī)模的相應(yīng)比例帶寬進(jìn)行交換，隨后（如有需要）再廣播至更大 TP 組的額外 GPU 中；廣播發(fā)生在 Scale-Up 擴(kuò)展域內(nèi)，可與接收激活操作重疊，實際上不產(chǎn)生額外開銷。在 NTP 和 NTP-PW 的模擬中，作者計入了以較低跨階段帶寬 Send/Recive 激活的開銷，并將所有必要的廣播視為完全重疊操作。

5.2 性能建模和預(yù)估

作者在現(xiàn)有系統(tǒng)上進(jìn)行了概念驗證設(shè)計的評估，但 NTP 的主要應(yīng)用場景針對 B200-NVL72 等具備更大 Scale-Up 能力的系統(tǒng)。鑒于此類系統(tǒng)尚未廣泛普及，作者采用高性能模擬器來評估 NTP 的優(yōu)勢——該模擬器能精準(zhǔn)建模大規(guī)模多節(jié)點(diǎn) GPU 系統(tǒng)的運(yùn)行表現(xiàn)。該模擬器具備高度復(fù)雜性，其特點(diǎn)包括：

• 對底層 GPU 微架構(gòu)的精細(xì)化建模。
• LLM 應(yīng)用并行映射的精確模擬。
• 通信/網(wǎng)絡(luò)行為的真實再現(xiàn)。

考慮到 LLM 中每個 GPU 的工作負(fù)載具有高度一致性，模擬器會根據(jù)所采用的并行配置策略，以單 GPU 為單位進(jìn)行任務(wù)劃分。在后續(xù)的結(jié)果分析中，通過對比模擬器預(yù)測性能與實際系統(tǒng)測量數(shù)據(jù)的相關(guān)性研究，證實了模擬器的保真度，從而為預(yù)測數(shù)值提供可靠的理論支撐。

具體實現(xiàn)層面：

• 將 LLM 建模為可分割的計算圖結(jié)構(gòu)，基于并行策略劃分并嵌入相應(yīng)的通信操作。
• 綜合考量 GPU 微架構(gòu)特性和系統(tǒng)規(guī)模，對單 GPU 上的計算與通信操作進(jìn)行性能建模。
• 支持不同計算-通信重疊策略的仿真模擬。

該模擬器具備雙重擴(kuò)展能力：既可模擬 NVL72（及以上規(guī)格）的大規(guī)模 Scale-Up 系統(tǒng)，也能仿真 Scale-Out 的大型計算集群。除應(yīng)用程序的大規(guī)模性能預(yù)估外，其功能還包括：

• 系統(tǒng)功耗估算。
• 通過類電源管理機(jī)制提升設(shè)備性能表現(xiàn)。

六、實驗 & 結(jié)果

6.1 原型評估

作者構(gòu)建了一個系統(tǒng)原型作為概念驗證，并針對 NTP 開銷進(jìn)行了一系列敏感性研究。該原型在 2×DGX-A100 計算節(jié)點(diǎn)上進(jìn)行評估，每個節(jié)點(diǎn)配備 8 個 80GB 顯存的 A100 GPU（NVLink 帶寬600GBps）、8 個 200Gbps IB NIC。實驗對象為隱層維度 12288 與 6144 的 LLM 訓(xùn)練過程，其 Attention Head 維度為 128，F(xiàn)FN 維度為隱層維度的 4 倍，序列長度介于 4K 至 16K 之間。

為精準(zhǔn)測量本方法中 DP 與 TP 的協(xié)同開銷（PP 開銷可忽略不計），采用1 個 PP Stage 配置，2 個 DP 副本（每個節(jié)點(diǎn) 1 個）。每個 DP 副本在參數(shù)同步前，基于 1-2 個樣本的 Local Batch 對 1-3 個網(wǎng)絡(luò)層進(jìn)行訓(xùn)練。通過 NTP 技術(shù)動態(tài)調(diào)整每個副本的 TP 規(guī)模，以量化分析并行策略的開銷特性。

作者在大規(guī)模不同 TDP 條件下驗證了性能模型的準(zhǔn)確性。針對這些實驗，作者在 DGX-H100 平臺上進(jìn)行了訓(xùn)練過程的性能剖析與模擬。

6.2 大規(guī)模模擬 & 敏感性研究

為了在大規(guī)模、大 NVL 域環(huán)境下評估 NTP 與 NTP-PW 性能，作者采用自主研發(fā)的仿真系統(tǒng)進(jìn)行實驗。實驗?zāi)M了基于 32,000 個 B200 GPU（單卡顯存 189 GB）的訓(xùn)練集群，其中 NVL 域規(guī)模為 32 個 GPU（單卡帶寬 1.8TB/s），每個 GPU 配備 1 個 800Gbps IB NIC。訓(xùn)練模型參數(shù)量達(dá) 480B，具體架構(gòu)為：隱含層維度 20,480，Attention Head 個數(shù)為 128，F(xiàn)FN 維度為隱含層的 4 倍，共 100 個 Layer。訓(xùn)練設(shè)置包括 16K 的序列長度，每個 Batch 1600 萬 Token。

如前所述，包含部分失效設(shè)備的 DP 副本組必須通過以下兩種方式避免對健康 DP 副本組形成性能瓶頸：降低 Local Batch 規(guī)模（DP-DROP）或?qū)Σ糠质У?NVL 域?qū)嵤┕β侍嵘１緦嶒炘?TP32 配置下運(yùn)行，同時支持降級至 TP30 和 TP28 模式。針對降級后的 TP 配置，通過仿真計算得出兩種解決方案的關(guān)鍵參數(shù)：非功率提升模式下的最大 Local Batch 規(guī)模，以及功率提升模式下的最低運(yùn)行功耗，確保降級 TP 副本的迭代時間不超過健康副本的基準(zhǔn)值。如下圖 Table 1 完整列出了所有配置方案及其對應(yīng)的性能指標(biāo)。

6.3 主要結(jié)果

如下圖 Figure 6 所示，根據(jù) Figure 4 觀測到的故障比例，沿 x 軸調(diào)整 GPU 故障比例參數(shù)。針對每個故障比例值，分別計算各容錯方法的吞吐量損失值。由于 GPU 故障分布模式會影響最終吞吐量，作者對每個比例參數(shù)進(jìn)行大量故障場景采樣，并繪制各場景的均值曲線。實驗數(shù)據(jù)顯示： 

• DP-DROP 方案的最大吞吐量降幅達(dá) 12%；
• NTP 方案將降幅控制在 3% 以內(nèi)；
• NTP-PW 方案在故障比例 ≤4×10?3 時仍能保持 <1% 的吞吐量損失。

與之對應(yīng)，如下圖 Figure 7 將 Mini-Batch 規(guī)模設(shè)為固定參數(shù)。當(dāng)故障導(dǎo)致實際 Mini-Batch 規(guī)模低于目標(biāo)值時，訓(xùn)練進(jìn)程將暫停直至足夠多故障 GPU 恢復(fù)運(yùn)行。本實驗采用與 LLaMA 3.1 觀測值相同的故障率參數(shù)，硬件故障恢復(fù)時間設(shè)為 5 天。通過增加備用 NVL 域數(shù)量，繪制單 GPU 吞吐量變化曲線。結(jié)果表明：

• NTP 方案僅需 16 個備用 NVL 域即可實現(xiàn)連續(xù)訓(xùn)練（無暫停），這是因為該訓(xùn)練任務(wù)每個 DP 副本使用 8 個 NVL 域，且基礎(chǔ) NTP 方案（無備用域）的吞吐量降幅始終不超過等效丟失 2 個 DP 副本的水平。
• DP-DROP 方案需要 90 個備用 NVL 域才能實現(xiàn)不間斷訓(xùn)練，這反而會降低單 GPU 吞吐量——因為要達(dá)到與 16 個備用 NVL 域的 NTP 方案相同吞吐量，需要額外配置 90 個NVL 域。
• NTP-PW 方案在此故障場景下無需備用域即可實現(xiàn)：

• 訓(xùn)練不中斷；
• 相較無故障基準(zhǔn)的吞吐量損失 <1%。

6.4 原型系統(tǒng)評估

作者進(jìn)一步對原型系統(tǒng)進(jìn)行性能評估，以量化 NTP 機(jī)制的開銷。由于 pre-sync 重分片操作與最終 Backward 過程存在重疊，重點(diǎn)測量其對 Backward 階段的性能影響。如下圖 Figure 8 所示，作者測試了不同隱藏層維度和序列長度的多種工作負(fù)載。實驗采用一個 TP8 并行度、Local Batch 為 2 的 DP 副本，與另一個降低 TP 并行度（Local Batch 為 1）的副本進(jìn)行對比訓(xùn)練，并在必須執(zhí)行重分片的 TP8 副本上測量 Backward 時延（基準(zhǔn)為：2 副本均采用 TP8，Local Batch 2 的訓(xùn)練配置）。

作者提出假設(shè)：性能下降主要受兩個參數(shù)影響：

1. 總 Backward 計算量——計算量越大，重分片操作獲得重疊優(yōu)化的機(jī)會越多；
2. 單 GPU 在重分片過程中的最大數(shù)據(jù)傳輸量——該值直接決定重分片耗時。

其中 1）受模型規(guī)模和序列長度影響，2）則與模型規(guī)模及降低后的 TP 并行度相關(guān)（TP 并行度降幅越大，重分片引發(fā)的通信開銷越高）。通過將 2）除以 1）計算通信計算比，并將其作為 Figure 8 橫坐標(biāo)，縱坐標(biāo)表示 Backward 時延增幅。

實驗證實：

• 通信計算比與性能降幅呈顯著線性相關(guān)，這意味著模型規(guī)模越大或序列越長，性能降幅越小。
• TP 并行度降幅越大（即初始 TP 并行度失效比例越高），性能降幅越顯著。這是因為盡管實現(xiàn)方案將大部分重分片操作與 Backward 重疊，仍存在部分無法隱藏的重分片操作；當(dāng) TP 降幅增大（需要更多重分片數(shù)據(jù)傳輸）時，這些未隱藏操作的出現(xiàn)概率也隨之增加。
• 在所有工作負(fù)載和設(shè)置中，最終 Backward 的減速幅度最多僅為4%。

6.5 仿真驗證

在如下圖 Figure 11b 中，在 DGX-H100 節(jié)點(diǎn)上采用 FP8 與 BF16 兩種精度格式，針對多種模型規(guī)模（8B 至 175B）、不同序列長度（2K 至 8K）以及在多個計算規(guī)模（8 至 512 個 GPU）條件下展開訓(xùn)練實驗。通過系統(tǒng)性地搜索各工作負(fù)載的最佳并行配置方案，并將實際硬件平臺觀測到的吞吐量與模擬器預(yù)測值進(jìn)行對比分析。可以看出，模擬器的預(yù)測結(jié)果與實際性能表現(xiàn)高度吻合。

如下圖 Figure 11a 展示了采用 DGX-H100 GPU 集群訓(xùn)練兩種規(guī)模模型（15B 和 340B）時，在不同單 GPU 功耗限制條件下的實驗結(jié)果。將實測訓(xùn)練性能數(shù)據(jù)與模擬器預(yù)測值進(jìn)行對比繪制，結(jié)果表明模擬器的預(yù)測輸出與實測數(shù)據(jù)間具有極高的相關(guān)性。

6.5 敏感性研究

功率提升機(jī)制允許部分失效的 NVL 域以更高吞吐量運(yùn)行（以避免對健康域形成瓶頸）。理論上也可通過提升健康域功率來進(jìn)一步提高其吞吐量。但功率提升會帶來性能/watt 比下降及數(shù)據(jù)中心能耗需求增加的代價。如 Table 1 中 TP30 配置所示，1.1 倍基準(zhǔn)功率時性能/watt 降低 2.8%，1.2 倍時降幅達(dá) 6.5%。鑒于 NTP-PW 方案僅對異常機(jī)架實施功率提升，最壞情況下僅 12% 的 NVL 域會承受這種效率損失。且該機(jī)制通過重新分配故障 GPU 的功率資源，不會額外增加供電需求。

實驗設(shè)計中通常假設(shè)單 GPU 故障僅影響所在 NVL 域的一個 GPU。[2503.11901] Characterizing GPU Resilience and Impact on AI/HPC Systems [7] 中表明：91% 的故障為不可控的內(nèi)存錯誤或 MMU 錯誤（僅影響單個 GPU），5% 為可能擴(kuò)散的 NVLink 錯誤。但在 GB200-NVL72 等架構(gòu)中，整個節(jié)點(diǎn)棄置（含 36 CPU + 72 GPU）比部分 GPU 隔離更易實施。如下圖 Figure 10 展示了故障影響范圍（單 GPU 故障導(dǎo)致的連帶失效 GPU 數(shù)量）對 NTP 方案的影響：

• DP-DROP 因需丟棄整個 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閑談