亚马逊网络服务并非第一家自研计算引擎的超大规模云厂商，但步伐紧随谷歌，后者早在2015年便通过自研TPU加&速器承载AI工作负载。

AWS自2017年推出“Nitro”DPU后，便下定决心要在计算引擎的各个维度全面自研，以便在服务器基础设施创新上不受制于人。如今，我们已经看到了多代Nitro DPU、四代Gra viton Arm服务器CPU、两代Inferentia AI推理加&速器，以及刚刚亮相的第二代Trainium AI训练加&速器。在拉斯维加斯举办的re:Invent 2023大会上，Trainium2芯片与Gra viton4服务器CPU一同登场，我们花了些功夫去摸清这个AI训练新引擎的底细，以及它在Inferentia系列中的定位。

AWS目前并未公布太多关于这些AI计算引擎的详细规格。好在我们与AWS Annapurna Labs部门的业务开发高级总监Gadi Hutt聊了聊，他正好主管计算引擎的设计、代工与落地。通过深入了解AWS系统，结合对技术文档的细致挖掘，我们试图梳理出Inferentia与Trainium之间的关系，特别是Trainium2的预期表现——在信息有限的条件下，尽量填补空白。

在进入正题之前，不妨先做点数学功课，了解一下AWS AI计算引擎的数据来源与速度背景。

在AWS首席执行官Adam Selipsky的re:Invent主题演讲中，Nvidia联合创始人兼CEO黄仁勋作为惊喜嘉宾登场，透露了一个关键数字：在“Ampere”A100和“Hopper”H100时代，AWS累计采购了200万颗这类芯片。

有传言称，AWS在2023年交付了约5万颗H100，我们推测2024年还有约2万颗订单。按每颗3万美元计算——面对供不应求的局面，Nvidia几乎不打折，最近几个季度的净利润率远超数据中心收入的50%——这笔订单价值约21亿美元。剩下193万颗A100，按2020年至今约1.25万美元的平均售价估算，总计约2413亿美元。

面对如此庞大的投资支出，性价比曲线显然还有优化的空间。AWS自研Titan模型，既要支撑母公司Amazon及其数万企业客户，也要让包括Anthropic的Claude 2在内的其他模型运行在自研的Inferentia和Trainium之上。

我们认为，这条曲线与AWS借助Gra viton服务器CPU走过的路径不会相差太多。AWS依然乐意销售Intel和AMD的CPU，但自研方案的性价比往往高出“传统芯片”30%到40%。通过Gra viton甩开中间商，AWS能以更低价格提供Arm实例，这对越来越多客户产生吸引力。可以预见，Nvidia、AMD GPU与AWS自研的Inferentia、Trainium之间，也将保持同样的传统定价差距。

撇开这些数学题，我们先看看Inferentia1，再概览Inferentia2和Trainium1，这样才能理清Trainium2的来龙去脉——它正与当前Hopper H100加&速器上的工作负载正面竞争。你懂的，H100贵得几乎像黄金（目前SXM5版本每盎司价格约为黄金的一半），像稀土一样难求，却支撑着AI经济的半边天。

推断TRAINIUM的架构

所有计算引擎本质上都是计算单元、存储单元和网络互联的层级结构，不同架构的抽象层可能差异显著，尤其是在架构新颖、工作负载快速变化的领域。

Inferentia1芯片由Annapurna Labs工程师打造，2018年11月首次发布，一年后全面上市，成为AWS AI引擎的基础。以下就是它的架构图：

该芯片包含四个核心，每个核心有四个不同的计算单元，搭配片上SRAM存储器和片外DDR4主存，这与当今多数AI芯片厂商的做法相似。AWS没有公布Inferentia1的SRAM容量、缓存大小、时钟速度，甚至每个NeuronCore内部单元的数量。不过，在Inferentia和Trainium芯片的Neuron SDK中，AWS确实提到了Inferentia2和Trainium1使用的NeuronCore-V2核心架构——我们正好以此为基础反推Inferentia1，再推测Trainium2的可能形态。

无论哪一代，NeuronCore都包含一个处理标量计算的ScalarEngine，以及一个处理多种精度整数与浮点数据向量计算的VectorEngine——大致相当于Nvidia GPU中的CUDA核心。根据Neuron SDK，NeuronCore-v1的ScalarEngine每周期处理512个浮点运算，VectorEngine每周期处理256个浮点运算（我们理解，这里是指ScalarEngine每周期512位、VectorEngine每周期256位，然后以特定格式泵送数据完成各自的计算）。

NeuronCore架构还包含一个用于加&速矩阵运算的TensorEngine，这超越了VectorEngine推动代数矩阵的范畴。矩阵运算对HPC和AI工作负载至关重要，往往承担大量计算并决定最大吞吐量。TensorEngine大致对应Nvidia GPU中的TensorCore，在NeuronCore-v1核心中，FP16/BF16精度下可提供16 teraflops。

Trainium1芯片于2020年12月发布，以Trn1和Trn1n两个实例发货。我们当时已对其做过详尽分析，坦率讲，AWS并未提供太多关于自研AI计算引擎的公开数据。Trainium1采用了新的NeuronCore-v2核心，结构元素相同，但核心数量更少，如下图所示：

Trainium1芯片增加了32 GB HBM堆叠DRAM内存以提升带宽，转向PCI-Express 5.0外形与I/O插槽，并增加了NeuronLink芯片间互连的链路数量，带宽提升了2至4倍。

