有了Trainium4,AWS將把除時鐘外的所有功能都調高

全球的人工智能模型開發者已經等待了一年多,終于能拿到Trainium3 XPU,這些XPU專門為訓練和推理設計,是英偉達“Blackwell”B200和B300 GPU以及谷歌“Trillium” TPU v6e和“Ironwood” TPU v7p加速器的可信替代品。
但當亞馬遜云服務首席執行官Matt Garmin開始談論預計可能在2026年底或2027年初交付的未來Tranium4 XPU時,所有排隊購買基于Trainium3的EC2容量模塊的人都做好了買家后悔的準備。因為盡管Trainium3相比公司Annapurna Labs芯片設計部門之前的Inferentia和Trainium XPU表現優異,Trainium4看起來將徹底突破插槽大小,不僅是一款非常強大的設備,還將擁有更具可擴展性的UltraServer集群,更適合運行多種專家級、思維鏈推理模型的混合。
在深入探討Trainium4可能是什么之前,讓我們先花點時間真正回顧一下Trainium3的XPU,尤其是因為該芯片及其前身Tranium2的許多技術規格直到最近才公開。讓我們先從Garmin在本周拉斯維加斯re:Invent 2025大會開幕式主題演講中談到的Tranium3規格表開始:

這是AWS去年此時展示的幻燈片的更新,預覽了它愿意談論的三個數據點。Tranium3采用了臺灣半導體制造股份有限公司的3納米工藝節點,這是一種對5納米技術的縮小,而大多數人認為Annapurna實驗室用于Tranium2芯片時正是如此。Tranium3預計能實現2倍的計算能力(這可能意味著很多不同的事情),并且能效提升40%(但這并不是一個有價值的指標,因為沒人知道Tranium1、Trainium2或Trainium3的功耗是多少)。但顯然,工藝縮小更多是為了削減功耗,而不是把新功能塞進芯片,而且插槽被做得更大以提升性能,每瓦凈性能提升了40%。
亞馬遜銷售基于Tranium2的常規服務器實例,以及擁有16個Trainium3套接字的UltraServer集群配置,但目前僅交付了單一內存域中64個Tranium3的UltraServers。

相比之下,Trn3 UltraServers在AWS上調用實例時計算能力提升了4.4倍,這與擁有四倍XPU數量的Trn2 UltraServers相比是合理的。據Garmin稱,最新的UltraServer擁有Trn2 UltraServer的3.9倍總HBM內存帶寬,更重要的是,對于那些擔心推理成本——這是生成式AI商業化的門檻因素——的人來說,每兆瓦能生成5倍的代幣數量。
以下是Garmin為該性能聲明分享的帕累托曲線,展示了每兆瓦輸出代幣在Y軸上的相互作用,輸出的交互性以每位用戶每秒代幣數表示:

將這個曲線向上推移,是2025年及以后獲得推理業務的全部策略。這組圖表比較了Trn2 UltraServer集群與運行OpenAI的GPT-OSS 120B模型的Trn3 UltraServer集群。
這張圖表還顯示了Garmin未提及的,如果對推理工作量重要,你可以在相同能量下獲得大約一個數量級更多的交互性。
在 Trainium2 實例云端擴展過程中,AWS 更新了該 XPU 的規格,我們也找到了一些 Tranium3 的規格,消除了一些謎團,填補了關于組件如何堆疊在 Trainium 套接字中以形成每個后續 XPU 的空白。
我們先從神經核開始,逐步推理。
所有NeuronCore設計都將四種不同的計算方式放入核心,就像CPU核心早已混合整數(標量)和向量單元,偶爾(Intel Xeon 5和6以及IBM Power10和Power11)也包含張量單元一樣。從Trainium系列開始,Annapurna Labs在架構中增加了集體通信核心(CC-Cores),以處理高性能計算和人工智能工作負載中常見的集體作的特殊處理,這樣實際上已經有五個核心了。
在僅用于Inferentia1芯片的NeuronCore-v1架構中,有用于整數計算的標量引擎(兩個整數輸入和一個整數輸出)、用于向量數學的向量引擎(兩個浮點輸入,一個浮點輸出)、一個用于張量數學的張量引擎(多個矩陣浮點輸入和一個矩陣浮點輸出)。
根據AWS文檔,NeuronCore-v1中的標量引擎每個時鐘周期可處理512次浮點運算,并可處理FP16、BF16、FP32、INT8、INT16和INT32數據類型。(我們認為AWS的意思是它處理512位數據。)文檔還說向量引擎每個周期能處理256次浮點運算(我們認為是256位數據),并且支持FP16、BF16、FP32、INT8、INT16和INT32數據格式。你可以根據數據寬度和每個單元能裝入多少個周期來計算每個周期的作。
TensorEngine for NeuronCore-v1 的維度從未公開,但我們知道它處理了 FP16、BF16 和 INT8 輸入,以及 FP32 和 INT32 輸出,并提供了 16 萬億次 FP16 或 BF16 張量處理。
在討論了第一個NeuronCore-v1設計后,讓我們把它們并排擺放,看看我們認為Trainium4可能的方向:

