• ?標(biāo)題：UB-Mesh: a Hierarchically Localized nD-FullMesh Datacenter Network Architecture
• 作者：Heng Liao, Bingyang Liu, Xianping Chen, Zhigang Guo, Chuanning Cheng, Jianbing Wang, Xiangyu Chen, Peng Dong, Rui Meng, Wenjie Liu, Zhe Zhou, Ziyang Zhang, Yuhang Gai, Cunle Qian, Yi Xiong, Zhongwu Cheng, Jing Xia, Yuli Ma, Xi Chen, Wenhua Du, Shizhong Xiao, Chungang Li, Yong Qin, Liudong Xiong, Zhou Yu, Lv Chen, Lei Chen, Buyun Wang, Pei Wu, Junen Gao, Xiaochu Li, Jian He, Shizhuan Yan, Bill McColl
• 鏈接：https://arxiv.org/abs/2503.20377v1
• 時(shí)間：26 Mar 2025

摘要

隨著大規(guī)模語(yǔ)言模型（LLM）的規(guī)模持續(xù)擴(kuò)大，其所需的計(jì)算能力和帶寬需求也在不斷提升。為應(yīng)對(duì)這一挑戰(zhàn)，我們提出了UB-Mesh，一種新穎的AI數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)，旨在提升可擴(kuò)展性、性能、成本效益以及可用性。與傳統(tǒng)數(shù)據(jù)中心提供對(duì)稱的節(jié)點(diǎn)間帶寬不同，UB-Mesh采用了一種分層局部化的n維全互連（nD-FullMesh）網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)。該設(shè)計(jì)充分利用了LLM訓(xùn)練中的數(shù)據(jù)局部性，優(yōu)先采用短距離直接互連，以最小化數(shù)據(jù)移動(dòng)距離并減少交換機(jī)的使用。

盡管UB-Mesh的n維全互連拓?fù)湓诶碚撋暇哂卸囗?xiàng)優(yōu)勢(shì)，但其具體的架構(gòu)設(shè)計(jì)、物理實(shí)現(xiàn)以及網(wǎng)絡(luò)系統(tǒng)優(yōu)化帶來(lái)了新的挑戰(zhàn)。為實(shí)現(xiàn)UB-Mesh的實(shí)際構(gòu)建，我們首先設(shè)計(jì)了基于4維全互連拓?fù)涞腢B-Mesh-Pod架構(gòu)。UB-Mesh-Pod通過(guò)一系列硬件組件實(shí)現(xiàn)，這些組件構(gòu)成了基礎(chǔ)構(gòu)建模塊，包括專門設(shè)計(jì)的NPU、CPU、低基數(shù)交換機(jī)（LRS）、高基數(shù)交換機(jī)（HRS）、網(wǎng)卡（NIC）等。這些組件通過(guò)一種新穎的統(tǒng)一總線（Unified Bus, UB）技術(shù)實(shí)現(xiàn)互連，該技術(shù)支持靈活的IO帶寬分配和硬件資源池化。在網(wǎng)絡(luò)系統(tǒng)優(yōu)化方面，我們提出了名為全路徑路由（All-Path-Routing, APR）的先進(jìn)路由機(jī)制，以高效管理數(shù)據(jù)流量。這些優(yōu)化措施結(jié)合拓?fù)涓兄男阅茉鰪?qiáng)和諸如64+1備份設(shè)計(jì)的穩(wěn)健可靠性機(jī)制，使得UB-Mesh相比傳統(tǒng)的Clos架構(gòu)實(shí)現(xiàn)了2.04倍的成本效益提升、7.2%的網(wǎng)絡(luò)可用性提升，并在多種LLM訓(xùn)練任務(wù)中實(shí)現(xiàn)了95%以上的線性擴(kuò)展能力。。

1.引言

新興的大規(guī)模語(yǔ)言模型（LLM）[1, 4, 6, 8, 10, 18, 22, 26, 27, 31, 32]正在改變AI產(chǎn)業(yè)和人類社會(huì)。根據(jù)規(guī)模定律（Scaling Laws）[5, 13]，LLM通過(guò)增加模型參數(shù)和訓(xùn)練數(shù)據(jù)量持續(xù)提升其理解、生成和推理能力。然而，這一趨勢(shì)對(duì)底層訓(xùn)練系統(tǒng)和基礎(chǔ)設(shè)施提出了日益嚴(yán)峻的挑戰(zhàn)，迫使下一代AI數(shù)據(jù)中心滿足以下要求：

圖1：傳統(tǒng)Clos數(shù)據(jù)中心架構(gòu)與UB-Mesh的比較

R1：大規(guī)模。隨著模型規(guī)模和訓(xùn)練數(shù)據(jù)量的增加，需要越來(lái)越多的NPU（神經(jīng)網(wǎng)絡(luò)處理單元）或GPU在合理的時(shí)間內(nèi)完成訓(xùn)練。例如，Llama-3的預(yù)訓(xùn)練使用1.6萬(wàn)個(gè)GPU需要54天[8]。領(lǐng)先公司最近的公告顯示，已成功部署了擁有10萬(wàn)GPU的AI訓(xùn)練系統(tǒng)[30]。可擴(kuò)展的基礎(chǔ)設(shè)施對(duì)于支持LLM技術(shù)的持續(xù)發(fā)展至關(guān)重要。

R2：高帶寬。在LLM訓(xùn)練系統(tǒng)中，AI計(jì)算節(jié)點(diǎn)（NPU/GPU）需要的互連帶寬超過(guò)每節(jié)點(diǎn)3.2 Tbps[12, 17]，大約是當(dāng)代數(shù)據(jù)中心中典型CPU節(jié)點(diǎn)互連帶寬的10倍。因此，最先進(jìn)的AI訓(xùn)練系統(tǒng)的總帶寬比當(dāng)前基于CPU的基礎(chǔ)設(shè)施即服務(wù)（IaaS）系統(tǒng)高出10到100倍。

R3：成本效益。構(gòu)建大規(guī)模AI數(shù)據(jù)中心需要巨額硬件投資，資本支出（CapEx）往往達(dá)到數(shù)十億美元。為實(shí)現(xiàn)所需的10到100倍總互連帶寬增長(zhǎng)，若使用傳統(tǒng)的對(duì)稱Clos數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)（見(jiàn)圖1-(a)），互連成本也將增加10到100倍。優(yōu)化網(wǎng)絡(luò)基礎(chǔ)設(shè)施為提升成本效益提供了重要機(jī)會(huì)。此外，降低運(yùn)營(yíng)成本（OpEx），包括能耗和維護(hù)費(fèi)用，對(duì)于確保整體成本效益同樣至關(guān)重要。

R4：高可用性。擁有10萬(wàn)計(jì)算節(jié)點(diǎn)和約100萬(wàn)個(gè)光模塊的大型LLM訓(xùn)練集群面臨顯著的可用性挑戰(zhàn)。當(dāng)前統(tǒng)計(jì)數(shù)據(jù)顯示，即使每條鏈路的平均故障間隔時(shí)間（MTBF）為5年，整個(gè)10萬(wàn)GPU的AI集群的原始MTBF仍下降至不到30分鐘[24]。為此，網(wǎng)絡(luò)架構(gòu)設(shè)計(jì)不僅需提升硬件可靠性，還需融入容錯(cuò)機(jī)制，以應(yīng)對(duì)互連、計(jì)算資源、控制系統(tǒng)和存儲(chǔ)的故障。

同時(shí)實(shí)現(xiàn)這些目標(biāo)極具挑戰(zhàn)性，需要向先進(jìn)的下一代數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)設(shè)計(jì)轉(zhuǎn)變。我們認(rèn)為，設(shè)計(jì)下一代AI數(shù)據(jù)中心應(yīng)以以下三個(gè)原則為核心：

P1：流量模式驅(qū)動(dòng)的網(wǎng)絡(luò)拓?fù)?。與傳統(tǒng)數(shù)據(jù)中心工作負(fù)載通常產(chǎn)生均勻且隨機(jī)的流量不同，大規(guī)模語(yǔ)言模型（LLM）的訓(xùn)練流量具有確定性并表現(xiàn)出強(qiáng)烈的數(shù)據(jù)局部性。例如，張量并行（Tensor Parallelism）[25]操作所需的集合通信通常占總數(shù)據(jù)流量的50%以上，且通常發(fā)生在8-64個(gè)相鄰NPU的集群內(nèi)。相比之下，數(shù)據(jù)并行（Data Parallelism）產(chǎn)生的集合通信占總流量的不到2%，但通常需要長(zhǎng)距離傳輸。因此，層次化、局部化的網(wǎng)絡(luò)架構(gòu)對(duì)于適配這些流量模式至關(guān)重要。

