大模型訓練如何應對GPU萬卡互聯難題

近日，字節跳動聯合北京大學的研究團隊發表了一篇論文《將大型語言模型訓練擴展至超過10，000塊GPU》，提出一個用于訓練大語言模型的生產系統，解決在萬卡集群上訓練大模型時面臨的效率和穩定性挑戰。

該論文介紹了系統的設計、實現和部署。此外，文中還提到了萬卡以上的集群規模遇到的問題及其解決方案。

01、萬卡集群的兩大挑戰

大模型時代，算力就是生產力。大模型的背后意味著巨大的計算資源，模型大小和訓練數據大小是決定模型能力的關鍵因素。市場的主力玩家們利用數萬個GPU構建大型人工智能集群，以訓練LLM。但當GPU集群達到萬卡規模，如何實現高效率、高穩定的訓練？

第一個挑戰是實現大規模的高效率訓練。模型浮點運算利用率（MFU）是實際吞吐量與假設最大吞吐量之比，是評估模型訓練效率的通用指標，可以直接反映端到端的訓練速度。為了訓練LLM，需要將模型分布為多個GPU上，并且GPU之間需進行大量通信以推動進展。除了通信之外，如操作符優化、數據預處理和GPU內存消耗等因素對MFU也有著顯著影響。

第二個挑戰是在大規模上實現訓練的高穩定性，即在整個過程中保持高效率訓練。在大模型訓練中，穩定性十分重要。失敗和延遲雖是大模型訓練中的常態，但其故障成本非常高。如何縮短故障恢復時間至關重要，一個掉隊者不僅會影響自己的工作，還會拖慢數萬個GPU的整個作業。

為了應對這些挑戰，字節跳動提出MegaScale（超大規模）系統，并已部署在自家的數據中心。那么字節是如何解決上述問題的呢？

02、如何實現大模型的高效訓練？

想要在不損害模型準確性的情況下處理急劇增加的計算需求，需要采用最先進的算法優化、通信策略、數據流水線管理以及網絡性能調優技術。下文深入探討了用于優化大型模型訓練的方法，以實現大規模的高效率訓練。

算法優化

在算法層面進行了一些優化，在不影響準確性的前提下，提高訓練效率。主要包括并行transformer塊、滑動窗口注意力（SWA）和LAMB優化器。

并行transformer塊：采用transformer塊的并行版本，代替標準的序列化公式。這種方法使得注意力塊和MLP塊的計算可以并行執行，從而減少了計算時間。先前的研究表明，這種修改不會降低具有數千億參數的模型的質量。

滑動窗口注意力（SWA）是一種稀疏注意力機制，它在輸入序列中的每個標記周圍使用固定大小的窗口，比全自注意力更高效。通過堆疊此類窗口注意力層，模型能夠有效地捕獲輸入數據中廣泛的上下文信息，同時創建大感受野，從而在不影響準確性的情況下加快訓練速度。

LAMB優化器：大規模的高效訓練通常受到批量大小限制的阻礙。特別是，增加批量大小可能會對模型收斂產生不利影響。LAMB優化器能夠使BERT的訓練批量大小擴展到64K，而不影響準確性。

3D并行中的通信重疊

3D并行指張量并行、流水線并行和數據并行。

在數據并行中有兩個主要通信操作：all-gather操作和reduce-scatter操作。在3D并行中，單個設備可能承載多個模型塊。重疊是基于模型塊實現的，以最大化帶寬利用。all-gather操作在模型塊的前向傳遞之前觸發，reduce-scatter操作在它的后向傳遞之后開始。這導致第一個all-gather操作和最后一個reduce-scatter操作無法隱藏。受到PyTorch FSDP的啟發，初始的all-gather操作在每次迭代的開始時被預取，允許它與數據加載操作重疊，有效地將減少了通信時間。

在流水線并行中，MegaScale使用交錯1F1B調度方法，以實現通信的重疊。在熱身階段，前向傳遞僅依賴于其先前的接收。我們解耦了通常一起實現的發送和接收，通過打破這種依賴關系，使得發送操作能夠與計算重疊。在張量/序列并行中，介紹了融合通信和計算等優化策略，以及將GEMM內核分成小塊并與通信進行流水線執行。

高效操作符

盡管在MegatronLM中已經對GEMM操作符進行了優化，但其他操作符中還有進一步增強的機會。注意力部分采用了FlashAttention-2，改進了不同線程塊和warp之間的工作分配。LayerNorm和GeLU由先前實現中的細粒度內核組成。通過將這些內核融合在一起，減少了與啟動多個內核相關的開銷，并有助于優化內存訪問模式，從而實現更好的性能。

數據流水線優化

數據預處理和加載經常被忽視。然而，這些操作在每個訓練步驟開始時會產生不可忽視的GPU空閑時間。優化這些操作對于訓練過程的效率至關重要。

異步數據預處理。數據預處理不在關鍵路徑上。因此，當GPU工作器在每個訓練步驟結束同步梯度時，可以開始后續步驟的數據預處理，這就隱藏了預處理的開銷。

