AWS Trainium2 架構與網絡技術

逍遙設計自動化

引言
Amazon 在人工智能加速器領域正在通過其 Trainium2 芯片架構取得進展，旨在與 NVIDIA 在人工智能訓練和推理領域展開競爭。本文將詳細探討 Trainium2 的架構、網絡功能和成本等關鍵方面[1] 。
2 q: t% p7 a+ }, r+ ~8 E


圖1：展示 AWS 硬件執(zhí)行 AI 工作負載的藝術表現(xiàn)，象征著 AWS 對 AI 加速的投入。

$ F- p/ E; L+ ~核心架構概述
Trainium2 架構相比前代產品有顯著提升，每個芯片提供 650 TFLOP/s 的密集 BF16 性能，配備 96GByte HBM3e 內存。每個 Trainium2 芯片包含兩個計算 chiplet 和四組 HBM3e 內存，通過 CoWoS-S/R 封裝連接。

圖2：NeuronCore-v2 的詳細架構，展示了張量引擎、向量引擎、標量引擎和 GPSIMD 引擎與 HBM 內存的集成。

) ]: Z, S* z1 z% F2 A1 h3 Z神經核心架構包含四個主要引擎：
3 \" @$ m3 n! i' E0 c: F張量引擎：128×128 脈動陣列，用于矩陣運算
向量引擎：處理向量運算和歸一化
標量引擎：管理按元素運算
3 q7 ^- R; g# V7 m6 c& OGPSIMD 引擎：執(zhí)行任意 C++ 運算

服務器架構與部署
2 J# o) Y" }2 `Trainium2 提供兩種主要規(guī)格：
% E) ^- m: u# E9 K7 F) o" v# x- N2 A$ hTrainium2 (Trn2)：每服務器16個芯片配置
Trainium2-Ultra (Trn2-Ultra)：跨四臺服務器的64個芯片