P2：拓?fù)涓兄挠?jì)算與通信。在分層局部化的數(shù)據(jù)中心網(wǎng)絡(luò)中，有效運(yùn)行LLM訓(xùn)練是另一重大挑戰(zhàn)。如果訓(xùn)練任務(wù)未在計(jì)算資源上優(yōu)化分布，或網(wǎng)絡(luò)系統(tǒng)未得到充分優(yōu)化，AI集群可能因流量擁塞或帶寬利用不足而性能低下。為解決這一問(wèn)題，并行策略選擇、路由、集合通信、負(fù)載均衡等因素必須與網(wǎng)絡(luò)拓?fù)渚_對(duì)齊。

P3：容錯(cuò)自愈系統(tǒng)。LLM訓(xùn)練系統(tǒng)必須具備自愈能力以確保魯棒性。鏈路故障時(shí)，路由系統(tǒng)應(yīng)自動(dòng)切換至備用路徑；同樣，若NPU發(fā)生故障，應(yīng)有機(jī)制無(wú)縫激活備用NPU，以維持系統(tǒng)完整性和LLM訓(xùn)練的連續(xù)性。

為滿足要求R1-R4并遵循原則P1-P3，我們提出了創(chuàng)新的UB-Mesh架構(gòu)。如圖1-(b)所示，UB-Mesh采用n維全互連（nD-FullMesh）網(wǎng)絡(luò)拓?fù)?，通過(guò)遞歸構(gòu)建全互連拓?fù)洹獜陌迳舷噜廚PU間的1維連接開(kāi)始，擴(kuò)展到機(jī)架內(nèi)相鄰1維網(wǎng)格間的2維連接，再通過(guò)互連更廣范圍的相鄰高維網(wǎng)格擴(kuò)展至3維及以上。這種分層局部化的網(wǎng)絡(luò)架構(gòu)最小化傳輸跳數(shù)并優(yōu)化每跳距離，優(yōu)先采用直接互連而非間接長(zhǎng)距離交換，從而減少對(duì)交換機(jī)和光模塊的依賴，滿足R1和R3。此外，它根據(jù)LLM訓(xùn)練的傳輸需求實(shí)現(xiàn)層次化帶寬分配，為短距離通信提供高帶寬，為長(zhǎng)距離通信提供較低帶寬，滿足R2和R3并遵循P1。

遵循P2，我們深入研究了先進(jìn)的網(wǎng)絡(luò)和系統(tǒng)優(yōu)化機(jī)制以增強(qiáng)UB-Mesh架構(gòu)。具體而言，我們引入了全路徑路由（All-Path Routing, APR）技術(shù)，以充分利用直接連接鏈路的帶寬。APR結(jié)合源路由、結(jié)構(gòu)化尋址與線性表查找以及無(wú)死鎖流量控制機(jī)制，實(shí)現(xiàn)自適應(yīng)路由、最小化轉(zhuǎn)發(fā)開(kāi)銷并避免死鎖。此外，我們加入了拓?fù)涓兄目焖俟收匣謴?fù)機(jī)制以提升可靠性。為進(jìn)一步優(yōu)化性能，我們提出了拓?fù)涓兄募贤ㄐ藕筒⑿袃?yōu)化，以提升訓(xùn)練期間的帶寬利用率。

為滿足R4并遵循P3，UB-Mesh采用64+1高可用性設(shè)計(jì)：每個(gè)機(jī)架包含一個(gè)額外的備用NPU。當(dāng)系統(tǒng)中NPU發(fā)生意外故障時(shí)，備用NPU被激活以恢復(fù)功能，確保LLM訓(xùn)練任務(wù)不間斷繼續(xù)。此外，路由系統(tǒng)通過(guò)一種新穎的直接通知技術(shù)在鏈路故障時(shí)實(shí)現(xiàn)快速恢復(fù)。

我們精心設(shè)計(jì)了UB-Mesh的硬件和系統(tǒng)堆棧，綜合考慮了各種工程約束和權(quán)衡。具體實(shí)現(xiàn)中，UB-Mesh-Pod采用4維全互連拓?fù)洌筓B-Mesh可無(wú)縫擴(kuò)展至8000個(gè)NPU，形成支持下一代AI數(shù)據(jù)中心構(gòu)建的高帶寬域。為此，我們開(kāi)發(fā)了一系列硬件組件作為基礎(chǔ)構(gòu)建模塊，包括NPU、CPU、低基數(shù)交換機(jī)（LRS）、高基數(shù)交換機(jī)（HRS）、網(wǎng)卡（NIC）等。

與采用多種互連技術(shù)（如PCIe、NVLINK、IB和RoCE）的基線系統(tǒng)不同，UB-Mesh使用新穎的統(tǒng)一總線（Unified Bus, UB）技術(shù)實(shí)現(xiàn)所有組件互連。這種統(tǒng)一方法提升了IO資源分配的靈活性，其點(diǎn)對(duì)點(diǎn)通信能力支持高效的硬件資源池化，UB還為無(wú)縫跨層優(yōu)化提供了機(jī)會(huì)。

全面評(píng)估表明，與非超訂閱的Clos網(wǎng)絡(luò)相比，UB-Mesh將高基數(shù)交換機(jī)使用量減少98%，光模塊使用量減少93%，實(shí)現(xiàn)系統(tǒng)級(jí)成本效益提升2.04倍。在多個(gè)LLM訓(xùn)練任務(wù)中的實(shí)驗(yàn)還顯示，與成本高昂的Clos網(wǎng)絡(luò)相比，UB-Mesh的性能降幅僅在7%以內(nèi)。這種低成本與高性能的結(jié)合不僅滿足當(dāng)前LLM訓(xùn)練需求，還使其能夠有效應(yīng)對(duì)未來(lái)的可擴(kuò)展性挑戰(zhàn)。

2.背景與動(dòng)機(jī)

2.1 LLM訓(xùn)練中的“通信墻”

大規(guī)模語(yǔ)言模型（LLM）訓(xùn)練是有史以來(lái)規(guī)模最大、計(jì)算和通信需求最為密集的并行計(jì)算應(yīng)用[11, 18, 20–22, 25, 27, 32]。遵循所謂的規(guī)模定律（Scaling Laws）[3, 5, 13]，LLM通過(guò)增加模型參數(shù)和訓(xùn)練數(shù)據(jù)量來(lái)提升性能。因此，LLM需要越來(lái)越多的AI加速器以在合理的時(shí)間內(nèi)完成訓(xùn)練。例如，開(kāi)源的Llama-3.1模型在1.6萬(wàn)個(gè)GPU上進(jìn)行訓(xùn)練[8]，而下一代LLM模型已使用10萬(wàn)GPU進(jìn)行訓(xùn)練[24, 29]。

標(biāo)準(zhǔn)的訓(xùn)練過(guò)程包括重復(fù)的訓(xùn)練迭代，每個(gè)迭代包含前向傳播以計(jì)算損失、反向傳播以確定梯度，以及優(yōu)化器步驟以調(diào)整模型參數(shù)。為充分利用分布式計(jì)算能力，LLM訓(xùn)練通過(guò)多種并行策略將數(shù)據(jù)、模型和激活值分割到數(shù)萬(wàn)個(gè)NPU上。在每個(gè)迭代中，NPU之間頻繁交換數(shù)據(jù)以分發(fā)輸入數(shù)據(jù)、同步激活值和梯度等。隨著訓(xùn)練系統(tǒng)規(guī)模的擴(kuò)大，數(shù)據(jù)移動(dòng)成為系統(tǒng)中成本最高的部分[9, 19]。若缺乏強(qiáng)大的NPU間通信能力，訓(xùn)練過(guò)程很容易受限于“通信墻”（Communication Wall）。

2.2 LLM訓(xùn)練中數(shù)據(jù)流量的局部性

圖2：LLM訓(xùn)練中的并行性

如圖2所示，LLM訓(xùn)練通常涉及多種并行技術(shù)，以下逐一介紹：

張量并行（Tensor Parallelism, TP）：TP以行或列的方式分割模型層，并將子層放置在多個(gè)NPU上并行計(jì)算[25]。其主要涉及AllReduce操作以合并分布式的部分結(jié)果。

