AWS擲出算力核彈，Trainium3深度解析

本文由半導(dǎo)體產(chǎn)業(yè)縱橫（ID：ICVIEWS）編譯自semianalysis

云廠商自研芯片終極形態(tài)？

亞馬遜在年度AWS re:Invent大會(huì)上正式發(fā)布了Trainium3 (Trn3)，并宣布推出Trainium4 (Trn4)。亞馬遜在數(shù)據(jù)中心定制芯片領(lǐng)域擁有最悠久、最廣泛的歷史。

今天，本文將深入探討Trainium3 芯片的飛躍式性能提升，包括微架構(gòu)、系統(tǒng)和機(jī)架架構(gòu)、擴(kuò)展性、性能分析器、軟件平臺(tái)以及數(shù)據(jù)中心加速等。

Amazon Basics GB200

借助Trainium3，AWS 始終專注于優(yōu)化總體擁有成本 (TCO) 下的性能。他們的硬件核心理念很簡(jiǎn)單：以最低的 TCO 實(shí)現(xiàn)最快的上市速度。AWS 不拘泥于任何單一的架構(gòu)設(shè)計(jì)，而是最大限度地提高運(yùn)營(yíng)靈活性。這體現(xiàn)在他們與多家合作伙伴在定制芯片方面的合作，以及對(duì)自身供應(yīng)鏈的管理，包括從多家組件供應(yīng)商采購(gòu)。

在系統(tǒng)和網(wǎng)絡(luò)方面，AWS 遵循“亞馬遜基礎(chǔ)架構(gòu)”理念，以性能與總體擁有成本 (TCO) 為導(dǎo)向進(jìn)行優(yōu)化。諸如使用 12.8T、25.6T 還是 51.2T 帶寬橫向擴(kuò)展交換機(jī)，或者選擇液冷還是風(fēng)冷等設(shè)計(jì)選擇，都只是實(shí)現(xiàn)特定客戶和數(shù)據(jù)中心最佳總體擁有成本的手段。

對(duì)于橫向擴(kuò)展網(wǎng)絡(luò)，雖然Trn2 僅支持 4x4x4 三維環(huán)形網(wǎng)格橫向擴(kuò)展拓?fù)?，?Trainium3 添加了一種獨(dú)特的交換式架構(gòu)，該架構(gòu)與 GB200 NVL36x2 拓?fù)浣Y(jié)構(gòu)有些類似，但也存在一些關(guān)鍵差異。添加這種交換式架構(gòu)的原因是，對(duì)于前沿的混合專家 (MoE) 模型架構(gòu)而言，交換式橫向擴(kuò)展拓?fù)渚哂懈叩慕^對(duì)性能和更高的總擁有成本 (TCO) 性能。

即使對(duì)于這種擴(kuò)展架構(gòu)中使用的交換機(jī)，AWS 也決定不做決定：在 Trainium3 的生命周期內(nèi)，他們將采用三種不同的擴(kuò)展交換機(jī)解決方案，首先是 160 通道、20 端口的 PCIe 交換機(jī)，以便快速上市，因?yàn)槟壳案咄ǖ篮投丝跀?shù)的 PCIe 交換機(jī)的可用性有限；之后切換到 320 通道 PCIe 交換機(jī)；最終切換到更大的 UALink，以實(shí)現(xiàn)最佳性能。

亞馬遜的軟件

在軟件方面，AWS 的 North Star 擴(kuò)展并開放了其軟件堆棧，以面向大眾，不再僅僅針對(duì)內(nèi)部 Bedrock 工作負(fù)載（例如運(yùn)行 vLLM v1 私有分支的 DeepSeek/Qwen 等）和 Anthropic 的訓(xùn)練和推理工作負(fù)載（運(yùn)行自定義推理引擎和所有自定義 NKI 內(nèi)核）優(yōu)化每 TCO 的性能。

事實(shí)上，他們正在進(jìn)行一項(xiàng)大規(guī)模、多階段的軟件戰(zhàn)略轉(zhuǎn)型。第一階段是發(fā)布并開源一個(gè)新的原生PyTorch 后端。他們還將開源其內(nèi)核語(yǔ)言“NKI”（神經(jīng)元內(nèi)核接口）的編譯器，以及內(nèi)核和通信庫(kù) matmul 和 ML ops（類似于 NCCL、cuBLAS、cuDNN 和 Aten Ops）。第二階段包括開源其 XLA 圖編譯器和 JAX 軟件棧。

通過(guò)開源其大部分軟件棧，AWS 將有助于擴(kuò)大 CUDA 的應(yīng)用范圍，并啟動(dòng)一個(gè)開放的開發(fā)者生態(tài)系統(tǒng)。我們認(rèn)為，CUDA 的護(hù)城河并非由建造城堡的英偉達(dá)工程師構(gòu)筑，而是由數(shù)百萬(wàn)外部開發(fā)者通過(guò)為 CUDA 生態(tài)系統(tǒng)做出貢獻(xiàn)而挖掘的。AWS 已經(jīng)意識(shí)到這一點(diǎn)，并正在推行同樣的戰(zhàn)略。

Trainium3 最初僅支持邏輯神經(jīng)元核心 (LNC) = 1 或 LNC = 2。LNC = 1 或 LNC = 2 是亞馬遜/Anthropic 的頂級(jí) L337 內(nèi)核工程師所期望的，但 LNC = 8 才是更廣泛的機(jī)器學(xué)習(xí)研究科學(xué)家群體在廣泛采用 Trainium 之前所更傾向的選擇。遺憾的是，AWS 計(jì)劃在 2026 年年中之前不會(huì)支持 LNC = 8。下文將詳細(xì)闡述 LNC 的概念以及不同模式對(duì)于研究科學(xué)家采用 Trainium 的重要性。

Trainium3 的上市又開辟了另一個(gè)戰(zhàn)線，Jensen 現(xiàn)在除了要面對(duì)另外兩個(gè)戰(zhàn)場(chǎng)——性能/總擁有成本 (TCO) 極高的 Google TPUv7和性能/總擁有成本 (TCO) 強(qiáng)勁的AMD MI450X UALoE72 （尤其是在 OpenAI 獲得“股權(quán)返還”后，AMD 最多可持有 10% 的股份）——之外，還必須應(yīng)對(duì)這個(gè)戰(zhàn)線。

我們?nèi)匀幌嘈?，只要英偉達(dá)繼續(xù)加快研發(fā)步伐，以光速前進(jìn)，他們就能繼續(xù)保持行業(yè)霸主的地位。詹森需要比過(guò)去四個(gè)月的速度更快。就像英特爾在CPU領(lǐng)域安于現(xiàn)狀，而AMD和ARM等競(jìng)爭(zhēng)對(duì)手卻一路領(lǐng)先一樣，如果英偉達(dá)也安于現(xiàn)狀，他們將會(huì)更快地失去領(lǐng)先地位。

在本文中，我們將討論支持可切換式擴(kuò)展機(jī)架的兩款Trainium3 機(jī)架 SKU：

• 空冷式Trainium3 NL32x2 開關(guān)機(jī)（代號(hào)“Teton3 PDS”）
• 液冷式Trainium3 NL72x2 開關(guān)式（代號(hào)“Teton3 MAX”）

我們將首先簡(jiǎn)要回顧Trainium 2架構(gòu)，并解釋Trainium 3引入的變更。文章前半部分將重點(diǎn)介紹各種Trainium 3機(jī)架式SKU的規(guī)格、芯片設(shè)計(jì)、機(jī)架架構(gòu)、物料清單（BOM）和功耗預(yù)算，然后再探討其縱向擴(kuò)展和橫向擴(kuò)展網(wǎng)絡(luò)架構(gòu)。文章后半部分將重點(diǎn)討論Trainium 3微架構(gòu)，并進(jìn)一步闡述亞馬遜的軟件戰(zhàn)略。最后，我們將討論亞馬遜和Anthropic的AI數(shù)據(jù)中心，并結(jié)合總體擁有成本（TCO）和每TCO性能分析，總結(jié)全文。

Trainium3 服務(wù)器類型和規(guī)格概述

Trainium2 和 Trainium3 之間共有四個(gè)不同的服務(wù)器 SKU，供應(yīng)鏈通常使用它們的代號(hào)來(lái)指代它們，這與 AWS 的品牌名稱不同。

讀者可能會(huì)發(fā)現(xiàn)，要理清各種不同的代際和機(jī)架尺寸組合，并在AWS品牌和ODM/供應(yīng)鏈?zhǔn)褂玫拇?hào)之間來(lái)回切換，實(shí)在令人困惑。我們懇請(qǐng)AWS：負(fù)責(zé)產(chǎn)品營(yíng)銷和命名的人員應(yīng)該停止使用這些令人費(fèi)解的名稱。理想情況下，他們應(yīng)該效仿Nvidia和AMD的命名規(guī)則，采用產(chǎn)品名稱后半部分分別表示擴(kuò)展技術(shù)和GPU數(shù)量，例如GB200中的NVL72指的是NVLink，GPU數(shù)量為72。

下表旨在為讀者提供清晰的解釋，幫助他們理解不同群體使用的各種命名規(guī)則，避免混淆：

來(lái)源：SemiAnalysis、AWS

Trainium3 在規(guī)格方面實(shí)現(xiàn)了幾項(xiàng)顯著的代際升級(jí)。

OCP MXFP8 的 FLOPs 吞吐量翻倍，并且增加了對(duì) OCP MXFP4 的支持，但性能與 OCP MXFP8 相同。有趣的是，對(duì)于更高精度的數(shù)字格式（例如 FP16 和 FP32），其性能與 Trn2 保持一致。在微架構(gòu)部分，我們將描述這些權(quán)衡的影響。

來(lái)源：SemiAnalysis、AWS

Trainium3 的 HBM3E 升級(jí)至 12 層，使每個(gè)芯片的內(nèi)存容量達(dá)到 144GB。盡管 Trn2 仍然使用 4 層 HBM3E，但 AWS 通過(guò)將引腳速度從低于平均水平的 5.7Gbps 提升至 Trn3 的 9.6Gbps，實(shí)現(xiàn)了 70% 的內(nèi)存帶寬提升，這是我們迄今為止見過(guò)的最高HBM3E引腳速度。

事實(shí)上，Trn2 中使用的 5.7Gbps 引腳速度更接近 HBM3 的速度，但由于它使用 24Gb 的芯片，在 8 層堆疊中每個(gè)堆疊提供 24GB 的容量，因此仍被歸類為 HBM3E。速度不足是由于使用了三星提供的內(nèi)存，其 HBM3E 的性能明顯低于海力士或美光。為了在 Trainium3 中實(shí)現(xiàn)更快的速度，AWS 正在切換到海力士和美光的 HBM。

與Trn2 相比，Trainium3 芯片的擴(kuò)展帶寬翻了一番，這得益于其采用 PCIe Gen 6，每條通道的帶寬為 64Gbps（單向），而 PCIe Gen 5 每條通道的帶寬為 32Gbps。Trainium3 使用 144 條 PCIe 活動(dòng)通道進(jìn)行擴(kuò)展，這意味著在 Gen6 上，每個(gè) Trainium3 芯片支持 1.2 TB/s（單向）的擴(kuò)展帶寬。

橫向擴(kuò)展帶寬支持翻倍至最大400 Gb/s，但大多數(shù)生產(chǎn)的 Trainium3 機(jī)架仍將沿用 Trn2 使用的每 XPU 200Gb/s 的橫向擴(kuò)展速度。

對(duì)于Trainium4，亞馬遜將使用 8 個(gè) HBM4 堆棧，與 Trainium3 相比，內(nèi)存帶寬提高 4 倍，容量提高 2 倍。

Trainium3 機(jī)架架構(gòu)

從機(jī)架解決方案層面來(lái)看，AWS 在 re:Invent 大會(huì)上發(fā)布了 Trainium3 (Gen1) UltraServer 和 Trainium3 (Gen2) UltraServer，分別對(duì)應(yīng) Trainium3 NL32x2 Switched 和 Trainium3 NL72x2 Switched。Trainium3 NL32x2 Switched 和 Trainium3 NL72x2 Switched 的主要區(qū)別在于其橫向擴(kuò)展網(wǎng)絡(luò)拓?fù)浜蜋C(jī)架架構(gòu)——本節(jié)將介紹這兩個(gè) SKU 在拓?fù)浜图軜?gòu)上的差異，并討論每種架構(gòu)最適合和優(yōu)化的 AI 工作負(fù)載。

