InfiniteHBD團隊 投稿
量子位 | 公眾號 QbitAI
隨著大模型的參數規模不斷擴大,分布式訓練已成為人工智能發展的中心技術路徑。
如此一來,高帶寬域的設計對提升大模型訓練效率至關重要。
然而,現有的HBD架構在可擴展性、成本和容錯能力等方面存在根本性限制:
以交換機為中心的HBD(如NVIDIA NVL-72)成本高昂、不易擴展規模;以GPU為中心的HBD(如 Google TPUv3和Tesla Dojo)存在嚴重的故障傳播問題;TPUv4等交換機-GPU混合HBD采用折中方案,但在成本和容錯方面仍然不甚理想。
為解決上述問題,北京大學、階躍星辰和曦智科技的研究團隊提出了InfiniteHBD,這是一種以光交換模組為中心的高帶寬域架構。
InfiniteHBD通過將低成本光交換(OCS,Optical Circuit Switching)能力嵌入每個光電轉換模組,實現了數據中心規模的低成本可擴展性和節點級故障隔離能力。
InfiniteHBD的單位成本僅為NVL-72的31%,GPU浪費率接近零(比NVL-72和TPUv4低一個數量級以上),且與NVIDIA DGX(每節點8個GPU)相比,MFU最高提升3.37 倍。
該項目論文已被SIGCOMM2025接收。
現有大模型訓練的HBD架構的三類關鍵組件
大模型的分布式訓練涉及多種并行策略,每種策略對應不同的通信需求。
數據并行(DP,Data Parallelism)、流水線并行(PP,Pipeline Parallelism)、上下文并行(CP,Context Parallelism)和序列并行(SP,Sequence Parallelism)等策略通信開銷較低,通常可通過傳統數據中心網絡(如Fat-Tree或Rail-Optimized架構)提供的200–800 Gbps帶寬完成。
而張量并行和專家并行則通信密集,需依賴高帶寬域(HBD,High-Bandwidth Domain)提供Tbps級帶寬支持,因此HBD成為影響訓練效率的關鍵因素。
現有用于大模型訓練的HBD架構,可根據其提供連接的關鍵組件分為三類。
第一類是以交換機為中心的HBD,如NVIDIA的DGX和GB200 NVL系列。
這類架構通過交換機(如NVLink Switch)互連GPU,能夠實現高性能的任意節點通信(any-to-any communication)。
然而,以交換機為中心的HBD存在幾個明顯的局限:
- 大量使用高帶寬交換機和鏈路,導致互連成本大幅上升,限制了系統的擴展能力,并引發嚴重的資源碎片化。例如,在NVL-36上運行TP-16作業時,即使無故障發生,仍有約1/9的GPU無法利用,造成資源浪費。
- 存在交換機級的故障爆炸半徑問題,即單個交換機故障可能導致其下所有GPU的帶寬受損,顯著影響整體訓練性能。
第二類是以GPU為中心的HBD架構,如SiP-Ring、Google TPUv3和Tesla Dojo。
這類架構通常構建Ring或Mesh拓撲的GPU間直連,顯著降低了互連成本并提升了擴展性。
但與此同時,它們也面臨HBD級別的故障爆炸半徑問題:單個節點故障會導致相鄰一組節點的帶寬降級,并且破壞整個拓撲結構。
例如,在SiP-Ring中,任一GPU故障都會中斷環形連接,將原本的環狀拓撲退化為一條線性拓撲,嚴重影響通信性能。
第三類是交換機-GPU混合型HBD架構,這類設計結合了GPU間直連和交換機鏈路。
一個典型代表是TPUv4,它將TPU組織成4×4×4的立方體結構,并通過光交換機(OCS)互連各個立方體,支持規模擴展至4096個TPU。
TPUv4在互連成本與擴展性之間取得了一定平衡,同時能夠將節點故障隔離在各自的立方體內部。
然而,TPUv4仍存在一些顯著不足:
由于OCS交換機故障,仍可能出現交換機級別的故障爆炸半徑問題。
立方體級故障爆炸半徑,即立方體內任一節點故障,可能導致整個立方體通信性能下降,影響訓練效率。
如下圖所示表格中所總結的,現有的HBD架構在可擴展性、互連成本、容錯性和碎片化方面存在根本性的限制。
為了指導更優的設計,研究人員分析了現有訓練工作負載,并總結出理想HBD應具備的三大關鍵屬性:
- 隨著集群規模和模型規模的擴大,最大化MFU(Model FLOPs Utilization)所需的最優TP組大小也不斷增長。這凸顯了HBD需要支持大規模、動態可重配置TP大小的重要性。因此,理想的HBD應兼具低成本、高擴展性和靈活重構能力。
- 由于EP存在負載不均的問題,MoE模型在采用大規模純TP訓練時,依然能夠相比EP保持較高的效率。這表明,只需針對TP的Ring-AllReduce通信進行優化,便可覆蓋大部分主流訓練場景,同時大幅簡化拓撲設計的復雜性。
- 此外,HBD還應具備出色的容錯能力和高GPU資源利用率,確保在節點故障情況下訓練任務依然能夠高效運行。
基于以上分析,研究團隊提出設計一種面向大規模訓練、支持動態重構、低成本且高容錯的HBD架構,專門針對TP Ring-AllReduce通信進行優化,助力下一代大模型訓練。
InfiniteHBD包含的三項關鍵創新
InfiniteHBD提出了一種以光交換模組為核心的HBD架構,采用了OCS技術。
通過在光電轉換模組中嵌入OCS能力,InfiniteHBD實現了動態可重構的點對多點連接,具備節點級故障隔離和低資源碎片化的能力,在可擴展性和成本上全面優于現有方案。
InfiniteHBD的設計包含三項關鍵創新:
- 基于硅光子技術的OCS光電轉換模組(OCSTrx)
- 可重配置的K-Hop Ring拓撲。
- HBD-DCN編排算法。
基于硅光子技術的OCS光電轉換模組(OCSTrx)
OCSTrx將基于MZI(Mach-Zehnder Interferometer)交換矩陣的OCS集成進商用QSFP-DD 800Gbps光電轉換模組。
硅光子技術大幅簡化了器件結構,降低了成本和功耗,顯著提升了InfiniteHBD的性價比和規模可擴展性。
每個OCSTrx連接兩塊GPU,并提供三種通信路徑:
- 回環路徑實現現節點內GPU間直連(Path3)。
- 外部鏈路分別連接到不同的外部節點(Path1&2)。
所有路徑采用時分復用設計,同一時刻僅激活一條通信路徑,獨占所有GPU帶寬,且切換延遲低于1ms,實現了動態故障切換和靈活拓撲構建。
可重配置的K-Hop Ring拓撲
通過OCSTrx,節點可以與所有K跳以內的節點直接互連。
在2-Hop Ring中,節點N3連接至N1、N2、N4和N5。
節點內拓撲通過激活回環路徑,可在任意位置動態構建任意大小的GPU粒度環,靈活支持大規模、可變尺寸的TP組。
例如,N1和N2通過OCSTrx的不同路徑激活,在N1和N2的GPU1-4之間形成一個完整環路。
節點間容錯:當某節點故障時,鄰居節點動態激活備用路徑,快速繞過故障節點,實現節點級故障隔離。
例如,若N2故障,N1和N3的外部路徑自動連接,GPU通信環路得以修復。
HBD-DCN編排算法
TP的節點放置方案直接影響數據中心網絡(DCN網絡,如Roce網絡)中的并行通信流量(如DP流量)。
不合理的 TP 分布會導致大量的跨架頂式交換機(ToR,Top of Rack)通信,增加網絡擁塞風險。
為此,InfiniteHBD設計了兩階段編排機制:
- 部署階段:在集群布線時優化DCN流量局部性,例如在Fat-Tree架構中,布線確保TP組內通信盡量在同一ToR內完成,減少高層交換機負載。
- 運行時階段:根據作業規模、并行策略、實時故障模式和DCN流量模式,動態計算最優TP放置方案,在最大化GPU利用率的同時,最小化跨ToR流量。
成本相比NVL-72降低69%,GPU浪費率接近零
在大規模仿真中,該項目采用配備4顆NVIDIA H100 GPU的節點作為仿真選型。
對比評估的HBD架構包括:
- Big-Switch(理想模型,所有節點通過一臺大型交換機互連)
- InfiniteHBD K-Hop Ring(K=2和K=3配置下的InfiniteHBD)
- NVL-36/72/576
- TPUv4
- SiP-Ring
所有HBD架構的單GPU帶寬均設置為6.4Tbps,數據中心網絡(DCN)采用傳統Fat-Tree拓撲,每顆GPU配備400 Gbps帶寬。
故障彈性評估基于兩種故障模式進行:
一是采集自真實10000 GPU規模生產環境的348天的故障追蹤數據,二是基于故障概率模型生成的仿真數據。
首先,研究人員評估了不同HBD架構的故障彈性表現。
具體來說,將“浪費的GPU”定義為因故障擴散或資源碎片化而無法參與計算的健康GPU。
GPU浪費率成為衡量HBD故障彈性的重要指標。
下圖展示了基于生產環境故障追蹤數據,不同TP規模下各HBD架構隨時間變化的GPU浪費率。
下圖則基于故障概率模型。
它描繪了在不同節點故障率下,各HBD架構GPU浪費率的變化趨勢。
在兩種故障模式下,InfiniteHBD均實現了近乎零的GPU浪費率,較NVL-36、NVL-72、TPUv4和SiP-Ring低一個數量級。
盡管NVL-576因其更大規模的HBD設計表現出一定程度的故障彈性,但其互連成本極高,幾乎無法接受。
接下來,研究者進行了模型訓練性能的端到端評估,探索了在Llama3.1-405B和GPT-MoE 1.1T訓練中最大化MFU的最優并行策略,進一步驗證了前文動機部分提出的關鍵觀點。
在Llama3.1實驗中,結果表明,訓練過程中需要采用TP-16、TP-32、TP-64甚至更大規模的TP組。
與傳統的8-GPU HBD架構(如NVIDIA DGX系統)相比,InfiniteHBD最高可將MFU提升至3.37倍。
在GPT-MoE實驗中,結果顯示,訓練MoE模型時,最優的并行策略并不依賴專家并行,通過采用大規模TP同樣可以實現高效訓練。
整體實驗結果表明,InfiniteHBD能夠有效滿足大規模LLM訓練對計算效率與通信性能的雙重需求。
隨后,團隊評估了HBD-DCN編排算法在通信效率優化方面的效果。
Baseline方法為一種貪婪算法,即隨機選擇節點,并采用第一個滿足作業需求的排列方案。
如下圖(a)所示,優化算法在不同集群規模下表現穩定,跨ToR流量幾乎無明顯波動,表明其對集群規模變化不敏感。
下圖(b)展示了作業規模比(Job Scale Ratio,作業占集群總計算資源的比例)對跨ToR流量的影響(節點故障率固定為5%)。
Baseline方法始終維持約10%的跨ToR流量,而優化算法即便在作業規模比達90%時,仍將跨ToR流量降低5.8倍,顯示出優異的高負載優化能力。
下圖(c)進一步分析了節點故障對算法性能的影響(作業規模比固定為85%)。
隨著故障率上升,基線方法的跨ToR流量線性增長,而優化算法在節點故障率低于7%時,持續保持近乎零的跨ToR流量,展現了出色的韌性和容錯性。
最后,團隊還分析了不同HBD架構的互連成本與能耗。
結果表明,InfiniteHBD在這兩方面均具有明顯優勢,其互連成本僅為NVL-72的31%、TPUv4的63%,同時在能耗方面也處于最低水平,僅為NVL-72的75%,并且與TPUv4持平。
團隊介紹
一作壽晨宸,北京大學信息科學技術學院本科生,當前主要研究方向為機器學習系統和人工智能基礎設施。
劉古月,本項目通訊作者之一。
她北京大學計算機學院助理教授、博士生導師,國家級青年人才。曾任卡耐基梅隆大學博士后,獲喬治華盛頓大學博士學位。長期擔任 SIGCOMM、NSDI、ASPLOS 專家組成員,并為首位 SIGCOMM Artifact 委員會亞洲共同主席。
在推動高帶寬互連技術發展的過程中,作者團隊與多方合作伙伴緊密合作。
階躍星辰:階躍星辰是行業領先的通用大模型創業公司,堅定探索實現通用人工智能的道路。公司于2023年4月成立,聚集人工智能領域的頂尖人才,已對外發布Step系列通用大模型矩陣,覆蓋了從語言、多模態到推理的全面能力。
曦智科技:曦智科技成立于2017年,是全球領先的光電混合算力提供商。公司秉持“馭光突破算力邊界”的愿景,致力于在算力需求大爆發的時代,通過光電混合算力新范式,為客戶提供一系列算力提升解決方案,共建更智能、更可持續的世界。
據悉圍繞OCS在高帶寬互連領域的應用,上海智能算力科技有限公司已經在籌備相關的集群建設,推動該領域的應用和實踐。作者團隊感謝上海智能算力科技有限公司對于科技創新和探索的支持。
arXiv地址:
https://arxiv.org/abs/2502.03885
劉古月個人主頁:
https://grace-liu.github.io/index.html
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.