?

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

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

01、萬卡集群的兩大挑戰(zhàn) 大模型時代，算力就是生產(chǎn)力。大模型的背后意味著巨大的計算資源，模型大小和訓(xùn)練數(shù)據(jù)大小是決定模型能力的關(guān)鍵因素。市場的主力玩家們利用數(shù)萬個GPU構(gòu)建大型人工智能集群，以訓(xùn)練LLM。但當GPU集群達到萬卡規(guī)模，如何實現(xiàn)高效率、高穩(wěn)定的訓(xùn)練？ 第一個挑戰(zhàn)是實現(xiàn)大規(guī)模的高效率訓(xùn)練。模型浮點運算利用率 (MFU)是實際吞吐量與假設(shè)最大吞吐量之比，是評估模型訓(xùn)練效率的通用指標，可以直接反映端到端的訓(xùn)練速度。為了訓(xùn)練LLM，需要將模型分布為多個GPU上，并且GPU之間需進行大量通信以推動進展。除了通信之外，如操作符優(yōu)化、數(shù)據(jù)預(yù)處理和GPU內(nèi)存消耗等因素對MFU也有著顯著影響。 第二個挑戰(zhàn)是在大規(guī)模上實現(xiàn)訓(xùn)練的高穩(wěn)定性，即在整個過程中保持高效率訓(xùn)練。在大模型訓(xùn)練中，穩(wěn)定性十分重要。失敗和延遲雖是大模型訓(xùn)練中的常態(tài)，但其故障成本非常高。如何縮短故障恢復(fù)時間至關(guān)重要，一個掉隊者不僅會影響自己的工作，還會拖慢數(shù)萬個GPU的整個作業(yè)。 為了應(yīng)對這些挑戰(zhàn)，字節(jié)跳動提出MegaScale（超大規(guī)模）系統(tǒng)，并已部署在自家的數(shù)據(jù)中心。那么字節(jié)是如何解決上述問題的呢？ ?

02、如何實現(xiàn)大模型的高效訓(xùn)練？ 想要在不損害模型準確性的情況下處理急劇增加的計算需求，需要采用最先進的算法優(yōu)化、通信策略、數(shù)據(jù)流水線管理以及網(wǎng)絡(luò)性能調(diào)優(yōu)技術(shù)。下文深入探討了用于優(yōu)化大型模型訓(xùn)練的方法，以實現(xiàn)大規(guī)模的高效率訓(xùn)練。 算法優(yōu)化 在算法層面進行了一些優(yōu)化，在不影響準確性的前提下，提高訓(xùn)練效率。主要包括并行transformer塊、滑動窗口注意力（SWA）和LAMB優(yōu)化器。 并行transformer塊：采用transformer塊的并行版本，代替標準的序列化公式。這種方法使得注意力塊和MLP塊的計算可以并行執(zhí)行，從而減少了計算時間。先前的研究表明，這種修改不會降低具有數(shù)千億參數(shù)的模型的質(zhì)量。 滑動窗口注意力（SWA）是一種稀疏注意力機制，它在輸入序列中的每個標記周圍使用固定大小的窗口，比全自注意力更高效。通過堆疊此類窗口注意力層，模型能夠有效地捕獲輸入數(shù)據(jù)中廣泛的上下文信息，同時創(chuàng)建大感受野，從而在不影響準確性的情況下加快訓(xùn)練速度。

LAMB優(yōu)化器：大規(guī)模的高效訓(xùn)練通常受到批量大小限制的阻礙。特別是，增加批量大小可能會對模型收斂產(chǎn)生不利影響。LAMB優(yōu)化器能夠使BERT的訓(xùn)練批量大小擴展到64K，而不影響準確性。 3D并行中的通信重疊 3D并行指張量并行、流水線并行和數(shù)據(jù)并行。 在數(shù)據(jù)并行中有兩個主要通信操作：all-gather操作和reduce-scatter操作。在3D并行中，單個設(shè)備可能承載多個模型塊。重疊是基于模型塊實現(xiàn)的，以最大化帶寬利用。all-gather操作在模型塊的前向傳遞之前觸發(fā)，reduce-scatter操作在它的后向傳遞之后開始。這導(dǎo)致第一個all-gather操作和最后一個reduce-scatter操作無法隱藏。受到PyTorch FSDP的啟發(fā)，初始的all-gather操作在每次迭代的開始時被預(yù)取，允許它與數(shù)據(jù)加載操作重疊，有效地將減少了通信時間。

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

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

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

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

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