消除冗余數據加載器。在分布式訓練的典型數據加載階段，每個GPU工作器都配備了自己的數據加載器，負責將訓練數據讀入CPU內存，然后轉發到GPU。這導致工作線程之間為爭奪磁盤讀取帶寬，因此產生了瓶頸。我們觀察到，在LLM訓練設置中，同一臺機器內的GPU工作器處于相同的張量并行組。因此，它們每次迭代的輸入本質上是相同的。基于這一觀察，我們采用了兩層樹狀的方法，在每臺機器上使用一個專用的數據加載器將訓練數據讀入共享內存。隨后，每個GPU工作器負責將必要的數據復制到自己的GPU內存中。這就消除了冗余讀取，并顯著提高了數據傳輸的效率。

集體通信群初始化

在分布式訓練中，初始化階段涉及在GPU工作器之間建立NVIDIA集體通信庫（NCCL）通信組。由于這種開銷在小規模場景中相對較小，因此默認使用torch.distributed。隨著GPU數量擴展到超過一萬個，naive實現引入的開銷變得無法忍受。

torch.distributed初始化時間過長有兩個原因。第一個問題在于同步步驟，其中每個進程在初始化特定通信組結束時參與了一個屏障操作。這個屏障使用TCPStore，以單線程、阻塞的讀寫方式操作。可以用非阻塞和異步的Redis替換TCPStore。第二個問題與全局屏障的不慎使用有關。每個進程在初始化其相應的通信組后執行一個全局屏障。我們精心設計了通信組的初始化順序，以最小化全局屏障的需求，降低了時間復雜度。

在未經優化的情況下，2048張GPU的集群初始化時間是1047秒，優化后可降至5秒以下；萬卡GPU集群的初始化時間則可降至30秒以下。

網絡性能調優

分析了3D并行中機器間的流量，并設計了技術方案來提高網絡性能。包括網絡拓撲設計、減少ECMP哈希沖突、擁塞控制和重傳超時設置。

網絡拓撲。我們的數據中心網絡是基于Broadcom Tomahawk 4芯片構建的高性能交換機。每個Tomahawk芯片的總帶寬為25.6Tbps，具有64×400Gbps端口。三層交換機以CLOS類似的拓撲連接，以連接超過10000個GPU。每層交換機的下行鏈路和上行鏈路的帶寬比為1:1。也就是說，32個端口用于下行，32個端口用于上行。該網絡以較小的直徑提供了高帶寬，每個節點都可以在有限的跳數內與其他節點通信。

減少ECMP哈希沖突。我們精心設計了網絡拓撲，并調度網絡流量以減少ECMP哈希沖突。首先，在機架ToR交換機上把上行與下行鏈路分開，一個 400G 下行鏈路端口通過特定的 AOC 電纜分為兩個 200G 下行鏈路端口，有效降低沖突率。

擁塞控制。在分布式訓練中大規模使用默認的DCQCN協議時，all-to-all通信可能會導致擁塞和PFC級別的提高。過度使用PFC可能會導致頭部阻塞（HoL），從而降低網絡吞吐量。為了緩解這些問題，我們開發了一個結合了Swift和DCQCN原理的算法，該算法將往返時間（RTT）的精確測量與顯式擁塞通知（ECN）的快速擁塞響應能力相結合。這種方法顯著提高了吞吐量，并最小化了與PFC相關的擁塞。

重傳超時設置。NCCL中的參數可以設置以控制重傳定時器和重試次數。我們調整這些參數以在鏈路抖動時快速恢復。為了進一步減少恢復時間，我們在NIC上啟用了adap_retrans功能。此功能支持在較短的時間間隔內進行重傳，當鏈路抖動周期較短時，有助于更快地恢復傳輸過程。

03、容錯性

隨著訓練集群擴展到超過數萬個GPU，軟件和硬件故障幾乎是不可避免的。我們為LLM訓練設計了一個健壯的訓練框架，實現了自動故障識別和快速恢復，在最小的人為干預和對正在進行的訓練任務最小影響的情況下實現容錯性。

如上圖所示，在接收到訓練任務后，驅動程序進程會與自定義的Kubernetes接口進行交互，以便分配計算資源并為每個執行器啟動相應的Pod。一個執行器管理一個節點。執行器完成初始化任務后將在每個GPU上創建訓練進程，并啟動一個健壯的訓練守護進程，定期向驅動程序發送heartbeat以便實時檢測異常并預警。當檢測到異常狀態或在預定時間內未收到狀態報告時，會觸發故障恢復程序，將暫停所有正在進行的訓練任務，并命令它們自我檢查診斷。

一旦識別出問題節點，驅動程序將向Kubernetes提交要被封鎖的節點的IP地址，以及在這些節點上運行的Pod信息，Kubernetes將驅逐故障節點，并用健康節點替換。此外，還有一個用戶界面可以手動刪除問題節點。恢復過程完成后，驅動程序會從最新的checkpoint恢復訓練。我們優化了checkpoint和恢復過程，以最小化訓練進度的損失。

為了增強對訓練穩定性和性能的監控，開發了一個精度達到毫秒級的監控系統。采用不同級別的監控來跟蹤各種指標。此外，文中還講述了如何實現checkpoint快速恢復、訓練故障排除，以及MegaScale部署和運營的經驗，感興趣的可下載論文查閱。

04、結論

本文深入探討了MegaScale的設計、實現和部署。通過算法-系統協同設計，MegaScale優化了訓練效率。在12288個GPU上訓練一個175B LLM模型時，MegaScale實現了55.2%的MFU，比Megatron-LM提高了1.34倍。