AWS Trainium2 架構(gòu)與網(wǎng)絡(luò)技術(shù)

逍遙設(shè)計(jì)自動(dòng)化

引言
1 p6 S/ X% B* B  PAmazon 在人工智能加速器領(lǐng)域正在通過其 Trainium2 芯片架構(gòu)取得進(jìn)展，旨在與 NVIDIA 在人工智能訓(xùn)練和推理領(lǐng)域展開競(jìng)爭(zhēng)。本文將詳細(xì)探討 Trainium2 的架構(gòu)、網(wǎng)絡(luò)功能和成本等關(guān)鍵方面[1] 。
, O0 F# {3 {: j


圖1：展示 AWS 硬件執(zhí)行 AI 工作負(fù)載的藝術(shù)表現(xiàn)，象征著 AWS 對(duì) AI 加速的投入。
1 [+ h  l% S. G# n
  L) d7 N! @0 K1 n  E/ j1 V核心架構(gòu)概述
3 S) X/ P; ?- W# I$ P8 JTrainium2 架構(gòu)相比前代產(chǎn)品有顯著提升，每個(gè)芯片提供 650 TFLOP/s 的密集 BF16 性能，配備 96GByte HBM3e 內(nèi)存。每個(gè) Trainium2 芯片包含兩個(gè)計(jì)算 chiplet 和四組 HBM3e 內(nèi)存，通過 CoWoS-S/R 封裝連接。
% }. U. W) I; }# }( }7 A+ r3 U' s! n
5 d2 t# ]9 T' z1 z; s  \+ G/ j圖2：NeuronCore-v2 的詳細(xì)架構(gòu)，展示了張量引擎、向量引擎、標(biāo)量引擎和 GPSIMD 引擎與 HBM 內(nèi)存的集成。
# P' I/ \: |% c  A3 I
7 b4 ?7 j  }. e  }% o; u: n神經(jīng)核心架構(gòu)包含四個(gè)主要引擎：
張量引擎：128×128 脈動(dòng)陣列，用于矩陣運(yùn)算
向量引擎：處理向量運(yùn)算和歸一化
; }- ]0 Y* Z0 K- f  @8 g4 V: e% p標(biāo)量引擎：管理按元素運(yùn)算
5 O& i2 }7 H9 c1 N/ K, c, k5 Y7 @$ HGPSIMD 引擎：執(zhí)行任意 C++ 運(yùn)算