集體通信群初始化 在分布式訓(xùn)練中，初始化階段涉及在GPU工作器之間建立NVIDIA集體通信庫（NCCL）通信組。由于這種開銷在小規(guī)模場景中相對較小，因此默認使用torch.distributed。隨著GPU數(shù)量擴展到超過一萬個，naive實現(xiàn)引入的開銷變得無法忍受。 torch.distributed初始化時間過長有兩個原因。第一個問題在于同步步驟，其中每個進程在初始化特定通信組結(jié)束時參與了一個屏障操作。這個屏障使用TCPStore，以單線程、阻塞的讀寫方式操作。可以用非阻塞和異步的Redis替換TCPStore。第二個問題與全局屏障的不慎使用有關(guān)。每個進程在初始化其相應(yīng)的通信組后執(zhí)行一個全局屏障。我們精心設(shè)計了通信組的初始化順序，以最小化全局屏障的需求，降低了時間復(fù)雜度。 在未經(jīng)優(yōu)化的情況下，2048張GPU的集群初始化時間是1047秒，優(yōu)化后可降至5秒以下；萬卡GPU集群的初始化時間則可降至30秒以下。

網(wǎng)絡(luò)性能調(diào)優(yōu) 分析了3D并行中機器間的流量，并設(shè)計了技術(shù)方案來提高網(wǎng)絡(luò)性能。包括網(wǎng)絡(luò)拓撲設(shè)計、減少ECMP哈希沖突、擁塞控制和重傳超時設(shè)置。

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

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

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

重傳超時設(shè)置。NCCL中的參數(shù)可以設(shè)置以控制重傳定時器和重試次數(shù)。我們調(diào)整這些參數(shù)以在鏈路抖動時快速恢復(fù)。為了進一步減少恢復(fù)時間，我們在NIC上啟用了adap_retrans功能。此功能支持在較短的時間間隔內(nèi)進行重傳，當鏈路抖動周期較短時，有助于更快地恢復(fù)傳輸過程。 ? 03、容錯性 隨著訓(xùn)練集群擴展到超過數(shù)萬個GPU，軟件和硬件故障幾乎是不可避免的。我們?yōu)長LM訓(xùn)練設(shè)計了一個健壯的訓(xùn)練框架，實現(xiàn)了自動故障識別和快速恢復(fù)，在最小的人為干預(yù)和對正在進行的訓(xùn)練任務(wù)最小影響的情況下實現(xiàn)容錯性。

如上圖所示，在接收到訓(xùn)練任務(wù)后，驅(qū)動程序進程會與自定義的Kubernetes接口進行交互，以便分配計算資源并為每個執(zhí)行器啟動相應(yīng)的Pod。一個執(zhí)行器管理一個節(jié)點。執(zhí)行器完成初始化任務(wù)后將在每個GPU上創(chuàng)建訓(xùn)練進程，并啟動一個健壯的訓(xùn)練守護進程，定期向驅(qū)動程序發(fā)送heartbeat以便實時檢測異常并預(yù)警。當檢測到異常狀態(tài)或在預(yù)定時間內(nèi)未收到狀態(tài)報告時，會觸發(fā)故障恢復(fù)程序，將暫停所有正在進行的訓(xùn)練任務(wù)，并命令它們自我檢查診斷。 一旦識別出問題節(jié)點，驅(qū)動程序?qū)⑾騅ubernetes提交要被封鎖的節(jié)點的IP地址，以及在這些節(jié)點上運行的Pod信息，Kubernetes將驅(qū)逐故障節(jié)點，并用健康節(jié)點替換。此外，還有一個用戶界面可以手動刪除問題節(jié)點?；謴?fù)過程完成后，驅(qū)動程序會從最新的checkpoint恢復(fù)訓(xùn)練。我們優(yōu)化了checkpoint和恢復(fù)過程，以最小化訓(xùn)練進度的損失。 為了增強對訓(xùn)練穩(wěn)定性和性能的監(jiān)控，開發(fā)了一個精度達到毫秒級的監(jiān)控系統(tǒng)。采用不同級別的監(jiān)控來跟蹤各種指標。此外，文中還講述了如何實現(xiàn)checkpoint快速恢復(fù)、訓(xùn)練故障排除，以及MegaScale部署和運營的經(jīng)驗，感興趣的可下載論文查閱。 ?

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

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

審核編輯：黃飛

?