前段時(shí)間發(fā)布的Llama 3.1 405B可謂是LLM界的良心開源。不僅公開了模型權(quán)重，而且在發(fā)布的論文中詳細(xì)介紹了所用的算法和工程方法，比如模型架構(gòu)、指令微調(diào)等等。

論文地址：https://ai.meta.com/blog/meta-llama-3-1/

此外，論文還難得地披露了訓(xùn)練基礎(chǔ)設(shè)施的各方面細(xì)節(jié)，比如4D并行、集群通信、故障率和可靠性等等。

其中，關(guān)于集群意外中斷及其歸因統(tǒng)計(jì)更是讓我們了解到，即使能用上最先進(jìn)的H100 GPU，也要面對(duì)如此頻繁的硬件故障。

1.6萬塊H100訓(xùn)Llama 3.1，每3小時(shí)故障1次！罪魁禍?zhǔn)拙故荊PU和HBM3顯存

但畢竟是1.6萬塊GPU組成的超大集群，工程量可想而知，即使Llama 3.1論文的篇幅有洋洋灑灑92頁，也很難深入、詳細(xì)地描述其構(gòu)建過程。

因此，Meta的工程師們最近又發(fā)表了一篇論文，專門介紹如何大規(guī)模設(shè)計(jì)、實(shí)施和運(yùn)營這個(gè)龐大的AI訓(xùn)練網(wǎng)絡(luò)。

論文地址：https://dl.acm.org/doi/10.1145/3651890.3672233

這篇論文集結(jié)了Meta多個(gè)團(tuán)隊(duì)的智慧結(jié)晶，包括AI生產(chǎn)工程（AI Production Engineering）、AI與系統(tǒng)協(xié)同設(shè)計(jì)、AI硬件系統(tǒng)等。

選擇RoCE的原因

分布式網(wǎng)絡(luò)中的GPU間通信主要包括兩個(gè)階段。首先是在單個(gè)訓(xùn)練節(jié)點(diǎn)上的4～8個(gè)GPU之間進(jìn)行「節(jié)點(diǎn)內(nèi)通信」，使用NVLink高速互聯(lián)方式。

如果訓(xùn)練工作需要額外GPU參與，就需要引入「節(jié)點(diǎn)間通信」，對(duì)此，業(yè)界通常有兩種設(shè)計(jì)方式。

- 標(biāo)準(zhǔn)的TCP/IP網(wǎng)絡(luò)或?qū)ζ渲械牟寮M(jìn)行修飾，比如fastsocket

- 專有的互連技術(shù)，比如InfiniBand、NVSwitch、Elastic Fabric Adaptor、Inter-rack等

由于CPU開銷和延遲的增加，前者容易造成集群性能的下降；后者盡管能提供較好的性能，但由于是專有技術(shù)，很難靈活部署。

因此，當(dāng)Meta引入基于GPU的分布式訓(xùn)練時(shí)，工程師們決定為其量身定制數(shù)據(jù)中心網(wǎng)絡(luò)，最后選擇RoCEv2（RDMA over Converged EtherNet v2）作為主要的節(jié)點(diǎn)間通信機(jī)制。

RDMA全稱Remote Direct Memory Access，可以在無需CPU參與的情況下實(shí)現(xiàn)互連GPU的內(nèi)存共享。

基于TCP/IP的通信機(jī)制中，數(shù)據(jù)包必須先發(fā)送到內(nèi)核才能拷貝至內(nèi)存，而RDMA則繞過內(nèi)核，信息可以直接到達(dá)或發(fā)送至應(yīng)用內(nèi)存。

RoCEv2則是實(shí)施RDMA的一種具體協(xié)議，規(guī)定使用以太網(wǎng)傳輸，數(shù)據(jù)包采用UDP格式，讀寫信息的封裝和解封都由RDMA NIC硬件處理。

之所以選擇RoCE，出于以下三方面的動(dòng)機(jī)：

- RoCE與訓(xùn)練工作負(fù)載常用的RDMA一脈相承，確保已有設(shè)施的無縫銜接