我們先來(lái)看一下每種服務(wù)器類型的物理布局。下表列出了每個(gè)機(jī)架式SKU 的主要規(guī)格：

來(lái)源：SemiAnalysis、AWS

Trainium2 僅提供前兩種機(jī)架 SKU 類型——即 Trn2 NL16 2D Torus 服務(wù)器和 Trn2 NL32x2 2D Torus 服務(wù)器，而 Trainium3 將提供所有四種機(jī)架 SKU 類型，其中大部分 Trainium3 將于 2026 年以 Trainium3 NL32x2 交換機(jī) SKU 的形式交付。我們預(yù)計(jì)，在 Trainium3 的整個(gè)生命周期內(nèi)，大部分 Trainium3 將部署在 Trainium3 NL32x2 交換機(jī) SKU 和 Trainium3 NL72x2 交換機(jī) SKU 中。

Trainium3 制程和封裝

Trainium3 的計(jì)算任務(wù)將從用于 Trn2 的 N5 節(jié)點(diǎn)遷移到 N3P 節(jié)點(diǎn)。Trainium3 將成為首批采用 N3P 的服務(wù)器之一，其他采用者還包括 Vera Rubin 和 MI450X 的主動(dòng)中介層芯片 (AID)。N3P 存在一些問(wèn)題需要解決，這可能會(huì)延后項(xiàng)目進(jìn)度。

來(lái)源：AWS、SemiAnalysis

我們認(rèn)為臺(tái)積電的N3P是3nm平臺(tái)上的“高性能計(jì)算旋鈕”，它在N3E的基礎(chǔ)上實(shí)現(xiàn)了雖小但意義重大的提升，無(wú)需新的設(shè)計(jì)規(guī)則即可提高頻率或降低功耗。公開數(shù)據(jù)顯示，N3P保留了N3E的規(guī)則和IP，但在相同漏電情況下速度提升約5%，或在相同頻率下功耗降低5%—10%，此外，在混合邏輯/SRAM/模擬設(shè)計(jì)中有效密度提升約4%。這正是超大規(guī)模集成電路制造商在大型AI ASIC中所需要的漸進(jìn)式、低摩擦的增益。

Trainium3 就是一個(gè)很好的例子，它說(shuō)明了為什么這類產(chǎn)品適合采用這種制程工藝。Trainium3 也充分體現(xiàn)了定制加速器為何會(huì)占據(jù) 3 納米高性能計(jì)算 (HPC) 產(chǎn)能的大部分：高密度矩陣引擎、大容量 SRAM 切片以及超長(zhǎng)的片上互連，這些都受益于器件延遲和漏電的每一次微小降低。

N3P 的底層技術(shù)與其說(shuō)是一項(xiàng)單一的突破，不如說(shuō)是多項(xiàng)設(shè)計(jì)-技術(shù)協(xié)同優(yōu)化 (DTCO) 調(diào)整疊加的結(jié)果。N3 代 FinFlex 庫(kù)允許設(shè)計(jì)人員在模塊內(nèi)混合使用更寬和更窄的鰭片，從而以精細(xì)的粒度權(quán)衡驅(qū)動(dòng)強(qiáng)度、面積和漏電。臺(tái)積電還改進(jìn)了 N3P 底層金屬堆疊中的襯線和阻擋層工藝，與之前的 3nm 工藝相比，降低了線路和過(guò)孔電阻。這些改進(jìn)共同提升了裕量，從而支持更高的時(shí)鐘頻率或更低的低電壓，尤其是在長(zhǎng)全局路徑上。

挑戰(zhàn)在于，N3P工藝在推進(jìn)互連線尺寸縮小和圖形化的同時(shí)，幾乎達(dá)到了當(dāng)前EUV光刻工具所能達(dá)到的極限。最小金屬間距在20納米左右，高縱橫比通孔以及更小的光學(xué)尺寸縮小都會(huì)加劇后端工藝的變異性和RC值。諸如通孔輪廓控制、蝕刻不足和介質(zhì)損傷等問(wèn)題都成為首要的時(shí)序問(wèn)題。對(duì)于臺(tái)積電而言，這意味著更脆弱的工藝窗口、更復(fù)雜的在線監(jiān)控以及更頻繁地使用DTCO反饋回路，以確保設(shè)計(jì)規(guī)則與生產(chǎn)線大規(guī)模生產(chǎn)的能力保持一致。目前，N3P工藝的缺陷密度改善速度低于預(yù)期，這導(dǎo)致芯片設(shè)計(jì)人員要么需要重新旋印以提高良率，要么只能等待工藝改進(jìn)。

能夠解讀芯片標(biāo)識(shí)的讀者會(huì)發(fā)現(xiàn)，上圖所示的封裝正是Trn2，而我們也正是采用了這款芯片，因?yàn)樗姆庋b布局與Trainium3完全相同。該封裝由兩個(gè)CoWoS-R組件構(gòu)成，而非一個(gè)大型中介層。兩個(gè)計(jì)算芯片通過(guò)基板相互連接。

Trainium3 將繼續(xù)采用臺(tái)積電的 CoWoS-R 平臺(tái)，該平臺(tái)在保持成本競(jìng)爭(zhēng)力的同時(shí)，突破了功耗和延遲的極限。與前代產(chǎn)品 Trainium2 不同，Trainium3 沒(méi)有采用全硅中介層，而是沿用了有機(jī)薄膜中介層，該中介層由六層銅 RDL 層構(gòu)成，覆蓋了光罩級(jí)尺寸，成本更低，機(jī)械柔順性也優(yōu)于硅中介層。它仍然支持芯片與中介層之間數(shù)十微米的精細(xì)布線和微凸點(diǎn)間距，這對(duì)于高密度芯片組結(jié)構(gòu)和 HBM 接口至關(guān)重要。其下方是二十層 ABF 基板，該基板將電源和 XSR 信號(hào)扇出到模塊邊界處的 130 至 150 微米 C4 凸點(diǎn)，MCM 在此處與電路板連接。

在CoWoS-R架構(gòu)上，超過(guò)六層的多層RDL并非硬性限制，而是一種有意為之的折衷方案。純有機(jī)中介層成本低廉且符合規(guī)范，但當(dāng)我們嘗試以32Gbps或更高的速率集成更多通道時(shí)，其性能終將達(dá)到極限。IPD（集成無(wú)源器件）通過(guò)在必要位置將小型硅無(wú)源元件嵌入有機(jī)層中來(lái)彌補(bǔ)這一不足。每個(gè)RDL中介層中數(shù)千個(gè)IPD能夠?qū)崿F(xiàn)亞微米級(jí)布線密度、極小的微凸點(diǎn)間距，并在芯片噪聲最大的區(qū)域（例如HBM PHY環(huán)和核心結(jié)構(gòu)）實(shí)現(xiàn)強(qiáng)大的去耦。

該芯片的前端由Annapurna設(shè)計(jì)，PCIe SerDes部分則獲得了Synopsys的授權(quán)。后端物理設(shè)計(jì)和封裝設(shè)計(jì)由Alchip負(fù)責(zé)。我們認(rèn)為Trainium3可能繼承了Marvell設(shè)計(jì)的Trainium2的部分接口IP，但就內(nèi)容而言意義不大。Marvell還委托其他第三方廠商進(jìn)行封裝設(shè)計(jì)。

有趣的是，有兩個(gè)流片項(xiàng)目，一個(gè)由Alchip擁有（稱為“Anita”），另一個(gè)則由Annapurna直接擁有（稱為“Mariana”）。在Anita項(xiàng)目中，Alchip直接從臺(tái)積電采購(gòu)芯片組件；而在Mariana項(xiàng)目中，Annapurna也直接采購(gòu)芯片組件。大部分產(chǎn)量將用于Mariana項(xiàng)目。雖然Alchip在Mariana項(xiàng)目中的設(shè)計(jì)參與度與Anita項(xiàng)目類似，但他們從Mariana項(xiàng)目獲得的收入應(yīng)該會(huì)低于Anita項(xiàng)目。亞馬遜和Annapurna都非常注重成本控制，并對(duì)供應(yīng)商施加了很大的壓力。與博通的ASIC交易相比，Trainium項(xiàng)目留給芯片設(shè)計(jì)合作伙伴Alchip和Marvell的利潤(rùn)空間要小得多。就每總擁有成本（TCO）的性能而言，Annapurna更加重視降低TCO。

Marvell最終成了這場(chǎng)競(jìng)爭(zhēng)的最大輸家。盡管Trainium2是由他們?cè)O(shè)計(jì)的，但他們?cè)谶@一代芯片的設(shè)計(jì)競(jìng)賽中輸給了Alchip。Marvell的Trainium3版本采用了基于芯片組（chiplet）的設(shè)計(jì)，將I/O功能放在一個(gè)獨(dú)立的芯片組上，而不是像Trainium2和即將推出的Trainium3那樣，將計(jì)算功能集成在一個(gè)單芯片上。

Marvell 因 Trainium2 的糟糕執(zhí)行而失去了這個(gè)插槽。開發(fā)周期過(guò)長(zhǎng)。Marvell 在設(shè)計(jì)該封裝的 RDL 中介層時(shí)也遇到了問(wèn)題，最終 Alchip 不得不介入，幫助其交付一個(gè)可行的方案。

Trainium4路線圖預(yù)估

Trainium4將由多家設(shè)計(jì)公司參與，分別基于不同的擴(kuò)展協(xié)議，分兩條不同的軌道進(jìn)行開發(fā)。我們?cè)缭?7 個(gè)月前的 5 月份就詳細(xì)介紹了 Trainium 4 加速器模型中 UALink/NVLink 的劃分。與 Trainium3 一樣，Alchip 將繼續(xù)主導(dǎo)兩條軌道的后端設(shè)計(jì)。

• 第一條軌道將采用UALink 224G。
• 第二條軌道將使用英偉達(dá)的NVLink 448G BiDi協(xié)議。

英偉達(dá)的VR芯片NVL144與Trainium4等NVLink Fusion產(chǎn)品之間很可能存在顯著的時(shí)間差。對(duì)于NVLink Fusion產(chǎn)品線而言，時(shí)間線可能會(huì)進(jìn)一步延后，因?yàn)镕usion芯片引入了額外的集成和驗(yàn)證要求，而且大多數(shù)英偉達(dá)混合信號(hào)工程師都將把精力集中在Nvidia VR芯片NVL144新產(chǎn)品的推出上。

雖然搭載NVLink Fusion 的 Trainium4 可能短期內(nèi)不會(huì)面世，但我們相信 AWS 已獲得有利的商業(yè)條款，不太可能支付英偉達(dá)通常約 75% 的毛利率。英偉達(dá)有強(qiáng)烈的戰(zhàn)略動(dòng)機(jī)來(lái)實(shí)現(xiàn)與 Trainium4 的互操作性，因?yàn)樵试S AWS 使用 NVLink 有助于英偉達(dá)維持其系統(tǒng)級(jí)鎖定。因此，英偉達(dá)可能會(huì)提供比其標(biāo)準(zhǔn)毛利率結(jié)構(gòu)下更具吸引力的價(jià)格。

與僅限于固定72封裝NVLink域的VR NVL144不同，Trainium4可通過(guò)跨機(jī)架AEC擴(kuò)展NVLink的規(guī)模，從而實(shí)現(xiàn)更大的144+相干域。NVLink 6采用400G雙向SerDes，允許在同一根導(dǎo)線上同時(shí)進(jìn)行200G接收和200G發(fā)送。這種400G雙向信號(hào)傳輸已經(jīng)接近銅纜的實(shí)際極限，盡管一些廠商可能會(huì)嘗試向前邁半代，推出600G雙向信號(hào)傳輸。

來(lái)源：英偉達(dá)

Trn2/3 NL16 2D 環(huán)面和 NL64 3D 環(huán)面

Trainium2 NL16 2D Torus 和 Trainium2 NL64 3D Torus SKU 分別被命名為 Trainium2 Server 和 Trainium2 UltraServer，并在 re:Invent 2024 大會(huì)上發(fā)布。我們?cè)?Trainium2 深度解析中介紹了這兩種架構(gòu)。

