有了Trainium4,AWS将把除时钟外的所有功能都调高

全球的人工智能模型开发者已经等待了一年多,终于能拿到Trainium3 XPU,这些XPU专门为训练和推理设计,是英伟达“Blackwell”B200和B300 GPU以及谷歌“Trillium” TPU v6e和“Ironwood” TPU v7p加速器的可信替代品。
但当亚马逊云服务首席执行官Matt Garmin开始谈论预计可能在2026年底或2027年初交付的未来Tranium4 XPU时,所有排队购买基于Trainium3的EC2容量模块的人都做好了买家后悔的准备。因为尽管Trainium3相比公司Annapurna Labs芯片设计部门之前的Inferentia和Trainium XPU表现优异,Trainium4看起来将彻底突破插槽大小,不仅是一款非常强大的设备,还将拥有更具可扩展性的UltraServer集群,更适合运行多种专家级、思维链推理模型的混合。
在深入探讨Trainium4可能是什么之前,让我们先花点时间真正回顾一下Trainium3的XPU,尤其是因为该芯片及其前身Tranium2的许多技术规格直到最近才公开。让我们先从Garmin在本周拉斯维加斯re:Invent 2025大会开幕式主题演讲中谈到的Tranium3规格表开始:

这是AWS去年此时展示的幻灯片的更新,预览了它愿意谈论的三个数据点。Tranium3采用了台湾半导体制造股份有限公司的3纳米工艺节点,这是一种对5纳米技术的缩小,而大多数人认为Annapurna实验室用于Tranium2芯片时正是如此。Tranium3预计能实现2倍的计算能力(这可能意味着很多不同的事情),并且能效提升40%(但这并不是一个有价值的指标,因为没人知道Tranium1、Trainium2或Trainium3的功耗是多少)。但显然,工艺缩小更多是为了削减功耗,而不是把新功能塞进芯片,而且插槽被做得更大以提升性能,每瓦净性能提升了40%。
亚马逊销售基于Tranium2的常规服务器实例,以及拥有16个Trainium3套接字的UltraServer集群配置,但目前仅交付了单一内存域中64个Tranium3的UltraServers。

相比之下,Trn3 UltraServers在AWS上调用实例时计算能力提升了4.4倍,这与拥有四倍XPU数量的Trn2 UltraServers相比是合理的。据Garmin称,最新的UltraServer拥有Trn2 UltraServer的3.9倍总HBM内存带宽,更重要的是,对于那些担心推理成本——这是生成式AI商业化的门槛因素——的人来说,每兆瓦能生成5倍的代币数量。
以下是Garmin为该性能声明分享的帕累托曲线,展示了每兆瓦输出代币在Y轴上的相互作用,输出的交互性以每位用户每秒代币数表示:

将这个曲线向上推移,是2025年及以后获得推理业务的全部策略。这组图表比较了Trn2 UltraServer集群与运行OpenAI的GPT-OSS 120B模型的Trn3 UltraServer集群。
这张图表还显示了Garmin未提及的,如果对推理工作量重要,你可以在相同能量下获得大约一个数量级更多的交互性。
在 Trainium2 实例云端扩展过程中,AWS 更新了该 XPU 的规格,我们也找到了一些 Tranium3 的规格,消除了一些谜团,填补了关于组件如何堆叠在 Trainium 套接字中以形成每个后续 XPU 的空白。
我们先从神经核开始,逐步推理。
所有NeuronCore设计都将四种不同的计算方式放入核心,就像CPU核心早已混合整数(标量)和向量单元,偶尔(Intel Xeon 5和6以及IBM Power10和Power11)也包含张量单元一样。从Trainium系列开始,Annapurna Labs在架构中增加了集体通信核心(CC-Cores),以处理高性能计算和人工智能工作负载中常见的集体作的特殊处理,这样实际上已经有五个核心了。
在仅用于Inferentia1芯片的NeuronCore-v1架构中,有用于整数计算的标量引擎(两个整数输入和一个整数输出)、用于向量数学的向量引擎(两个浮点输入,一个浮点输出)、一个用于张量数学的张量引擎(多个矩阵浮点输入和一个矩阵浮点输出)。
根据AWS文档,NeuronCore-v1中的标量引擎每个时钟周期可处理512次浮点运算,并可处理FP16、BF16、FP32、INT8、INT16和INT32数据类型。(我们认为AWS的意思是它处理512位数据。)文档还说向量引擎每个周期能处理256次浮点运算(我们认为是256位数据),并且支持FP16、BF16、FP32、INT8、INT16和INT32数据格式。你可以根据数据宽度和每个单元能装入多少个周期来计算每个周期的作。
TensorEngine for NeuronCore-v1 的维度从未公开,但我们知道它处理了 FP16、BF16 和 INT8 输入,以及 FP32 和 INT32 输出,并提供了 16 万亿次 FP16 或 BF16 张量处理。
在讨论了第一个NeuronCore-v1设计后,让我们把它们并排摆放,看看我们认为Trainium4可能的方向:

AWS于2020年12月在re:Invent上开始谈论Trainium1,并花了两年时间才全面推广,这可以理解,因为这是亚马逊首个自家开发的数据中心级培训加速器。我们认为Trainium1是通过台积电7纳米工艺蚀刻的;我们知道它有550亿个晶体管,频率为3 GHz。该芯片采用了与后续于 2023 年 4 月上市的 Inferentia2 相同的 NeuronCore-v2 架构,工艺规模缩小至 5 纳米,晶体管数量大致相同,但针对推理特定工作负载做了一些调整,比如将 NeuronLink 芯片互连端口减半。
随着 Trainium2 于2023年11月发布,2024年12月批量发货,AWS 转向 NeuronCore-v3 架构,停止生产 Inferentia 芯片,因为推理开始变得更像训练。Trainium2使每个套接字核心数增加了四倍,单个内存域中的神经核总数增加了16倍,每个实例的套接字数也增加了四倍。据我们所知,AWS还通过将Trainium2的频率从7纳米缩小到5纳米,提升了Trainium2的频率。有趣的是,每个NeuronCore的峰值标量和矢量性能在v3下降了约60%,峰值张量核心吞吐量下降了12%。但AWS为芯片增加了1:4稀疏度支持以支持张量运算,加上更高的核心数量,使Trainium2在FP16或BF16精度下的有效吞吐量提升了3.5倍。事实上,NeuronCore-v3支持多种稀疏度模式:4:16、4:12、4:8、2:8、2:4、1:4和1:2。
NeuronCore-v3的SRAM内存(由三个计算单元共享)提升至每核28 MB,但具体提升的幅度尚不清楚。HBM内存最终提升至96 GB,提升了3倍,带宽提升3.5倍至2.9 TB/秒。这可以说是首个竞争性Trainium芯片,Anthropic使用这些Trainium2设备进行模型开发和推断,且AWS基岩模型服务的大部分推断均由Trainium完成,这绝非巧合。我们怀疑Garmin在主题演讲中提到的数百万台Trainium大多是Trainium2设备。