序列并行（Sequence Parallelism, SP）：SP（在某些論文中也稱為上下文并行）通常用于將序列分割到多個(gè)NPU以實(shí)現(xiàn)并行處理。SP依賴RingAttention[15]技術(shù)，并可采用AllGather操作收集不同NPU的部分結(jié)果。

專家并行（Expert Parallelism, EP）：對(duì)于采用專家混合（Mix-of-Experts, MoE）技術(shù)的LLM模型，密集MLP層被MoE層替代，每個(gè)MoE層包含若干“專家”。執(zhí)行時(shí)專家被稀疏激活。EP將專家分布在不同NPU上，輸入標(biāo)記通過(guò)All2All通信動(dòng)態(tài)發(fā)送至目標(biāo)專家。

流水線并行（Pipeline Parallelism, PP）：與TP分割每個(gè)模型層不同，PP將層分布到多個(gè)設(shè)備上，并以流水線方式執(zhí)行前向和反向傳播。PP涉及低開(kāi)銷的P2P通信以傳輸層間激活值，但需要高效的調(diào)度算法[11, 16]以最小化流水線中的空泡。

數(shù)據(jù)并行（Data Parallelism, DP）：DP在多個(gè)NPU上復(fù)制模型和優(yōu)化器狀態(tài)，每個(gè)副本并行處理一部分輸入批次。訓(xùn)練過(guò)程中需通過(guò)AllReduce操作同步梯度。

表1：LLM訓(xùn)練中數(shù)據(jù)流量的分析

這些并行技術(shù)共同將訓(xùn)練任務(wù)分布到數(shù)千個(gè)NPU上。需要注意的是，并非所有并行技術(shù)產(chǎn)生的數(shù)據(jù)流量均等。根據(jù)我們基于內(nèi)部MoE-2T模型的分析（如表1詳述），通信強(qiáng)度呈層次化并表現(xiàn)出強(qiáng)烈的局部性。具體而言，TP和SP約占總流量的97%，而其余并行技術(shù)通常產(chǎn)生不到2%的流量。其他模型架構(gòu)可能表現(xiàn)出略有不同的數(shù)據(jù)流量分布，但同樣具有較強(qiáng)的局部性[28]。因此，架構(gòu)設(shè)計(jì)應(yīng)優(yōu)先采用層次化的網(wǎng)絡(luò)帶寬分配方式。

2.3 數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)

圖3：傳統(tǒng)數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)

設(shè)計(jì)大規(guī)模LLM訓(xùn)練集群涉及一個(gè)關(guān)鍵問(wèn)題：如何將眾多NPU資源組織成一個(gè)協(xié)調(diào)的分布式系統(tǒng)？這需要確定合適的數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)。傳統(tǒng)數(shù)據(jù)中心和超算系統(tǒng)探索了多種拓?fù)浣Y(jié)構(gòu)，但這些拓?fù)湮幢剡m用于大規(guī)模LLM訓(xùn)練。圖3展示了一些常見(jiàn)拓?fù)洌?/p>

• Clos：傳統(tǒng)數(shù)據(jù)中心通常采用Clos架構(gòu)，通過(guò)兩到三層交換機(jī)對(duì)稱連接NPU/CPU，提供高性能并適應(yīng)多種流量模式。然而，其高成本是一個(gè)缺點(diǎn)，主要源于對(duì)高性能交換機(jī)和光模塊的廣泛使用。
• 3D Torus：與依賴交換機(jī)互連的Clos不同，3D Torus直接連接相鄰NPU。雖然降低了成本，但3D Torus的NPU間帶寬較低，且難以適應(yīng)復(fù)雜的通信模式，例如MoE模型中常見(jiàn)的All-to-All通信[20]。
• Dragon-Fly：Dragon-Fly拓?fù)鋄14]將交換機(jī)組織為互連組，每組內(nèi)交換機(jī)彼此直接相連，組間通過(guò)長(zhǎng)距離電纜連接，數(shù)據(jù)包通過(guò)間接路由傳輸。該架構(gòu)在傳統(tǒng)數(shù)據(jù)中心和高性能計(jì)算（HPC）環(huán)境中降低了維度。Dragon-Fly比Clos成本更低，但因需滿足完整的NPU-交換機(jī)帶寬需求仍較昂貴，且其架構(gòu)在LLM訓(xùn)練流量（特別是P2P和AllReduce場(chǎng)景）中表現(xiàn)不佳。
• Fugaku Tofu：Fugaku Tofu[2]提出了一種獨(dú)特的6D Torus拓?fù)溆糜贖PC應(yīng)用。然而，與3D Torus類似，其NPU間帶寬較低，且難以適應(yīng)復(fù)雜的通信模式。
• 綜上所述，如何為大規(guī)模LLM訓(xùn)練設(shè)計(jì)一個(gè)高性能且成本效益高的數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)仍是一個(gè)開(kāi)放性問(wèn)題。

3.UB-Mesh架構(gòu)

為滿足第1節(jié)中提出的要求R1-R4并遵循原則P1-P3，我們提出了UB-Mesh架構(gòu)。UB-Mesh采用了一種新穎的n維全互連（nD-FullMesh）拓?fù)浣Y(jié)構(gòu)，最大限度地利用直接電線互連，從而減少對(duì)昂貴高帶寬交換機(jī)和光模塊的需求。接下來(lái)，我們將介紹一系列專為構(gòu)建UB-Mesh設(shè)計(jì)的硬件模塊。最后，我們將詳細(xì)討論UB-Mesh-Pod的架構(gòu)設(shè)計(jì)，即UB-Mesh的具體4維全互連實(shí)現(xiàn)，同時(shí)充分考慮了若干工程約束和權(quán)衡。

3.1 nD-FullMesh拓?fù)?/h4>

圖4：UB-Mesh的n維全互連拓?fù)浼翱赡艿奈锢碛成?/p>

如圖4所示，UB-Mesh引入了一種新穎的n維全互連拓?fù)?，?qiáng)調(diào)跨所有網(wǎng)絡(luò)層級(jí)的直接互連。該拓?fù)鋸南噜徆?jié)點(diǎn)形成1維全互連開(kāi)始，其中同一層級(jí)的每個(gè)節(jié)點(diǎn)與該層內(nèi)其他所有節(jié)點(diǎn)相連。基于這一概念，相鄰1維全互連的節(jié)點(diǎn)之間再進(jìn)行直接互連，形成2維全互連。這一過(guò)程遞歸進(jìn)行，相鄰的2維全互連形成3維全互連，依此類推，最終根據(jù)需要擴(kuò)展至n維全互連。這種靈活的虛擬拓?fù)淇梢詿o(wú)縫映射到各種物理NPU組織形式：?jiǎn)伟迳系?維全互連、機(jī)架內(nèi)的2維全互連、機(jī)架行間的3維全互連、樓層機(jī)架組內(nèi)的4維全互連，甚至跨越整棟建筑的5維全互連及更高維度。

圖5：高維拓?fù)渑c靈活的帶寬分配

這種分層局部化的拓?fù)湓跇?gòu)建下一代LLM訓(xùn)練數(shù)據(jù)中心時(shí)具有多項(xiàng)優(yōu)勢(shì)。首先，由于nD-FullMesh在每個(gè)網(wǎng)絡(luò)維度上形成緊密耦合的直接連接域，并可提供逐層超訂閱的帶寬，它充分利用了LLM訓(xùn)練固有的數(shù)據(jù)局部性和從密集到稀疏的流量模式（見(jiàn)第2.2節(jié)）。此外，我們可以靈活調(diào)整每個(gè)維度的每節(jié)點(diǎn)帶寬分配，以滿足未來(lái)LLM訓(xùn)練任務(wù)的具體需求。例如，如圖5-(a)所示，假設(shè)存在6維全互連拓?fù)?，我們可以通過(guò)在NPU的IO模塊中分配不同數(shù)量的互連鏈路來(lái)調(diào)整??????和??????維度的帶寬，如圖5-(b)所示。

其次，與需要大量交換的Clos架構(gòu)不同，這種直接連接的拓?fù)湓跍p少傳輸距離方面具有巨大潛力。結(jié)合數(shù)據(jù)放置和集合通信優(yōu)化（見(jiàn)第5.2節(jié)），大多數(shù)傳輸（通常涉及TP或SP，如表1所示）可在0-2跳內(nèi)完成，大幅降低數(shù)據(jù)移動(dòng)開(kāi)銷。

表2：不同類型鏈路的使用估算

