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

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

01、萬卡集群的兩大挑戰(zhàn)

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

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

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

為了應對這些挑戰(zhàn)，字節(jié)跳動提出MegaScale（超大規(guī)模）系統(tǒng)，并已部署在自家的數(shù)據(jù)中心。那么字節(jié)是如何解決上述問題的呢？

02、如何實現(xiàn)大模型的高效訓練？

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

算法優(yōu)化

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

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

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

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

3D并行中的通信重疊

3D并行指張量并行、流水線并行和數(shù)據(jù)并行。

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

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

高效操作符

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

數(shù)據(jù)流水線優(yōu)化

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

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

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

集體通信群初始化

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

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

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

網絡性能調優(yōu)

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

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

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

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

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

03、容錯性

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

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

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

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

04、結論

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

我們強調在整個訓練過程中需要容錯，并實現(xiàn)了一個定制的健壯訓練框架，以自動定位和修復故障。此外，還提供了一套全面的監(jiān)控工具，用于深入觀察系統(tǒng)組件和事件，便于復雜異常的根本原因識別。我們相信，我們的工作不僅為那些從事LLM訓練的人提供了實用的見解，也為這個快速發(fā)展的領域的未來研究鋪平了道路。

審核編輯：黃飛