AWS于2020年12月在re:Invent上開始談論Trainium1,并花了兩年時間才全面推廣,這可以理解,因為這是亞馬遜首個自家開發的數據中心級培訓加速器。我們認為Trainium1是通過臺積電7納米工藝蝕刻的;我們知道它有550億個晶體管,頻率為3 GHz。該芯片采用了與后續于 2023 年 4 月上市的 Inferentia2 相同的 NeuronCore-v2 架構,工藝規模縮小至 5 納米,晶體管數量大致相同,但針對推理特定工作負載做了一些調整,比如將 NeuronLink 芯片互連端口減半。
隨著 Trainium2 于2023年11月發布,2024年12月批量發貨,AWS 轉向 NeuronCore-v3 架構,停止生產 Inferentia 芯片,因為推理開始變得更像訓練。Trainium2使每個套接字核心數增加了四倍,單個內存域中的神經核總數增加了16倍,每個實例的套接字數也增加了四倍。據我們所知,AWS還通過將Trainium2的頻率從7納米縮小到5納米,提升了Trainium2的頻率。有趣的是,每個NeuronCore的峰值標量和矢量性能在v3下降了約60%,峰值張量核心吞吐量下降了12%。但AWS為芯片增加了1:4稀疏度支持以支持張量運算,加上更高的核心數量,使Trainium2在FP16或BF16精度下的有效吞吐量提升了3.5倍。事實上,NeuronCore-v3支持多種稀疏度模式:4:16、4:12、4:8、2:8、2:4、1:4和1:2。
NeuronCore-v3的SRAM內存(由三個計算單元共享)提升至每核28 MB,但具體提升的幅度尚不清楚。HBM內存最終提升至96 GB,提升了3倍,帶寬提升3.5倍至2.9 TB/秒。這可以說是首個競爭性Trainium芯片,Anthropic使用這些Trainium2設備進行模型開發和推斷,且AWS基巖模型服務的大部分推斷均由Trainium完成,這絕非巧合。我們懷疑Garmin在主題演講中提到的數百萬臺Trainium大多是Trainium2設備。