第三，與Clos等其他流行拓?fù)湎啾?，nD-FullMesh拓?fù)錅p輕了對(duì)高帶寬交換機(jī)和長(zhǎng)距離光互連的依賴。相反，它最大限度地利用短距離直接互連。例如，根據(jù)表2的估算，短距離無(wú)源電線占據(jù)了總電纜消耗的86.7%。這一設(shè)計(jì)顯著降低了交換機(jī)和光模塊的成本，同時(shí)提升了系統(tǒng)可靠性，因?yàn)殡娋€和連接器比光模塊更穩(wěn)定。

最后，與不適合復(fù)雜集合通信操作（如All-to-All）的3D-Torus拓?fù)洳煌?，全互連拓?fù)淠軌蚋咝еС诌@些操作，具體將在第5節(jié)中討論。

3.2 UB-Mesh的構(gòu)建模塊

所提出的nD-FullMesh拓?fù)湓诶碚撋暇哂卸囗?xiàng)優(yōu)勢(shì)，但其實(shí)踐實(shí)現(xiàn)和系統(tǒng)性能的提升帶來(lái)了新的挑戰(zhàn)，需要全面考慮各種工程約束和權(quán)衡。為此，我們?cè)O(shè)計(jì)并制造了一系列硬件模塊，作為UB-Mesh的基礎(chǔ)構(gòu)建塊。這些模塊通過(guò)一種名為統(tǒng)一總線（Unified Bus, UB）的新型互連機(jī)制實(shí)現(xiàn)連接。具體細(xì)節(jié)如下：

3.2.1 硬件模塊

表3：UB-Mesh的主要構(gòu)建模塊

表3列舉了UB-Mesh的主要組件。其核心是作為AI計(jì)算單元的NPU。NPU配備兩個(gè)UB IO控制器，總計(jì)提供x72條UB通道。負(fù)責(zé)執(zhí)行主機(jī)程序的CPU也配備一個(gè)UB IO控制器，提供x32的UB IO。

盡管UB-Mesh優(yōu)先采用直接互連，但交換能力對(duì)于滿足特定需求仍不可或缺。例如，在圖5中聚合??和??????維度的NPU IO有助于實(shí)現(xiàn)高效的機(jī)架間互連。此外，我們期望CPU-NPU比例及其綁定關(guān)系能夠靈活調(diào)整，以實(shí)現(xiàn)NPU和內(nèi)存資源的高效池化，這可以通過(guò)交換實(shí)現(xiàn)。

為在UB-Mesh-Pod內(nèi)提供成本效益高的交換能力，我們開(kāi)發(fā)了輕量級(jí)低基數(shù)交換機(jī)（LRS）。其制造成本低，可高效聚合機(jī)架間IO帶寬并支持CPU-NPU通信。

盡管UB-Mesh僅使用低基數(shù)交換機(jī)即可實(shí)現(xiàn)，但實(shí)際網(wǎng)絡(luò)約束通常將其規(guī)模限制在一定范圍內(nèi)（UB-Mesh-Pod）。因此，UB-Mesh還包括高基數(shù)交換機(jī)（HRS），提供x512的UB IO容量以支持Pod級(jí)交換。UB-Mesh的詳細(xì)架構(gòu)設(shè)計(jì)和考慮將在第3.3節(jié)中介紹。

3.2.2 統(tǒng)一總線（UB）互連

圖6：統(tǒng)一總線互連與傳統(tǒng)混合互連的對(duì)比

連接這些異構(gòu)硬件模塊需要特定的互連技術(shù)。如圖6-(a)所示，傳統(tǒng)的基于GPU的LLM系統(tǒng)通常采用多種互連方式，而UB-Mesh通過(guò)統(tǒng)一總線（UB）技術(shù)實(shí)現(xiàn)所有組件的互連，其優(yōu)勢(shì)包括：  

• 靈活的IO資源分配：UB互連與特定用例解耦，支持芯片內(nèi)不同類型IO的靈活資源分配，如圖5所示。NPU間帶寬和CPU-NPU帶寬也可根據(jù)具體需求靈活調(diào)整，因?yàn)樗鼈兪褂孟嗤腢B鏈路。
• 硬件資源池化：UB的點(diǎn)對(duì)點(diǎn)通信能力支持包括DDR DRAM、CPU、NPU和NIC在內(nèi)的高效硬件資源池化。例如，在圖6-(b)中，CPU和NPU通過(guò)UB互連實(shí)現(xiàn)池化，以提升資源利用率。
• 系統(tǒng)優(yōu)化收益：統(tǒng)一互連消除了協(xié)議轉(zhuǎn)換的需要，大幅降低開(kāi)銷并簡(jiǎn)化驅(qū)動(dòng)程序、通信庫(kù)和框架的設(shè)計(jì)與優(yōu)化。

3.3 架構(gòu)概覽

圖7：UB-Mesh的架構(gòu)設(shè)計(jì)

通過(guò)組合表3中列出的硬件構(gòu)建模塊并通過(guò)統(tǒng)一總線（Unified Bus, UB）進(jìn)行互連，我們提出了UB-Mesh的整體架構(gòu)設(shè)計(jì)。如圖7-(a)所示，我們按照所提出的nD-FullMesh拓?fù)鋵?shí)現(xiàn)了一個(gè)UB-Mesh-Pod。具體而言，在該架構(gòu)中，我們?cè)诿總€(gè)機(jī)架內(nèi)創(chuàng)建了一個(gè)2D全互連拓?fù)?，并將其擴(kuò)展到機(jī)架之外的另一個(gè)2D全互連拓?fù)?，從而形成一個(gè)4D全互連拓?fù)?。我們選擇不在這一代產(chǎn)品中將UB-Mesh擴(kuò)展至5D全互連或更高維度，以在成本效益和靈活性之間實(shí)現(xiàn)工程平衡，在4D全互連之外仍采用Clos拓?fù)洹Ｒ韵聻樵敿?xì)介紹：

3.3.1 在機(jī)架內(nèi)實(shí)現(xiàn)2D全互連

圖8：UB-Mesh機(jī)架的硬件實(shí)現(xiàn)

如圖7-(b)所示，每個(gè)機(jī)架配置為2D全互連拓?fù)?，集成了多個(gè)組件以確保高效通信和資源利用。機(jī)架的核心由八個(gè)NPU板組成，每個(gè)板包含八個(gè)NPU，如圖8-(a)所示。這64個(gè)NPU在機(jī)架內(nèi)互連，形成一個(gè)2D全互連網(wǎng)絡(luò)，確保高NPU間帶寬。需要注意的是，由于UB IO控制器也具備路由能力，每個(gè)NPU在此架構(gòu)中既充當(dāng)計(jì)算單元又充當(dāng)路由器，支持間接路由。

除NPU外，機(jī)架還包括專用CPU板，如圖8-(c)所示。與傳統(tǒng)設(shè)置中CPU和NPU位于同一板不同，此處它們是分離的。CPU通過(guò)交換機(jī)連接到NPU，支持靈活的CPU/NPU比例調(diào)整，并實(shí)現(xiàn)CPU/NPU/DDR資源池化以提升資源利用率。

機(jī)架配備多個(gè)背板交換機(jī)，管理機(jī)架內(nèi)和機(jī)架間的連接。這些交換機(jī)采用低基數(shù)設(shè)計(jì)（稱為L(zhǎng)RS），以降低成本，同時(shí)確保設(shè)備間的無(wú)阻塞通信?？傮w而言，這些背板交換機(jī)輸出四個(gè)UB x256 IO。

3.3.2 64+1高可用性設(shè)計(jì)

正如第1節(jié)所述，在大規(guī)模AI數(shù)據(jù)中心中確保系統(tǒng)可用性是一個(gè)關(guān)鍵挑戰(zhàn)。系統(tǒng)中單個(gè)NPU可能會(huì)發(fā)生意外故障。雖然軟件系統(tǒng)可以檢測(cè)這些故障并使用剩余健康NPU重啟訓(xùn)練，但由于計(jì)算能力和系統(tǒng)帶寬的減少，系統(tǒng)性能將受到顯著影響。

圖9：通過(guò)啟用備用NPU實(shí)現(xiàn)容錯(cuò)