8 o4 \. x4 P3 h服務(wù)器架構(gòu)與部署
! D& U" }( ?7 T1 KTrainium2 提供兩種主要規(guī)格：
8 ^1 ~$ ?, ~( K3 bTrainium2 (Trn2)：每服務(wù)器16個(gè)芯片配置
! g* a- p* x, \" D$ VTrainium2-Ultra (Trn2-Ultra)：跨四臺(tái)服務(wù)器的64個(gè)芯片
/ u( a+ [' `- B% V$ y. O! ~7 L% D

( W% o( O2 ?' E( F! r( B圖3：Trainium2 服務(wù)器架構(gòu)，展示了計(jì)算托盤和 CPU 頭托盤在18機(jī)架單元中的排列。
4 h. V& `) p' [
物理服務(wù)器架構(gòu)占用18個(gè)機(jī)架單元，包括：
1 D* v" W8 c! ^. [  U! _一個(gè) 2U CPU 頭托盤
5 @3 }% V9 Q# t3 L7 Y) e八個(gè) 2U 計(jì)算托盤
每個(gè)計(jì)算托盤包含兩個(gè) Trainium2 芯片
計(jì)算托盤中沒有 CPU（JBOG - "僅GPU組"）

' E. X) o6 h' P! _% V網(wǎng)絡(luò)功能
* S* e' W& i  j# v! @; jTrainium2 的網(wǎng)絡(luò)基礎(chǔ)設(shè)施整合了多項(xiàng)技術(shù)：
NeuronLinkv3（擴(kuò)展上行網(wǎng)絡(luò)）
% O" ]0 x' w: w+ M- m彈性 Fabric 適配器 EFAv3（擴(kuò)展出口網(wǎng)絡(luò)）
- u: W  `2 t. _+ ?. K7 U4 o前端和存儲(chǔ)網(wǎng)絡(luò)
帶外管理網(wǎng)絡(luò)
, f2 I$ `0 e* J  d" @4 x0 V+ k9 P

0 R) g' y9 X6 i' M圖4：4x4x4 三維環(huán)形網(wǎng)絡(luò)拓?fù)�，顯示節(jié)點(diǎn)間連接和帶寬規(guī)格。
8 N: q& x) J$ B! D. c- G1 ^$ j/ C7 K
' @; b7 m- n& f# G供電創(chuàng)新
% e: d+ ]1 B# D' v0 p; t+ x0 dAmazon 在 Trainium2 中實(shí)施了垂直供電，這在芯片供電管理方面是重大進(jìn)展。
) Z1 o! A3 L/ n! R7 t  c
# [( O# D7 N- m+ W圖5：傳統(tǒng)與垂直供電系統(tǒng)的比較，展示了改進(jìn)的供電效率和散熱管理。

3 ?$ u/ s; y. ?9 G$ R7 _2 E成本分析與性能
與 NVIDIA 的 H100 相比，Trainium2 的總擁有成本顯示出有利的經(jīng)濟(jì)性：

圖6：Trainium2 和 H100 配置的詳細(xì)成本比較，突顯了 AWS 解決方案的經(jīng)濟(jì)優(yōu)勢(shì)。
0 g" Z" |$ n. H" N
9 l- d4 L3 {, }% u( m& N$ N+ j主要經(jīng)濟(jì)因素包括：
較低的前期資本成本（每芯片4,000美元，相比H100的23,000美元）
由于更好的能效而降低運(yùn)營(yíng)成本
8 j. @& s' Y& I7 x, G在部署生命周期內(nèi)更有利的總擁有成本

& d4 v( ~2 u$ t* W5 EProject Rainier 實(shí)施
AWS 目前正在為 Anthropic 部署一個(gè)名為"Project Rainier"的大型集群，配備40萬個(gè) Trainium2 芯片。該部署展示了 Trainium2 可以有效運(yùn)行的規(guī)模。
/ H. i9 L4 ^3 n
$ P# I2 I, C8 f6 ~5 J2 S+ G9 B圖7：AWS 印第安納數(shù)據(jù)中心園區(qū)的空中視圖，Project Rainier 正在此處部署。

軟件棧與開發(fā)工具
Trainium2 的軟件生態(tài)系統(tǒng)包括：
/ {: s" A5 M9 q8 o5 A- PNeuronX 集體通信庫(kù)
通過 TorchDynamo 的 PyTorch 集成
# Y. e, y: E" NJAX 測(cè)試版支持
2 P1 Q4 v1 K' k/ ?. [用于低級(jí)優(yōu)化的 Neuron 內(nèi)核語言（NKI）
- w& S& {* u) k/ B  K7 _, ?6 x5 ?, z

$ b6 a+ U( N/ i! U0 u, x3 ~2 z: B圖8：Neuron 分布式事件追蹤界面，顯示全面的調(diào)試和分析功能。

未來發(fā)展
Trainium2 代表了 Amazon 在 AI 加速器領(lǐng)域的重要投入。憑借具有競(jìng)爭(zhēng)力的價(jià)格和性能特征，對(duì) NVIDIA 在 AI 訓(xùn)練市場(chǎng)的主導(dǎo)地位構(gòu)成了顯著挑戰(zhàn)。Project Rainier 的成功和 Anthropic 的采用將是 Trainium2 實(shí)際效果的重要指標(biāo)。

圖9：下一代供電架構(gòu)，展示了 AWS 對(duì)未來可擴(kuò)展性和效率的投入。
' x$ [9 E$ b9 D! P4 {* y
通過對(duì)供電、網(wǎng)絡(luò)功能和軟件集成的嚴(yán)密關(guān)注，AWS 創(chuàng)建了一個(gè)有競(jìng)爭(zhēng)力的平臺(tái)，解決了現(xiàn)代 AI 加速中的許多挑戰(zhàn)。該平臺(tái)的發(fā)展將可能影響 AI 硬件開發(fā)和云服務(wù)供應(yīng)商策略的未來方向。
) `% m9 Z1 h( g: t6 v
. v2 v# ?, D/ @$ f9 {/ T" s" P$ ?參考文獻(xiàn)
2 c8 p6 Q( Q. P[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/
, K/ K- ^9 e1 Z: L/ n
END
; _6 H/ G5 ^! R
/ {$ H5 v4 Q/ T2 H" c
軟件申請(qǐng)我們歡迎化合物/硅基光電子芯片的研究人員和工程師申請(qǐng)?bào)w驗(yàn)免費(fèi)版PIC Studio軟件。無論是研究還是商業(yè)應(yīng)用，PIC Studio都可提升您的工作效能。
點(diǎn)擊左下角"閱讀原文"馬上申請(qǐng)

歡迎轉(zhuǎn)載
& [6 s2 Y0 L7 g' V; q
轉(zhuǎn)載請(qǐng)注明出處，請(qǐng)勿修改內(nèi)容和刪除作者信息！
2 r- v" t& t  I: @4 l5 [

. @2 W# I% D% u. Y. B+ f% W4 @5 ^


' x" j1 t6 B: e2 p關(guān)注我們
( W, @- Q) s- s8 [1 l* O
: |/ Q" D6 p3 ]2 @

0 \3 \: t3 e" B1 o" Z 1 x0 d$ R  N( Z/ o, M

' z3 ?  d; T" L2 @

                     
) w; |2 T& ?$ _1 q$ V; }
  {' |" S2 k. s% D% G% F

1 Y+ X$ @) R" H" G! t  }  [- S' S關(guān)于我們：
深圳逍遙科技有限公司（Latitude Design Automation Inc.）是一家專注于半導(dǎo)體芯片設(shè)計(jì)自動(dòng)化（EDA）的高科技軟件公司。我們自主開發(fā)特色工藝芯片設(shè)計(jì)和仿真軟件，提供成熟的設(shè)計(jì)解決方案如PIC Studio、MEMS Studio和Meta Studio，分別針對(duì)光電芯片、微機(jī)電系統(tǒng)、超透鏡的設(shè)計(jì)與仿真。我們提供特色工藝的半導(dǎo)體芯片集成電路版圖、IP和PDK工程服務(wù)，廣泛服務(wù)于光通訊、光計(jì)算、光量子通信和微納光子器件領(lǐng)域的頭部客戶。逍遙科技與國(guó)內(nèi)外晶圓代工廠及硅光/MEMS中試線合作，推動(dòng)特色工藝半導(dǎo)體產(chǎn)業(yè)鏈發(fā)展，致力于為客戶提供前沿技術(shù)與服務(wù)。
5 Q' ^& }8 Y3 {5 r7 o
4 @$ H. A, ?& P) `$ _6 {$ Thttp://www.latitudeda.com/
0 k9 I9 d: W2 k6 S  ?（點(diǎn)擊上方名片關(guān)注我們，發(fā)現(xiàn)更多精彩內(nèi)容）