簡(jiǎn)要回顧一下Trainium2 SKU——Trainium2 NL16 2D Torus 和 Trainium2 NL64 3D Torus SKU 的主要區(qū)別在于擴(kuò)展世界的大小。Trainium2 NL16 2D Torus 在整個(gè)擴(kuò)展世界大小下占用半個(gè)服務(wù)器機(jī)架，該世界大小包含 16 個(gè) Trainium2 設(shè)備，構(gòu)成一個(gè) 4x4 網(wǎng)格的 2D 環(huán)面。而 Trainium2 NL32x2 3D Torus 則由四個(gè) Trainium2 NL16 2D Torus 半機(jī)架服務(wù)器連接而成，總共占用兩個(gè)機(jī)架。這四個(gè) Trainium2 NL16 2D Torus 半機(jī)架服務(wù)器通過(guò) AEC 連接，從而創(chuàng)建一個(gè)包含 64 個(gè) Trainium2 設(shè)備的 4x4x4 3D 環(huán)面。

請(qǐng)注意，在上圖中，Trainium2 NL16 2D 環(huán)形圖表示一個(gè)完整的機(jī)架，但包括圖中所示機(jī)架內(nèi)的兩個(gè)放大世界。

如前所述，Trainium2 目前僅提供 NL16 2D Torus 和 NL32x2 3D Torus 機(jī)架式 SKU，分別采用 2D 和 3D Torus 拓?fù)浣Y(jié)構(gòu)，不提供其他拓?fù)浣Y(jié)構(gòu)。由于 Anthropic 的 Rainier 項(xiàng)目是 Trainium2 的主要需求來(lái)源，因此 Trainium2 的部署將主要采用 NL64 3D Torus 外形尺寸，生產(chǎn)也將遵循其對(duì) NL64 3D Torus 的偏好。這是因?yàn)?Anthropic 的推理模型需要更大規(guī)模的拓?fù)浣Y(jié)構(gòu)。

可切換機(jī)架規(guī)模架構(gòu)

當(dāng)英偉達(dá)推出采用全對(duì)全擴(kuò)展拓?fù)浣Y(jié)構(gòu)、擴(kuò)展規(guī)?？蛇_(dá)72顆芯片的Oberon架構(gòu)（GB200 NVL72）時(shí)，許多ASIC和GPU制造商調(diào)整了未來(lái)的機(jī)架式設(shè)計(jì)路線圖，以模仿英偉達(dá)的Oberon架構(gòu)。雖然AMD率先發(fā)布了類似Oberon架構(gòu)的MI400 Helios機(jī)架式服務(wù)器，但AWS將成為除英偉達(dá)之外首家實(shí)際交付并部署類似全對(duì)全交換式擴(kuò)展架構(gòu)的廠商，其產(chǎn)品采用Trainium3 NL32x2交換式和Trainium3 NL72x2交換式封裝。AMD的首款機(jī)架式擴(kuò)展設(shè)計(jì)MI450X UALoE72將比Trainium3的機(jī)架式擴(kuò)展設(shè)計(jì)晚一年發(fā)布，預(yù)計(jì)在年底上市。Meta也將在AMD的MI450X之前推出其首款交換式架構(gòu)產(chǎn)品。我們將在本文的“3D 環(huán)面架構(gòu)與切換架構(gòu)”部分解釋為什么切換式擴(kuò)展架構(gòu)優(yōu)于環(huán)面架構(gòu)。

在re:Invent 大會(huì)上展出的 Trainium3 UltraServer 對(duì)應(yīng)于 Trainium3 NL72x2 Switched 機(jī)架式服務(wù)器，但 Trainium3 NL72x2 Switched 只是 Trainium3 全交換機(jī)機(jī)架式服務(wù)器的兩種型號(hào)之一——另一種型號(hào)是 Trainium3 NL32x2 Switched。與 Trainium3 NL72x2 Switched 一樣，Trainium3 NL32x2 Switched 也是全交換機(jī)式服務(wù)器，但它與 Trainium3 NL72x2 Switched 的區(qū)別在于，Trainium3 NL32x2 Switched 采用風(fēng)冷散熱，因此其擴(kuò)展規(guī)模較小，功率密度也較低。

Trainium3 NL32x2 開關(guān)式（Teton3 PDS）

機(jī)架架構(gòu)

Trainium3 NL32x2 Switched 的機(jī)架布局與 Trainium NL32x2 3D Torus 非常相似。兩者每個(gè)機(jī)架都包含 16 個(gè) JBOG（GPU 集群）托架和兩個(gè)主機(jī) CPU 托架。每個(gè) JBOG 托架包含兩個(gè) Trainium3 加速器，因此每個(gè)機(jī)架總共有 32 個(gè) Trainium3 芯片。一個(gè)完整的 Trainium3 NL32x2 Switched 擴(kuò)展世界由兩個(gè)機(jī)架組成，每個(gè)機(jī)架包含 32 個(gè) Trainium3 芯片，總共可構(gòu)建一個(gè)包含 64 個(gè) Trainium3 芯片的世界。

Trainium NL32x2 3D Torus 和 Trainium NL32x2 Switched 的主要區(qū)別在于，Trainium NL32x2 Switched 在機(jī)架中間增加了可擴(kuò)展的 NeuronLink 交換機(jī)托架，從而實(shí)現(xiàn)了全對(duì)全交換網(wǎng)絡(luò)。NeuronLink 交換機(jī)托架之所以放置在機(jī)架中間，原因與 Nvidia Oberon 機(jī)架中 NVLink 交換機(jī)托架的放置位置相同，都是為了最大限度地縮短加速器和可擴(kuò)展交換機(jī)之間最長(zhǎng) SerDes 傳輸距離。此外，Trainium NL32x2 Switched 還將 CPU 托架、電源架、電池備份單元 (BBU) 托架和機(jī)架頂部 (ToR) 交換機(jī)的位置也進(jìn)行了調(diào)整，從始終位于各自 8 個(gè) JBOG 托架組的頂部，改為位于機(jī)架的頂部和底部，從而縮短了 16 個(gè) JBOG 托架和 4 個(gè) NeuronLink 交換機(jī)托架之間的距離。此外，還將推出配備五個(gè) NeuronLink 交換機(jī)托架的設(shè)計(jì)，支持熱插拔交換機(jī)托架，無(wú)需停機(jī)。這與英偉達(dá)的 GB200/300 NVL72/VR NVL144 形成鮮明對(duì)比，后者的操作員必須先清空機(jī)架上的所有工作負(fù)載才能更換交換機(jī)托架。亞馬遜的理念始終以現(xiàn)場(chǎng)可維護(hù)性和可靠性為核心，因?yàn)樗麄冊(cè)诓渴鸷凸芾碇邪缪葜陵P(guān)重要的角色。而英偉達(dá)則為了追求性能而忽略其他方面，因?yàn)檫@正是他們所追求的。

其他主要區(qū)別在于NeuronLink 拓?fù)浣Y(jié)構(gòu)和相關(guān)的連接組件。Trainium NL32x2 交換式架構(gòu)將通過(guò)跨機(jī)架 AEC 連接兩個(gè)機(jī)架，機(jī)架 A 中的一顆芯片可直接連接到機(jī)架 B 中的另一顆芯片。我們將在本文后面更詳細(xì)地討論這種擴(kuò)展網(wǎng)絡(luò)。

JBOG/計(jì)算托盤

Trainium3 NL32x2 Switched 的 JBOG 托盤與 Trainium NL32x2 3D Torus 的非常相似。每個(gè) JBOG 將包含兩顆 Trainium3 芯片。Trainium3 NL32x2 Switched 采用基于 PCIe 6.0 的連接方式，相比使用 PCIe 5.0 連接的 Trainium2 NL16 2D Torus 和 Trainium2 NL32x2 3D Torus JBOG 而言，這是一項(xiàng)升級(jí)。因此，PCB 材料也需要從 M8 級(jí)覆銅板 (CCL)（具體而言，是低 DK2 玻璃纖維布 + HVLP2 銅箔）升級(jí)到 M8.5 級(jí)覆銅板 (CCL)（低 DK2 玻璃纖維布 + HVLP4 銅箔）。

迄今為止，所有Trainium 機(jī)架均采用無(wú)電纜設(shè)計(jì)理念以提高組裝效率，因此所有信號(hào)均通過(guò) PCB 走線傳輸。信號(hào)在 PCB 上傳輸時(shí)的損耗遠(yuǎn)高于通過(guò)架空電纜傳輸時(shí)的損耗，因此必須在 JBOG 板中間放置四個(gè) PCIe 6.0 x16 重定時(shí)器，以補(bǔ)償在前置 I/O 端口和兩個(gè) Trainium3 封裝之間通過(guò) PCB 傳輸?shù)男盘?hào)。

Trainium3 的網(wǎng)卡也位于 JBOG 托盤內(nèi)。對(duì)于 AWS 的后端網(wǎng)絡(luò) EFAv4，Trainium3 NL32x2 交換式網(wǎng)卡有兩種配置可供選擇：

• 方案一：每個(gè)裝有兩塊Trainium3 芯片的 JBOG 托架配備一個(gè) Nitro-v6 (2*200G) 400Gbps 網(wǎng)卡模塊：每塊 Trainium3 芯片提供 200Gbps 的 EFA 帶寬。
• 方案二：每個(gè)JBOG 托架配備兩個(gè) Nitro-v6 (2*200G) 400Gbps Nitro NIC 模塊，每個(gè) Trainium3 芯片提供 400Gbps 的 EFA 帶寬。

絕大多數(shù)基于Trainium3 的服務(wù)器將采用方方案一。每?jī)蓧K Trainium3 芯片配備一塊 Nitro-v6 400G 網(wǎng)卡，這意味著每塊 Trainium3 芯片擁有 200Gbps 的橫向擴(kuò)展帶寬。AWS 認(rèn)為，即使對(duì)于目前規(guī)模最大的生產(chǎn)級(jí)推理模型，每塊 GPU 200Gbps 的帶寬也足以滿足預(yù)填充實(shí)例和解碼實(shí)例之間鍵值緩存?zhèn)鬏數(shù)闹丿B需求。對(duì)于訓(xùn)練模型，AWS 的理念是，像 Anthropic 這樣擁有精英程序員的小型公司會(huì)使用流水線并行 (PP) 來(lái)降低網(wǎng)絡(luò)流量，而不僅僅依賴于 FSDP/TP/Ctx 并行/DP。然而，需要注意的是，流水線并行對(duì)于大規(guī)模訓(xùn)練來(lái)說(shuō)是絕對(duì)必要的，但維護(hù)和調(diào)試使用 PP 的代碼庫(kù)卻非常麻煩。

對(duì)于AWS 的前端網(wǎng)絡(luò) ENA，CPU 托架內(nèi)配備了一個(gè)專用的 Nitro-v6 (2*200) 400Gbps 網(wǎng)卡模塊。為了將每個(gè) JBOG 托架連接到 CPU 托架，服務(wù)器正面鋪設(shè)了一條專用的 PCIe 6.0 x16 DAC 線纜（128Gbyte/s 單向）。Trainium2 NL16 2D Torus 也以相同的方式連接 CPU 托架和 JBOG 托架。

Trainium3 NL32x2 Switched 的 CPU 托架布局與 Trainium2 NL16 2D Torus 的布局非常相似，在Trainium2 文章中有詳細(xì)說(shuō)明。

Trainium3 NL32x2 Switched 是 Trainium3 最新推出的產(chǎn)品，支持可切換式擴(kuò)展架構(gòu)。由于采用風(fēng)冷機(jī)架設(shè)計(jì)，其機(jī)架功率密度較低。其功率密度與 Trainium3 NL32x2 3D Torus 基本相同，主要區(qū)別在于增加了擴(kuò)展交換機(jī)托架。此外，Trainium3 NL32x2 Switched 也是唯一一款支持可切換式擴(kuò)展架構(gòu)且可在非液冷數(shù)據(jù)中心部署的產(chǎn)品。

由于液冷數(shù)據(jù)中心的部署目前是部署的關(guān)鍵瓶頸，因此采用風(fēng)冷散熱的Trainium3芯片相比其他競(jìng)爭(zhēng)對(duì)手的液冷式可擴(kuò)展交換加速器，具有上市時(shí)間優(yōu)勢(shì)。當(dāng)運(yùn)營(yíng)商試圖將液冷機(jī)架強(qiáng)行部署到風(fēng)冷數(shù)據(jù)中心時(shí)，他們必須使用效率低下的液冷-風(fēng)冷側(cè)掛式冷卻器。因此，我們預(yù)計(jì)到2026年部署的Trainium3芯片中，大部分將是Trainium3 NL32x2交換式SKU。

Trainium3 NL72x2 開關(guān)式（Teton3 Max）