- 使用以太網(wǎng)可以保留原數(shù)據(jù)中心相當(dāng)比例的組件和工具，并能繼續(xù)使用基于Clos的設(shè)計(jì)

- 整個(gè)技術(shù)棧都以開放標(biāo)準(zhǔn)為基礎(chǔ)，確保網(wǎng)絡(luò)基礎(chǔ)設(shè)施的兼容和靈活

擴(kuò)展后的RoCE網(wǎng)絡(luò)中，每個(gè)集群可容納數(shù)千甚至數(shù)萬個(gè)GPU，而且可用于支持生產(chǎn)場(chǎng)景下的各種GPU工作任務(wù)，比如排名、內(nèi)容推薦、內(nèi)容理解、NLP和GenAI模型訓(xùn)練等。

拓?fù)浣Y(jié)構(gòu)

前后端分離

訓(xùn)練集群主要依靠兩個(gè)互相獨(dú)立的網(wǎng)絡(luò)：前端網(wǎng)絡(luò)（Frontend Network, FE）負(fù)責(zé)數(shù)據(jù)攝取、檢查點(diǎn)和日志記錄，后端網(wǎng)絡(luò)（Backend Network, BE）用于訓(xùn)練。AI訓(xùn)練機(jī)架分別連接到FE和BE。

FE的網(wǎng)絡(luò)層級(jí)中包含機(jī)架交換機(jī) (RSW)、結(jié)構(gòu)交換機(jī) (FSW) 以及更高層的存儲(chǔ)倉庫，為GPU提供訓(xùn)練所需的輸入數(shù)據(jù)。

BE是一種專用結(jié)構(gòu)，以非阻塞架構(gòu)連接所有RDMA NIC，在集群中的任意兩個(gè)GPU之間，無論物理距離如何，都能提供高帶寬、低延遲的無損傳輸。后端結(jié)構(gòu)使用的協(xié)議即為RoCEv2。

FE和BE的分離是部署RoCE的早期做出的一項(xiàng)主要決策，主要是希望兩個(gè)網(wǎng)絡(luò)能夠獨(dú)立進(jìn)化、互不干擾。此外，將用于AI訓(xùn)練的網(wǎng)絡(luò)分隔出來，也能簡(jiǎn)化并加速路由和通信機(jī)制的迭代。

BE網(wǎng)絡(luò)曾被多次修改。最初的GPU集群使用簡(jiǎn)單的星形拓?fù)?，讓一些AI機(jī)架連接到中央以太交換機(jī)（運(yùn)行不可路由的RoCEv1協(xié)議）。

這種設(shè)計(jì)在GPU規(guī)模和交換機(jī)冗余方面有明顯的限制。因此，后來迅速過渡到基于結(jié)構(gòu)（fabric-based）的架構(gòu)，以實(shí)現(xiàn)更好的擴(kuò)展性和可用性。

AI Zone

AI機(jī)架的設(shè)計(jì)含有兩層Clos拓?fù)洌环Q為AI Zone。

RTSW：Rack Training Switch，機(jī)架訓(xùn)練交換機(jī)

CTSW：Cluster Training Switch，集群訓(xùn)練交換機(jī) 

ATSW：Aggregator Training Switch，聚合訓(xùn)練交換機(jī)

RTSW作為葉交換機(jī)，為機(jī)架內(nèi)的 GPU 提供縱向擴(kuò)展連接。主干層由模塊化的CTSW組成，在集群中的所有機(jī)架之間提供橫向擴(kuò)展連接。

AI Zone旨在以非阻塞方式支持大量GPU互連，然而，Llama這種大型模型需要的GPU規(guī)模大于單個(gè)AI Zone。為了適應(yīng)這一點(diǎn)，又加入了額外的ATSW層，用于連接數(shù)據(jù)中心內(nèi)的CTSW，將RoCE域擴(kuò)展到單個(gè)AI Zone之外。

跨AI Zone的連接在設(shè)計(jì)上做到了oversubscription，并使用ECMP平衡網(wǎng)絡(luò)流量。