这就引出了Trainium3,它现在已经在UltraServer实例中批量发布了。Trainium3 设备的核心是 NeuronCore-v4 架构——是的,如果核心名称与设备名称同步会更好——有一些重大变化。首先,矢量引擎经过调整,能够进行快速的指数函数评估,其性能是标量引擎的4倍,而标量引擎是生成式AI模型自注意算法的一部分。其次,FP16和BF16数据格式可以量化成MXFP8格式,AWS表示这对生成式AI模型中多层感知器(MLP)层之间的数据量化非常有用。NeuronCore-v3 设计还将 SRAM 提升至 32 MB 的每核。Trainium2 和 Trainium3 之间的时钟频率似乎名义上有所变化,但变化不大,设备最大的变化是 NeuronLink-v4 XPU 互连端口带宽翻倍至 2.5 TB/sec,HBM 内存容量提升 1.5 倍至 144 GB,HBM 带宽提升 1.7 倍至 4.9 TB/秒。
我们认为 Trainium3 设计的变更旨在让计算、内存和互联重新达到更好的平衡,从而提升 Trainium3 套接字的理论峰值性能,而非理论上的峰值性能。Trn3 Gen1 UltraServer的内存域保持在64台设备,与Trainium2相同,但随着Trn3 Gen2 UltraServer的发货,域大小增加到144个套接字。这会使AI训练或推理任务中可投入的核心数量增加2.25倍。
这就引出了预计明年这个时候开始推出的Trainium4。
我们推测将采用的 NeuronCore-v5 架构,AWS 将为 Trainium 处理添加适当的 FP4 支持,而不仅仅是把 MXP4 塞进 FP8 的张量槽,留下大量空白。Garmin在主旨演讲中表示,采用FP4原生格式后,Tranium4的性能将是Tranium3的6倍,这意味着FP8的处理能力将提升3倍。Garmin进一步表示,Trainium4的HBM内存容量是Trainium4的2倍,HBM带宽是Trainium4的4倍。
在上面的怪兽表中,我们尝试推测了Trainium4可能的样子,以及如何为耦合的Trainium4设备进一步扩展内存域。
实现这一目标的方式有很多种,我们认为 AWS 充其量会转向 2 纳米工艺以节省一些功耗,或者坚持 3 纳米工艺,节省一些成本,制造更大更热的 XPU。这是一个艰难的决定,但我们认为Trainium4会倾向于采用2纳米蚀刻技术。
如果你看上面Garmin的图表,你不会发现它说Tranium4会同时支持NVLink和UALink接口——Nvidia曾大肆宣传AWS将采用NVLink技术,但我们有预感,AWS会推出带有NVlink接口的Graviton系列芯片变体,并获得Nvidia一直不愿透露的内容: 它能够将定制CPU和定制XPU粘合到一个巨大的共享内存域中,使用NVLink端口和NVSwitch内存织物交换机。到目前为止,英伟达乐于让客户拥有连接英伟达GPU的定制CPU或连接英伟达GPU的定制XPU,但并未允许第三种选项。
我们认为AWS购买的GPU数量足够多,能够以合理的价格提出并获得这样的需求。我们还认为AWS将支持Nvidia的NVFP4数据格式以及用于FP4处理的MXFP4格式,这很可能是交换的一部分,以便将Tranium4芯片的工作转移到Nvidia的“Blackwell”和“Rubin”GPU上。当然,这些只是直觉。我们也认为AWS希望能够将这些功能接入自家机架,这些机架本质上将是英伟达机架的克隆。
但有趣的是,UALink也出现在上面的幻灯片里,来自Garmin。AWS保持开放选择,毫无疑问希望为Tranium4封装设计芯片组架构,使其能够将NVLink端口替换为UALink端口,并采用机架设计,允许NVSwitch交换机在明年晚些时候上市时替换为UALink交换机。有可能 NeuronLink-v5 经过调整以兼容 UALink 2.0,这些交换机将来自 Annapurna Labs,而非 Astera Labs、Upscale AI、Marvell、Cisco Systems 或其他提供可扩展互连 ASIC 的厂商。
在相同或稍小的热包内实现3倍性能的最简单方法是将核心数量增加三倍,同时保持时钟频率不变,同时采用2纳米工艺。如果晶体管收缩较大(台积电的1.6纳米A16工艺),那么热量可以降低一点,或者把时钟调高一点点。我们的建议是利用散热优势保持不变,就像AWS在Trainium2和Trainium3之间所做的那样,只增加3倍的核心。
如果你把核心数增加3倍,每个插槽24个核心,可能分布在四个芯片组上,那么在恒定精度下能达到3倍;如果从FP8缩小到FP4,每个插槽的效能提升6倍。
接下来事情变得有趣了。如果你将设备数量翻倍到每台288个(与英伟达的做法一致),Trainium4 UltraServer集群中可以拥有6912个Neuro核心,全部集中在一个内存域中,配备1944TB的HBM内存。
当然,这远不及谷歌在单一内存域中能提供的9,612个Ironwood TPU v7p XPU的数量......但它比目前销售的Trn2 Gen2 UltraServer集群好13.5倍。










评论