Trainium3 NL32x2 Switched 和 Trainium3 NL72x2 Switched 都采用了全交換架構(gòu)，但 Trainium3 NL72x2 Switched 的機(jī)架式架構(gòu)與 Nvidia 的 GB200 NVL72 Oberon 架構(gòu)最為相似。除了 Oberon 和 Trainium3 NL72x2 Switched 都采用液冷散熱外，Trainium3 NL72x2 Switched 還將 CPU 集成到計(jì)算托架中，就像 Nvidia 將 Grace 和 Vera 與 GPU 集成在同一個(gè)計(jì)算托架上一樣。相比之下，Trainium NL32x2 Switched 仍然使用分離式 CPU 節(jié)點(diǎn)。與 Oberon 一樣，Trainium NL72x2 Switched 也使用冷板對(duì) Trainium3 加速器和 Graviton 4 CPU 進(jìn)行液冷散熱。Trainium NL72x2 Switched 與 Oberon 架構(gòu)的主要區(qū)別在于它采用了跨機(jī)架連接，從而擴(kuò)展了規(guī)模，使其能夠跨越兩個(gè)機(jī)架。

機(jī)架架構(gòu)

Trainium3 NL72x2 Switched 使用兩個(gè)機(jī)架來(lái)實(shí)現(xiàn) 144 個(gè) XPU 的全球規(guī)模，每個(gè)機(jī)架包含 18 個(gè)計(jì)算托架和 10 個(gè)位于中間的 NeuronLink 交換機(jī)托架。每個(gè)計(jì)算托架包含四個(gè) Trainium3 和一個(gè) Graviton4 CPU，因此 Trainium3 NL72x2 Switched 的全球規(guī)模由兩個(gè)機(jī)架組成，共包含 144 個(gè) Trainium3 和 36 個(gè) Graviton4。與 Trainium3 NL32x2 Switched 一樣，Trainium3 NL72x2 Switched 也使用母線供電。背板采用了 TE 和 Amphenol 的混合連接器，因此我們?cè)诰W(wǎng)絡(luò)模型中采用了不同的連接器。

計(jì)算機(jī)托盤

Trainium3 NL72x2 Switched 的高計(jì)算密度和功率密度始于其計(jì)算托架，每個(gè)計(jì)算托架包含四個(gè) Trainium3 芯片。Trainium3 NL72x2 Switched 的連接主要基于 PCIe 6.0，因此其 PCB 材料與上述 Trainium3 NL32x2 Switched 相同。為了增加各個(gè) Trainium3 芯片與前置 I/O 端口之間的信號(hào)傳輸距離，使用了六個(gè) PCIe 6.0 x16 重定時(shí)器。需要注意的是，由于其采用無(wú)電纜設(shè)計(jì)以提高生產(chǎn)速度，因此其上市設(shè)計(jì)中包含了一些低成本的重定時(shí)器，以降低設(shè)計(jì)風(fēng)險(xiǎn)。在成功完成初始生產(chǎn)部署后，AWS 可以考慮優(yōu)化設(shè)計(jì)，并有可能移除部分重定時(shí)器。

如上所述，Trainium3 NL32x2 Switched 和 Trainium3 NL72x2 Switched 的主要區(qū)別在于散熱方式：Trainium3 NL72x2 Switched 采用液冷散熱，而 Trainium3 NL32x2 Switched 采用風(fēng)冷散熱。液冷散熱用于冷卻 Trainium3 模塊、NeuronLinkv4 x32 通道 PCIe 6.0 交換機(jī)和 Graviton4 CPU。計(jì)算托架中的其他組件，包括 PCIe 6.0 x16 重定時(shí)器、Nitro-v6 網(wǎng)卡、PCIe 6.0 x16 AEC 硬盤籠、DIMM 內(nèi)存條和 2 塊 8TB 本地 NVMe 硬盤，則采用風(fēng)扇進(jìn)行風(fēng)冷散熱。

就主機(jī)CPU而言，Graviton4將是Trainium3 NL72x2 Switched目前唯一可用的CPU選項(xiàng)。在Trainium3的生命周期內(nèi)，用戶可以將CPU升級(jí)到下一代Graviton處理器。理論上，x86 CPU也受支持，因?yàn)樗鼈円部梢酝ㄟ^(guò)PCIe與其他組件連接，但我們認(rèn)為廠商不會(huì)推出x86版本的Trainium3 NL72x2 Switched，而只會(huì)提供x86版本的Trainium NL32x2 Switched。由于Trainium3使用PCIe 6.0，而Graviton4使用PCIe 5.0，因此需要在CPU旁邊放置兩個(gè)PCIe轉(zhuǎn)換器，以實(shí)現(xiàn)CPU和GPU之間的PCIe 6.0到PCIe 5.0的轉(zhuǎn)換。CPU內(nèi)存方面，CPU旁邊設(shè)有12個(gè)DDR5 DIMM插槽，主流版本將使用64GB和128GB容量的DDR5 DIMM內(nèi)存條。每個(gè)計(jì)算托架將使用兩塊 8TB 本地 NVMe 硬盤作為本地存儲(chǔ)。

Trainium3 NL72x2 交換式橫向擴(kuò)展網(wǎng)絡(luò)

Trainium3 NL72x2 Switched 將具有與 Trainium3 NL32x2 Switched 相同的橫向擴(kuò)展網(wǎng)絡(luò)配置，即每個(gè) Trainium3 芯片可以選擇 400G 或 200G 的橫向擴(kuò)展帶寬：

• 方案一：每個(gè)裝有四顆Trainium3 芯片的 JBOG 托架配備兩個(gè) Nitro-V6 (2*200G) 400Gbps 網(wǎng)卡模塊：每顆 Trainium3 芯片提供 200Gbps 的 EFA 帶寬。
• 方案二：每個(gè)JBOG 托盤包含四個(gè) Trainium3 芯片，每個(gè)芯片配備四個(gè) Nitro-V6 (2*200G) 400Gbps Nitro NIC 模塊：每個(gè) Trainium3 芯片提供 400Gbps 的 EFA 帶寬。

與Trainium3 NL32x2 交換式交換機(jī)一樣，絕大多數(shù) Trainium3 NL72x2 產(chǎn)品將采用方案一。每?jī)蓚€(gè) Trainium3 芯片配備一個(gè) Nitro-v6 400G 網(wǎng)卡，每個(gè) Trainium3 芯片的橫向擴(kuò)展帶寬為 200Gbps。

不過(guò)，Trainium3 NL72x2 交換機(jī)版與普通版的區(qū)別在于，主機(jī) CPU 位于計(jì)算托架上，而專用于 CPU 的 Nitro-V6 (2*200) 400Gbps 網(wǎng)卡模塊也位于該托架內(nèi)。此外，借助 PCIe 交換機(jī)，CPU 還可以使用 Trainium3 專用網(wǎng)卡與外部網(wǎng)絡(luò)通信。

Trainium3 NL72x2 Switched 是 AWS 針對(duì) Nvidia Oberon 機(jī)架式架構(gòu)推出的解決方案。Trainium3 NL72x2 Switched 架構(gòu)的功率密度遠(yuǎn)高于其前代產(chǎn)品。由于其高功率密度以及對(duì)液冷數(shù)據(jù)中心的需求，我們預(yù)計(jì) Trainium3 的部分產(chǎn)品將采用 Trainium3 NL72x2 Switched SKU，但大部分產(chǎn)品仍將采用 Trainium3 NL32x2 Switched。了解了機(jī)架布局和計(jì)算托架布局/拓?fù)浣Y(jié)構(gòu)之后，現(xiàn)在是時(shí)候深入了解 Trainium3 NL32x2 Switched 和 Trainium3 NL72x2 Switched 的真正核心優(yōu)勢(shì)——交換式縱向擴(kuò)展網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)。

擴(kuò)展網(wǎng)絡(luò)架構(gòu)

2D/3D 環(huán)形網(wǎng)絡(luò)與交換式擴(kuò)展網(wǎng)絡(luò)

在詳細(xì)介紹新型交換式橫向擴(kuò)展網(wǎng)絡(luò)的具體拓?fù)浣Y(jié)構(gòu)之前，我們先來(lái)解釋一下AWS 為何選擇從 2D/3D 環(huán)面架構(gòu)轉(zhuǎn)向交換式架構(gòu)。Trainium2 NL26 2D 環(huán)面服務(wù)器和 Trainium2 NL32x2 3D 環(huán)面服務(wù)器的 NeuronLinkv3 橫向擴(kuò)展拓?fù)浣Y(jié)構(gòu)分別是 2D 網(wǎng)狀環(huán)面和 3D 網(wǎng)狀環(huán)面。然而，對(duì)于需要全連接集合的MoE模型而言，保持環(huán)面拓?fù)浣Y(jié)構(gòu)并非最優(yōu)。相比之下，密集模型并不大量使用全連接集合，這意味著交換式架構(gòu)在與密集模型配合使用時(shí)性能優(yōu)勢(shì)并不明顯，但總體擁有成本 (TCO) 卻更高。

采用3D Torus 架構(gòu)時(shí)，由于擴(kuò)展域內(nèi)芯片間的過(guò)載，當(dāng)消息大小從 16KB 增長(zhǎng)到 1MB（即批處理大小增加）時(shí)，擴(kuò)展網(wǎng)絡(luò)會(huì)因過(guò)載而突然面臨帶寬瓶頸。相比之下，Trainium3 獨(dú)特的交換式拓?fù)浣Y(jié)構(gòu)，即使其第一代交換式網(wǎng)絡(luò)并非扁平的單層交換式拓?fù)?，也不?huì)出現(xiàn)過(guò)載現(xiàn)象。

對(duì)于預(yù)填充而言，更大的Trainium3 NL72x2 Switched擴(kuò)展拓?fù)浣Y(jié)構(gòu)并不能帶來(lái)顯著優(yōu)勢(shì)，因?yàn)轭A(yù)填充通常受限于計(jì)算能力，而更大的拓?fù)浣Y(jié)構(gòu)主要適用于解碼階段的大規(guī)模專家并行計(jì)算。對(duì)于參數(shù)總數(shù)為2-3的前沿MoE模型，Trainium3 NL32x2 Switched的擴(kuò)展網(wǎng)絡(luò)足以滿足解碼階段的需求；但對(duì)于參數(shù)總數(shù)超過(guò)4萬(wàn)億的前沿MoE模型，將這些模型部署在更大的Trainium3 NL72x2 Switched擴(kuò)展網(wǎng)絡(luò)上則能帶來(lái)顯著的性能提升。

Trainium3 NL32x2 交換式交換機(jī)和 Trainium3 NL72x2 交換式交換機(jī)旨在最終提供全對(duì)全的縱向擴(kuò)展交換解決方案，但 AWS 秉持著以低 TCO 實(shí)現(xiàn)快速上市的理念，決定構(gòu)建一種與當(dāng)時(shí)可用的不同代縱向擴(kuò)展交換機(jī)兼容的網(wǎng)絡(luò)架構(gòu)。

在深入探討不同的機(jī)架架構(gòu)和交換機(jī)世代之前，我們將首先分解Trainium3 的 NeuronLinkv4/芯片間互連 (ICI) 帶寬的組成部分。

用于Trainium3 的 NeuronLinkv4 服務(wù)器內(nèi)擴(kuò)展網(wǎng)絡(luò)使用三種不同的連接介質(zhì)連接 XPU：通過(guò) PCB、背板和跨機(jī)架連接。我們將詳細(xì)介紹擴(kuò)展網(wǎng)絡(luò)連接，但為了便于參考，我們?cè)诒竟?jié)末尾提供了一個(gè)簡(jiǎn)要表格進(jìn)行總結(jié)。

Trainium3 在所有三種介質(zhì)上共提供 160 條 PCIe NeuronLinkv4 連接通道，其中 144 條為活動(dòng)通道，16 條為背板上的冗余通道。每臺(tái) Trainium3 的 160 條通道分布如下：背板：共有 80 條通道，其中 64 條處于活動(dòng)狀態(tài)，16 條為冗余通道。每條通道通過(guò)一個(gè) Strada Whisper 背板連接器連接到背板，每個(gè) Trainium3 芯片支持 160 對(duì)差分信號(hào)（DP，即 80 個(gè) Tx DP 和 80 個(gè) Rx DP）。AWS 利用冗余通道來(lái)應(yīng)對(duì)背板電纜故障、交換機(jī)托盤級(jí)故障和端口托盤級(jí)故障，從而實(shí)現(xiàn)容錯(cuò)。

