亚马逊推出Trainium3人工智能加速器，在FP8性能上直接与布莱克威尔Ultra竞争——新的Trn3 Gen2 UltraServer从英伟达的剧本中获取垂直扩容笔记

亚马逊网络服务本周推出了用于AI训练和推理的下一代Trainium3加速器。正如AWS所说，新处理器的速度是其前身的两倍，效率提高了四倍。这使其成为AI训练和推理成本方面的最佳解决方案之一。从绝对数量来看，Trainium3提供了高达2517个MXFP8 TFLOPS，与英伟达的Blackwell Ultra相比，这一数字低了近两倍。然而，AWS的Trn3 UltraServer每个机架打包144个Trainium3芯片，并提供0.36 ExaFLOPS的FP8性能，因此与英伟达的NVL72 GB300的性能相匹配。这是一件非常大的事情，因为很少有公司能够挑战英伟达的机架规模AIsystems.AWSTrainium3 AWS Trainium3是一款双芯片AI加速器，配备了144 GB的HBM3E内存，使用四个堆栈，可提供高达4.9 TB/s的峰值内存带宽。据称，每个计算芯片由台积电使用其3nm级制造工艺制造，包含四个NeuronCore-v4内核（与前代相比，具有扩展的ISA）并连接两个HBM3E内存堆栈。这两个芯片使用专有的高带宽接口连接，并共享128个独立的硬件数据移动引擎（这是Traite架构的关键）、协调芯片之间流量的集体通信内核和四个用于横向扩展连接的NeuronLink v4接口。NeuronCore-v4集成了四个执行块：张量引擎、向量引擎、标量引擎、GPSIMD块和32 MB本地SRAM，由编译器显式管理，而不是缓存控制。从软件开发的角度来看，内核是围绕软件定义的数据流模型构建的，在该模型中，数据由DMA引擎暂存到SRAM中，由执行单元处理，然后写回，因为近内存积累使DMA能够在单个事务中执行读-加-写操作。SRAM在内核之间不一致，用于平铺、暂存和积累，而不是一般缓存。（图片来源：AWS）（图片来源：AWS）张量引擎是用于GEMM、卷积、转置和点积操作的收缩式矩阵处理器，支持MXFP4、MXFP8、FP16、BF16、TF32和带有BF16或FP32输出的FP32输入。每个内核，它在MXFP8/MXFP4中提供315个TFLOPS，在BF16/FP16/TF32中提供79个TFLOPS，在FP32中提供20个TFLOPS，并且它使用M： N模式（例如4:16、4:12、4:8、2:8、2:4、1:4和1:2）实现结构化稀疏加速，允许在支持的稀疏工作负载上实现相同的315个TFLOPS峰值。用于矢量转换的矢量引擎提供大约1.2 TFLOPS FP32、硬件转换为MXFP格式，以及一个快速指数单元，其吞吐量是注意力工作负载标量指数路径的四倍。该单元支持各种数据类型，包括FP8、FP16、BF16、TF32、FP32、INT8、INT16和INT32。标量引擎还提供了大约1.2 TFLOPS FP32，用于跨FP8到FP32和整数数据类型的控制逻辑和小型操作。也许NeuronCore-v4最有趣的组件是GPSIMD块，它集成了八个完全可编程的512位矢量处理器，可以在访问本地SRAM的同时执行用C/C++编写的通用代码。GPSIMD被集成到NeuronCore中，因为并非真正的人工智能模型中的所有内容都清晰地映射到张量引擎。现代人工智能工作负载包含大量用于异常数据布局、后处理逻辑、索引和特定模型数学的代码。这些很难或效率低下地表达为矩阵运算，在主机CPU上运行它们会引入延迟和昂贵的数据传输。GPSIMD通过在内核内提供真正的通用可编程矢量单元来解决这个问题，因此这种逻辑以全速直接在张量旁边运行，并使用相同的本地SRAM.In简而言之，NeuronCore-v4作为一个紧密耦合的数据流引擎运行，其中张量数学、矢量变换、标量控制和自定义代码都共享一个本地32MB暂存板，并由Neuron编译器而不是Nvidia硬件上使用的扭曲调度器进行编排。

就性能而言，Trainium3在FP8计算（嗯，MXFP8）中的性能比其直接前身高出近两倍，每个包达到2.517 PFLOPS（明显领先于英伟达的H100/H200，但落后于布莱克威尔B200/B300），并增加了MXFP4支持。然而，Trainium3的BF16、TF32和FP32性能仍与Trainium2持平，这清楚地表明AWS押注MXFP8进行未来的训练和推理。