這就引出了Trainium3,它現在已經在UltraServer實例中批量發布了。Trainium3 設備的核心是 NeuronCore-v4 架構——是的,如果核心名稱與設備名稱同步會更好——有一些重大變化。首先,矢量引擎經過調整,能夠進行快速的指數函數評估,其性能是標量引擎的4倍,而標量引擎是生成式AI模型自注意算法的一部分。其次,FP16和BF16數據格式可以量化成MXFP8格式,AWS表示這對生成式AI模型中多層感知器(MLP)層之間的數據量化非常有用。NeuronCore-v3 設計還將 SRAM 提升至 32 MB 的每核。Trainium2 和 Trainium3 之間的時鐘頻率似乎名義上有所變化,但變化不大,設備最大的變化是 NeuronLink-v4 XPU 互連端口帶寬翻倍至 2.5 TB/sec,HBM 內存容量提升 1.5 倍至 144 GB,HBM 帶寬提升 1.7 倍至 4.9 TB/秒。
我們認為 Trainium3 設計的變更旨在讓計算、內存和互聯重新達到更好的平衡,從而提升 Trainium3 套接字的理論峰值性能,而非理論上的峰值性能。Trn3 Gen1 UltraServer的內存域保持在64臺設備,與Trainium2相同,但隨著Trn3 Gen2 UltraServer的發貨,域大小增加到144個套接字。這會使AI訓練或推理任務中可投入的核心數量增加2.25倍。
這就引出了預計明年這個時候開始推出的Trainium4。
我們推測將采用的 NeuronCore-v5 架構,AWS 將為 Trainium 處理添加適當的 FP4 支持,而不僅僅是把 MXP4 塞進 FP8 的張量槽,留下大量空白。Garmin在主旨演講中表示,采用FP4原生格式后,Tranium4的性能將是Tranium3的6倍,這意味著FP8的處理能力將提升3倍。Garmin進一步表示,Trainium4的HBM內存容量是Trainium4的2倍,HBM帶寬是Trainium4的4倍。
在上面的怪獸表中,我們嘗試推測了Trainium4可能的樣子,以及如何為耦合的Trainium4設備進一步擴展內存域。
實現這一目標的方式有很多種,我們認為 AWS 充其量會轉向 2 納米工藝以節省一些功耗,或者堅持 3 納米工藝,節省一些成本,制造更大更熱的 XPU。這是一個艱難的決定,但我們認為Trainium4會傾向于采用2納米蝕刻技術。
如果你看上面Garmin的圖表,你不會發現它說Tranium4會同時支持NVLink和UALink接口——Nvidia曾大肆宣傳AWS將采用NVLink技術,但我們有預感,AWS會推出帶有NVlink接口的Graviton系列芯片變體,并獲得Nvidia一直不愿透露的內容: 它能夠將定制CPU和定制XPU粘合到一個巨大的共享內存域中,使用NVLink端口和NVSwitch內存織物交換機。到目前為止,英偉達樂于讓客戶擁有連接英偉達GPU的定制CPU或連接英偉達GPU的定制XPU,但并未允許第三種選項。
我們認為AWS購買的GPU數量足夠多,能夠以合理的價格提出并獲得這樣的需求。我們還認為AWS將支持Nvidia的NVFP4數據格式以及用于FP4處理的MXFP4格式,這很可能是交換的一部分,以便將Tranium4芯片的工作轉移到Nvidia的“Blackwell”和“Rubin”GPU上。當然,這些只是直覺。我們也認為AWS希望能夠將這些功能接入自家機架,這些機架本質上將是英偉達機架的克隆。
但有趣的是,UALink也出現在上面的幻燈片里,來自Garmin。AWS保持開放選擇,毫無疑問希望為Tranium4封裝設計芯片組架構,使其能夠將NVLink端口替換為UALink端口,并采用機架設計,允許NVSwitch交換機在明年晚些時候上市時替換為UALink交換機。有可能 NeuronLink-v5 經過調整以兼容 UALink 2.0,這些交換機將來自 Annapurna Labs,而非 Astera Labs、Upscale AI、Marvell、Cisco Systems 或其他提供可擴展互連 ASIC 的廠商。
在相同或稍小的熱包內實現3倍性能的最簡單方法是將核心數量增加三倍,同時保持時鐘頻率不變,同時采用2納米工藝。如果晶體管收縮較大(臺積電的1.6納米A16工藝),那么熱量可以降低一點,或者把時鐘調高一點點。我們的建議是利用散熱優勢保持不變,就像AWS在Trainium2和Trainium3之間所做的那樣,只增加3倍的核心。
如果你把核心數增加3倍,每個插槽24個核心,可能分布在四個芯片組上,那么在恒定精度下能達到3倍;如果從FP8縮小到FP4,每個插槽的效能提升6倍。
接下來事情變得有趣了。如果你將設備數量翻倍到每臺288個(與英偉達的做法一致),Trainium4 UltraServer集群中可以擁有6912個Neuro核心,全部集中在一個內存域中,配備1944TB的HBM內存。
當然,這遠不及谷歌在單一內存域中能提供的9,612個Ironwood TPU v7p XPU的數量......但它比目前銷售的Trn2 Gen2 UltraServer集群好13.5倍。










評論