他們沒(méi)有將這16 條額外的通道用作額外的帶寬，主要原因如下：

對(duì)于解碼等延遲密集型工作負(fù)載，增加通道數(shù)并不會(huì)帶來(lái)任何好處。這就像使用更大的管道并不會(huì)影響一滴水從管道一端流到另一端的速度一樣。

即使啟用全部80 條通道，像訓(xùn)練這樣通信密集型的工作負(fù)載也無(wú)法獲得任何顯著的性能提升，這是因?yàn)榇嬖凇暗絷?duì)效應(yīng)”。在任何大型訓(xùn)練運(yùn)行中，至少會(huì)有幾個(gè)機(jī)架的某些通道出現(xiàn)故障。如果只有一個(gè) Trainium3 機(jī)架出現(xiàn)故障通道，那么整個(gè)訓(xùn)練作業(yè)將只能有效使用 80 條通道中的 64 條，因?yàn)樗衅渌麢C(jī)架都會(huì)等待速度最慢的機(jī)架。

PCB：通過(guò) NeuronLinkv4 PCB 走線，為相鄰的 Trainium3 芯片提供 64 條通道。對(duì)于 Trainium3 NL32x2 Switched 芯片，PCB 直接連接到相鄰的 Trainium3 芯片；而對(duì)于 Trainium3 NL72x2 Switched 芯片，PCB 連接則通過(guò)八個(gè) PCIe 6.0 32 通道交換機(jī)（或四個(gè) 64 通道交換機(jī)或兩個(gè) 128 通道交換機(jī)）實(shí)現(xiàn)。這種設(shè)計(jì)的優(yōu)勢(shì)在于，AWS 可以在制造時(shí)選擇每通道成本最低的方案。由于 PCB 芯片間互連 (ICI) 故障率遠(yuǎn)低于背板，因此無(wú)需在 PCB 上設(shè)置冗余通道。

跨機(jī)架：16 條通道從每個(gè) Trainium3 通過(guò) PCB 連接到 OSFP-XD 籠，然后通過(guò) PCIe 有源電纜 (AEC) 連接到相鄰的機(jī)架。

NeuronLink PCIe PHY 和 UALink 交換機(jī)

AWS 的NeuronLink目標(biāo)是在盡可能實(shí)現(xiàn)供應(yīng)鏈多樣性的前提下，以最快的速度和最低的總體擁有成本將產(chǎn)品推向市場(chǎng)。其設(shè)計(jì)橫向擴(kuò)展網(wǎng)絡(luò)架構(gòu)的方法也遵循同樣的原則。在 Trainium3 的生命周期內(nèi)，共推出了三代橫向擴(kuò)展交換機(jī)：首先是 160 通道的 Scorpio X PCIe 6.0 交換機(jī)，然后是 320 通道的 Scorpio-X PCIe 6.0 交換機(jī)，最后還可以選擇升級(jí)到更高基數(shù)的 72+ 端口 UALink 交換機(jī)。160 通道的 Scorpio-X 交換機(jī)可以實(shí)現(xiàn)快速上市，但缺點(diǎn)是它強(qiáng)制采用一種非理想的橫向擴(kuò)展網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)，并非全連接，并且在同一 Trainium3 NL72x2 交換機(jī)架內(nèi)的兩個(gè) Trainium3 芯片之間最多需要三個(gè)跳才能連接。然而，使用 320 通道的 Scorpio-X 交換機(jī)或 UALink 交換機(jī)可以改善這種情況。

前兩代交換機(jī)托架采用多平面擴(kuò)展交換架構(gòu)，理論上無(wú)法在不經(jīng)過(guò)多次跳轉(zhuǎn)的情況下實(shí)現(xiàn)機(jī)架內(nèi)所有設(shè)備間的完全通信。第一代交換機(jī)很快就會(huì)被更高帶寬和更高基數(shù)的交換機(jī)所取代。下表展示了兩種機(jī)架SKU 和三種不同交換機(jī)托架代數(shù)下六種不同組合的擴(kuò)展特性。

交換機(jī)托架

Trainium3 NL32x2 交換型交換機(jī)將首先在第一代交換機(jī)托架中使用 160 通道 PCIe 交換機(jī)來(lái)實(shí)現(xiàn)擴(kuò)展拓?fù)浣Y(jié)構(gòu)。每個(gè)機(jī)架包含兩個(gè)交換平面，每個(gè)平面由 8 臺(tái)交換機(jī)組成。由于每臺(tái) PCIe 交換機(jī)的端口數(shù)量有限，每臺(tái) Trainium3 只能連接到同一平面上的 8 臺(tái) PCIe 交換機(jī)，而不能連接到所有其他 PCIe 交換機(jī)。因此，并非所有 Trainium3 都能直接與其他 Trainium3 通信，而是需要經(jīng)過(guò)多個(gè)躍點(diǎn)。

例如，節(jié)點(diǎn)1 中的 Trainium3 A 可以通過(guò)一個(gè)交換機(jī)跳與所有其他節(jié)點(diǎn)中的 Trainium3 A 通信。節(jié)點(diǎn) 1 中的 Trainium3 B 也同樣可以通過(guò)一個(gè)交換機(jī)跳與所有其他節(jié)點(diǎn)中的 Trainium3 B 通信。

然而，考慮以下三種情況：位于不同交換平面或機(jī)架上的Trainium 芯片進(jìn)行通信。在每種情況下，數(shù)據(jù)都可以通過(guò)多條路徑在兩個(gè) Trainium3 芯片之間傳輸，以下列出一些可能的路徑：

• 路徑A：從機(jī)架 A 的節(jié)點(diǎn) 1 中的 Trainium3 A 到機(jī)架 A 的節(jié)點(diǎn) 16 中的 Trainium3 A，總共 1 跳。
• 路徑B：從機(jī)架 A 節(jié)點(diǎn) 1 中的 Trainium3 A 到機(jī)架 A 節(jié)點(diǎn) 2 中的 Trainium3 B，總共需要 2 跳。
• 路徑C：從機(jī)架 A 節(jié)點(diǎn) 1 中的 Trainium3 A 到機(jī)架 B 節(jié)點(diǎn) 2 中的 Trainium3 B，總共需要 3 跳。

由于Trainium3 芯片具備自動(dòng)轉(zhuǎn)發(fā)功能，且集合延遲基于 SBUF 到 SBUF 的數(shù)據(jù)傳輸，AWS 聲稱這種多跳傳輸不會(huì)造成延遲問(wèn)題。我們認(rèn)為，與 Nvidia GPU 相比，這種延遲可以忽略不計(jì)，因?yàn)?Nvidia GPU 的集合傳輸需要從 HBM 開始或結(jié)束。但最終，性能機(jī)器學(xué)習(xí)工程師需要根據(jù)縱向擴(kuò)展拓?fù)浣Y(jié)構(gòu)，協(xié)同優(yōu)化模型并行性，以盡可能減少通信跳數(shù)，同時(shí)還要考慮到跨機(jī)架連接的帶寬不到機(jī)架內(nèi)背板帶寬的 10%。

最直接的方法是在機(jī)架內(nèi)部署專家并行（EP），然后通過(guò)AEC將機(jī)架上的Trainium3實(shí)例“配對(duì)”，并在配對(duì)的機(jī)架之間使用張量并行（TP）。另一種顯而易見的并行策略是，在機(jī)架內(nèi)使用專家并行，并在機(jī)架對(duì)之間使用上下文并行。

另一種潛在的并行策略是在兩個(gè)機(jī)架上都使用完整的EP（擴(kuò)展并行），但要考慮到額外的躍點(diǎn)。這種策略對(duì)于非常稀疏的模型可能非常有效，因?yàn)檫@類模型由于d_model維度太小而無(wú)法跨機(jī)架實(shí)現(xiàn)TP（并行并行）。因此，即使通過(guò)直接連接的Trainium3進(jìn)行額外躍點(diǎn)帶來(lái)的延遲也是值得的。

接下來(lái)是Trainium3 NL72x2 交換機(jī) SKU，其擴(kuò)展拓?fù)浣Y(jié)構(gòu)更為復(fù)雜。它包含四個(gè)平面，每個(gè)平面配備 10 個(gè) 160 通道 PCIe 交換機(jī)，每個(gè)機(jī)架的交換機(jī)托架上共有 40 個(gè)交換機(jī)；此外，18 個(gè)計(jì)算托架上每個(gè)托架還配備 8 個(gè) 32 通道 PCIe 交換機(jī)，這意味著每個(gè)機(jī)架的計(jì)算托架上總共有 144 個(gè)較小的 PCIe 交換機(jī)。因此，每個(gè)機(jī)架總共有 184 個(gè)擴(kuò)展交換機(jī)，或者說(shuō)，144 臺(tái) Trainium3 擴(kuò)展機(jī)架共計(jì) 368 個(gè)擴(kuò)展交換機(jī)。為了便于跟蹤交換機(jī)配置，我們?cè)诖酥噩F(xiàn)之前的匯總表：

與同樣在交換機(jī)托架中采用160 通道 PCIe 交換機(jī)的 Trainium3 NL32x2 交換式設(shè)計(jì)類似，這種擴(kuò)展設(shè)計(jì)也受到每個(gè)交換機(jī)托架上 PCIe 交換機(jī)僅有 20 個(gè)端口的限制，這意味著每個(gè)交換機(jī)每個(gè)節(jié)點(diǎn)只能連接到四個(gè) Trainium3 芯片中的一個(gè)（這 20 個(gè)端口中有兩個(gè)未使用或用于管理）。在同一交換平面內(nèi)，每個(gè) Trainium3 芯片之間僅相隔一個(gè)交換機(jī)跳數(shù)。

與每個(gè)JBOG 僅包含兩個(gè) Trainium3 的 Trainium3 NL32x2 交換式設(shè)計(jì)不同，Trainium3 NL72x2 交換式設(shè)計(jì)在同一計(jì)算托架板上集成了四個(gè) Trainium3。同一板上不同交換平面上的每個(gè) Trainium3 芯片通過(guò) 8 個(gè) 32 通道 Scorpio-P PCIe 交換機(jī)進(jìn)行通信，這意味著位于不同交換平面上的 Trainium3 芯片之間的芯片間通信需要額外的交換機(jī)跳轉(zhuǎn)。

當(dāng)Trainium3 設(shè)備不在同一交換平面上時(shí)，交換跳數(shù)大于 1。請(qǐng)考慮以下三種不同的情況：

• 路徑A：從機(jī)架 A 的節(jié)點(diǎn) 1 中的 Trainium3 A 到機(jī)架 A 的節(jié)點(diǎn) 2 中的 Trainium3 A，總共需要 1 跳。
• 路徑B：從機(jī)架 A 節(jié)點(diǎn) 1 中的 Trainium3 A 到機(jī)架 A 節(jié)點(diǎn) 2 中的 Trainium3 C，總共需要 3 跳。
• 路徑C：從機(jī)架 A 節(jié)點(diǎn) 1 中的 Trainium3 A 到機(jī)架 B 節(jié)點(diǎn) 2 中的 Trainium3 C，總共需要 4 跳。

Trainium3 NL72x2 已開啟 Gen1 備用交換機(jī)托架

AWS 也針對(duì)各種場(chǎng)景進(jìn)行了配置。如果 Scorpio X 160 通道 PCIe 交換機(jī)不可用，則可以使用 Broadcom PEX90144 交換機(jī)作為備用方案，該交換機(jī)具有 144 條通道和 72 個(gè)最大可用端口。然而，這種具有更高基數(shù)（每個(gè)端口 2 條通道，最大端口數(shù)為 72）的替代方案并不意味著可擴(kuò)展交換平面的數(shù)量會(huì)減少。

由于潛在的高串行延遲，從Trainium3 延伸出的 ICI 通道可能不適合每個(gè)端口拆分成兩條通道。這意味著，對(duì)于 144 通道的 PEX90144 擴(kuò)展交換機(jī)，AWS 將使用每個(gè)交換機(jī) 36 個(gè)端口（每個(gè)端口 8 條通道）或 18 個(gè)端口（每個(gè)端口 4 條通道）作為備份。下圖展示了一種這樣的擴(kuò)展拓?fù)浣Y(jié)構(gòu)，其中每個(gè) PCIe 交換機(jī)有 18 個(gè)端口，每個(gè)端口使用 8 條通道的 PEX90144。

第二代交換機(jī)托架

