AWS Trainium2 架构与网络技术

逍遥设计自动化

引言
+ i4 O9 t) A, M. VAmazon 在人工智能加速器领域正在通过其 Trainium2 芯片架构取得进展，旨在与 NVIDIA 在人工智能训练和推理领域展开竞争。本文将详细探讨 Trainium2 的架构、网络功能和成本等关键方面[1] 。
9 K! k3 E5 S, `

4xbsgsoswjt64071050008.png

8 |: F8 f8 O9 E' `3 `
6 X2 x" k8 z8 h8 W* Z* b

speir54thtv64071050109.png

图1：展示 AWS 硬件执行 AI 工作负载的艺术表现，象征着 AWS 对 AI 加速的投入。
0 h  e+ s3 l6 K0 i4 h! t2 B6 c
核心架构概述
# `) f7 }- G: r( pTrainium2 架构相比前代产品有显著提升，每个芯片提供 650 TFLOP/s 的密集 BF16 性能，配备 96GByte HBM3e 内存。每个 Trainium2 芯片包含两个计算 chiplet 和四组 HBM3e 内存，通过 CoWoS-S/R 封装连接。
$ j* t1 i9 U9 H! {" W8 z

gcluvrifk1w64071050210.png

; h% b# Z. G( z: J图2：NeuronCore-v2 的详细架构，展示了张量引擎、向量引擎、标量引擎和 GPSIMD 引擎与 HBM 内存的集成。
1 _: I. H7 K7 }' T0 N- c
5 j' X0 K- F- b' Z  @神经核心架构包含四个主要引擎：
% b, ?& O2 n  T' D( a' @( M3 G张量引擎：128×128 脉动阵列，用于矩阵运算
向量引擎：处理向量运算和归一化
: R% l& T$ W: g' E" X, {/ {标量引擎：管理按元素运算
GPSIMD 引擎：执行任意 C++ 运算
1 d' g  J$ B$ j& C
服务器架构与部署
Trainium2 提供两种主要规格：
Trainium2 (Trn2)：每服务器16个芯片配置
Trainium2-Ultra (Trn2-Ultra)：跨四台服务器的64个芯片
6 F: i5 @! Q4 w! `
/ E$ u6 y" U$ i- j- A6 o# p

lzvbysrl5nc64071050310.png

% B& O+ y- o* {% G) {" u图3：Trainium2 服务器架构，展示了计算托盘和 CPU 头托盘在18机架单元中的排列。

& W# w: I) t& E6 \; K! W; X物理服务器架构占用18个机架单元，包括：
一个 2U CPU 头托盘
八个 2U 计算托盘
9 e4 n3 x' K  m* ~每个计算托盘包含两个 Trainium2 芯片
. N1 ?" Y! f, a1 {) z- W- @! L计算托盘中没有 CPU（JBOG - "仅GPU组"）

: G  j) L3 ~- t, r网络功能
: U: R; @" s" u. B' n( b3 |Trainium2 的网络基础设施整合了多项技术：
NeuronLinkv3（扩展上行网络）
) i# i. q0 _$ X3 G弹性 Fabric 适配器 EFAv3（扩展出口网络）
前端和存储网络
带外管理网络
; G: r9 T4 ^  V7 ?, z. d0 Q
! H% f3 A* Z% z- W  w

opmhs2rj1my64071050410.png

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

供电创新
Amazon 在 Trainium2 中实施了垂直供电，这在芯片供电管理方面是重大进展。
  D# W: K, Z; [1 o9 n5 h

jpgwtrojuah64071050511.png

图5：传统与垂直供电系统的比较，展示了改进的供电效率和散热管理。

成本分析与性能
! l1 H  m* R5 y9 r9 o# n与 NVIDIA 的 H100 相比，Trainium2 的总拥有成本显示出有利的经济性：

pj03qewwrbu64071050611.png

* f3 `+ H+ D7 G$ |# h$ R- v% ^图6：Trainium2 和 H100 配置的详细成本比较，突显了 AWS 解决方案的经济优势。
3 O9 s4 R. P5 B; P. K8 p# V  L/ O
: r  y- ~' j. x9 H; b$ U1 J: I. }% [主要经济因素包括：
较低的前期资本成本（每芯片4,000美元，相比H100的23,000美元）
由于更好的能效而降低运营成本
在部署生命周期内更有利的总拥有成本

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

hv5dm52fngf64071050712.png

图7：AWS 印第安纳数据中心园区的空中视图，Project Rainier 正在此处部署。

软件栈与开发工具
Trainium2 的软件生态系统包括：
NeuronX 集体通信库
+ S8 D1 p/ D! ]- I) g0 I0 {. _# M通过 TorchDynamo 的 PyTorch 集成
JAX 测试版支持
& s$ [; |! y: S1 o1 q. m) p8 }+ T( z用于低级优化的 Neuron 内核语言（NKI）
0 V1 V. ~' y3 T; F9 C2 i3 E0 e  y# _
+ f6 u: N0 [5 c7 ~8 k6 y, F/ @- N

gyfcekujxgb64071050812.png

图8：Neuron 分布式事件追踪界面，显示全面的调试和分析功能。

未来发展
; F% \* v  T- UTrainium2 代表了 Amazon 在 AI 加速器领域的重要投入。凭借具有竞争力的价格和性能特征，对 NVIDIA 在 AI 训练市场的主导地位构成了显著挑战。Project Rainier 的成功和 Anthropic 的采用将是 Trainium2 实际效果的重要指标。
4 c1 t) r" t  T( m, `! C

gerd3esh4cn64071050913.png

图9：下一代供电架构，展示了 AWS 对未来可扩展性和效率的投入。

通过对供电、网络功能和软件集成的严密关注，AWS 创建了一个有竞争力的平台，解决了现代 AI 加速中的许多挑战。该平台的发展将可能影响 AI 硬件开发和云服务供应商策略的未来方向。

参考文献
[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/
  \5 U2 s: Z0 X% A8 D
2 i/ B. a; o* yEND


软件申请我们欢迎化合物/硅基光电子芯片的研究人员和工程师申请体验免费版PIC Studio软件。无论是研究还是商业应用，PIC Studio都可提升您的工作效能。
, V8 E# G( f& n& B- c% Y点击左下角"阅读原文"马上申请

欢迎转载
! P$ |0 \, ?7 v5 {
: V" F, o9 v  P, x2 d6 p3 F转载请注明出处，请勿修改内容和删除作者信息！
1 `; t3 f8 x/ I' k( w
. b' [( \) m- E

nlj50ijxcp464071051013.gif

5 ~: ]1 l1 w" W9 D9 C% X& V, F( v
, [' W+ Q; h" A! c关注我们
: `- a! s3 K5 r, r1 p
6 p# D' z& }( t/ ]% _% P

: G4 Y' M- S% K5 G& K2 H2 q 1g21mosmfw364071051113.png

' j5 Y. O2 X3 |# a exohpa2ha4m64071051213.png ( G, Y' R* g7 A8 F4 P  \

1 m% ]' W# E! a w0egxzsvpbt64071051313.png

                     


- x& C6 w& ~! F- p
# S8 k4 M2 n4 Y" G# C5 Z: O关于我们：
% A' v. A- l0 o; u! m0 ^深圳逍遥科技有限公司（Latitude Design Automation Inc.）是一家专注于半导体芯片设计自动化（EDA）的高科技软件公司。我们自主开发特色工艺芯片设计和仿真软件，提供成熟的设计解决方案如PIC Studio、MEMS Studio和Meta Studio，分别针对光电芯片、微机电系统、超透镜的设计与仿真。我们提供特色工艺的半导体芯片集成电路版图、IP和PDK工程服务，广泛服务于光通讯、光计算、光量子通信和微纳光子器件领域的头部客户。逍遥科技与国内外晶圆代工厂及硅光/MEMS中试线合作，推动特色工艺半导体产业链发展，致力于为客户提供前沿技术与服务。

/ [" X3 b' H7 ?/ B/ x( H0 j8 Mhttp://www.latitudeda.com/
（点击上方名片关注我们，发现更多精彩内容）