虽然没有直接证据，但我们推测，Trainium1将每颗芯片的NeuronCore数量比Inferentia1减少了一半（从4个降至2个），然后将每个核心内的标量、矢量和张量引擎数量翻了一番。实际的变化是缓存与内存层次的抽象化，使得NeuronCore能在每种计算单元类型中实现多线程。

更有趣的是，Trainium芯片首用的NeuronCore-v2还包含称为GPSIMD引擎的单元——一组8个512位宽的通用处理器（这确实非常有意思）。这些设备可以直接用C和C++编程，访问片上SRAM和内核上的其他三种引擎，适合实现其他引擎格式无法直接支持的自定义操作。（我们曾查阅Flynn分类法试图理清GPSIMD引擎的定位，但从文档来看，仍不明确它属于阵列处理器、流水线处理器还是关联处理器。）

用于推理的Inferentia2芯片本质上是半配的Trainium1——NeuronLink-v2互连端口只有一半，部分芯片单元可能未激活，时钟速度和性能大致相当，HBM内存容量与带宽也基本一致。

“Inferentia2和Trainium1的芯片架构几乎相同，”Hutt告诉The Next Platform。“我们为Inferentia2保留了HBM带宽，这对推理至关重要，而不只是面向训练。LLM推理实际受内存限制，而非计算限制。因此，我们采用类似的硅架构并尽可能降低成本——比如不需要那么多NeuronLink。推理场景中，从一个加&速器到另一个时，只需环状链路生成token。而训练则需要完整的网状连接，以最大限度减少服务器内部各加&速器之间的跳数。当然，访问训练服务器时还需要大量网络带宽。”

使用16个Trainium1芯片的服务器互连采用2D环面或3D超立方体——按Hutt的说法，两者是同一事物，只是表述不同。示意图如下：

下表汇集了我们对Inferentia1、Inferentia2和Trainium1已知的馈送与速度信息，以及我们对Trainium2的预测（以粗体红色斜体标注）。Neuron SDK目前尚未更新Trainium2的相关数据。

根据 Trainium2 芯片封装图片，我们推测它本质上是两个 Trainium1 芯片互连而成——要么是单片形式，要么是通过某种高速互连连接的两个小芯片插槽。

Hutt 没有透露 Trainium 芯片的具体数据与速度，但确认 Trainium2 拥有更多核心和更高内存带宽。他补充说，芯片有效性能将提升至 4 倍——并称这是个保守数字，可能已超越现实世界 AI 训练工作负载对 Trainium1 的预期。原因在于这些工作负载更多受内存容量与带宽限制，而非计算上限。

我们认为 Trainium2 芯片将有 32 个核心，制程从 Trainium1 的 7nm 缩至 4nm，这样核心数量翻倍的同时，功耗维持在 Trainium1 相同或略高水平。我们还推测，时钟速度将从 Trainium1 的 3 GHz（AWS 已公开）适度提升到 3.4 GHz。同时，基于直觉判断，Trainium2 的总 NeuronLink 带宽将增加 33%，达到每端口 256 GB/秒，总带宽 2 TB/秒，并且仍支持 2D 环面互连。每个芯片上的 NeuronLink 端口数量可能会增加，以扩展环面维度、减少数据共享时设备间的跳数。2D 环面意味着集群中任意两个 NeuronCore 芯片之间有固定的两跳。网络似乎不太可能扩展到全面配置，但这并非没有可能——SGI 曾在其超级计算机芯片组中做到过这一点。

鉴于 AWS 希望通过 Trainium2 将 UltraCluster 扩展到 10 万个设备，我们猜测实际的核心数量将少于上表所示的 64 个。

具体数字很难确定，假设约有 10% 的核心可能因良率问题不可用，但因此芯片的整体良率会更高。可以预见，AWS 会将那些所有核心均能正常运行的 Trainium2 芯片保留下来——很可能用于每个核心都至关重要的内部工作负载。由此，Trainium2 可能拥有 56 个或 58 个核心，甚至接近 56 个，且所有这些核心的带宽均可充分利用。AWS 承诺的 96 GB 可能只对应设备上四个内存堆栈中的三个，实际 HBM 内存可能达到 128 GB。我们强烈怀疑这将是 HBM3 内存，但 Hutt 并未证实。

不过他多次强调，性能是由内存驱动的，而不是指望峰值理论算力快于内存带宽的增长。如果我们的判断正确，Trainium1 到 Trainium2 之间，内存带宽将增长 5 倍。

以下是使用 Inferentia 和 Trainium 芯片可用的实例：

目前 Trn2 实例在价格或性能方面的具体比较仍是未知数。但根据暗示与直觉，我们坚信 Trainium2 将提供约 2 倍于 Nvidia H100 的性能，这意味着在许多模型中它可以与 Nvidia 刚刚发布的、配备更大更快 HBM3e 内存的 H200 旗鼓相当。当我们建议 AWS 可能将基于 Trainium2 的 EC2 实例定价在类似其 Gra viton CPU 与 AMD、Intel X86 处理器之间的性价比差距——即 30% 到 40% 之间——Hutt 并没有泼冷水。他虽未做出任何承诺，仅表示性价比肯定会更好。

这并非空谈——当 10 万个设备以 FP16 精度连接在一起，在没有任何稀疏技巧的情况下，以真正的 FP16 分辨率达到 65 exaflops，这完全有潜力成为世界上最大的人工智能集群。