提升大模型訓練 MFU:字節“拖后腿”現象分析和歸因
一、背景
在之前的系列文章中,筆者已經系統性地介紹過大規模 LLM 訓練面臨的各種挑戰以及可能涉及的問題和解決方案。在對大規模任務進行 Profiling 分析的時候,面對成千上萬的 kernel 也經常苦不堪言,想要通過統計分析來診斷相應的問題,并為優化提供更多的可能性。碰巧看到了字節跳動 Seed 的這篇文章,雖然社區內沒有看到太多討論,不過其確實與我們的一些思路不謀而合,這里進行簡單介紹。
其實文章中的大部分結論性內容筆者在之前的文章中都已總結過,比如:
- 大規模 GPU 集群運維實踐:假裝萬卡 GPU 集群經驗中,我們提到過硬件(GPU 降頻,PCIe 鏈路降級)導致的訓練任務降速;以及軟件(Python 垃圾回收)導致的周期性降速等。
- 大模型訓練中的負載均衡:挑戰與解決方案解析中,我們也提到了一些因為負載不均衡導致無法充分發揮算力的問題,包括:
流水線并行(PP) Stage 的負載不均衡導致的問題和解決方案。
MoE 中專家負載不均衡導致的問題和方案。
長序列負載不均衡問題和優化方案。
多模態和長序列場景中樣本序列長度不均衡導致的問題和改進方案。
一些碎碎念:最近停更了將近半個月,一方面是確實太忙,沒有太多精力;另一方面是技術迭代太快,看了很多東西都沒來得及記錄。比如,筆者不久之前還推測“不久的未來一定會有 30B 規模 Dense 模型 Reasoning 能力接近或超越 DeepSeek R1”,沒想到這個“不久”這么快到來;在介紹 Reasoning 相關工作時也提到,混合 Reasoning 和非 Reasoning 模型必定會出現,沒想到現在基本已是標配;最后,LLaMA-4、Qwen3 等模型相繼推出,Google Gemini 也屢創新高,期待 DeepSeek R2 能再創輝煌。
二、摘要
LLM 訓練是當前對分布式計算要求最高的任務之一,通常需要成千上萬的 GPU,并且需要頻繁進行跨機器同步。這種工作模式容易受到 Straggler 的影響(即,訓練可能因少數慢 Worker 而變慢)。在字節跳動的場景中,作者發現 Straggler 并非總是簡單的硬件故障引起,可能源于多種復雜因素。
本文研究旨在通過對字節跳動 LLM 訓練集群 5 個月的追蹤數據,全面探討 LLM 訓練中的 Straggler 問題。核心方法是采用假設分析,模擬無 Straggler 的理想場景,并與實際情況對比。通過這種方式作者探討了以下問題:
- Straggler 對于訓練任務的影響頻率及對訓練任務的具體影響。
- Straggler 能否表現出時間或空間上的規律性。
- Straggler 現在的潛在根源是什么。
三、引言
3.1 問題背景
LLM 訓練中,Straggler 的影響與集群中訓練任務采用的分布式并行策略密切相關。典型的 LLM 訓練會采用混合并行策略,包括 DP(Data Parallelism)、PP(Pipeline Parallelism)、TP(Tensor Parallelism);在 MoE 模型中,也通常會采用 EP(Expert Parallelism);在長序列場景還會采用 SP(Sequence Parallelism)或 CP(Context Parallelism)。所有這些策略都需要頻繁協調以在 GPU 和服務器間同步結果,因而容易受到 Straggler 的影響。
為了緩解 Straggler 問題,常見的手段如下所示,通常比較少采用:
- 使用冗余的備份節點,資源浪費較嚴重。
- 異步 SGD,數學不等價,精度有損。
- 丟棄慢 Worker 的更新,數學不等價,精度有損。
鑒于 LLM 訓練任務容易受 Straggler 影響,作者分析了字節跳動 LLM 訓練集群 2024 年 1 月到 5 月期間收集的追蹤數據,以便進行分析。
3.2 混合并行訓練和 Straggler
Straggler 定義:若所有 Worker 完成分配任務所需時間相同,則采用混合并行技術的 LLM 訓練任務被視為無 Straggler 現象。這可以最大限度減少了同步所需時間,實現理想的無 Straggler 場景,達到最佳性能。基于這一定義,任何導致 Worker 進度滯后的問題均被視為 Straggle 現象。不僅包括針對單一 Worker 的硬件故障,也包括影響所有 Worker 的不可預測停滯(比如垃圾回收),以及數據不均等問題導致 Workload 不均。
DP/ZeRO/FSDP:將訓練數據分割至多個 Worker,每個 Worker 均持有整個模型的副本。
- 采用 DP 時,每個 Step 各 Worker 被分配一個訓練 Batch。該 Worker 執行其 Batch 的 Forward 與 Backward 計算,隨后所有 Worker 在進入下一 Step 訓練前進行梯度 AllReduce 操作。此梯度 AllReduce 步驟要求 Worker 間同步,任一 Worker 的延遲都可能導致整體停滯,形成拖尾效應。
- ZeRO 和 FSDP 對 DP 進行了擴展,通過跨 Worker 切分優化器狀態、參數和/或梯度來降低單 GPU 內存需求。與梯度 AllReduce 不同,ZeRO 和 FSDP 需執行梯度計算的 ReduceScatter 步驟、Device-Local 參數更新步驟,以及參數 AllGather 步驟來完成每次訓練迭代。其中 ReduceScatter 與 AllGather 操作均需跨 Worker 同步,因此同樣易受 Straggler Worker 影響。
PP:將模型切分至多個 Worker,每個 Worker 持有模型連續層的不相交子集,稱為 PP Stage。PP 降低了模型權重和激活值對單塊 GPU 內存的需求。訓練過程中,一個 Batch 的數據被劃分為若干 Micro-Batch,通過 PP Stage 進行流水線式訓練。已有多種 Micro-Batch 調度策略,如 GPipe、1F1B 及 VPP(Virtual PP)。這些調度策略均假設計算在各流水線階段間均勻分配,旨在最小化 PP Bubble —— 即某 Stage 因等待前一 Stage 數據而空閑的時間。若 PP Stage 分配不均,最慢 Stage 將阻礙其他階段,形成性能瓶頸。因此,PP 易因 Stage 間計算分配不均而受拖累。
TP 和 CP:除 PP 外,還可采用 TP 和 CP,進一步降低單 GPU 內存需求。TP 在各 Worker 間劃分每層權重,CP 則在節點間分割序列 Token。二者均需在每層 Transformer 后執行同步步驟,以聚合 TP 或 CP 組內所有 Worker 的部分計算結果。由于同步時慢速設備會拖累整體進度,TP 與 CP 同樣易受 Straggler Worker 影響。本文不分析 TP 和 CP 組內的系統性落后 Worker 問題。
在實際訓練中,通常會采用混合并行策略,其性能優于任何單一策略。采用混合并行策略時,Worker 可組織為一個超立方結構,每個維度對應一種并行策略。如下圖 Figure 1 所示,展示了一種 DP-PP-TP 的混合并行策略。
- 一個速度較慢的 TP Worker 會拖慢所屬的 PP Rank,進而產生 PP 性能瓶頸。
- 這個 PP 瓶頸又會進一步減緩該 Worker 所屬的整個 DP Rank,延后梯度同步,并阻礙其他 DP Rank 的進度。
四、方法
4.1 LLM 訓練任務追蹤
集群配置:收集追蹤數據的集群專用于訓練,并由多個團隊內部共享。集群中的機器采用與 NVIDIA DGX 服務器類似的硬件配置:每臺服務器 8 個 GPU,通過 NVLink 或 PCIe 鏈路互連,包括 4 或 8 個百 Gbps NIC、一個專用于存儲和管理的獨立 NIC、數百個 CPU Core 及數 TB 內存。服務器間通過高性能交換機以三層 CLOS 拓撲結構互聯。網絡無收斂,并且經過精心調優,確保不會因網絡擁塞導致性能下降。
任務調度:在集群上可同時調度多個任務,每個任務在其執行期間獨占分配的 GPU 資源。調度器確保每項任務使用相同類型的 GPU,以實現硬件配置的一致性。此外,調度器會以最優方式分配 GPU,保證任務所需 GPU 在網絡拓撲中位置相鄰。由于大型任務以 8 的倍數申請 GPU,不同大規模任務不會共用同一臺服務器。加之網絡無擁塞,這意味著即便任務運行于同一集群,也不會因資源爭搶出現 Straggler 現象。
數據采集:本研究所用追蹤數據采集自 2024 年 1 月 1 日至 2024 年 5 月 31 日期間提交的 LLM 預訓練任務。鑒于研究聚焦大規模任務,僅分析至少使用 128 塊 GPU 的任務,并根據規則剔除無效或不適合分析的數據,最終得到 3079 個有效任務樣本。這些任務包含采用 Dense 架構和 MoE 架構的模型,分別配置了短序列或長序列上下文訓練。其中 31.7% 的任務使用 ≥256 塊GPU,18.3% 使用 ≥ 512 塊 GPU,3.6% 使用 ≥ 5000 塊 GPU。總體而言,這些分析樣本覆蓋了 LLM 訓練任務總 GPU 時長的一半左右。
其中所有任務均采用開源 Megatron-LM 的定制版本完成。使用的 Megatron-LM 框架已經集成自研的性能分析工具 NDTimeline(GitHub - volcengine/veScale: A PyTorch Native LLM Training Framework [3]),默認情況下,該工具會對任務中 10% 的訓練 Step 進行采樣分析。對于每個被分析的 Step,工具會記錄一系列重要操作的開始和結束時間,這些操作包括 Forward 和 Backward 計算以及通信操作。詳細如下圖 Table 1 所示。
為了降低追蹤大量小型 Kernel 的成本,分析中的 Forward 和 Backward 操作通常包括多個 GPU Kernel。此外,NDTimeline 會定期同步任務中所有機器的時鐘,這使得能夠對齊不同機器上的相關操作,以便進行假設分析。
在追蹤記錄中的每個操作,其日志包括操作類型、開始和結束時間戳,以及一組元數據,如訓練 Step ID、Micro-Batch ID、PP Rank 和 DP Rank。這些元數據可以幫助重建操作間的依賴關系,對于模擬無 Straggler 情況下的替代 Timeline 至關重要。
4.2 用于假設分析的模擬器
假設分析的目標是通過回答以下問題來評估 Straggler 的影響:
- 如果所有 Straggler 都不存在,任務會耗時多久?
- 如果除特定一組 Straggler 外其他都不存在,任務又會耗時多久?
為了解答這些問題,作者模擬了一個不存在 Straggler 的替代 Timeline。核心觀點在于,在沒有 Straggler 的情況下,相似操作的耗時應當是相同的。基于這一觀點,作者首先嘗試估算在無 Straggler 的理想情境下每個操作的持續時間。然后,模擬這樣一個替代時間線:操作按其依賴關系啟動,并在估算的理想時長內完成。通過比較模擬作業完成時間(JCT)與實際追蹤記錄中的時間,便能評估 Straggler 現象對整體效率的影響。
估計無 Straggle 的理想操作時長。從概念上講,可以將追蹤到的操作組織成一個四維張量,維度分別為:訓練 Step、Micro-Batch、PP Rank 和 DP Rank。作者將此張量稱為 OpDuration 張量。如上 Table 1 中的每種操作類型都有一個這樣的張量。
- 對于計算操作,張量中的元素直接對應追蹤到的操作時長。
- 對于通信操作,計算追蹤時長中的一個子部分,稱為傳輸時長。由于通信是作為集合操作(或 P2P)的一部分進行的,單個操作的追蹤時長受兩個因素影響:(1)數據傳輸到另一 Rank 所需的時間,即“傳輸時長”;(2)等待同一集合操作(或 P2P)中其他操作開始的時間,即“阻塞時長”。在這兩個因素中,“傳輸時長”是集合操作固有的,而“阻塞時長”則由操作調度決定。因此,在 OpDuration 張量中,作者僅為通信操作類型(如 Params-Sync、Forward-Send 等)存儲“傳輸時長”。為了估算一個操作的“傳輸時長”,作者取同一集合操作(或 P2P 操作)中所有對等操作的最大開始時間,并從該操作的結束時間中減去這一最大開始時間。
提取操作依賴關系。模擬器需要兩項輸入:理想化操作時長與操作依賴模型。接下來,闡述模擬器的依賴模型,該模型源自當前使用的基于 Megatron-LM 的訓練系統。
在此依賴模型中,每個 Worker 運行若干“Stream”以執行其操作。同一 Stream 上安排的所有操作按序執行,而跨 Stream 的操作只要滿足依賴關系便可并發執行。具體而言,每個 Worker 擁有:
- 一個執行全部 Forward 與 Backward 操作的 Stream
- 一個執行所有 DP 專用通信操作的 Stream
- 四個各執行不同類型 PP 專用通信操作的流:
Forward-recv(RF)
Backward-recv(RB)
Forward-send(SF)
Backward-send(SB)
各操作間的依賴關系如下圖 Figure 2 所示:
- 相同 Stream 的依賴關系:Stream 內的操作根據其在追蹤記錄中的啟動時間進行排序。假設相鄰操作之間存在隱式依賴。
- DP 通信與計算依賴關系:在每個 PP Stage:
首個 Micro-Batch 的 Forward 操作需在完成對應的 params-sync 集合通信(以獲取該 Stage 參數)之后進行,如下圖 Figure 2 所示(Syncparams → CF, mid=1)。
這些參數會被本地緩存,供后續 Micro-Batch 使用。不同 Micro-Batch 計算出的梯度會在本地累積,然后在 DP Rank 間進行聚合。因此,最后一個 Micro-Batch 的 Backward 應在執行 grads-sync 集體通信之前完成,以便跨 DP Rank 聚合該 PP Stage 的梯度,如下圖 Figure 2 所示(CB,mid=8 → Syncgrads)。
- PP 通信與計算依賴關系:
除首個 PP Rank 外,任何一個 Micro-Batch 在 PP Rank p 上的 Forward 與 Backward 操作,必須等待同一 GPU 上該 Micro-Batch 的 Forward-receive(RF) 與 Backward-receive(RB) 通信操作完成后才能啟動,如下圖 Figure 2 所示(例如,RF,mid=1 → CF,mid=1,RB,mid=7 → CB,mid=7)。
同理,除最后一個 PP Rank 外,任何 Micro-Batch 在 PP Rank p 上的 Forward-send(SF) 與 Backward-send(SB) 操作,必須待同一 GPU 上該 Micro-Batch 的 Forward 與 Backward 操作結束后方可開始,如下圖 Figure 2 所示(例如,CF,mid=1 → SF,mid=1,CB,mid=7 → SB,mid=7)。
- 跨 Rank 通信依賴關系:對于給定的 Micro-Batch,其 DP 通信操作(即 params-sync 和 grads-sync)在所有具有相同 PP Rank 的 DP 進程間形成一個集合通信組。同理,該 Micro-Batch 的 PP send 與 receive 操作則在相鄰的、具有相同 DP Rank 的 PP 進程之間形成配對關系。此類集合(或 P2P)操作組的依賴模型規定,任一單獨操作的數據傳輸均需在所有操作啟動后方能開始。
模擬一個替代 Timeline。根據依賴模型和理想時長,可以通過以下規則模擬另一種執行 Timeline:
- 模擬器在所有依賴操作完成后立即啟動下一項操作。換言之,操作的開始時間為其依賴的上一個操作的最晚結束時間。
- 計算操作一經啟動即視為完成,其結束時間為開始時間 + OpDuration 中對應的操作時長。
- 對于每個通信操作,模擬器會等待同一集合通信組(或 P2P)中的所有對等操作都啟動。操作的結束時間為該組中最晚啟動時間 + OpDuration 中存儲的相應“傳輸時長”。
4.3 與 Straggler 相關的降速和 GPU 浪費的指標
模擬器可以估算在沒有 Straggler 干擾的替代 Timeline 的 JCT,那么應該計算哪些指標來量化 Straggler 的影響呢?作者用 Tideal 表示無 Straggler 的 JCT。為了考慮模擬過程中引入的誤差,也使用未修改的操作時長對原始 Timeline 進行模擬,并將得到的 JCT 記為 T。模擬誤差相對較小。此外,部分追蹤數據存在較大的模擬誤差,為了確保分析的準確性,作者剔除了模擬誤差 >= 5% 的追蹤數據。為了量化與 Straggle 任務相關的性能下降程度,作者計算慢速指標 S,即比例:
除了與整體 Straggle 相關的性能下降外,還希望量化由不同類型操作中的 Straggle 導致的性能下降(如 Table 1)。為此,作者首先為每種操作類型 t 計算其操作類型降速值 St,具體方法如下:
其中,Tideal 是按上述方法計算出的理想 JCT,而 T?tideal 表示當操作類型 t 的 OpDuration 中的元素未被固定時的 JCT。
在作者的集群中,一個作業在其整個運行期間獨占 GPU 資源。因此,JCT 的增加(或下降)可以直接轉化為該作業所浪費的 GPU 小時數。具體來說,通過以下公式估算作業的資源浪費比例,即浪費的 GPU 小時百分比:
同樣地,可以通過計算 1 ? 1/St 來估算因不同操作類型導致的資源浪費情況。
五、Straggler 的影響
5.1 Straggler 普遍存在,并導致不可忽視的資源浪費
如下圖 Figure 3 展示了所有作業中資源浪費百分比的累積分布函數(CDF)。可以看出,追蹤的訓練任務中,42.5% 存在 Straggler。更為嚴重的是,由于 Straggler 的存在,超過 10% 的作業至少浪費了分配到的 GPU 時間的 21.3%,而約 1% 的作業浪費了分配的 GPU 資源的 45% 以上。整體數據顯示,由于 Straggler,追蹤的所有作業中有 10.4% 的 GPU 時間被浪費。
此外,作者還研究了那些執行速度大幅下降(S > 3)的作業,發現這些均為大型作業,且問題往往源于不到 3% 的 Worker。在大多數情況下,速度緩慢的操作集中在計算環節,而非通信過程。基于此,作者認為服務器問題(可能是硬件故障或配置錯誤)通常是導致這些問題的罪魁禍首。
5.2 Step 在落后任務中表現出類似的降速現象
接下來,作者探究了任務降速究竟是少數極慢步驟還是多數步驟共同作用。作者將顯著降速比 S > 1.1 的任務定義為 Straggle 任務。為此,作者將單個 Step 的降速比定義為該 Step 執行時間與理想 Step 執行時間(例如,Tideal/n 對應含 n 個訓練步驟的任務)之比。
如下圖 Figure 4 展示了經過任務整體降速比歸一化后的單 Step 降速比 CDF。數據顯示,中位數 Step 的歸一化降速比為 1.0,即使 90 分位 Step 的歸一化降速比也僅為 1.06。表明大多數步驟的降速程度與任務整體降速比相當,說明 Straggle 現象并非由臨時環境因素導致,而是源于更持續性的問題。這一結果同時意味著,僅需采樣少量訓練 Step 來分析 Straggle 任務就足以實現成本效益最大化。
5.3 哪些操作導致了 Straggle
如下圖 Figure 5 展示了追蹤數據中所有作業按操作類型劃分的資源浪費比例。與 FALCON 的研究結論不同,作者發現多數減速源于計算操作而非通信環節。這得益于作者集群中充足的網絡帶寬、專用 LLM 訓練集群的使用以及多項內部網絡優化與調參措施。
Figure 5 同時顯示,PP-level 通信的影響略高于 DP-level 通信操作。與預期相符,因為 DP-level 通信存在大量 Overlap,相比 PP-level 通信 Kernel(多發生于訓練 Step 關鍵路徑上的預熱和冷卻階段)能容忍更多降速。
5.4 作業規模與 Straggle 有何關聯
令人驚訝的是,并未觀察到任務降速與任務規模之間存在明顯的正相關關系。這表明,任務規模并非導致滯后的決定性因素,而模型類型或人為因素等其他因素可能扮演著更為關鍵的角色。例如,超大規模任務通常由值班團隊精心維護,并能比其他任務獲得更好的優化,因此它們的降速情況未必比小型任務更糟。再比如,長上下文任務更容易觀察到受到 Straggle 影響,但這些任務通常規模較小,當與其他模型混合時,這種偏差反而使結果呈現出任務規模與降速之間的反向關聯。
六、根本原因
作者重點針對那些執行速度明顯減慢的作業進行研究,即速度降幅 S ≥ 1.1 的情況。由于確定作業降速的根本原因需手動檢查,因此作者并沒有探究所有可能的成因,而是聚焦于常見誘因。
6.1 是否應該歸因到個別 Worker
作者首先分析了有多少 Straggle 是由于 Worker 的硬件或軟件問題所致。由于作者執行了健康檢查,預計僅有少數節點容易出現問題。因此,若某項任務因少數 Worker 的問題導致降速,則修正這些問題 Worker 上的操作的執行時間(將其設置為理想執行時間),足以優化整個任務的完成時間。
作者基于這一觀察來衡量問題 Worker 對落后任務的影響。采用與之前相同的方法,通過估算操作降速百分比來計算每個 Worker 造成的降速,具體而言,定義 Worker w 的降速 Sw 為:
其中,T?wideal 是僅修復非 Worker w 上執行的操作時的 JCT(由模擬器得出),而 Tideal 是所有操作均被修復時的 JCT。
接下來,針對每項任務,篩選出速度減慢程度 Sw 位居該任務前 3% 的 Worker 集合 W。若少數 Worker 因硬件(或軟件)問題表現不佳,那么 W 將包含這些 Worker。隨后,計算這些 Worker 對任務整體降速所貢獻的比例 MW。
其中,T 代表模擬的原始 Step 時長(未修正任何 Straggle),TWideal 為僅修正選定 Worker 上運行的操作時的模擬 Step 時長,而 Tideal 則是所有 Straggle 被修正后的理想模擬 Step 時長。
對于大規模作業(涉及數千個 Worker),計算 MW 成本很高,每個 Worker w 需要運行數千次模擬以計算 Sw。因此,作者采用一種近似方法來擴展分析:不單獨計算各 Worker 的降速程度,而是測量整個 DP Rank 和 PP Rank 的降速情況。隨后,為每個 Worker 賦予其所屬 DP 和 PP Rank 中較小的(即最小的)降速值(任何 Worker 必屬于一個 DP Rank 和一個 PP Rank)。這一方法將所需模擬次數從 DP-Degree × PP-Degree 減少到 DP-Degree + PP-Degree,從而大幅降低計算復雜度。
如下圖 Figure 6 展示了 Straggle 作業的 MW CDF。可以看出,僅有 1.7% 的 Straggle 作業中,Worker 問題導致了超過 50% 的顯著時延。據此可以得出結論:在追蹤的數據中,問題 Worker 并非大多數 Straggle 作業的主導因素。作者進一步調查了少數由問題 Worker 主導作業降速的情況,發現這些情況造成的降速程度顯著更大:有問題 Worker 的作業降速比為 3.04,而平均降速比僅為 1.28。
6.2 PP Stage 不均衡
在作者的追蹤分析中,發現 PP 最后一個 Stage 相比其他 Stage 的計算不均衡問題是導致 Straggle 現象的常見原因。最后一個 Stage 包含 LM Head,還有 Loss 計算,因此在模型切分時經常無意間造成最后一個 Stage 計算量突增的情況,進而引發性能瓶頸。
為了驗證損失層的耗時問題,作者運行了一個包含 4 個 PP Stage 的任務:前 3 個 Stage 各 9 個 Transformer Layer,最后一個 Stage 額外多一個損失層(包括 LM Head)。最后一層的邏輯運算耗時是普通 Transformer 層的 9 倍以上,導致最后一個 Stage 的 Forward(Backward)速度比平均水平慢 2.07x(1.41x)。
采用與之前類似的量化方法,作者獲得一些觀測結果,如下圖 Figure 7 所示,展示了不同任務中 MS(最后一個 PP Stage 的影響系數)的 CDF,可以看出,39.3% 的任務中主要降速(MS >= 0.5)源自最后一個 PP Stage。
接下來,作者進一步探討了能否緩解這一問題。采用與 LLaMA 3 類似的方法:給最后一個 PP Stage 少分配 ε 個層數。然而,手動調整 ε 值存在多重挑戰:
- 首先,劃分模型時,必須將完整的 Transformer 層分配給 PP Stage,限制了各 Stage 間的計算分配靈活性;
- 其次,當詞表增大或最大序列長度/隱藏層尺寸減小時,損失層耗時(相對于 Transformer 層)占比會上升。隨著詞表增加,這種現象將愈發普遍,這會導致需要在前置階段執行更多 Transformer 層以平衡最后一個 PP Stage 的耗時。如此一來,可用 PP Stage 數量受限,每個 Virtual Stage 包含的層數超出理想值,制約了模型并行效率。因此,即便 ε 取值合理,仍可能導致性能未達最優。
針對前述包含 9 個 Transformer 層的任務,作者也嘗試手動調整各 PP Stage 的層數分配,發現通過人工切分可獲得 9.9% 的加速效果。但即便經過手動調優,各 Stage 計算負載仍不完全均衡 —— 例如調優后最后一個 PP Stage 的 Forward 計算耗時仍為其他 Stage 的 1.55x。
6.3 序列長度不均衡
作者又進一步分析了 Worker 速度慢、PP Stage 負載不均無法解釋的任務。深入研究發現,對于長上下文任務,訓練數據間序列長度的差異是導致任務降速的重要原因。如下圖 Figure 10 展示了某個最大序列長度為 32K 的訓練任務的序列長度分布情況,可以看出有比較明顯的長尾現象:
序列長度的差異之所以成為問題,主要是 SelfAttention 的算法復雜度呈二次方關系。在訓練中通常會采用 Sample Packing 的方式盡可能降低序列不均導致的問題,但是 SelfAttention 部分的不均衡依然無法避免,比如一個 32K 序列的 SelfAttention 計算量是 32 個 1K 序列計算量的 32 倍(預訓練階段通常也會采用 Sample Packing,但不會添加 Attention Mask,可以保障負載比較均衡,而長度擴展或 SFT 等階段會添加 Attention Mask,因此會存在這個問題)。
不同 Micro-Batch 間計算時間的差異會導致 PP Stage 出現 Bubble,并使得 DP 的各個 Worker 完成時間不同步,如下圖 Figure 8 所示,這兩種情況共同造成了 Straggle 問題。
上述分析基于計算時間與 O(Σsi2) 成正比的觀察。如下圖 Figure 9 中通過實驗驗證了這一假設:針對典型任務的前幾個訓練 Step,繪制了每個 Micro-Batch 的處理時長與 Σsi2 的關系圖,結果證實二者確實存在正比關系。
接下來,作者量化了因序列長度不平衡而導致作業降速的比例。由于追蹤數據未能提供足夠信息來修正序列長度不平衡的問題,故無法沿用之前的度量標準。因此采用一種 Forward-Backward 相關性指標來衡量序列長度失衡引發的降速現象。該指標基于以下觀察:若某 Micro-Batch 的 Forward 計算因序列長度不均而變慢,則其 Backward 同樣會以相近程度放緩,且 Forward-Backward 計算時間應存在關聯性。如下圖 Figure 11 展示了每個 Straggle 作業(即 S ≥ 1.1 的作業)中某 PP Stage 的皮爾遜相關系數。經驗表明,相關系數 ≥ 0.9 的作業極可能因序列長度失衡而降速。以此閾值為界,發現 21.4% 的作業受到序列長度不平衡影響,其平均降速程度達 1.34 倍。
此外,如下圖 Figure 12 中,作者分析了最大序列長度的變化如何影響作業降速。隨著最大序列長度的增長,序列長度不平衡的影響更為顯著。然而,上下文長度也在不斷增加,因此解決這一可擴展性挑戰變得愈發重要。
解決這一問題的思路也比較簡單,因為通過序列長度可以準確預測 Micro-Batch 的計算時間(如 Figure 9),因此可以在一個 Global Batch 進行數據的重排(數學等價),讓不同的 DP Rank 中的計算盡可能均衡,這一目標可以通過貪心算法實現。在一個最大序列長度為 32K 的典型任務上測試此方法后,觀察到吞吐量提升了 23.9%。
6.4 Python 自動垃圾回收
Python 的垃圾回收(GC)是導致降速的另一個主要原因。集群中的任務都使用 Python 運行,其 Runtime 在認為必要時會觸發 GC。一旦觸發,GC 可能耗費數百毫秒,在此期間,用戶程序會被暫停,新 Kernel 無法啟動,進而阻塞 Forward 計算。需要注意的是,Backward 不受影響,因為它們由 C++ 發起。
由于不同 Python 進程在不同時間觸發 GC,該問題進一步惡化,如下圖 Figure 13 所示,當 GC 暫停單個 Worker 進程時,整個訓練作業都會阻塞,隨著 Worker 進程數量的增加,此類暫停次數也隨之增加。因此,隨著模型規模的擴大,由 GC 引發的 Straggle 愈發嚴重。
因此,作者工程團隊實施了一項有計劃的 GC 優化方案,旨在減少因 GC 導致的 Straggle 問題。該優化措施關閉了 Python 的自動垃圾回收機制,改為根據用戶設定的訓練 Step 間隔手動觸發 GC,確保所有 Worker 同步執行 GC。此措施顯著降低了 GC 相關的 Straggle:在使用 128 個 DP Rank 的任務中,通過設定每 500 Step 執行一次 GC,性能提升了 12.6%。
然而如何確定合適的 GC 間隔也很有挑戰,間隔過長可能導致內存耗盡而任務崩潰,間隔過短可能影響性能。由于不同任務的內存分配速率不同,合適的 GC 與具體任務密切相關,因此作者默認不開啟,而讓用戶自主選擇。
此外,作者還觀察到,隨著任務推進,GC 導致的暫停時間會逐漸延長,造成訓練吞吐持續下降。推測原因是內存泄露導致堆內存持續增長,進而延長了 GC 暫停時間。需要說明的是,計劃性 GC 能有效掩蓋此類泄露的影響,維持訓練吞吐的穩定。
6.5 其他根本原因
除了上述問題外,作者還探討了其他可能得原因:
- CUDA 內存碎片化。偶發情況下,發現內存碎片化會顯著拖慢 PyTorch 的 CUDA 內存分配性能,致使 cudaFree 與 cudaMalloc 調用頻次激增,進而造成 Forward 或 Backward 計算操作異常緩慢。在數據追蹤中,表現為同一 TP Group 內不同 TP Rank 上運行的通信 Kernel 啟動時間參差不齊但卻近乎同時完成,暗示部分通信 Kernel 存在啟動延遲。通過調取 PyTorch 運行日志確認相關任務確實頻繁調用內存管理接口后,作者啟用了內存分配器追蹤功能進行復現測試。重跑過程中觀察到大量 segment_alloc 與 segment_free 調用,證實 PyTorch 內存分配器為癥結所在。
- 虛假 Kernel 依賴問題。在初期訓練大型 MoE 模型時,發現用于梯度同步的 ReduceScatter Kernel 會阻塞其他無依賴關系的 Kernel 啟動,導致任務嚴重降速。推測成因在于無關 Kernel 共享同一 CUDA 硬件隊列引發的虛假依賴。實踐證明調高 CUDA_DEVICE_MAX_CONNECTIONS 參數可緩解此現象。但值得注意的是,該問題會隨模型框架迭代時隱時現,作者仍在持續探究其深層機制。
6.6 總結:觀察和啟發
基于之前的根因分析可以得到如下結論:
- 極少數的 Straggle 任務是機器問題(無論硬件還是軟件)引起,這意味著傳統的健康指標不太可能有助于檢測或預防大多數 Straggle 問題。
- 導致 Straggle 的最常見原因包括:PP Stage 分配不均、每個 Micro-Batch 中序列長度的不均衡,以及 Python GC 引起的暫停。
- 作者還發現兩種不太常見的 Straggle 原因:PyTorch 內存碎片化和虛假內存依賴。
七、在線檢測 Straggler
作者還開發了一個名為 SMon 的 Online 服務。該服務在每次 NdTimeline 性能分析會話(記錄數十個訓練 Step)后自動運行。SMon 能夠提供熱力圖來呈現 Worker 節點的降速情況:
- 每個單元格表示一個 Worker,其 x 和 y 軸對應 DP 和 PP 的 Rank,顏色深淺表示降速程度。
- 可以簡化識別 Straggle Worker 的過長;也可以幫助識別降速的根源。
- 如下圖 Figure 14 所示,Worker 問題、PP Stage 劃分不均衡以及序列長度不均導致的降速展現出了各自獨特的模式。
八、參考鏈接