為提升系統(tǒng)可用性，我們引入了特殊的64+1設(shè)計(jì)：如圖7-(b)和圖8-(b)所示，系統(tǒng)包含64個(gè)常規(guī)NPU以及一個(gè)額外的備用NPU。該備用NPU通過(guò)LRS與64個(gè)常規(guī)NPU相連。當(dāng)任何NPU發(fā)生故障時(shí)，備用NPU將被策略性地啟用。如圖9所示，當(dāng)NPU-3發(fā)生故障時(shí)，管理系統(tǒng)激活備用NPU（圖中節(jié)點(diǎn)B）以替換NPU-3。與NPU-3相關(guān)的原始直接連接鏈路將被重定向。例如，路徑5-3被重定向?yàn)槁窂?-LRS-B。盡管此策略將原始直接連接變?yōu)閱翁酚?，略微增加了傳輸延遲，但相比簡(jiǎn)單屏蔽NPU-3并在剩余七個(gè)NPU上運(yùn)行任務(wù)，其效果要優(yōu)越得多。在機(jī)架內(nèi)帶寬為主要關(guān)注點(diǎn)的LLM訓(xùn)練場(chǎng)景中，增加的延遲對(duì)整體訓(xùn)練性能影響微乎其微。

3.3.3 在UB-Mesh-Pod中擴(kuò)展至4D全互連

在機(jī)架內(nèi)，我們將64個(gè)NPU組織為2D全互連拓?fù)洹Ｔ跈C(jī)架之外，我們實(shí)現(xiàn)另一個(gè)2D全互連來(lái)組織機(jī)架，形成名為UB-Mesh-Pod的4D全互連架構(gòu)。

如圖7-(a)和(c)所示，一行中四個(gè)相鄰機(jī)架被組織為緊密耦合的1D全互連，機(jī)架通過(guò)每個(gè)機(jī)架中的LRS端口直接互連。在每個(gè)機(jī)架列中，機(jī)架也直接相連，形成一個(gè)16機(jī)架的2D全互連拓?fù)?。此配置中的每條鏈路代表一個(gè)UB x128 IO，如圖8-(d)所示。我們沿兩個(gè)維度連接四個(gè)相鄰機(jī)架以構(gòu)建機(jī)架間全互連，這是考慮到有源電線的覆蓋范圍得出的最佳點(diǎn)。由于每個(gè)機(jī)架有64個(gè)NPU，每個(gè)Pod有16個(gè)機(jī)架，一個(gè)4D全互連的UB-Mesh-Pod總共包含1024個(gè)NPU。

3.3.4 UB-Mesh-SuperPod及更高規(guī)模

基于1000規(guī)模的UB-Mesh-Pod，我們進(jìn)一步構(gòu)建了UB-Mesh-SuperPod，可容納多個(gè)UB-Mesh-Pod?？紤]到當(dāng)前云場(chǎng)景中小規(guī)?；蛑幸?guī)模LLM訓(xùn)練負(fù)載可能不會(huì)消耗整個(gè)SuperPod，我們選擇在Pod級(jí)互連中采用對(duì)稱的Clos拓?fù)洌抢^續(xù)使用全互連設(shè)計(jì)。這種設(shè)計(jì)允許云管理者根據(jù)用戶需求靈活劃分SuperPod，并保證每個(gè)劃分域內(nèi)的全帶寬。如圖7-(c)所示，我們使用高基數(shù)Pod交換機(jī)（HRS）連接SuperPod中的每個(gè)機(jī)架，可擴(kuò)展至8000個(gè)NPU。

最后，SuperPod中的機(jī)架還通過(guò)UB交換機(jī)（方案(a)）或位于CPU板上的NIC（方案(b)）連接到大規(guī)模數(shù)據(jù)中心網(wǎng)絡(luò)（DCN）。DCN域通常支持大規(guī)模數(shù)據(jù)并行訓(xùn)練。DCN交換機(jī)采用Clos拓?fù)浣M織。DCN域可擴(kuò)展至10萬(wàn)NPU或更多。

4.實(shí)現(xiàn)架構(gòu)

在UB-Mesh-SuperPod中，4D全互連與Clos混合拓?fù)淦惹行枰咝У幕旌下酚芍С郑撼送ㄟ^(guò)LRS和HRS進(jìn)行的基于交換機(jī)的路由外，NPU本身通過(guò)UB控制器也具備路由能力。UB-Mesh中大量路由器和層次化拓?fù)錇槁酚上到y(tǒng)引入了新的復(fù)雜性。我們認(rèn)為，路由系統(tǒng)必須滿足以下五個(gè)關(guān)鍵要求：  

• 支持混合拓?fù)洌郝酚上到y(tǒng)必須滿足UB-Mesh的4D全互連與Clos混合拓?fù)渌鶐?lái)的需求。
• 高效轉(zhuǎn)發(fā)：由于每個(gè)NPU同時(shí)也是路由器，且整個(gè)系統(tǒng)包含大量NPU，為節(jié)省NPU硬件資源，路由系統(tǒng)必須高效處理路由表查詢和轉(zhuǎn)發(fā)操作。
• 支持非最短路徑：在nD全互連拓?fù)渲?，兩個(gè)端點(diǎn)之間存在多個(gè)路徑，距離各異。系統(tǒng)應(yīng)支持使用非最短路徑，以最大化網(wǎng)絡(luò)帶寬利用率。
• 快速故障恢復(fù)：為確?？煽啃院涂捎眯裕酚上到y(tǒng)必須快速?gòu)墓收现谢謴?fù)，減少對(duì)訓(xùn)練過(guò)程的影響。
• 無(wú)死鎖：最后，整個(gè)網(wǎng)絡(luò)系統(tǒng)必須在無(wú)死鎖風(fēng)險(xiǎn)下運(yùn)行，確保數(shù)據(jù)流暢通無(wú)阻。

表4：路由系統(tǒng)比較

然而，表4中列出的現(xiàn)有路由技術(shù)，如Clos、Torus和DragonFly拓?fù)渖鲜褂玫拈L(zhǎng)前綴匹配（LPM）、基于主機(jī)的路由和維度順序路由（DOR），均無(wú)法完全滿足我們的需求。因此，我們提出了全路徑路由（All-Path-Routing, APR）技術(shù)和基于直接通知的快速故障恢復(fù)技術(shù)，以滿足上述要求。

4.1 全路徑路由（APR）

圖10：最短路徑路由與全路徑路由

UB-Mesh的全互連網(wǎng)絡(luò)架構(gòu)為任意兩個(gè)NPU之間提供了多條路徑。傳統(tǒng)的路由策略，如最短路徑優(yōu)先（SPF）路由（見(jiàn)圖10-(a)），往往無(wú)法充分利用網(wǎng)絡(luò)帶寬，且易受鏈路故障影響。為提升網(wǎng)絡(luò)效率和魯棒性，我們?cè)赨B-Mesh中提出了全路徑路由（APR）機(jī)制。

如圖10-(b)所示，APR利用源節(jié)點(diǎn)與目標(biāo)節(jié)點(diǎn)之間的所有可用路徑。這種靈活性允許動(dòng)態(tài)切換路徑，以應(yīng)對(duì)故障或擁塞，從而提升網(wǎng)絡(luò)魯棒性。為高效實(shí)現(xiàn)這一目標(biāo)，APR依賴以下三種底層機(jī)制：(1) 源路由（SR）機(jī)制實(shí)現(xiàn)自適應(yīng)路由；(2) 結(jié)構(gòu)化尋址與線性表查詢技術(shù)，減少路由表查詢和轉(zhuǎn)發(fā)開(kāi)銷；(3) 拓?fù)涓兄獰o(wú)死鎖流量控制，避免死鎖。

4.1.1 源路由

圖11：源路由頭部格式

為充分利用APR提供的路徑，一種實(shí)用方法是采用源路由（SR）機(jī)制。如圖11所示，在原始數(shù)據(jù)包頭部添加一個(gè)8字節(jié)的SR頭部，包含轉(zhuǎn)發(fā)指令。每個(gè)路由器使用一個(gè)4位指針（ptr）指示12位位圖字段中的位偏移。第??位的值指定第??跳的轉(zhuǎn)發(fā)方式（即，位值為1表示該跳采用SR轉(zhuǎn)發(fā)，0表示傳統(tǒng)轉(zhuǎn)發(fā)）。在SR轉(zhuǎn)發(fā)的情況下，位圖字段還用于定位六個(gè)指令字段之一，指示如何轉(zhuǎn)發(fā)該數(shù)據(jù)包。SR信息高度壓縮，因此開(kāi)銷較低。

4.1.2 結(jié)構(gòu)化尋址與線性表查詢