2 W% |3 j4 O' _# ~- j8 _) h4 |
+ {1 q4 x" I9 q- e; W+ c圖3：Trainium2 服務器架構，展示了計算托盤和 CPU 頭托盤在18機架單元中的排列。
8 n# A8 H5 F' }8 o
物理服務器架構占用18個機架單元，包括：
! s% O5 d, w$ R  x: j一個 2U CPU 頭托盤
八個 2U 計算托盤
每個計算托盤包含兩個 Trainium2 芯片
計算托盤中沒有 CPU（JBOG - "僅GPU組"）
0 K( \5 K3 R9 Y6 u. h/ w0 T
6 z' Q& y6 g8 R8 K: C$ C網絡功能
Trainium2 的網絡基礎設施整合了多項技術：
% I/ y/ J9 [' X. J; Q# F- tNeuronLinkv3（擴展上行網絡）
! T) A& d2 G8 F彈性 Fabric 適配器 EFAv3（擴展出口網絡）
前端和存儲網絡
帶外管理網絡

3 U8 M& R7 C: P3 F2 ]$ S
圖4：4x4x4 三維環(huán)形網絡拓撲，顯示節(jié)點間連接和帶寬規(guī)格。
' N6 a$ p* r' b7 A( s
供電創(chuàng)新
Amazon 在 Trainium2 中實施了垂直供電，這在芯片供電管理方面是重大進展。

8 X& N3 x8 ~6 G, _8 B3 i圖5：傳統(tǒng)與垂直供電系統(tǒng)的比較，展示了改進的供電效率和散熱管理。

成本分析與性能
4 c1 _/ K# u+ p1 \; U與 NVIDIA 的 H100 相比，Trainium2 的總擁有成本顯示出有利的經濟性：

, _( s/ a8 I  a  v" _: b' U圖6：Trainium2 和 H100 配置的詳細成本比較，突顯了 AWS 解決方案的經濟優(yōu)勢。

主要經濟因素包括：
較低的前期資本成本（每芯片4,000美元，相比H100的23,000美元）
由于更好的能效而降低運營成本
在部署生命周期內更有利的總擁有成本
1 ^! @- Q5 K$ V
; }" M# _( i0 a: A" u- r6 X! nProject Rainier 實施
AWS 目前正在為 Anthropic 部署一個名為"Project Rainier"的大型集群，配備40萬個 Trainium2 芯片。該部署展示了 Trainium2 可以有效運行的規(guī)模。

圖7：AWS 印第安納數(shù)據中心園區(qū)的空中視圖，Project Rainier 正在此處部署。

軟件棧與開發(fā)工具
Trainium2 的軟件生態(tài)系統(tǒng)包括：
NeuronX 集體通信庫
7 u/ z( M+ S6 B通過 TorchDynamo 的 PyTorch 集成
8 ^9 s! O# Y' U( n8 m7 p3 L$ h" X: ]JAX 測試版支持
3 E7 O* {' s) l; P2 E用于低級優(yōu)化的 Neuron 內核語言（NKI）

& s3 y- k; C7 I" K
. {( z* i/ q3 w' x0 u) a! w  y) [圖8：Neuron 分布式事件追蹤界面，顯示全面的調試和分析功能。

8 {$ K( m' @! d  w& T1 f未來發(fā)展
; @" ]  T, C; DTrainium2 代表了 Amazon 在 AI 加速器領域的重要投入。憑借具有競爭力的價格和性能特征，對 NVIDIA 在 AI 訓練市場的主導地位構成了顯著挑戰(zhàn)。Project Rainier 的成功和 Anthropic 的采用將是 Trainium2 實際效果的重要指標。
* |( G) R3 m1 K/ s6 g
& @: u$ X5 a3 W: ^: h) O: w7 F圖9：下一代供電架構，展示了 AWS 對未來可擴展性和效率的投入。
1 p6 ~5 |' r8 {! G1 u* F! C
通過對供電、網絡功能和軟件集成的嚴密關注，AWS 創(chuàng)建了一個有競爭力的平臺，解決了現(xiàn)代 AI 加速中的許多挑戰(zhàn)。該平臺的發(fā)展將可能影響 AI 硬件開發(fā)和云服務供應商策略的未來方向。
0 u! J4 y" }+ }) `  u4 ^# m9 y
參考文獻
3 e% b- p% \7 H  a[1] D. Patel, D. Nishball, and R. Knuhtsen, "Amazon's AI Self Sufficiency | Trainium2 Architecture & Networking," SemiAnalysis, Dec. 3, 2024. [Online]. Available: https://semianalysis.com/2024/12/03/amazons-ai-self-sufficiency-trainium2-architecture-networking/

END
) h5 R) C4 N/ k% s
& U  ~" _& f4 ]
* N$ @# _% G0 u軟件申請我們歡迎化合物/硅基光電子芯片的研究人員和工程師申請體驗免費版PIC Studio軟件。無論是研究還是商業(yè)應用，PIC Studio都可提升您的工作效能。
點擊左下角"閱讀原文"馬上申請
( ]: ]4 f& t+ |
歡迎轉載
0 {; y5 l. @& |
轉載請注明出處，請勿修改內容和刪除作者信息！
$ H4 J' L, [$ d
  U, \- F  Y7 N8 p( d: l



關注我們

0 l7 ^/ o" F- V3 Y4 \( \9 L

: z7 V; I% e! z" Z

?7 [5 ]/ V2 }1 u1 D

& e0 T1 ^- c2 u

                     
7 C$ h' \! K* u/ k8 k( q/ L


* `5 I% q9 K( w7 W( y# y; z關于我們：
4 F) b) U: C& C深圳逍遙科技有限公司（Latitude Design Automation Inc.）是一家專注于半導體芯片設計自動化（EDA）的高科技軟件公司。我們自主開發(fā)特色工藝芯片設計和仿真軟件，提供成熟的設計解決方案如PIC Studio、MEMS Studio和Meta Studio，分別針對光電芯片、微機電系統(tǒng)、超透鏡的設計與仿真。我們提供特色工藝的半導體芯片集成電路版圖、IP和PDK工程服務，廣泛服務于光通訊、光計算、光量子通信和微納光子器件領域的頭部客戶。逍遙科技與國內外晶圓代工廠及硅光/MEMS中試線合作，推動特色工藝半導體產業(yè)鏈發(fā)展，致力于為客戶提供前沿技術與服務。

http://www.latitudeda.com/
- P. L6 H" B" `) x（點擊上方名片關注我們，發(fā)現(xiàn)更多精彩內容）