為了緩解跨AI Zone流量的性能瓶頸，調(diào)度程序也進(jìn)行了改進(jìn)，能學(xué)習(xí)到GPU服務(wù)器在網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)中的位置。將訓(xùn)練節(jié)點(diǎn)劃分到不同的AI Zone時(shí)，調(diào)度程序可以找到「最小切割」（minimum cut），從而減少流量。

路由

上面討論的計(jì)算能力和網(wǎng)絡(luò)拓?fù)涞臄U(kuò)展引發(fā)了大量訓(xùn)練流量的路由和負(fù)載均衡問題。具體來說，AI訓(xùn)練的工作負(fù)載體現(xiàn)出以下3個(gè)具有挑戰(zhàn)性的特征：

- 低熵：與傳統(tǒng)數(shù)據(jù)中心相比，AI工作負(fù)載的流的數(shù)量和多樣性要小得多，流的模式通常是重復(fù)的和可預(yù)測(cè)的

- 突發(fā)性：在時(shí)間維度上，流量通常在毫秒級(jí)的時(shí)間粒度上出現(xiàn)或結(jié)束

- 「大象流」：每次流量爆發(fā)時(shí)，強(qiáng)度可以與NIC的線路速率相當(dāng)

ECMP和路徑固定

關(guān)于路由機(jī)制，最初考慮的是廣泛采用的ECMP機(jī)制（Equal-Cost Multi-Path，等成本多路徑），根據(jù)五元組的哈希值隨機(jī)路由，其中五元組包括源和目標(biāo)IP、源和目標(biāo)UDP端口以及協(xié)議。

然而，正如預(yù)期的那樣，由于低熵特征，ECMP的表現(xiàn)不佳。

另一種方案是在最初部署時(shí)設(shè)計(jì)的「路徑固定」（path-pinning），根據(jù)目的地「切片」（即RTSW下行鏈路的索引）將數(shù)據(jù)包路由到特定路徑。

如果每個(gè)機(jī)架都分配完全相同的作業(yè)且網(wǎng)絡(luò)中沒有故障，這種方法就會(huì)有不錯(cuò)的效果，但事實(shí)并非如此。

如圖5所示，機(jī)架中的兩臺(tái)主機(jī)中只有一臺(tái)使用上行鏈路帶寬，這種不均衡的碎片化工作安排導(dǎo)致了特定RTSW上行鏈路的流量分布不均，使訓(xùn)練性能下降超過30%。

此外，上行鏈路或CTSW上可能發(fā)生故障，導(dǎo)致受影響的流通過ECMP被重新分配到其他CTSW，造成流量不均甚至與現(xiàn)有流的沖突，拉低了整體的訓(xùn)練效率。

要采用短期的緩解措施，可以將RTSW上行鏈路的帶寬升級(jí)2倍，可以減輕性能影響但過于昂貴。

隊(duì)列對(duì)擴(kuò)展

接下來，團(tuán)隊(duì)對(duì)ECMP機(jī)制進(jìn)行了進(jìn)一步審視，利用NCCL等庫中的隊(duì)列對(duì)擴(kuò)展（Queue Pair Scaling）功能，以實(shí)現(xiàn)路由機(jī)制的演進(jìn)。

此外，將交換機(jī)配置為E-ECMP (Enhanced ECMP)，使用其UDF功能對(duì)RoCE數(shù)據(jù)包的目標(biāo)QP字段進(jìn)行額外哈希處理。

對(duì)于QP擴(kuò)展，論文評(píng)估了兩種策略。一是將每條消息拆分為通過單個(gè)QP發(fā)布而非多個(gè)QP（split），這會(huì)同時(shí)降低消息容量并產(chǎn)生多個(gè)ACK；另一種方法則是以循環(huán)方式將每條消息發(fā)布到不同的隊(duì)列（round-robin）。

如圖7所示，在AllReduce基準(zhǔn)上，E-ECMP和QP擴(kuò)展搭配使用時(shí)，性能相比基線可提升40%。在兩種QP擴(kuò)展策略中，round-robin的表現(xiàn)更好。

但值得注意的是，哈希操作的潛在概率性質(zhì)是該方案的本質(zhì)缺陷。此外，需要根據(jù)具體的工作負(fù)載定制QP擴(kuò)展策略和參數(shù)，雖然短期內(nèi)可行，但長期來看會(huì)提升操作復(fù)雜性。