為減少每個(gè)NPU的UB IO控制器中的路由表查詢和轉(zhuǎn)發(fā)開(kāi)銷，APR路由系統(tǒng)利用UB-Mesh拓?fù)涞慕Y(jié)構(gòu)，采用結(jié)構(gòu)化尋址和線性表查詢機(jī)制。具體而言，地址空間根據(jù)網(wǎng)絡(luò)元素的物理位置（如Pod、機(jī)架和板）劃分為多個(gè)段。由于同一段內(nèi)的NPU共享相同前綴，僅需存儲(chǔ)短段地址，NPU可通過(guò)相對(duì)于段地址的線性偏移進(jìn)行尋址。此設(shè)計(jì)顯著減少表空間，加速路由表生成和分發(fā)，提升收斂速度，并支持快速響應(yīng)狀態(tài)變化，例如故障恢復(fù)操作。

4.1.3 拓?fù)涓兄獰o(wú)死鎖流量控制

鑒于UB-Mesh的nD全互連拓?fù)浒瑥?fù)雜的環(huán)結(jié)構(gòu)，且APR機(jī)制支持多跳路由，在有限的虛擬通道（VL）資源下實(shí)現(xiàn)無(wú)死鎖流量控制具有挑戰(zhàn)性。為應(yīng)對(duì)這一挑戰(zhàn)，我們提出了TFC（拓?fù)涓兄獰o(wú)死鎖流量控制）算法。該算法在僅使用2個(gè)VL資源的情況下，最小化VL資源使用，同時(shí)實(shí)現(xiàn)無(wú)死鎖的全路徑路由。

TFC算法通過(guò)通道依賴圖（CDG）建模死鎖，并將UB-Mesh拓?fù)鋭澐譃樽訄D。在每個(gè)子圖內(nèi)，拓?fù)鋵?dǎo)向規(guī)則和VL約束被統(tǒng)一為單一集合。算法應(yīng)用N維跨維度斷環(huán)原則，將集合分解為單VL無(wú)環(huán)子集，隨后采用同維度斷環(huán)原則計(jì)算冪集元素的排列和笛卡爾積，生成所有路徑組合和VL映射，確保無(wú)死鎖運(yùn)行。由于篇幅限制，詳細(xì)說(shuō)明在此省略。

4.2 基于直接通知的快速故障恢復(fù)

圖12：從逐跳通知到直接通知

當(dāng)網(wǎng)絡(luò)中發(fā)生鏈路故障時(shí)，傳統(tǒng)路由系統(tǒng)通常采用逐跳通知方式。如圖12左側(cè)所示，當(dāng)鏈路1-3發(fā)生故障時(shí)，該信息由節(jié)點(diǎn)1和3逐跳傳播，通常耗時(shí)較長(zhǎng)。在UB-Mesh中，由于每個(gè)節(jié)點(diǎn)具有確定的通信目標(biāo)集，我們可以通過(guò)直接通知受影響的節(jié)點(diǎn)加速路由收斂。如圖12右側(cè)所示，當(dāng)鏈路1-3發(fā)生故障時(shí)，根據(jù)預(yù)計(jì)算的路由關(guān)系，該信息直接發(fā)送至節(jié)點(diǎn)6。通過(guò)這種拓?fù)涓兄闹苯油ㄖ绞剑纱蠓鶞p少控制平面的開(kāi)銷。

5.性能最大化

盡管UB-Mesh的硬件架構(gòu)設(shè)計(jì)與LLM訓(xùn)練的流量模式相匹配，但在這種層次化的AI集群上運(yùn)行工作負(fù)載時(shí)，如果訓(xùn)練任務(wù)未能在計(jì)算資源上有效分布，可能會(huì)導(dǎo)致低利用率。此外，集合通信的性能對(duì)整體訓(xùn)練效率至關(guān)重要。為確保LLM訓(xùn)練期間的系統(tǒng)性能最優(yōu)，我們引入了若干拓?fù)涓兄獌?yōu)化策略，以進(jìn)一步提升系統(tǒng)性能。

5.1 拓?fù)涓兄贤ㄐ?/h4>

為優(yōu)化UB-Mesh上的集合通信并充分利用層次化直接互連的帶寬，我們提出了利用UB-Mesh全路徑路由能力（見(jiàn)第4.1節(jié)）的拓?fù)涓兄贤ㄐ潘惴?。我們以All-Reduce和All-to-All為例闡述我們的概念：

圖13：AllReduce的多環(huán)算法

All-Reduce：我們提出了一個(gè)多環(huán)（Multi-Ring）算法，以在UB-Mesh上高效實(shí)現(xiàn)All-Reduce。我們首先通過(guò)統(tǒng)一的抽象模型對(duì)網(wǎng)絡(luò)拓?fù)溥M(jìn)行建模，考慮節(jié)點(diǎn)數(shù)量、節(jié)點(diǎn)間連接、鏈路帶寬和延遲等因素。接下來(lái)，我們將集合通信與路徑映射相結(jié)合，采用邏輯多環(huán)算法，確保路徑獨(dú)占使用，避免流量沖突。如圖13-(a)所示，原始路徑表示默認(rèn)映射。未包含在這些路徑中的空閑鏈路通過(guò)APR機(jī)制被利用，以提升帶寬。最后，如圖13-(b)所示，我們優(yōu)化了多路徑上的流量分區(qū)，以緩解瓶頸并通過(guò)APR最大化借用帶寬的收益。

All-to-All：在All-to-All通信場(chǎng)景中，我們考慮了兩種代表性用例：

圖14：多路徑與層次化All-to-All

(1) 通用All-to-All：如圖14-(a)所示，當(dāng)UB-Mesh中的源節(jié)點(diǎn)（為簡(jiǎn)化起見(jiàn)，圖中展示2D網(wǎng)格，提出的技術(shù)可擴(kuò)展至更高維度）向不同目標(biāo)節(jié)點(diǎn)發(fā)送不同數(shù)據(jù)時(shí)，我們采用多路徑All2All優(yōu)化。具體而言，每個(gè)元素（向量或張量）被分為兩個(gè)分區(qū)，分別沿X-FullMesh和Y-FullMesh互連同時(shí)傳輸，并最多使用一跳轉(zhuǎn)發(fā)到達(dá)目標(biāo)。這種策略保證了UB-Mesh nD全互連架構(gòu)中的高帶寬利用率。

(2) 廣播+歸約（Broadcast + Reduce）：對(duì)于涉及標(biāo)記分發(fā)和專家數(shù)據(jù)收集的All-to-All操作[7]，其語(yǔ)義等價(jià)于多個(gè)廣播和歸約操作的重疊。如圖14-(b)和(c)所示，我們可采用層次化廣播/歸約以節(jié)省帶寬使用，充分利用UB-Mesh的層次拓?fù)洹?/p>

5.2 拓?fù)涓兄⑿谢?/p>

圖15：拓?fù)涓兄⑿谢?/p>

給定訓(xùn)練任務(wù)和系統(tǒng)，我們確定最優(yōu)并行策略，以充分利用UB-Mesh的高帶寬局部互連。我們采用拓?fù)涓兄⑿谢瘷C(jī)制優(yōu)化LLM訓(xùn)練的模型和數(shù)據(jù)分割。如圖15所示，我們的方法包括：

步驟①：生成可行的并行配置并將其映射到UB-Mesh架構(gòu)上。

步驟②：利用拓?fù)涓兄ㄐ懦杀灸Ｐ驮u(píng)估通信成本。

步驟③：迭代最小化通信開(kāi)銷以找到最優(yōu)配置。

該機(jī)制有兩個(gè)重要要求：(1) 構(gòu)建適當(dāng)?shù)乃阉骺臻g以平衡效率和性能；(2) 確保成本模型盡可能準(zhǔn)確。對(duì)于要求(1)，我們使用基于優(yōu)先級(jí)的啟發(fā)式方法剪枝搜索空間：涉及高通信量的TP和SP（或CP）優(yōu)先分配到高帶寬域，而用于梯度更新的PP和DP優(yōu)先級(jí)最低。對(duì)于需要EP的MoE模型，我們強(qiáng)制SP*DP為EP的整數(shù)倍。對(duì)于要求(2)，我們?cè)赨B-Mesh拓?fù)渖蠝?zhǔn)確建模APR和拓?fù)涓兄贤ㄐ诺男袨椋⑹褂脙?nèi)部高精度仿真基礎(chǔ)設(shè)施校準(zhǔn)模型。

6.評(píng)估

