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

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

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

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

Amazon Basics GB200

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

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

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

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

亞馬遜的軟件

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

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

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

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

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

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

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

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

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

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

Trainium2 和 Trainium3 之間共有四個不同的服務器 SKU，供應鏈通常使用它們的代號來指代它們，這與 AWS 的品牌名稱不同。

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

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

來源：SemiAnalysis、AWS

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

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

來源：SemiAnalysis、AWS

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

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

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

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

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

Trainium3 機架架構(gòu)

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

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

來源：SemiAnalysis、AWS

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

Trainium3 制程和封裝

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

來源：AWS、SemiAnalysis

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

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

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

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

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

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

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

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

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

Marvell最終成了這場競爭的最大輸家。盡管Trainium2是由他們設計的，但他們在這一代芯片的設計競賽中輸給了Alchip。Marvell的Trainium3版本采用了基于芯片組（chiplet）的設計，將I/O功能放在一個獨立的芯片組上，而不是像Trainium2和即將推出的Trainium3那樣，將計算功能集成在一個單芯片上。

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

Trainium4路線圖預估

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

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

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

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

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

來源：英偉達

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 大會上發(fā)布。我們在 Trainium2 深度解析中介紹了這兩種架構(gòu)。

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

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

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

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

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

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

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

機架架構(gòu)

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

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

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

JBOG/計算托盤

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

迄今為止，所有Trainium 機架均采用無電纜設計理念以提高組裝效率，因此所有信號均通過 PCB 走線傳輸。信號在 PCB 上傳輸時的損耗遠高于通過架空電纜傳輸時的損耗，因此必須在 JBOG 板中間放置四個 PCIe 6.0 x16 重定時器，以補償在前置 I/O 端口和兩個 Trainium3 封裝之間通過 PCB 傳輸?shù)男盘枴?/p>

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

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

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

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

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

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

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

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

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

機架架構(gòu)

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

計算機托盤

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

即使啟用全部80 條通道，像訓練這樣通信密集型的工作負載也無法獲得任何顯著的性能提升，這是因為存在“掉隊效應”。在任何大型訓練運行中，至少會有幾個機架的某些通道出現(xiàn)故障。如果只有一個 Trainium3 機架出現(xiàn)故障通道，那么整個訓練作業(yè)將只能有效使用 80 條通道中的 64 條，因為所有其他機架都會等待速度最慢的機架。

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

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

NeuronLink PCIe PHY 和 UALink 交換機

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

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

交換機托架

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

例如，節(jié)點1 中的 Trainium3 A 可以通過一個交換機跳與所有其他節(jié)點中的 Trainium3 A 通信。節(jié)點 1 中的 Trainium3 B 也同樣可以通過一個交換機跳與所有其他節(jié)點中的 Trainium3 B 通信。

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

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

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

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

另一種潛在的并行策略是在兩個機架上都使用完整的EP（擴展并行），但要考慮到額外的躍點。這種策略對于非常稀疏的模型可能非常有效，因為這類模型由于d_model維度太小而無法跨機架實現(xiàn)TP（并行并行）。因此，即使通過直接連接的Trainium3進行額外躍點帶來的延遲也是值得的。

接下來是Trainium3 NL72x2 交換機 SKU，其擴展拓撲結(jié)構(gòu)更為復雜。它包含四個平面，每個平面配備 10 個 160 通道 PCIe 交換機，每個機架的交換機托架上共有 40 個交換機；此外，18 個計算托架上每個托架還配備 8 個 32 通道 PCIe 交換機，這意味著每個機架的計算托架上總共有 144 個較小的 PCIe 交換機。因此，每個機架總共有 184 個擴展交換機，或者說，144 臺 Trainium3 擴展機架共計 368 個擴展交換機。為了便于跟蹤交換機配置，我們在此重現(xiàn)之前的匯總表：

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

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

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

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

Trainium3 NL72x2 已開啟 Gen1 備用交換機托架

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

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

第二代交換機托架

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

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

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

最終目標：第三代交換機托架

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

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

對于Trainium3 NL72x2 交換式架構(gòu)，每個機架的擴展拓撲結(jié)構(gòu)將變?yōu)槿B接，因為每個 UALink 交換機都可以連接到機架內(nèi)的每個 Trainium3 芯片。通過 8 臺 32 通道 Scorpio P 交換機實現(xiàn)的本地計算托架連接現(xiàn)在體現(xiàn)了剩余帶寬。

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

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

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

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

Trainium3 機架功率預算和物料清單

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

Trainium3 NL72x2 開關(guān)系統(tǒng)的總系統(tǒng)功耗自然更高，因為它在兩個機架上部署了 144 個芯片，而 Trainium3 NL32x2 開關(guān)系統(tǒng)在兩個機架上僅部署了 64 個芯片。然而，一旦將功耗標準化為每個芯片的功耗，64 個芯片的 Trainium3 NL32x2 開關(guān)系統(tǒng)和 144 個芯片的 Trainium3 NL72x2 開關(guān)系統(tǒng)實際上每個芯片的功耗非常接近，因為 Trainium3 芯片的熱設計功耗 (TDP) 是影響整體功耗預算的最大因素。 Trainium3 NL72x2 Switched 的機架功率密度自然更高，因為它在一個機架中可以容納 64 個芯片，而 Trainium3 NL32x2 Switched 只能容納 32 個芯片。

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

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

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

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

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

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

英偉達應該注意產(chǎn)品變現(xiàn)時間，因為從芯片生產(chǎn)到客戶產(chǎn)生收益的變現(xiàn)時間，GB200 NVL72 已經(jīng)延長，而 Vera Rubin Kyber 機架式顯卡的變現(xiàn)時間還會更長。這將給 OEM/ODM 廠商和終端云平臺帶來巨大的營運資金壓力，增加其總體擁有成本 (TCO) 并降低盈利能力。

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

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

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

EFA 是一種網(wǎng)絡接口，它使用自身定制的可擴展可靠數(shù)據(jù)報 (SRD) 傳輸層來降低延遲，并提供擁塞控制和負載均衡。這些特性對于人工智能至關(guān)重要，因為缺少這些特性，大規(guī)模通信將無法擴展。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

來源：SemiAnalysis AI 網(wǎng)絡模型

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

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

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

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

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

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

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

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

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

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

Trainium3 微架構(gòu)

Trainium3 采用了與 Trainium2 和 Google 的 TPU 類似的設計思路，芯片由少量的大型 NeuronCore 構(gòu)成。這與 Nvidia 和 AMD 等 GPU 架構(gòu)形成鮮明對比，后者通常使用大量較小的 Tensor Core。大型核心通常更適合 GenAI 工作負載，因為它們的控制開銷更小。與 Trainium2 一樣，Trainium3 每個封裝包含 8 個 NeuronCore，每個 NeuronCore 包含以下四個引擎：

• 張量引擎
• 矢量引擎
• 標量引擎
• GPSIMD

來源：AWS

張量引擎

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

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

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

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

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

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

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

此外，Trainium3 的 MXFP4 性能與 MXFP8 相同，但與 AMD/Nvidia 的 GPU 相比，它在推理方面優(yōu)化得不夠好，因為它們能夠以略低的質(zhì)量換取更快的推理速度。

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

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

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

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

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

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