擁塞控制

將集群過渡到400G網(wǎng)絡(luò)部署時(shí)，團(tuán)隊(duì)嘗試調(diào)整原有的DCQCN算法以適應(yīng)新的網(wǎng)絡(luò)速度和拓?fù)浣Y(jié)構(gòu)，然而卻遇到了無法解決的問題，相比200G網(wǎng)絡(luò)出現(xiàn)了性能下降。

因此，團(tuán)隊(duì)選擇在沒有DCQCN的情況下繼續(xù)進(jìn)行400G部署，在一年多的時(shí)間中僅使用PFC，沒有任何其他傳輸級(jí)的擁塞控制機(jī)制。結(jié)果發(fā)現(xiàn)，訓(xùn)練集群表現(xiàn)穩(wěn)定，并沒有出現(xiàn)持續(xù)擁堵的情況。

接收方驅(qū)動(dòng)的流量準(zhǔn)入

為了緩解400G及以上網(wǎng)路的擁塞，論文對(duì)集群庫和RoCE傳輸進(jìn)行協(xié)同設(shè)計(jì)，強(qiáng)制執(zhí)行接收方驅(qū)動(dòng)的流量準(zhǔn)入。

圖14展示了這種GPU到GPU的通信架構(gòu)，包含兩個(gè)階段的拷貝，以及接收方發(fā)起的通過NCCL集群庫的通信。

每個(gè)GPU的HBM內(nèi)存維護(hù)多個(gè)通道，用于并行傳輸消息塊。

首先，發(fā)送方的GPU線程將數(shù)據(jù)從計(jì)算緩沖區(qū)（compute buffer）拷貝至可用的通道緩沖區(qū)（channel buffer）。

對(duì)于發(fā)送方的CPU代理線程，只有在收到來自接收方的CTS數(shù)據(jù)包（clear-to-send，包括大小和內(nèi)存信息）后，才能發(fā)出RDMA寫入請(qǐng)求。

之后，接收方的GPU線程將通道緩沖區(qū)中的內(nèi)容復(fù)制到目標(biāo)的計(jì)算緩沖區(qū)。

最后，雙方的CPU代理線程回收通道緩沖區(qū)。通道緩沖區(qū)回收完畢并準(zhǔn)備就緒后，接收方的CPU代理就會(huì)發(fā)送另一個(gè)CTS數(shù)據(jù)包，開啟下一輪通信。

擁塞開始時(shí)，這種機(jī)制可以有效限制網(wǎng)絡(luò)中的in-flight流量。然而，要實(shí)現(xiàn)正確的參數(shù)配置可能具有挑戰(zhàn)性，比如：

- 由于大量的并發(fā)計(jì)算爭(zhēng)奪GPU線程資源，通道數(shù)量受限

- 由于RoCE更粗粒度的流量控制，以及終端主機(jī)可能比較遲緩，設(shè)置通道緩沖區(qū)大小需要比Infiniband更仔細(xì)，在擁塞擴(kuò)散和帶寬利用率不足之間進(jìn)行權(quán)衡。

擁塞控制一直是RDMA網(wǎng)絡(luò)研究的焦點(diǎn)，其中DCQCN一直是以存儲(chǔ)為中心的網(wǎng)絡(luò)的黃金標(biāo)準(zhǔn)。然而，Meta團(tuán)隊(duì)所述的經(jīng)驗(yàn)為定制擁塞控制算法提供了不同的視角。

總結(jié)

除了介紹工程設(shè)計(jì)方面的考量，論文也描述了使用的集群觀測(cè)工具，以及一些故障排除案例。

通過分離FE和BE網(wǎng)絡(luò)、采用不同的路由方案并優(yōu)化集群流量模式，Meta團(tuán)隊(duì)構(gòu)建了高性能且可靠的網(wǎng)絡(luò)基礎(chǔ)設(shè)施，強(qiáng)調(diào)了深入理解訓(xùn)練工作負(fù)載的重要性，并對(duì)相關(guān)的網(wǎng)絡(luò)組件進(jìn)行了「量身定制」。