在本節(jié)中，我們探索了UB-Mesh的架構(gòu)設(shè)計(jì)空間，并分析其相較于基線Clos架構(gòu)的優(yōu)勢(shì)。為便于更詳細(xì)的比較，我們將分析分為兩個(gè)層級(jí)：機(jī)架內(nèi)架構(gòu)比較和機(jī)架間架構(gòu)比較。

6.1 實(shí)驗(yàn)設(shè)置

表5：基準(zhǔn)模型

表5列出了基準(zhǔn)工作負(fù)載：Llama-70B和GPT3-175B為密集模型，而GPT4-2T為采用MoE技術(shù)的稀疏模型[20]。需要注意的是，GPT4的架構(gòu)尚未正式發(fā)布，我們采用了推測(cè)參數(shù)[23]。為評(píng)估系統(tǒng)在更大模型上的性能，我們還包括了Dense-1T和MoE-10T模型。為探索UB-Mesh的架構(gòu)，我們構(gòu)建了一個(gè)內(nèi)部仿真基礎(chǔ)設(shè)施，并與真實(shí)PoC硬件對(duì)齊，以評(píng)估集群規(guī)模的LLM訓(xùn)練性能。

6.2 機(jī)架內(nèi)架構(gòu)探索

圖16：基線機(jī)架內(nèi)架構(gòu)

我們比較了不同的機(jī)架內(nèi)網(wǎng)絡(luò)架構(gòu)，如圖16所示：

(a) 2D-FM：UB-Mesh的架構(gòu)。64個(gè)NPU通過(guò)電線直接互連，形成2D全互連網(wǎng)絡(luò)拓?fù)洌╔-Fullmesh + Y-Fullmesh）。該設(shè)計(jì)成本低，采用48個(gè)LRS（見(jiàn)表3）聚合機(jī)架間帶寬并實(shí)現(xiàn)CPU-NPU互連（為簡(jiǎn)化起見(jiàn)，圖中省略了CPU）。

(b) 1D-FM-A：此替代架構(gòu)保留1D X-Fullmesh，即每個(gè)板上的8個(gè)NPU仍直接互連。然而，跨板通信通過(guò)36個(gè)LRS實(shí)現(xiàn)。每個(gè)NPU有一個(gè)UB x16 IO連接到LRS，另一個(gè)UB x16 IO連接到四個(gè)高基數(shù)交換機(jī)（HRS）用于機(jī)架間通信。

(c) 1D-FM-B：此架構(gòu)進(jìn)一步用HRS替換LRS。每個(gè)背板中的四個(gè)LRS用于NPU到CPU通信?？绨錘PU通信通過(guò)四個(gè)背板中的八個(gè)交換機(jī)實(shí)現(xiàn)。這些交換機(jī)還連接到機(jī)架間網(wǎng)絡(luò)，為每個(gè)NPU提供UB x32 IO用于機(jī)架間通信。

(d) Clos：此架構(gòu)不使用NPU間直接連接，而是將64個(gè)NPU的所有端口連接到72個(gè)LRS，形成對(duì)稱的Clos拓?fù)?。該架?gòu)提供最高靈活性，但需要大量交換資源。

圖17：不同機(jī)架內(nèi)拓?fù)涞男阅?/p>

性能比較：在圖17中，我們固定機(jī)架間架構(gòu)（2D-FM，見(jiàn)第6.3節(jié)），并比較不同機(jī)架內(nèi)架構(gòu)的訓(xùn)練吞吐量，相對(duì)于Clos基線。SuperPod規(guī)模為8000（128個(gè)機(jī)架），評(píng)估的序列長(zhǎng)度范圍為8000至1000萬(wàn)。我們計(jì)算了不同序列長(zhǎng)度的平均性能。

如圖17-(a)所示，與Clos架構(gòu)相比，2D-FM架構(gòu)實(shí)現(xiàn)了93.2%至95.9%的訓(xùn)練性能。1D-FM-A架構(gòu)顯示出較低的性能下降，對(duì)于Llama2-70B相比2D-FM提升了2.44%的性能。對(duì)于其他參數(shù)更多的模型，提升幅度小于1.6%。1D-FM-B架構(gòu)由于更高的機(jī)架間帶寬，相比2D-FM顯示出略高的性能提升，超過(guò)3%，但提升幅度仍較小。

圖17-(b)探索了不同序列長(zhǎng)度下的性能，相對(duì)于Clos架構(gòu)的基線場(chǎng)景（所有模型平均值）。對(duì)于8000至3.2萬(wàn)的序列長(zhǎng)度，2D-FM架構(gòu)實(shí)現(xiàn)了95.5%的性能，略低于1D-FM-A（98.1%）和1D-FM-B（99.2%）。對(duì)于6.4萬(wàn)至1000萬(wàn)的序列長(zhǎng)度，2D-FM架構(gòu)相比Clos架構(gòu)實(shí)現(xiàn)了95.0%的性能。

可以看出，與Clos相比，2D-FM架構(gòu)以遠(yuǎn)低于硬件成本提供了相似的訓(xùn)練性能（性能差距在7%以內(nèi)），這將在第6.4節(jié)中進(jìn)行評(píng)估。

6.3 機(jī)架間架構(gòu)探索

圖18：機(jī)架間網(wǎng)絡(luò)架構(gòu)

在機(jī)架內(nèi)2D-FM之外，UB-Mesh在UB-Mesh Pod內(nèi)采用另一個(gè)2D-FM，形成4D-FM網(wǎng)絡(luò)架構(gòu)。如圖18所示，我們比較了UB-Mesh的2D全互連機(jī)架間架構(gòu)與基線Clos架構(gòu)：

(a) 2D-FM：16個(gè)機(jī)架在水平和垂直方向直接互連，形成2D全互連直接連接。所有機(jī)架連接到HRS交換機(jī)，支持跨Pod互連。在UB-Mesh-Pod中，可支持三種路由策略：  

• 最短路徑（Shortest）：基線路由策略，僅選擇2D網(wǎng)格上的最短路徑用于P2P通信。
• 繞行（Detour）：根據(jù)全路徑路由機(jī)制，還使用其他路徑以最大化帶寬。
• 借用（Borrow）：由于所有機(jī)架也連接到交換機(jī)，我們?cè)试S機(jī)架“借用”交換機(jī)的帶寬。

(b) Clos：所有機(jī)架連接到HRS交換機(jī)，無(wú)直接連接。此架構(gòu)比2D-FM消耗更多交換機(jī)，但提供最高的All-to-All帶寬和靈活性。

圖19：不同機(jī)架間互連的端到端性能比較

性能比較：如圖19所示，我們比較了2D-FM（包括不同路由策略）和Clos架構(gòu)的性能?？梢钥闯?，與理想的Clos架構(gòu)相比，2D-FM與Clos之間的性能差距微乎其微，尤其是應(yīng)用Detour和Borrow策略時(shí)。GPT3-175B對(duì)機(jī)架間通信性能不敏感，而在GPT4-2T上，采用最短路徑路由的樸素2D-FM訓(xùn)練性能下降0.73%。采用Detour和Borrow路由策略后，性能差距縮小至僅0.46%。總體而言，2D-FM機(jī)架間互連表現(xiàn)出與昂貴的Clos架構(gòu)幾乎相同的性能。

圖20：機(jī)架間帶寬探索

機(jī)架間帶寬探索：圖20-(b)比較了8000 SuperPod在不同機(jī)架間帶寬條件下的吞吐性能。評(píng)估的機(jī)架間帶寬為每NPU的x4、x8、x16和x32 UB IO。對(duì)于8000至3.2萬(wàn)的序列長(zhǎng)度，最優(yōu)機(jī)架間帶寬為UB x16；而對(duì)于6.4萬(wàn)至1000萬(wàn)的序列長(zhǎng)度，最優(yōu)機(jī)架間帶寬為UB x32。從UB x8增加到UB x16在8000-3.2萬(wàn)序列長(zhǎng)度范圍內(nèi)的性能增益僅為0.44%，但從UB x16增加到UB x32在6.4萬(wàn)-1000萬(wàn)序列長(zhǎng)度范圍內(nèi)的性能增益更為顯著，達(dá)1.85%。在6.4萬(wàn)至1000萬(wàn)序列長(zhǎng)度的場(chǎng)景中，部分TP（張量并行）和SP（序列并行）流量不可避免地穿越機(jī)架間鏈路。更高的機(jī)架間帶寬顯著減少了TP和SP的通信時(shí)間，在這些場(chǎng)景中帶來(lái)更明顯的性能提升。這些數(shù)據(jù)強(qiáng)調(diào)了將機(jī)架間帶寬與不同模型場(chǎng)景的特定序列長(zhǎng)度需求相匹配的重要性，特別是在減少大規(guī)模模型的TP和SP通信延遲方面。