Trainium3 NL32x2 交換機(jī)版 SKU 也兼容 320 通道 PCIe 擴(kuò)展交換機(jī)，一旦 320 通道 PCIe 交換機(jī)上市，即可替換現(xiàn)有的 160 通道 PCIe 交換機(jī)。由于 320 通道 PCIe 交換機(jī)的端口數(shù)量是 160 通道交換機(jī)的兩倍，因此每個(gè)機(jī)架只需 8 臺(tái)交換機(jī)即可構(gòu)建擴(kuò)展網(wǎng)絡(luò)，機(jī)架內(nèi)的每個(gè) Trainium 芯片之間也僅需經(jīng)過(guò)一個(gè)交換機(jī)即可連接。由于擴(kuò)展拓?fù)浣Y(jié)構(gòu)已實(shí)現(xiàn)全連接，PCB 上相鄰 Trainium3 芯片之間的直接連接成為一項(xiàng)額外優(yōu)勢(shì)。

在此設(shè)計(jì)中，機(jī)架中任意Trainium3 芯片之間的最大跳距僅為一跳——與 Trainium3 NL32x2 交換式 SKU 的 Gen1 設(shè)計(jì)最壞情況下的兩跳距離相比，這是一個(gè)很大的改進(jìn)，當(dāng) SKU 升級(jí)到使用 Gen2 交換機(jī)托架時(shí)，可提供延遲優(yōu)勢(shì)。

對(duì)于Trainium3 NL72x2 交換式芯片，從 160 通道 Scorpio X PCIe 交換機(jī)升級(jí)到 320 通道 Scorpio X PCIe 交換機(jī)意味著交換平面數(shù)量從四個(gè)減少到兩個(gè)。JBOG 托架上相鄰的 Trainium3 芯片仍然需要通過(guò) Scorpio P 交換機(jī)進(jìn)行通信。

最終目標(biāo)：第三代交換機(jī)托架

UALink 上線后，即可在數(shù)據(jù)中心內(nèi)安裝 72+ 端口的 Scorpio X UALink 交換機(jī)，替換原有的 320 通道 Scorpio X 交換機(jī)（每臺(tái)交換機(jī) 40 個(gè)端口）。UALink 交換機(jī)的延遲低于基于 PCI 的交換機(jī)，并且將使用 UALink 協(xié)議。

下圖展示了安裝了基于Gen3 UALink 的交換機(jī)托架的 Trainium3 NL32x2 交換式 SKU 的拓?fù)浣Y(jié)構(gòu)。借助 UALink 交換機(jī)，Trainium3 NL32x2 交換式 SKU 將繼續(xù)支持全連接，就像 320 通道 Scorpio X PCIe 交換機(jī)一樣。每個(gè)邏輯端口的具體端口數(shù)量和通道數(shù)尚未確定，但每個(gè)機(jī)架的整體擴(kuò)展帶寬將保持不變。

對(duì)于Trainium3 NL72x2 交換式架構(gòu)，每個(gè)機(jī)架的擴(kuò)展拓?fù)浣Y(jié)構(gòu)將變?yōu)槿B接，因?yàn)槊總€(gè) UALink 交換機(jī)都可以連接到機(jī)架內(nèi)的每個(gè) Trainium3 芯片。通過(guò) 8 臺(tái) 32 通道 Scorpio P 交換機(jī)實(shí)現(xiàn)的本地計(jì)算托架連接現(xiàn)在體現(xiàn)了剩余帶寬。

所有這些加起來(lái)就意味著需要大量的PCIe 交換機(jī)，即使只考慮一代產(chǎn)品也是如此！

幸運(yùn)的是，亞馬遜與Astera Labs 建立了戰(zhàn)略合作伙伴關(guān)系。事實(shí)證明亞馬遜購(gòu)買得越多，節(jié)省得就越多！

如果AWS 達(dá)到其對(duì) ALAB PCIe 交換機(jī)和重定時(shí)器的采購(gòu)量承諾，它將獲得與這些產(chǎn)品采購(gòu)掛鉤的 ALAB 股票認(rèn)股權(quán)證。這些股票認(rèn)股權(quán)證會(huì)在 AWS 達(dá)到采購(gòu)里程碑時(shí)逐步生效，由于行權(quán)價(jià)格僅為 20.34 美元，任何高于該價(jià)格的市場(chǎng)價(jià)格都會(huì)立即為 AWS 帶來(lái)價(jià)值。這種結(jié)構(gòu)實(shí)際上為 AWS 提供了基于股權(quán)的組件采購(gòu)“返利”。在以下示例中，截至 9 月 25 日生效的股票認(rèn)股權(quán)證相當(dāng)于約 23% 的有效折扣。

由于不同SKU 的擴(kuò)展拓?fù)浣Y(jié)構(gòu)（交換機(jī)或網(wǎng)狀網(wǎng)絡(luò)）和 NeuronLink 通道數(shù)量不同，Trainium 各代產(chǎn)品的銅纜數(shù)量也各不相同。Trainium2 NL16 2D Torus 采用單背板和相對(duì)較少的 AEC 鏈路，而 Trainium NL32x2 3D Torus 則增加了通道數(shù)量，需要四個(gè) NeuronLink 背板以及約 6,100 根銅纜來(lái)支持更密集的 3D Torus 拓?fù)浣Y(jié)構(gòu)。Trainium3 NL32x2 Switched 保持了類似的背板數(shù)量和約 5,100 根銅纜，而 Trainium3 NL72x2 Switched 則進(jìn)一步擴(kuò)展了擴(kuò)展范圍，每個(gè)服務(wù)器組可容納 144 個(gè)芯片（Trainium3 NL32x2 為 64 個(gè)），銅纜數(shù)量也因此增加到 11,520 根。

Trainium3 機(jī)架功率預(yù)算和物料清單

我們已針對(duì)不同Trainium 系統(tǒng)中的主要組件組，編制了一份詳細(xì)的零件清單和主要功耗預(yù)算。訂閱我們AI 總擁有成本 (TCO) 模型和 AI 硬件物料清單 (BoM) 模型的用戶，可以查看有關(guān)數(shù)量、平均售價(jià) (ASP) 和系統(tǒng)總成本的詳細(xì)信息。

Trainium3 NL72x2 開關(guān)系統(tǒng)的總系統(tǒng)功耗自然更高，因?yàn)樗趦蓚€(gè)機(jī)架上部署了 144 個(gè)芯片，而 Trainium3 NL32x2 開關(guān)系統(tǒng)在兩個(gè)機(jī)架上僅部署了 64 個(gè)芯片。然而，一旦將功耗標(biāo)準(zhǔn)化為每個(gè)芯片的功耗，64 個(gè)芯片的 Trainium3 NL32x2 開關(guān)系統(tǒng)和 144 個(gè)芯片的 Trainium3 NL72x2 開關(guān)系統(tǒng)實(shí)際上每個(gè)芯片的功耗非常接近，因?yàn)?Trainium3 芯片的熱設(shè)計(jì)功耗 (TDP) 是影響整體功耗預(yù)算的最大因素。 Trainium3 NL72x2 Switched 的機(jī)架功率密度自然更高，因?yàn)樗谝粋€(gè)機(jī)架中可以容納 64 個(gè)芯片，而 Trainium3 NL32x2 Switched 只能容納 32 個(gè)芯片。

Trainium3 實(shí)現(xiàn)盈利的時(shí)間策略

考慮到其雄心勃勃的機(jī)架架構(gòu)，AWS 在機(jī)架設(shè)計(jì)中做出了戰(zhàn)略性決策，以優(yōu)化 Trainium3 機(jī)架的貨幣化時(shí)間。我們相信這是 Trainium3 的一項(xiàng)巨大優(yōu)勢(shì)，部署 Trainium3 的客戶將從中受益。貨幣化時(shí)間的優(yōu)化對(duì)領(lǐng)先的 AI 實(shí)驗(yàn)室的代幣經(jīng)濟(jì)投資回報(bào)率 (ROIC) 有著顯著影響。下文我們將探討 AWS 在 Trainium3 的設(shè)計(jì)和部署方面所做的創(chuàng)新和戰(zhàn)略決策，以優(yōu)化貨幣化時(shí)間。

在供應(yīng)鏈方面，AWS自2024年底以來(lái)，已通過(guò)Trn2項(xiàng)目提升了供應(yīng)鏈和產(chǎn)能一年多。盡管從晶圓廠出廠到2025年上半年機(jī)架交付之間存在較長(zhǎng)的延遲，但我們預(yù)計(jì)機(jī)架ODM和數(shù)據(jù)中心供應(yīng)鏈目前已做好準(zhǔn)備，能夠以更短的CoWoS到機(jī)架交付周期應(yīng)對(duì)Trainium3的產(chǎn)能提升。我們觀察到，平均交付周期已縮短至一個(gè)季度以內(nèi)，并且還在持續(xù)縮短。

除了為產(chǎn)能爬坡做好供應(yīng)鏈準(zhǔn)備外，AWS 還針對(duì) Trainium 機(jī)架架構(gòu)的設(shè)計(jì)做出了多項(xiàng)戰(zhàn)略決策。正如我們?cè)谖恼虑懊嫣岬降?，Trainium 服務(wù)器采用無(wú)線纜設(shè)計(jì)理念，所有信號(hào)都通過(guò) JBOG 或節(jié)點(diǎn)內(nèi)的 PCB 傳輸，以優(yōu)化組裝效率。雖然通過(guò)飛線傳輸信號(hào)性能更佳，但線纜在組裝過(guò)程中也可能成為潛在的故障點(diǎn)。GB200 組裝的一大挑戰(zhàn)在于內(nèi)部線纜數(shù)量龐大，因此，Nvidia 實(shí)際上正在效仿 Trainium，為 Vera Rubin 平臺(tái)采用無(wú)線纜計(jì)算托架設(shè)計(jì)，以提高生產(chǎn)效率。這種設(shè)計(jì)的缺點(diǎn)是需要額外的 PCIe 重定時(shí)器，但考慮到 AWS 每次購(gòu)買 PCIe 重定時(shí)器時(shí)都能從 ALAB 獲得有效的返利，這仍然是一個(gè)相對(duì)經(jīng)濟(jì)的解決方案，尤其對(duì)于 AWS 而言，能夠縮短盈利時(shí)間，因此非常值得。

這種理念的另一個(gè)例子是連接到背板的擴(kuò)展鏈路的冗余設(shè)計(jì)。如上所述，每個(gè)Trainium3 都配備了 80 條專用于背板的 NeuronLinkv4 通道，其中 16 條通道用于冗余。這樣做的目的是為了彌補(bǔ)背板潛在的不可靠性。鑒于英偉達(dá) GB200 的背板可靠性較差，需要花費(fèi)時(shí)間進(jìn)行調(diào)試和更換，Trainium3 中設(shè)計(jì)的冗余通道有望實(shí)現(xiàn)熱插拔擴(kuò)展交換機(jī)托架，而不會(huì)中斷整個(gè)機(jī)架的工作負(fù)載。

最后，AWS 靈活的硬件方案也使其能夠在諸如高密度機(jī)架液冷數(shù)據(jù)中心尚未就緒和 UALink 交換機(jī)不可用等限制條件下部署 Trainium3。同時(shí)擁有風(fēng)冷（Trainium NL32x2 交換式）和液冷（Trainium NL72x2 交換式）全交換式擴(kuò)展機(jī)架的優(yōu)勢(shì)在于，即使液冷數(shù)據(jù)中心尚未準(zhǔn)備就緒，AWS 也可以在其原有的低密度數(shù)據(jù)中心部署 Trainium NL32x2 交換式機(jī)架。這種靈活性避免了單個(gè)設(shè)施的延遲導(dǎo)致收入延遲，正如我們最近在 CoreWeave 的 Denton 數(shù)據(jù)中心所看到的那樣。關(guān)于擴(kuò)展型 Neuron 交換機(jī)，我們討論了 AWS 計(jì)劃如何部署低基數(shù)交換機(jī)以縮短產(chǎn)品上市時(shí)間。這再次展現(xiàn)了他們的靈活性和優(yōu)化盈利時(shí)間的決心。

英偉達(dá)應(yīng)該注意產(chǎn)品變現(xiàn)時(shí)間，因?yàn)閺男酒a(chǎn)到客戶產(chǎn)生收益的變現(xiàn)時(shí)間，GB200 NVL72 已經(jīng)延長(zhǎng)，而 Vera Rubin Kyber 機(jī)架式顯卡的變現(xiàn)時(shí)間還會(huì)更長(zhǎng)。這將給 OEM/ODM 廠商和終端云平臺(tái)帶來(lái)巨大的營(yíng)運(yùn)資金壓力，增加其總體擁有成本 (TCO) 并降低盈利能力。

