深度揭秘：Meta工程師如何構建超大規模AI訓練網絡？
前段時間發布的Llama 3.1 405B可謂是LLM界的良心開源。不僅公開了模型權重，而且在發布的論文中詳細介紹了所用的算法和工程方法，比如模型架構、指令微調等等。
論文地址：https://ai.meta.com/blog/meta-llama-3-1/
此外，論文還難得地披露了訓練基礎設施的各方面細節，比如4D并行、集群通信、故障率和可靠性等等。
其中，關于集群意外中斷及其歸因統計更是讓我們了解到，即使能用上最先進的H100 GPU，也要面對如此頻繁的硬件故障。
1.6萬塊H100訓Llama 3.1，每3小時故障1次！罪魁禍首竟是GPU和HBM3顯存
但畢竟是1.6萬塊GPU組成的超大集群，工程量可想而知，即使Llama 3.1論文的篇幅有洋洋灑灑92頁，也很難深入、詳細地描述其構建過程。
因此，Meta的工程師們最近又發表了一篇論文，專門介紹如何大規模設計、實施和運營這個龐大的AI訓練網絡。
論文地址：https://dl.acm.org/doi/10.1145/3651890.3672233
這篇論文集結了Meta多個團隊的智慧結晶，包括AI生產工程（AI Production Engineering）、AI與系統協同設計、AI硬件系統等。
選擇RoCE的原因
分布式網絡中的GPU間通信主要包括兩個階段。首先是在單個訓練節點上的4～8個GPU之間進行「節點內通信」，使用NVLink高速互聯方式。
如果訓練工作需要額外GPU參與，就需要引入「節點間通信」，對此，業界通常有兩種設計方式。
- 標準的TCP/IP網絡或對其中的插件進行修飾，比如fastsocket
- 專有的互連技術，比如InfiniBand、NVSwitch、Elastic Fabric Adaptor、Inter-rack等
由于CPU開銷和延遲的增加，前者容易造成集群性能的下降；后者盡管能提供較好的性能，但由于是專有技術，很難靈活部署。
因此，當Meta引入基于GPU的分布式訓練時，工程師們決定為其量身定制數據中心網絡，最后選擇RoCEv2（RDMA over Converged EtherNet v2）作為主要的節點間通信機制。
RDMA全稱Remote Direct Memory Access，可以在無需CPU參與的情況下實現互連GPU的內存共享。
基于TCP/IP的通信機制中，數據包必須先發送到內核才能拷貝至內存，而RDMA則繞過內核，信息可以直接到達或發送至應用內存。
RoCEv2則是實施RDMA的一種具體協議，規定使用以太網傳輸，數據包采用UDP格式，讀寫信息的封裝和解封都由RDMA NIC硬件處理。
之所以選擇RoCE，出于以下三方面的動機：
- RoCE與訓練工作負載常用的RDMA一脈相承，確保已有設施的無縫銜接
- 使用以太網可以保留原數據中心相當比例的組件和工具，并能繼續使用基于Clos的設計
- 整個技術棧都以開放標準為基礎，確保網絡基礎設施的兼容和靈活
擴展后的RoCE網絡中，每個集群可容納數千甚至數萬個GPU，而且可用于支持生產場景下的各種GPU工作任務，比如排名、內容推薦、內容理解、NLP和GenAI模型訓練等。
拓撲結構
前后端分離
訓練集群主要依靠兩個互相獨立的網絡：前端網絡（Frontend Network, FE）負責數據攝取、檢查點和日志記錄，后端網絡（Backend Network, BE）用于訓練。AI訓練機架分別連接到FE和BE。
FE的網絡層級中包含機架交換機 (RSW)、結構交換機 (FSW) 以及更高層的存儲倉庫，為GPU提供訓練所需的輸入數據。
BE是一種專用結構，以非阻塞架構連接所有RDMA NIC，在集群中的任意兩個GPU之間，無論物理距離如何，都能提供高帶寬、低延遲的無損傳輸。后端結構使用的協議即為RoCEv2。
FE和BE的分離是部署RoCE的早期做出的一項主要決策，主要是希望兩個網絡能夠獨立進化、互不干擾。此外，將用于AI訓練的網絡分隔出來，也能簡化并加速路由和通信機制的迭代。
BE網絡曾被多次修改。最初的GPU集群使用簡單的星形拓撲，讓一些AI機架連接到中央以太交換機（運行不可路由的RoCEv1協議）。
這種設計在GPU規模和交換機冗余方面有明顯的限制。因此，后來迅速過渡到基于結構（fabric-based）的架構，以實現更好的擴展性和可用性。
AI Zone
AI機架的設計含有兩層Clos拓撲，被稱為AI Zone。
RTSW：Rack Training Switch，機架訓練交換機
CTSW：Cluster Training Switch，集群訓練交換機
ATSW：Aggregator Training Switch，聚合訓練交換機
RTSW作為葉交換機，為機架內的 GPU 提供縱向擴展連接。主干層由模塊化的CTSW組成，在集群中的所有機架之間提供橫向擴展連接。
AI Zone旨在以非阻塞方式支持大量GPU互連，然而，Llama這種大型模型需要的GPU規模大于單個AI Zone。為了適應這一點，又加入了額外的ATSW層，用于連接數據中心內的CTSW，將RoCE域擴展到單個AI Zone之外。
跨AI Zone的連接在設計上做到了oversubscription，并使用ECMP平衡網絡流量。
為了緩解跨AI Zone流量的性能瓶頸，調度程序也進行了改進，能學習到GPU服務器在網絡拓撲結構中的位置。將訓練節點劃分到不同的AI Zone時，調度程序可以找到「最小切割」（minimum cut），從而減少流量。
路由
上面討論的計算能力和網絡拓撲的擴展引發了大量訓練流量的路由和負載均衡問題。具體來說，AI訓練的工作負載體現出以下3個具有挑戰性的特征：
- 低熵：與傳統數據中心相比，AI工作負載的流的數量和多樣性要小得多，流的模式通常是重復的和可預測的
- 突發性：在時間維度上，流量通常在毫秒級的時間粒度上出現或結束
- 「大象流」：每次流量爆發時，強度可以與NIC的線路速率相當
ECMP和路徑固定
關于路由機制，最初考慮的是廣泛采用的ECMP機制（Equal-Cost Multi-Path，等成本多路徑），根據五元組的哈希值隨機路由，其中五元組包括源和目標IP、源和目標UDP端口以及協議。
然而，正如預期的那樣，由于低熵特征，ECMP的表現不佳。
另一種方案是在最初部署時設計的「路徑固定」（path-pinning），根據目的地「切片」（即RTSW下行鏈路的索引）將數據包路由到特定路徑。
如果每個機架都分配完全相同的作業且網絡中沒有故障，這種方法就會有不錯的效果，但事實并非如此。
如圖5所示，機架中的兩臺主機中只有一臺使用上行鏈路帶寬，這種不均衡的碎片化工作安排導致了特定RTSW上行鏈路的流量分布不均，使訓練性能下降超過30%。
此外，上行鏈路或CTSW上可能發生故障，導致受影響的流通過ECMP被重新分配到其他CTSW，造成流量不均甚至與現有流的沖突，拉低了整體的訓練效率。
要采用短期的緩解措施，可以將RTSW上行鏈路的帶寬升級2倍，可以減輕性能影響但過于昂貴。
隊列對擴展
接下來，團隊對ECMP機制進行了進一步審視，利用NCCL等庫中的隊列對擴展（Queue Pair Scaling）功能，以實現路由機制的演進。
此外，將交換機配置為E-ECMP (Enhanced ECMP)，使用其UDF功能對RoCE數據包的目標QP字段進行額外哈希處理。
對于QP擴展，論文評估了兩種策略。一是將每條消息拆分為通過單個QP發布而非多個QP（split），這會同時降低消息容量并產生多個ACK；另一種方法則是以循環方式將每條消息發布到不同的隊列（round-robin）。
如圖7所示，在AllReduce基準上，E-ECMP和QP擴展搭配使用時，性能相比基線可提升40%。在兩種QP擴展策略中，round-robin的表現更好。
但值得注意的是，哈希操作的潛在概率性質是該方案的本質缺陷。此外，需要根據具體的工作負載定制QP擴展策略和參數，雖然短期內可行，但長期來看會提升操作復雜性。
擁塞控制
將集群過渡到400G網絡部署時，團隊嘗試調整原有的DCQCN算法以適應新的網絡速度和拓撲結構，然而卻遇到了無法解決的問題，相比200G網絡出現了性能下降。
因此，團隊選擇在沒有DCQCN的情況下繼續進行400G部署，在一年多的時間中僅使用PFC，沒有任何其他傳輸級的擁塞控制機制。結果發現，訓練集群表現穩定，并沒有出現持續擁堵的情況。
接收方驅動的流量準入
為了緩解400G及以上網路的擁塞，論文對集群庫和RoCE傳輸進行協同設計，強制執行接收方驅動的流量準入。
圖14展示了這種GPU到GPU的通信架構，包含兩個階段的拷貝，以及接收方發起的通過NCCL集群庫的通信。
每個GPU的HBM內存維護多個通道，用于并行傳輸消息塊。
首先，發送方的GPU線程將數據從計算緩沖區（compute buffer）拷貝至可用的通道緩沖區（channel buffer）。
對于發送方的CPU代理線程，只有在收到來自接收方的CTS數據包（clear-to-send，包括大小和內存信息）后，才能發出RDMA寫入請求。
之后，接收方的GPU線程將通道緩沖區中的內容復制到目標的計算緩沖區。
最后，雙方的CPU代理線程回收通道緩沖區。通道緩沖區回收完畢并準備就緒后，接收方的CPU代理就會發送另一個CTS數據包，開啟下一輪通信。
擁塞開始時，這種機制可以有效限制網絡中的in-flight流量。然而，要實現正確的參數配置可能具有挑戰性，比如：
- 由于大量的并發計算爭奪GPU線程資源，通道數量受限
- 由于RoCE更粗粒度的流量控制，以及終端主機可能比較遲緩，設置通道緩沖區大小需要比Infiniband更仔細，在擁塞擴散和帶寬利用率不足之間進行權衡。
擁塞控制一直是RDMA網絡研究的焦點，其中DCQCN一直是以存儲為中心的網絡的黃金標準。然而，Meta團隊所述的經驗為定制擁塞控制算法提供了不同的視角。
總結
除了介紹工程設計方面的考量，論文也描述了使用的集群觀測工具，以及一些故障排除案例。
通過分離FE和BE網絡、采用不同的路由方案并優化集群流量模式，Meta團隊構建了高性能且可靠的網絡基礎設施，強調了深入理解訓練工作負載的重要性，并對相關的網絡組件進行了「量身定制」。