UB-Mesh默認(rèn)分配每NPU UB x16 IO用于機(jī)架間通信，以在成本和性能之間實(shí)現(xiàn)平衡。我們還可調(diào)整機(jī)架內(nèi)/機(jī)架間帶寬比例，以匹配特定LLM訓(xùn)練工作負(fù)載的需求。

6.4 成本效益比較

圖21：CapEx比較

系統(tǒng)成本通常以總擁有成本（TCO）衡量，即TCO = CapEx + OpEx?？紤]到相對(duì)于基線的訓(xùn)練性能，我們定義系統(tǒng)的成本效益如下：

我們使用內(nèi)部數(shù)據(jù)估算CapEx成本，包括NPU、CPU、LRS、HRS、電纜和其他模塊的成本，并比較不同架構(gòu)。如圖21所示，UB-Mesh的4D-FM+Clos架構(gòu)相較于2D-FM+x16（表示每NPU UB x16 IO）Clos、1D-FM+x16 Clos和x64T Clos架構(gòu)分別實(shí)現(xiàn)了1.18倍、1.26倍、1.65倍和2.46倍的CapEx降低。與基線Clos架構(gòu)相比，UB-Mesh成功將系統(tǒng)中網(wǎng)絡(luò)基礎(chǔ)設(shè)施成本占比從67%降低至20%，這得益于節(jié)省了高性能交換機(jī)和長(zhǎng)距離光纜/模塊。根據(jù)我們的評(píng)估，與基線Clos架構(gòu)相比，98%的高基數(shù)交換機(jī)和93%的光模塊得以節(jié)省。

OpEx降低主要包括系統(tǒng)生命周期內(nèi)的電費(fèi)和維護(hù)成本。UB-Mesh相較于Clos減少了約35%的OpEx，這得益于其大幅減少了交換機(jī)和光模塊的使用。根據(jù)我們?cè)撇块T的AI系統(tǒng)估算，OpEx約占TCO的30%。最終，根據(jù)公式(1)，UB-Mesh實(shí)現(xiàn)了2.04倍更高的成本效益。

6.5 線性度評(píng)估

圖22：序列長(zhǎng)度256K時(shí)的線性度分析

AI集群的線性度指的是集群性能隨NPU數(shù)量增加的線性擴(kuò)展程度。具體而言，我們通過(guò)以下公式測(cè)量線性度：

圖22評(píng)估了UB-Mesh在不同集群規(guī)模下的線性度?；鶞?zhǔn)規(guī)模（即圖中的1×規(guī)模）因任務(wù)而異。具體而言，Llama2-70B使用128個(gè)NPU，GPT3-175B的基準(zhǔn)規(guī)模為512，Dense-1T和GPT4-2T使用1000個(gè)NPU。

可以看出，UB-Mesh在所有任務(wù)的1×至32×規(guī)模下線性度超過(guò)100%，這是因?yàn)橐?guī)模提升提供了更多高帶寬域，并釋放了搜索更優(yōu)并行策略以提升MFU（模型浮點(diǎn)利用率）的潛力。GPT4-2T和Dense-1T模型在64×規(guī)模（涉及6.4萬(wàn)個(gè)NPU）下線性度有所下降，但仍高于95%。

6.6 網(wǎng)絡(luò)可靠性分析

AI訓(xùn)練過(guò)程對(duì)硬件故障敏感。經(jīng)典的基于時(shí)間的指標(biāo)用于評(píng)估系統(tǒng)可靠性：

其中，MTBF表示平均故障間隔時(shí)間，MTTR表示平均修復(fù)時(shí)間。

表6：MTBF估算

如表6估算，優(yōu)先使用直接連接的電線（E-Cables）而非光纖和交換機(jī)，大幅降低了網(wǎng)絡(luò)模塊的年化故障率（AFR）。

然后，我們計(jì)算兩種架構(gòu)的MTBF

：基線Clos在8000個(gè)NPU集群中為13.8小時(shí)，而UB-Mesh達(dá)到98.5小時(shí)，提升了7.14倍。最終，根據(jù)公式(3)（我們根據(jù)現(xiàn)有統(tǒng)計(jì)數(shù)據(jù)假設(shè)MTTR為75分鐘），UB-Mesh的可用性為98.8%，顯著優(yōu)于Clos的91.6%（提升7.2%）。

為進(jìn)一步提升可用性，我們還精心設(shè)計(jì)了內(nèi)部網(wǎng)絡(luò)監(jiān)控工具，可在10分鐘內(nèi)快速識(shí)別和定位網(wǎng)絡(luò)故障，并在3分鐘內(nèi)觸發(fā)任務(wù)遷移，大幅減少M(fèi)TTR。通過(guò)這種MTTR優(yōu)化，根據(jù)我們的評(píng)估，UB-Mesh的可用性可進(jìn)一步提升至99.78%。

總結(jié)：與基線Clos架構(gòu)相比，UB-Mesh以微小的性能下降（7%以內(nèi)）為代價(jià)實(shí)現(xiàn)了2.04倍更高的成本效益。由于大幅減少了交換機(jī)和光模塊的使用，UB-Mesh將網(wǎng)絡(luò)可用性提升了7.2%。UB-Mesh還在多個(gè)LLM訓(xùn)練任務(wù)上實(shí)現(xiàn)了95%以上的線性度。

7.討論

集合通信協(xié)處理器

值得一提的是，在UB IO控制器中還配備了一個(gè)特殊的協(xié)處理器，稱為集合通信單元（Collective Communication Unit, CCU），用于卸載集合通信任務(wù)。具體而言，CCU執(zhí)行指令，主動(dòng)從HBM讀取或?qū)懭霐?shù)據(jù)，啟動(dòng)NPU間傳輸，并利用片上SRAM緩沖區(qū)執(zhí)行在線數(shù)據(jù)歸約操作。這種設(shè)計(jì)消除了從應(yīng)用程序內(nèi)存緩沖區(qū)到通信緩沖區(qū)的冗余數(shù)據(jù)拷貝，有效緩解了HBM帶寬的消耗，并通過(guò)基于檢查位的細(xì)粒度同步機(jī)制維持確定的歸約順序。CCU還能與計(jì)算核心無(wú)縫協(xié)作，實(shí)現(xiàn)高效的計(jì)算與通信重疊。

支持大規(guī)模專家模型

除了密集模型和常規(guī)MoE模型外，我們注意到人們正在積極探索包含大規(guī)模專家的MoE模型[6, 7]。這種設(shè)計(jì)通常需要大規(guī)模、細(xì)粒度的All-to-All通信。UB-Mesh通過(guò)多路徑和層次化的All-to-All優(yōu)化（見(jiàn)第5.1節(jié)）以及基于UB的Load/Store數(shù)據(jù)傳輸，能夠高效支持此類大規(guī)模專家模型。CCU還能有效卸載All-to-All操作，節(jié)省寶貴的計(jì)算核心使用，這與[7]中的預(yù)期一致。

8.結(jié)論

本文提出了UB-Mesh，一種為下一代大規(guī)模語(yǔ)言模型（LLM）訓(xùn)練設(shè)計(jì)的新型數(shù)據(jù)中心網(wǎng)絡(luò)架構(gòu)。UB-Mesh采用n維全互連（nD-FullMesh）網(wǎng)絡(luò)拓?fù)?，減少了交換機(jī)和光模塊的使用，同時(shí)適配LLM工作負(fù)載的分層局部化流量模式。我們深入考慮了架構(gòu)設(shè)計(jì)、物理實(shí)現(xiàn)和網(wǎng)絡(luò)系統(tǒng)優(yōu)化，并提出了多項(xiàng)技術(shù)以應(yīng)對(duì)各種挑戰(zhàn)。與傳統(tǒng)的Clos網(wǎng)絡(luò)相比，UB-Mesh在提供相似LLM訓(xùn)練性能的同時(shí)，實(shí)現(xiàn)了2.04倍的系統(tǒng)級(jí)成本效益提升和7.2%的可用性提升。此外，UB-Mesh在多項(xiàng)LLM訓(xùn)練任務(wù)上實(shí)現(xiàn)了95%以上的線性擴(kuò)展能力。

本文轉(zhuǎn)載自??Andy730??，作者：常華Andy