AWS Trainium2 架构与网络技术

逍遥设计自动化

引言
  i; ?6 ]9 J' S3 u- YAmazon 在人工智能加速器领域正在通过其 Trainium2 芯片架构取得进展，旨在与 NVIDIA 在人工智能训练和推理领域展开竞争。本文将详细探讨 Trainium2 的架构、网络功能和成本等关键方面[1] 。

4xbsgsoswjt64071050008.png

* Y3 O& r6 E+ W

speir54thtv64071050109.png

2 i" o5 l8 R, r图1：展示 AWS 硬件执行 AI 工作负载的艺术表现，象征着 AWS 对 AI 加速的投入。

核心架构概述
Trainium2 架构相比前代产品有显著提升，每个芯片提供 650 TFLOP/s 的密集 BF16 性能，配备 96GByte HBM3e 内存。每个 Trainium2 芯片包含两个计算 chiplet 和四组 HBM3e 内存，通过 CoWoS-S/R 封装连接。

gcluvrifk1w64071050210.png

图2：NeuronCore-v2 的详细架构，展示了张量引擎、向量引擎、标量引擎和 GPSIMD 引擎与 HBM 内存的集成。
/ X6 k  W" t$ H! o+ F8 U7 A
! R" c2 d5 a2 c! S  d5 D神经核心架构包含四个主要引擎：
! T4 x, U' M; ~/ z0 G张量引擎：128×128 脉动阵列，用于矩阵运算
% o4 Q0 O$ J7 N/ o  D% Y2 B/ o向量引擎：处理向量运算和归一化
标量引擎：管理按元素运算
! c0 V* }0 V+ P6 E2 d* i/ P* c1 _GPSIMD 引擎：执行任意 C++ 运算
) L- \0 G! B7 l2 Q' X4 h; X8 l
$ ^/ }6 C" W9 `& f; N服务器架构与部署
, y" J: D+ O7 @; z2 o2 f  p) Q9 fTrainium2 提供两种主要规格：
3 S* O& F& N+ d' Y& lTrainium2 (Trn2)：每服务器16个芯片配置
! ^; j. h+ S( P' ~Trainium2-Ultra (Trn2-Ultra)：跨四台服务器的64个芯片

" S8 J- c; v+ i; E. ~

lzvbysrl5nc64071050310.png

图3：Trainium2 服务器架构，展示了计算托盘和 CPU 头托盘在18机架单元中的排列。
$ J- T6 C# k6 L  q4 K
  j1 e* n8 W4 Z2 p( z物理服务器架构占用18个机架单元，包括：
& b: w- w4 {: I! ~( e, `, f4 K一个 2U CPU 头托盘
4 D/ T* A+ G, q八个 2U 计算托盘
每个计算托盘包含两个 Trainium2 芯片
计算托盘中没有 CPU（JBOG - "仅GPU组"）

网络功能
Trainium2 的网络基础设施整合了多项技术：
NeuronLinkv3（扩展上行网络）
* a% W2 ~8 F5 X$ e; p弹性 Fabric 适配器 EFAv3（扩展出口网络）
# P* e/ ^+ w7 _7 T9 J( e前端和存储网络
带外管理网络

$ w9 s. t( u9 d! {; I0 m2 T

opmhs2rj1my64071050410.png

图4：4x4x4 三维环形网络拓扑，显示节点间连接和带宽规格。

供电创新
Amazon 在 Trainium2 中实施了垂直供电，这在芯片供电管理方面是重大进展。

jpgwtrojuah64071050511.png

1 w$ h" j. ]( O' r图5：传统与垂直供电系统的比较，展示了改进的供电效率和散热管理。

成本分析与性能
3 m5 D! J; k+ M9 Z# {# K2 f0 Q与 NVIDIA 的 H100 相比，Trainium2 的总拥有成本显示出有利的经济性：

pj03qewwrbu64071050611.png

图6：Trainium2 和 H100 配置的详细成本比较，突显了 AWS 解决方案的经济优势。
4 d( ~% r2 H2 _
7 f, z3 y; G. g! @$ j: U主要经济因素包括：
较低的前期资本成本（每芯片4,000美元，相比H100的23,000美元）
由于更好的能效而降低运营成本
在部署生命周期内更有利的总拥有成本

- ^3 L7 ?. r7 e6 iProject Rainier 实施
AWS 目前正在为 Anthropic 部署一个名为"Project Rainier"的大型集群，配备40万个 Trainium2 芯片。该部署展示了 Trainium2 可以有效运行的规模。

hv5dm52fngf64071050712.png

# n/ i# U8 w3 [+ ]图7：AWS 印第安纳数据中心园区的空中视图，Project Rainier 正在此处部署。
5 k, `2 A; c& Q: f# z% J0 r3 g
* s0 M* ^! f$ {8 I软件栈与开发工具
Trainium2 的软件生态系统包括：
% v) G. _- i0 ~* i) E6 u, N( rNeuronX 集体通信库
通过 TorchDynamo 的 PyTorch 集成
JAX 测试版支持
6 n6 }( ^* ^& x! E5 g. [用于低级优化的 Neuron 内核语言（NKI）
) B3 V2 ^3 B2 t) p' `1 a9 `

gyfcekujxgb64071050812.png

# x6 ]/ O, X- O: E$ Q- J图8：Neuron 分布式事件追踪界面，显示全面的调试和分析功能。
1 l0 H3 I2 }8 @, c2 z( g  f- B
未来发展
* m# `' D9 Q/ a, @0 n) v/ KTrainium2 代表了 Amazon 在 AI 加速器领域的重要投入。凭借具有竞争力的价格和性能特征，对 NVIDIA 在 AI 训练市场的主导地位构成了显著挑战。Project Rainier 的成功和 Anthropic 的采用将是 Trainium2 实际效果的重要指标。

gerd3esh4cn64071050913.png

4 \/ o7 ?8 g9 {- |. ^9 r9 o图9：下一代供电架构，展示了 AWS 对未来可扩展性和效率的投入。
; {/ ^: t  k5 ~
% u/ K+ p9 z" f9 e1 Z通过对供电、网络功能和软件集成的严密关注，AWS 创建了一个有竞争力的平台，解决了现代 AI 加速中的许多挑战。该平台的发展将可能影响 AI 硬件开发和云服务供应商策略的未来方向。
' I/ o1 r( [4 M) b, p
' J- O2 n: m3 v; x3 y2 @参考文献
[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/
% o. n0 E: e* h! M. V- Y
. n" {0 `: u$ N+ p3 V3 N. aEND
0 \5 a0 P9 i* j1 L7 O$ ]4 N' Q& U

5 k% u7 Z/ H: q/ E8 [9 c% t软件申请我们欢迎化合物/硅基光电子芯片的研究人员和工程师申请体验免费版PIC Studio软件。无论是研究还是商业应用，PIC Studio都可提升您的工作效能。
  D( W! g1 O. c6 s6 e点击左下角"阅读原文"马上申请
' |. c; ^; V% E; o
欢迎转载
; [# m4 r; ^1 f2 l. e( {+ M+ _
转载请注明出处，请勿修改内容和删除作者信息！
8 `% s  i1 C! s

- i/ h" b8 F1 T  I2 b* ]% l: n

nlj50ijxcp464071051013.gif

4 F  ]9 Z/ q3 U7 V' W7 a" \2 a/ B
3 k2 D3 d7 @& b$ ?6 u0 l; G关注我们
1 J' R. D, }) s8 r3 j/ P  u* m

# O5 D+ p/ ?. [, a9 W 1g21mosmfw364071051113.png   B: Z$ k; J- q+ @, B7 `

c3 ^7 m4 c$ J  l exohpa2ha4m64071051213.png . ~$ S4 g* g+ K

w0egxzsvpbt64071051313.png ; X. O- ]  M: G  y5 j

                     
& I: u8 a. g2 E7 l. [8 s: s
% c0 x! S3 C$ v+ F" `

3 N5 M8 f! I$ P关于我们：
深圳逍遥科技有限公司（Latitude Design Automation Inc.）是一家专注于半导体芯片设计自动化（EDA）的高科技软件公司。我们自主开发特色工艺芯片设计和仿真软件，提供成熟的设计解决方案如PIC Studio、MEMS Studio和Meta Studio，分别针对光电芯片、微机电系统、超透镜的设计与仿真。我们提供特色工艺的半导体芯片集成电路版图、IP和PDK工程服务，广泛服务于光通讯、光计算、光量子通信和微纳光子器件领域的头部客户。逍遥科技与国内外晶圆代工厂及硅光/MEMS中试线合作，推动特色工艺半导体产业链发展，致力于为客户提供前沿技术与服务。
7 e$ q. V& L1 f; I  \/ X3 w# U
http://www.latitudeda.com/
0 |/ Z1 M8 d8 r0 l$ A- v（点击上方名片关注我们，发现更多精彩内容）