接之前的文章（AWS 不用英伟达 GPU，打造与众不同的超级计算机），AWS 推出的自研 AI 芯片 Trainium 及其升级版 Trainium 2，正在重塑云计算和 AI 训练领域的格局，我们开始来看这颗芯片的细节。
Trainium 2 以显著增强的性能和架构设计，填补了上一代芯片的不足，通过 650 TFLOP/s 的计算能力和 96GB 的 HBM3 内存支持，面向大规模生成式 AI 模型训练与推理。
随着 Rainier 项目中 40 万颗 Trainium 2 芯片的部署，AWS 正在全球范围内推动 ExaFLOPS 级超级计算集群的应用。
我们来分析 Trainium 系列芯片的技术演进及其对 AI 计算未来的影响，并展望其迭代方向和 AWS 的战略潜力。
Part 1.从 Trainium 到 Trainium 2：技术升级与核心改进
● 什么是 Trainium 处理器？
Trainium 是 AWS 于 2022 年推出的首款 AI 加速器，旨在支持深度学习模型的训练和推理。
作为 Inferentia 的升级版，Trainium 1 提供了一定的并行计算能力，但由于互连网络性能有限（NeuronLink-v2）、软件集成度不足，其在生成式 AI（GenAI）训练中的竞争力不强。
Trainium 1 在 GenAI 前沿模型训练或推理方面暴露出明显的短板，Trainium1 具备四个端口，而 Inferentia2 仅有两个，关键的问题在于其纵向和横向扩展网络缺乏竞争力，这一缺陷严重制约了它们在大规模 GenAI 训练任务中的表现。
众多软件错误的存在，也进一步干扰了客户工作负载的正常运行，使得这两款芯片难以在 GenAI 核心领域施展拳脚。
● Trainium 2 的架构与设计改进
Trainium2 应运而生，成为 AWS 在 AI 芯片领域战略调整的关键棋子，设计目标明确指向复杂的 GenAI LLM 推理和训练工作负载，旨在弥补前代产品的不足，并在与其他竞品的竞争中占据一席之地。
变革的驱动力主要源于市场对高效、强大的 AI 芯片的迫切需求，以及 AWS 自身在 AI 领域扩张战略的推动。
Trainium 2 采用了更高效的设计，每颗芯片包含两个计算芯片组和四个 HBM3 堆栈，支持 96GB 内存和高达 46TB/s 的带宽。
通过 NeuralLink-v3 网络，Trainium 2 实现了芯片间更快速的通信，能够形成 64 芯片的 3D 环面拓扑。
● 性能提升：
◎ 计算能力：650 TFLOP/s 的 BF16 性能（相比 Trainium 1 显著提升）。
◎ 扩展能力：支持 Trn2-Ultra 配置，可连接 64 颗芯片，形成大规模并行计算能力。
◎ 能效比：优化的互连网络设计显著降低了通信功耗，提升整体能效比。
Trainium2 是一款专为大规模 AI 模型训练和推理设计的高性能芯片，运行功耗约为 500W，提供 650 TFLOP/s 的 BF16 性能，并配备 96GByte HBM3e 内存。
其 NeuronLinkv3 扩展网络通过铜背板实现高效的服务器内芯片间连接，每个芯片利用 JBOG PCB 上的 PCB 走线及 32 个 PCIe Gen 5.0 通道（每条通道 32Gbit/s 单向），与其他三个服务器内芯片通信，形成 2x2x2x2 超立方体网格结构，等同于 4×4 2D 环面。
这种拓扑确保了低延迟、高带宽的数据交互，特别适用于多芯片协同处理任务，如大规模矩阵运算。
在服务器间，NeuronLinkv3 通过有源电缆连接四台物理服务器，构建一个 64 芯片的 4x4x4 3D 环面结构，Z 轴带宽 64GByte/s，X 和 Y 轴带宽 128GByte/s，支持高效张量并行和激活分片，提升集群计算能力。
Trainium2 采用少量大型 NeuronCore 设计理念，对比传统 GPU 的小型张量核心，更适应 GenAI 工作负载。
● 每个 NeuronCore 包含：
◎ 张量引擎：128×128 脉动阵列，承担主要计算任务。
◎ 矢量引擎：加速矢量运算，如 softmax 计算。
◎ 标量引擎：执行简单映射操作，如偏差添加。
◎ GPSIMD 引擎：允许自定义 C++代码执行，增强功能扩展性。
Trainium2 拥有专用集体通信核心，优化芯片间通信效率，避免资源争用，减少通信延迟对整体训练效率的影响。然而，预先确定的资源比例可能限制其对不同工作负载的适应性，某些情况下可能导致资源闲置或成为性能瓶颈。
每个 Trainium2 芯片由两个计算芯片组和四个 HBM3e 内存堆栈组成。计算芯片组通过 CoWoS - S / R 封装与相邻的 HBM3e 堆栈通信，芯片的两半则通过 ABF 基板相互连接，封装结构在保证芯片内部数据传输效率的同时，也面临着一些挑战，如当计算芯片组访问非直接相邻的 HBM 堆栈内存时，性能会略有下降！
Trainium 2 在硬件性能和扩展性上取得了重大突破，但其扩展网络的算术强度（225.9 BF16 FLOP/字节）仍低于 Google TPUv6e 和 Nvidia H100 的 300-560 BF16 FLOP/字节，NeuronLink 的拓扑规模（64 芯片）也小于 TPU 的 256 芯片世界规模。
AWS 通过优化软件堆栈和高效的扩展网络设计，弥补了部分差距，使 Trainium 2 成为一款兼具性能和成本优势的 AI 芯片。
Trainium2 和 Trainium2-Ultra 服务器采用独特设计，每个物理服务器占用 18 个机架单元（RU），由一个 2U CPU 托盘和八个 2U 计算托盘组成。这种架构通过无源铜背板以点对点方式连接计算托盘，形成 4×4 2D 环面结构，减少了传统交换机带来的延迟和带宽损耗。
每个计算托盘包含两个 Trainium 芯片，被称为“一堆 GPU”（JBOG），依赖 CPU 托盘进行控制和数据交互。普通 Trn2 实例的计算托盘配备 8 个 200G EFAv3 NIC，提供高达 800Gbit/s 的横向扩展带宽。
Trn2-Ultra SKU 则专注于构建 64 芯片的纵向扩展网络，横向扩展带宽为 200Gbit/s，通过有源电铜缆连接四台物理服务器，实现大规模集群扩展。CPU 托盘内的 PCIe 交换机连接计算托盘与本地 NVMe 磁盘，使 Trainium2 能通过 GPUDirect-Storage 直接访问存储，提高数据读取速度。
托盘还配备了 80Gbit/s 弹性块存储链路和 100Gbit/s Nitro 卡（ENA），用于与外部存储设备和网络连接。CPU 托盘内含两个 Intel Xeon Sapphire Rapids CPU 和最多 2TB 的 DDR5 内存，采用 48V DC 配电系统，确保稳定运行。
NeuronLinkv3 互连技术在 Trn2-Ultra 中将 64 个芯片连接成 4x4x4 3D 环面结构，提升了单服务器内及跨服务器间的通信效率。这种设计使得 Trainium2 系列服务器成为处理复杂 AI 工作负载的理想选择。
Part 2.Trainium 系列芯片的影响与未来趋势
Trainium2 在计算性能和内存容量之间取得了较好的平衡。其 650 TFLOP/s 的密集 BF16 性能与 96GByte 的 HBM3e 内存容量相匹配，能够在处理大规模数据时减少内存瓶颈对计算效率的影响。
例如，在 LLM 训练过程中，模型参数和中间数据的存储与读取能够与计算操作高效协同，避免因内存不足或带宽受限导致的计算停顿。
NeuronLinkv3 的扩展网络拓扑为 Trainium2 提供了强大的扩展能力。无论是服务器内的 2D 环面结构还是服务器间的 3D 环面结构，都能够根据不同的任务需求和集群规模进行灵活配置。
◎ 在小规模集群中，服务器内的高速连接可以满足快速的数据交互需求；
◎ 而在大规模集群中，服务器间的扩展网络能够实现跨服务器的高效协同计算，
◎ 如在构建超大规模的 LLM 推理集群时，能够将多个服务器的计算资源整合起来，提高整体推理速度。
张量引擎、矢量引擎、标量引擎和 GPSIMD 引擎的组合以及专用集体通信核心的设计，使得 Trainium2 能够针对不同类型的 AI 任务进行优化。
在 LLM 的矩阵运算、矢量运算、自定义操作以及芯片间通信等方面都能展现出较高的效率。例如，在处理自注意力机制中的复杂运算时，各个引擎能够协同工作，提高运算速度；而专用集体通信核心则在多芯片通信中减少延迟，提升整体性能。
● 高效扩展与灵活配置：Trainium 2 通过 NeuronLink 网络支持 2D/3D 环面拓扑，可实现高效的张量并行和激活分片方案，适应从小规模到超大规模的计算需求。
● 成本与能效优化：自研芯片的成本效益使得 AWS 能够以更低的费用支持复杂的 AI 工作负载。
● 软硬件协同：通过与 Neuron SDK 和 JAX 等框架的深度整合，AWS 大幅降低了开发者的使用门槛，提升了软硬件协同效率。
Trainium2 在硬件设计上取得显著进展，软件生态系统的完善仍是关键挑战。
目前，Pytorch XLA 与 Trainium2 的结合存在 API 不完善和代码路径维护的局限性。未来，AWS 需加强与 Meta 等公司的合作，优化 Pytorch 和 JAX 在 Trainium2 上的运行效率。
同时，持续开发和优化 Neuron Kernel Language (NKI)，通过与高校及研究机构的合作，扩大开发者社区，提供培训资源、开源示例代码和举办开发者竞赛，吸引更多开发者使用 NKI 进行内核开发，丰富 Trainium2 的软件生态系统。
网络性能方面，NeuronLinkv3 和 EFAv3 已具备一定优势，但随着 AI 模型规模扩大和计算速度要求提高，仍有提升空间。可以探索更高带宽连接技术或优化网络拓扑结构以减少数据传输延迟。
● 对行业趋势的推动
◎ AI 芯片自研成为主流，AWS 的成功案例将鼓励其他云服务商（如 Google、微软）加大 AI 芯片研发投入，从而推动整个行业向定制化加速器方向发展。
◎ 分布式 AI 超算的崛起，Rainier 项目展示了 AI 集群在大规模部署中的潜力，未来类似的分布式超算架构将进一步普及，为全球 AI 研发提供基础设施支持。
◎ AI 芯片性能和成本的优化，更多企业将能负担得起大规模生成式 AI 模型的训练与推理，推动其在各行业的广泛应用。
小结
AWS 通过 Trainium 2 展示了自研 AI 加速器的潜力与价值，为生成式 AI 的快速发展提供了强大的硬件支持，芯片技术的不断迭代和集群规模的扩展，A 自研芯片成为云服务商的核心竞争力。
END
作者：芝能芯芯
文章来源：芝能智芯
推荐阅读
更多汽车电子工程领域相关技术干货，请关注专栏芝能汽车电子设计 欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