橫向擴(kuò)展和橫向擴(kuò)展網(wǎng)絡(luò)

理解EFA 首先要了解彈性網(wǎng)絡(luò)附加 (ENA)。在 AWS 中，啟動(dòng)虛擬機(jī)時(shí)，系統(tǒng)會(huì)通過(guò) ENA 分配一定的網(wǎng)絡(luò)容量。ENA 用于集群中的實(shí)例之間通信，以及連接到其他資源，例如 S3 和 EFS 等存儲(chǔ)服務(wù)或負(fù)載均衡器等網(wǎng)絡(luò)服務(wù)。ENA 還可以通過(guò) Nitro 系統(tǒng)用于 EBS，并提供上行鏈路/廣域網(wǎng)連接以及互聯(lián)網(wǎng)連接。

ENA 為上述服務(wù)提供了足夠的容量，但眾所周知，AI 服務(wù)器需要來(lái)自無(wú)阻塞網(wǎng)絡(luò)拓?fù)浣Y(jié)構(gòu)的更大容量。這就是 EFA 的作用所在。它是后端網(wǎng)絡(luò)或“東西向”網(wǎng)絡(luò)，而 ENA 是前端網(wǎng)絡(luò)或“南北向”網(wǎng)絡(luò)。

EFA 是一種網(wǎng)絡(luò)接口，它使用自身定制的可擴(kuò)展可靠數(shù)據(jù)報(bào) (SRD) 傳輸層來(lái)降低延遲，并提供擁塞控制和負(fù)載均衡。這些特性對(duì)于人工智能至關(guān)重要，因?yàn)槿鄙龠@些特性，大規(guī)模通信將無(wú)法擴(kuò)展。

EFA并非以太網(wǎng)的直接替代品，因?yàn)樗鼧?gòu)建于以太網(wǎng)之上，分別負(fù)責(zé)第1層（物理層）和第2層（數(shù)據(jù)鏈路層），但它是RoCEv2的替代方案，RoCEv2是對(duì)以太網(wǎng)的擴(kuò)展。在許多方面，AWS聲稱EFA超越了RoCEv2和InfiniBand，因?yàn)樗€包含許多更高層的功能。

AWS 聲稱 EFA 具有以下優(yōu)勢(shì)：

安全性：在安全特性方面，EFA 構(gòu)建于亞馬遜 VPC 控制平面之上，這意味著它繼承了 VPC 的核心云安全特性。例如，Nitro 強(qiáng)制執(zhí)行實(shí)例隔離，用戶空間不允許一個(gè)租戶訪問(wèn)另一個(gè)租戶的內(nèi)存。EFA 還使用線速加密（AES-256），這意味著流量將進(jìn)行端到端完全加密。

可擴(kuò)展性：SRD 發(fā)送器支持多路徑并具備擁塞感知能力，它會(huì)將內(nèi)存數(shù)據(jù)包分散到網(wǎng)絡(luò)中的多條路徑上，同時(shí)避開擁塞熱點(diǎn)。AWS 聲稱，這種處理?yè)砣⒗眯侣窂蕉怀霈F(xiàn)瓶頸的能力，使得 AWS 能夠在無(wú)需大型緩沖區(qū)交換機(jī)的情況下，跨區(qū)域構(gòu)建大型網(wǎng)絡(luò)架構(gòu)。這與 Nvidia 的 Spectrum-XGS 和 OpenAI 的 MRC 協(xié)議類似，它們也聲稱無(wú)需大型緩沖區(qū)交換機(jī)即可實(shí)現(xiàn)跨區(qū)域傳輸。

通用性：Libfabric 應(yīng)用程序編程接口 (API) 將網(wǎng)卡 (NIC) 和系統(tǒng)參考數(shù)據(jù) (SRD) 暴露給 MPI 實(shí)現(xiàn)，例如 Nvidia 集體通信庫(kù) (NCCL)。通過(guò) Libfabric，EFA 的通用性得以提升，因?yàn)樵S多網(wǎng)卡都在開發(fā)中使用相同的 API 接口，這意味著更多網(wǎng)絡(luò)通過(guò) Libfabric 實(shí)現(xiàn)了 EFA 兼容。然而，由于 Nvidia 擁有許多常用的高級(jí)技術(shù)，這種通用性的說(shuō)法在實(shí)踐中并不成立。

在AWS Nvidia GPU 上，由于用戶體驗(yàn)持續(xù)不佳，我們?nèi)匀徽J(rèn)為 EFA 相較于 Spectrum-X、InfiniBand 或配備 Connect-X 的 Arista 交換機(jī)，性能并無(wú)提升。然而，Trainium 的用戶體驗(yàn)要好得多，這得益于 AWS 能夠控制 Trainium 的整個(gè)堆棧。

為了支持EFA，AWS 構(gòu)建了定制的網(wǎng)絡(luò)接口卡 (NIC)。下表顯示了每一代 EFA 與特定 EC2 服務(wù)器的映射關(guān)系：

與我們之前關(guān)于Trainium2 的說(shuō)明（其中假設(shè)前端和后端網(wǎng)絡(luò)是分開的）相反，AWS 和 Google 一樣，將兩種類型的流量匯聚到同一個(gè)網(wǎng)絡(luò)上。他們通過(guò)將 Trainium 托架和 CPU 托架上的 Nitro-v6 網(wǎng)卡連接到同一個(gè)機(jī)架頂部 (ToR) 交換機(jī)來(lái)實(shí)現(xiàn)這一點(diǎn)。

對(duì)于Trainium3 上的 EFAv4，有兩種橫向擴(kuò)展網(wǎng)絡(luò)速度可供選擇：每顆 Trainium3 芯片配備一個(gè) 400G Nitro-v6 網(wǎng)卡，實(shí)現(xiàn) 400G 網(wǎng)絡(luò)速度；或者兩個(gè) Trainium3 芯片共享一個(gè) 400G Nitro-v6 網(wǎng)卡；或者每顆 Trainium3 芯片實(shí)現(xiàn) 200G 網(wǎng)絡(luò)速度。大多數(shù)機(jī)架將采用每顆 Trainium3 芯片 200G 的方案，我們將在下文中討論并繪制此方案的示意圖。無(wú)論采用哪種方案，Nitro-v6 網(wǎng)卡都將連接到兩個(gè) 200G OSFP 插槽。

在每臺(tái)Trainium3 芯片支持 200G 帶寬的版本中，每塊 400G Nitro-v6 網(wǎng)卡將支持兩顆 Trainium3 芯片。AWS 采用雙機(jī)架頂部 (ToR) 設(shè)計(jì)，其中一塊 Nitro-v6 網(wǎng)卡支持兩條 200G 鏈路，分別連接到位于計(jì)算托架上方的兩臺(tái) ToR 交換機(jī)。網(wǎng)卡端的 56G SerDes 通道通過(guò)一根帶齒輪箱的 400G Y 型有源銅纜 (AEC) 轉(zhuǎn)換為 ToR 端的 112G SerDes 通道，用于 Trainium 托架；而兩臺(tái) CPU 托架則通過(guò)一根直連 AEC 或直接有源銅纜 (DAC) 連接到同一對(duì) ToR 交換機(jī)。

對(duì)亞馬遜來(lái)說(shuō)幸運(yùn)的是，由于獲得了Credo 的股票返利，他們?cè)?AEC 方面的交易比在 PCIe 交換機(jī)和重定時(shí)器方面更劃算。Credo 的股票返利結(jié)構(gòu)與 AWS 與 ALAB 的交易類似，但由于 AWS 在此協(xié)議中獲得的認(rèn)股權(quán)證數(shù)量遠(yuǎn)高于 ALAB，且此后 Credo 的股價(jià)大幅上漲，因此實(shí)際返利金額要高得多。這意味著亞馬遜獲得的 Credo 認(rèn)股權(quán)證的價(jià)值超過(guò)了授予這些認(rèn)股權(quán)證所需的支出。Credo 實(shí)際上是付錢給亞馬遜，才讓亞馬遜購(gòu)買了 AEC！

盡管英偉達(dá)InfiniBand 或 Spectrum Ethernet 參考網(wǎng)絡(luò)架構(gòu)采用軌道優(yōu)化的 Clos 拓?fù)浣Y(jié)構(gòu)來(lái)減少 GPU 之間的交換機(jī)跳數(shù)，但 AWS 仍盡可能優(yōu)先選擇 ToR 交換機(jī)作為第一層交換層。這樣，通過(guò)用銅纜代替芯片和第一層交換層之間的光纖鏈路，可以降低整體網(wǎng)絡(luò)成本。此外，如果 ToR 交換機(jī)的額外上行鏈路端口可用于上層鏈路的欠載，從而實(shí)現(xiàn)容錯(cuò)、虛擬軌道或與其他服務(wù)的連接，則還能提供更多選擇。AWS 認(rèn)為這種權(quán)衡是值得的。

大多數(shù)Neocloud和超大規(guī)模數(shù)據(jù)中心的默認(rèn)網(wǎng)絡(luò)配置采用每個(gè)邏輯端口400G或800G的帶寬，與網(wǎng)卡帶寬相匹配。例如，Nvidia針對(duì)2000個(gè)GPU的H100集群的參考架構(gòu)將使用25.6T QM9700 InfiniBand交換機(jī)，這些交換機(jī)配備64個(gè)400G的邏輯端口，與CX-7網(wǎng)卡提供的每個(gè)GPU 400G的帶寬相匹配。交換機(jī)的帶寬限制意味著，由64端口交換機(jī)構(gòu)建的兩層網(wǎng)絡(luò)最多只能服務(wù)2048個(gè)GPU。

為了增加網(wǎng)絡(luò)上的最大GPU數(shù)量，越來(lái)越多的網(wǎng)絡(luò)采用高基數(shù)網(wǎng)絡(luò)（即將鏈路拆分成更多更小的邏輯端口），這表明堅(jiān)持使用較大的默認(rèn)邏輯端口尺寸會(huì)浪費(fèi)大量的網(wǎng)絡(luò)優(yōu)化和成本節(jié)約空間。

在微軟的文章中，我們也討論了超大規(guī)模數(shù)據(jù)中心如何在大型人工智能實(shí)驗(yàn)室的推動(dòng)下，開始部署高基數(shù)網(wǎng)絡(luò)。下圖展示了一個(gè)例子——OpenAI 在 Oracle 的網(wǎng)絡(luò)，該網(wǎng)絡(luò)可以使用 100G 邏輯端口連接兩層共 131,072 個(gè) GPU。

AWS 的 AI 網(wǎng)絡(luò)方案已經(jīng)直奔主題，默認(rèn)使用 100G 邏輯端口。這有兩個(gè)主要優(yōu)勢(shì)：

1. AWS 僅需 12.8T 交換機(jī)即可構(gòu)建大型網(wǎng)絡(luò)。

如果我們沿用傳統(tǒng)方法，使用與GPU網(wǎng)卡規(guī)格相匹配的邏輯端口（通常為400G），那么僅使用12.8T交換機(jī)構(gòu)建的網(wǎng)絡(luò)規(guī)模將非常有限。對(duì)于完全由12.8T交換機(jī)組成的兩層網(wǎng)絡(luò)，我們最多只能連接512個(gè)GPU。然而，如果我們改用100G邏輯端口，兩層網(wǎng)絡(luò)可以連接8192個(gè)GPU，三層網(wǎng)絡(luò)可以連接524288個(gè)GPU——足以滿足目前最大規(guī)模的多建筑集群的需求。

來(lái)源：SemiAnalysis AI 網(wǎng)絡(luò)模型

但為什么要用12.8T 交換機(jī)構(gòu)建網(wǎng)絡(luò)呢？實(shí)際上，AWS 并沒(méi)有這樣的目標(biāo)。AWS 的宗旨是最大限度地降低總體擁有成本，并圍繞這一核心目標(biāo)靈活調(diào)整采購(gòu)決策。打個(gè)比方，對(duì) AWS 來(lái)說(shuō)，貓是黑是白并不重要，重要的是它能抓到老鼠。因此，只要能降低總體擁有成本，AWS 會(huì)選擇任何交換機(jī)方案，無(wú)論是 12.8T、25.6T 還是 51.2T，亦或是 400G DR4、800G DR8 光模塊。

2. 如果我們引入 25.6T 和 51.2T 交換機(jī)，AWS 僅需兩層即可實(shí)現(xiàn)更大的規(guī)模。

按照同樣的計(jì)算方法，如果AWS引入25.6T和51.2T的交換機(jī)，就能實(shí)現(xiàn)巨大的規(guī)?！绻覀兒?jiǎn)單地假設(shè)默認(rèn)的400G或800G邏輯端口大小，那么實(shí)際規(guī)模會(huì)被低估。下表顯示，對(duì)于雙層網(wǎng)絡(luò)，在51.2T交換機(jī)上使用100G端口可以連接的GPU數(shù)量是使用400G端口的16倍。對(duì)于三層網(wǎng)絡(luò)，這個(gè)比例會(huì)擴(kuò)展到GPU數(shù)量的64倍。

使用100G 端口的缺點(diǎn)在于連接極其復(fù)雜——運(yùn)營(yíng)商通常需要使用配線架、跳線板或笨重的八爪魚線纜才能將這些 100G 鏈路路由到正確的目的地，而且對(duì)布線錯(cuò)誤的容忍度很低。亞馬遜則使用定制的光纖網(wǎng)絡(luò) ViaPhoton 來(lái)支持這種布線，從而最大限度地降低了這種復(fù)雜性帶來(lái)的影響。

Trainium3 橫向擴(kuò)展網(wǎng)絡(luò)

每個(gè)Trainium pod 內(nèi)的 ToR 交換機(jī)在葉層和脊層被分割成多個(gè)平面，并針對(duì)特定網(wǎng)絡(luò)進(jìn)行了優(yōu)化。所有 Trainium pod（均為獨(dú)立的可擴(kuò)展單元）均通過(guò)脊層連接。在下圖所示的示例中，我們假設(shè)在一個(gè)三層網(wǎng)絡(luò)中，每個(gè)平面上配備 12.8T 帶寬的交換機(jī)，并以此為基礎(chǔ)計(jì)算最大集群規(guī)模。

將葉脊層交換機(jī)替換為25.6T 交換機(jī)，而非 12.8T 交換機(jī)，意味著在相同的三層網(wǎng)絡(luò)中可支持的 Trainium3 芯片數(shù)量將增加 4 倍，同時(shí) pod 數(shù)量和每個(gè) pod 的機(jī)架數(shù)量也將翻倍。如果將葉脊層交換機(jī)升級(jí)到 51.2T，則該網(wǎng)絡(luò)中的芯片數(shù)量將再次增加四倍。

這種橫向擴(kuò)展網(wǎng)絡(luò)可以跨越多棟建筑物。FR 光模塊可用于幾公里的距離，而 ZR 光模塊可用于數(shù)百公里的距離。秉承讓網(wǎng)卡和網(wǎng)絡(luò)架構(gòu)管理長(zhǎng)距離傳輸帶來(lái)的延遲的理念，AWS 將放棄使用深度緩沖交換機(jī)，直接將不同建筑物的骨干層連接在一起。

確切的跨尺度拓?fù)浣Y(jié)構(gòu)尚不清楚——但下圖表示了一種已用于各種超大規(guī)模數(shù)據(jù)中心之間跨尺度擴(kuò)展的拓?fù)浣Y(jié)構(gòu)。

最后，在橫向擴(kuò)展網(wǎng)絡(luò)設(shè)備采購(gòu)方面，由于OpenAI 在 AWS 上的集群完全沒(méi)有使用 EFA，許多人對(duì) AWS 為 Trainium 使用的網(wǎng)絡(luò)架構(gòu)感到困惑。相反，該集群使用 GB300 交換機(jī)，并運(yùn)行 CX-8 協(xié)議，該協(xié)議運(yùn)行 OpenAI 的自定義協(xié)議 MultiPath Reliable Connection (MRC)。OpenAI 甚至可能使用 OCS 來(lái)連接不同的集群。這可能導(dǎo)致整個(gè)供應(yīng)鏈對(duì) AWS 的網(wǎng)絡(luò)構(gòu)建方式產(chǎn)生一些困惑，我們希望本節(jié)內(nèi)容能夠幫助那些有興趣了解其核心原理的人厘清思路。

Trainium3 微架構(gòu)

Trainium3 采用了與 Trainium2 和 Google 的 TPU 類似的設(shè)計(jì)思路，芯片由少量的大型 NeuronCore 構(gòu)成。這與 Nvidia 和 AMD 等 GPU 架構(gòu)形成鮮明對(duì)比，后者通常使用大量較小的 Tensor Core。大型核心通常更適合 GenAI 工作負(fù)載，因?yàn)樗鼈兊目刂崎_銷更小。與 Trainium2 一樣，Trainium3 每個(gè)封裝包含 8 個(gè) NeuronCore，每個(gè) NeuronCore 包含以下四個(gè)引擎：

• 張量引擎
• 矢量引擎
• 標(biāo)量引擎
• GPSIMD

來(lái)源：AWS

張量引擎

張量引擎采用128x128 BF16 脈動(dòng)陣列和 512x128 MXFP8/MXFP4 脈動(dòng)陣列。Trainium3 上的 BF16 脈動(dòng)陣列大小與 Trn2 的 BF16 陣列大小相同，但在 FP8 平臺(tái)上，其大小是 Trnium2 的兩倍。

脈動(dòng)陣列從名為“SBUF”的SRAM緩沖區(qū)獲取輸入，并將結(jié)果輸出到名為“PSUM”的部分和SRAM緩沖區(qū)。張量引擎可以遍歷矩陣乘法（matmul）的K維，并將每個(gè)結(jié)果的部分和相加，從而得到最終結(jié)果?，F(xiàn)代LLM工作負(fù)載中超過(guò)80%的算力和浮點(diǎn)運(yùn)算能力（FLOPS）都將用于張量引擎/脈動(dòng)陣列。張量引擎還支持MXFP8 4:8和4:16結(jié)構(gòu)化稀疏性，其浮點(diǎn)運(yùn)算能力是同等稠密模式的4倍，但我們懷疑會(huì)有客戶使用此功能。

MXFP4/MXFP8 的 512x128 脈動(dòng)數(shù)組也可以拆分成 4 個(gè) (128x128) 脈動(dòng)數(shù)組，這樣每個(gè)周期就可以將四個(gè)結(jié)果推送到 PSUM 緩沖區(qū)。在某些 GEMM 指令集中，有一些優(yōu)化方法可以使 4 個(gè) (128x128) 脈動(dòng)數(shù)組在更高的 MFU 下獲得比使用 512x128 脈動(dòng)數(shù)組指令集更好的性能。

通常情況下，即使對(duì)于BF16/MFP8，GEMM 也完全以 FP32 格式累加（在 Nvidia Hopper 中僅為 FP22），但某些工作負(fù)載可以容忍略低的累加精度。對(duì)于 Trainium3 張量引擎，可以選擇以 FP32 格式累加 128 個(gè)元素，然后在最后將其向下轉(zhuǎn)換為 BF16 格式。

張量引擎支持的數(shù)字格式和每瓦性能優(yōu)化

Trainium3 團(tuán)隊(duì)通過(guò)專注于提升 MXFP8 的性能，同時(shí)保持 BF16 的性能不變，并結(jié)合其他物理優(yōu)化措施（例如采用 3nm 工藝、優(yōu)化布局設(shè)計(jì)以及使用定制單元庫(kù)），在相同的芯片面積和功耗預(yù)算下實(shí)現(xiàn)了 MXFP8 性能翻倍。為了獲得更高的 MXFP8 每瓦 FLOPS，他們還采用了與 Trn2 相比更新的垂直供電系統(tǒng)。許多關(guān)鍵的物理設(shè)計(jì)工作都是在內(nèi)部完成的，而不是外包給供應(yīng)商。為了將通常精度較高的主權(quán)重轉(zhuǎn)換為精度較低的計(jì)算權(quán)重，Trainium3 在芯片中集成了硬件加速單元，用于加速 MXFP8/MXFP4 的量化/去量化。

遺憾的是，只專注于MXFP8 的弊端在于 BF16 的性能并沒(méi)有提升。像 Anthropic 這樣的高級(jí) L337 用戶不需要 BF16 進(jìn)行訓(xùn)練，他們具備使用 MXFP8 進(jìn)行訓(xùn)練的技能，但普通機(jī)器學(xué)習(xí)訓(xùn)練師只知道如何使用 BF16 進(jìn)行訓(xùn)練。

此外，Trainium3 的 MXFP4 性能與 MXFP8 相同，但與 AMD/Nvidia 的 GPU 相比，它在推理方面優(yōu)化得不夠好，因?yàn)樗鼈兡軌蛞月缘偷馁|(zhì)量換取更快的推理速度。

然而，對(duì)于推理解碼等內(nèi)存密集型工作負(fù)載而言，這一點(diǎn)影響不大，因?yàn)閬嗰R遜/Anthropic 可以將權(quán)重存儲(chǔ)在自定義的 4 位塊大小的存儲(chǔ)格式中，同時(shí)使用 MXFP8 進(jìn)行計(jì)算。這種技術(shù)通常被稱為 W4A8。對(duì)于內(nèi)存密集型操作，使用 W4A8 可以實(shí)現(xiàn)從 HBM 加載和存儲(chǔ)數(shù)據(jù)的速度提升一倍，因?yàn)閺?HBM 到芯片的傳輸將以 4 位而非 8 位進(jìn)行，并且在輸入到張量引擎之前會(huì)在芯片上進(jìn)行反量化。

此外，Trainium3 不支持 NVFP4（塊大小 16，塊縮放格式 E4M3），僅硬件支持 OCP MXFP4（塊大小 32，塊縮放格式 E8M0）。這意味著與 Nvidia GPU 相比，Trainium3需要更高級(jí)的 QAT/PTQ 技術(shù)。E8M0 塊縮放比 E4M3 塊縮放更差的原因在于，E8M0 會(huì)將縮放因子精確到最近的 2^n，從而導(dǎo)致更嚴(yán)重的量化誤差。雖然 Trainium3 在技術(shù)上支持 NVFP4 作為存儲(chǔ)格式（或任何 4 位任意存儲(chǔ)格式），并且可以支持在線解量化到 OCP MXFP8，但它本身并不提供硬件加速的 NVFP4 到 OCP MXFP8 解量化支持，必須通過(guò)軟件驅(qū)動(dòng)的方式實(shí)現(xiàn)。

Trainium3 不支持 NVFP4 使得 4 位訓(xùn)練變得更加困難。英偉達(dá)研究院（以及英偉達(dá)市場(chǎng)部）最近發(fā)布了一篇關(guān)于 NVFP4 訓(xùn)練的研究論文，其中展示了一種實(shí)驗(yàn)性的訓(xùn)練方案，該方案在前向傳播和反向傳播中都使用了 4 位浮點(diǎn)運(yùn)算。我們認(rèn)為，在未來(lái) 12 個(gè)月內(nèi)，西方前沿實(shí)驗(yàn)室不太可能在前向傳播和反向傳播中都采用 4 位浮點(diǎn)訓(xùn)練，但我們認(rèn)為，隨著訓(xùn)練方案的成熟，他們最終可能會(huì)轉(zhuǎn)向 4 位訓(xùn)練。

盡管如此，一些西方前沿實(shí)驗(yàn)室已經(jīng)在前向訓(xùn)練階段采用了NVFP4格式，但反向訓(xùn)練階段目前仍使用更高精度的數(shù)字格式，而且這種方法似乎效果良好，沒(méi)有造成明顯的質(zhì)量損失。一些使用4位浮點(diǎn)數(shù)進(jìn)行前向訓(xùn)練的前沿實(shí)驗(yàn)室已經(jīng)將這些模型部署到生產(chǎn)環(huán)境中，擁有數(shù)百萬(wàn)活躍用戶。

AWS Trainium3 的缺點(diǎn)是，如果 4 位前向傳播訓(xùn)練繼續(xù)在高級(jí)用戶中得到廣泛應(yīng)用，那么 Trainium3 可能由于不支持 NVFP4（塊大小 16，塊縮放格式 E8M0）并且僅啟用與 OCP MXFP8 相同的 OCP MXFP4（塊大小 32，塊縮放格式 E4M3）計(jì)算能力而表現(xiàn)不佳。

*聲明：本文系原作者創(chuàng)作。文章內(nèi)容系其個(gè)人觀點(diǎn)，我方轉(zhuǎn)載僅為分享與討論，不代表我方贊成或認(rèn)同，如有異議，請(qǐng)聯(lián)系后臺(tái)。