AWS最强AI芯片，深度解读！

（原标题：AWS最强AI芯片，深度解读！）

要是您但愿不错经常碰面，接待标星保藏哦~

亚马逊面前正在进行全球最大的 AI 集群建造之一，部署了宽广 Hopper 和 Blackwell GPU。除了对基于 Nvidia 的集群进入宽广老本支拨外，AWS 还向 Trainium2 AI 集群进入了数十亿好意思元的老本支拨。AWS 面前正在为 Anthropic 部署一个名为“Rainier 名目”的包含 40 万个 Trainium2 芯片的集群。近一年来，咱们行业滥觞的加速器模子还是掌捏了单元数目、成本、规格以及与亚马逊与 Marvell 等公司合营的弘大鸿沟相关的很多其他细节。

到面前为止，由于硬件规格较弱和软件集成度较差，亚马逊基于 Trainium1 和 Inferentia2 的实例在 GenAI 前沿模子考试或推理方面尚无竞争力。跟着 Trainium2 的发布，亚马逊进行了紧要的道路调整，并最终在芯片、系统和软件编译器/框架级别上提供具有竞争力的定制硅片进行考试和推理。

需要明确的是，由于其里面模子（如 Titan 和 Olympus）失败，亚马逊仍处于危急状态。此外，诚然他们还是牢固地在定制 AI 芯片竞赛中位居第二，仅次于谷歌，但他们仍然严重依赖 Nvidia 的产能。亚马逊的 Trainium2 并非经过考证的“考试”芯片，大部分容量将用于 LLM 推理。亚马逊对 Anthropic 的新 40 亿好意思元投资实质上将回到 Project Rainier 400k Trainium2 集群中，而且面前还莫得其他主要客户。

AWS Trainium1 / Inferentia2 GenAI 瑕玷

2022 年，AWS 发布了 Trainium1 和 Inferentia2 芯片。Trainium1 芯片和 Inferentia2 芯片真的相通，仅仅 Inferentia2 芯片只消两个 Neuronlink-v2 互连端口，而 Trainium1 有四个端口。

由于短少竞争力的纵向和横向推广相聚，Trainium1/Inferentia2 在 GenAI 考试方面进展欠安，很多软件空虚也影响了客户的责任负载。因此，Trainium1/Inferentia2 被用于考试非复杂的非 GenAI 里面 Amazon 责任负载（举例信用卡诈骗检测 AI 模子）以及 Anthropic 和 Amazon 里面责任负载的推理。

讪笑的是，Trainium1 更恰当 GenAI 推理而不是考试。在里面，亚马逊也一直在使用 Inferentia2 进行推理，举例在 2024 年 Prime Day，卓绝 80,000 个 Inferentia2/Trainium1 芯片被用于为亚马逊的基于 ML 的助手提供复旧，匡助 Amazon.com Prime 会员。

AWS Trainium2 规格综合

跟着 Trainium2 的出现，这一切齐将发生调动，因为 AWS 当今顽强地对准复杂的 GenAI LLM 推理和考试责任负载，使用约 500W 的芯片，每个芯片具有 650 TFLOP/s 的密集 BF16 性能和 96GByte 的 HBM3e 内存容量。

Trainium2 的另一个迫切进步是其推广相聚。Nvidia 的推广相聚称为 NVLink，对于 H100，其每 GPU 的运行速率为 400Gbyte/s，而 InfiniBand 推广相聚的运行速率为 50Gbyte/s；Google TPU 的推广相聚称为 ICI，而 AWS 的推广相聚称为 NeuronLink。系数具有推广相聚的 AI 集群部署仍然使用后端推广相聚。推广相聚提供的高带宽域用于竣事需要高带宽和低延伸的并行决策，举例张量并行，而推广相聚的低带宽域用于其他相貌的并行，举例对延伸相对不太明锐的数据并行。Nvidia 通过其 GB200 NVL72 从 8-GPU NVLink 域转动为 72-GPU NVLink 域，这是将大型模子推理成本裁减约 14 倍的最迫切驱动成分之一，因为它不错竣事多种各样的并行决策，而这在微型推广相聚域中是不成能竣事的。

Trainium2 有两种 SKU，第一种是将每台行状器单元的 16 个 Trainium2 芯片集聚在一齐，形成一个 4×4 二维环面的单一推广全国大小，而第二种是将每台行状器单元的 64 个 Trainium2 芯片（跨两个机架）集聚在一齐，形成一个 4x4x4 三维环面的单一推广全国大小，称为 Trainium2-ultra。与仅具有 4×4 二维环面的 Trainium1 比较，Trainium2-ultra 当今提供了罕见的集聚维度。这个罕见的维度允许在系数这个词推广域中进行张量并行和激活分片。

Trainium2-ultra 将成为 GenAI 前沿模子考试和推理的主要 SKU，用于亚马逊里面责任负载和 Anthropic 的责任负载。咱们将在本文背面深刻探讨 NeuronLink 推广相聚。

Trainium2 与 TPUv6e/GB200/H100 的比较

Trainium2 的推广拓扑是 16 芯片 SKU/64 芯片 SKU 的 2D/3D 环面，这意味着 Trainium2 的推广相聚更接近于 TPU 类拓扑（除了 Trainium2 的全国大小要小得多），而不是 Nvidia NVLink 拓扑。关节辞别在于 Trainium 和 TPU 具有点对点集聚，而 NVLink 具有交换机并复旧系数到系数集聚。

Trainium2 与其他加速器的主要辞别在于，其算术强度低得多，为每字节 225.9 BF16 FLOP，而 TPUv6e/GB200/H100 的贪图是每字节 300 到 560 BF16 FLOP。算术强度是通过将 FLOP/s 除以 HBM 带宽（以字节/秒为单元）计较得出的，它暗示计较婉曲量与内存带宽的比率。分析很迫切，因为很多应用设施（举例推理）经常受到内存带宽的瓶颈为止，导致计较 FLOPS 未得到充分运用，因此不同的算术强度可能标明加速器更恰当特定任务或时间——稍后将详确先容这一丝。

想象具有较低算术强度的 Trainium2 可能是正确的罗致，因为由于 ML 辩论的进步，模子的算术强度增长较慢。隆起的例子包括非常流行的搀和众人 (MoE：Mixture of Experts)，它使用分组 GEMM。在分组 GEMM 中，每个 token 最多只会被路由到几个众人，因此加载权重所需的内存量与密集前馈相聚 (FFN：feedforward networks) 比较要大得多，在密集前馈相聚中，每个“看到”的 token 齐会对每个权重进行计较。

在按推广域中的芯片数目进行表率化后，咱们不错看到 Trainium2 仍然具有较低的算术强度，但是每个推广的全国鸿沟（Trainium2-Ultra 的情况为 64 个芯片）的团员峰值 FLOP/s 齐低得多，这是因为 Trainium2 的全国鸿沟小于全国鸿沟为 256 个芯片的 TPUv6e。

AWS Trainium2 的封装

每个 Trainium2 芯片由两个计较芯片组和四个 HBM3e 内存堆栈构成。每个计较芯片组通过 CoWoS-S / R 封装与其两个平直的 HBM3e 堆栈通讯，芯片的两半通过 ABF 基板互相集聚。当计较芯片组尝试拜谒非平直相邻的 HBM 堆栈的内存时，性能会略有下落，况兼可能需要 NUMA 感知编程才能竣事与 MI300X 芯片组肖似的峰值性能。还有两个无源结构硅片。

AWS Trainium2 微架构

和 Trainium1 和 Google TPU 一样，Trainium2 由极少大型 NeuronCore 构成。这与 GPU 形成对比，后者使用宽广较小的张量中枢。大型中枢经常更恰当 GenAI 责任负载，因为它们的规定开销较少。在 H100 SXM 上，有 528 个张量中枢，在 GB200 上，有大致 640 个张量中枢。每个 NeuronCore 中有四个引擎：

张量引擎

矢量引擎

标量引擎

GPSIMD

滥觞是张量引擎，它是一个 128×128 脉动阵列，它从名为“SBUF”的 SRAM 缓冲区采集输入，并将其效果输出到名为“PSUM”的部分和 SRAM 缓冲区中。张量引擎不错轮回矩阵乘法 (matmul) 的 K 维，并将每个效果的部分和相加以赢得圆善效果。张量引擎/脉动阵列是当代 LLM 责任负载的 80% 以上的功率和 FLOPS 进入的场合。

接下来是矢量引擎，它旨在加速矢量运算——每个输出元素齐依赖于多个输入元素的运算，举例在稳固层入网算 softmax 或在层/批处理表率化层入网算挪动平均值和方差时。NeuronCore 休养设施不错并行化，这么系数引擎齐不错同期责任。举例在稳固层中，矢量引擎不错同期计较刻下图块的 softmax，而脉动阵列不错计较 QxK^T matmul 或 AxV matmul。

第三个是标量引擎，它被想象用于履行 1：1 映射的操作，举例元素级操作（如 SeLU）或 Exp，或者在线性层末尾添加偏差。

临了，在 NeuronCore 里面，有多个图灵完备的GpSimd 引擎，不错运行放浪 C++，这么任何 C++ 开发东说念主员齐不错搪塞快速地运行自界说操作。举例，在自稳固力层中使用 GpSimd 引擎时，您需要应用三角掩码，以便刻下 token 不错看到任何畴昔的 token，况兼只可看到刻下和昔时的 token。但是，跟着Tri Dao 的 Flash-Attention和Horace He 的 FlexAttention引申的 Block Sparse Attention 的出现，跟着时期的推移，应用三角掩码可能会变得不那么迫切。

此外，与 Trainium1 一样，Trainium2 也有专用的集体通讯中枢，专门用于与其他芯片通讯。这是一项出色的翻新，因为它允许计较通讯重迭，而不会在计较资源和通讯资源之间产生任何争用。

比较之下，在 Nvidia 和 AMD GPU 上，通讯操作与计较操作在相通的中枢 (SM) 上运行。因此，最终用户需要仔细均衡运行通讯操作的 SM 与运行计较操作的 SM 的比例。在 GPU 上，这是使用“ NCCL_MIN_NCHANNELS ”环境象征完成的，实质上这是一个相当复杂的调整。

由于这种复杂性，只消首先进的用户才会进行通讯/计较 SM 比率调整。此外，履行通讯操作经常会裁减运行计较操作的 SM 的 L2 缓存教唆率。因此，Nvidia PTX 提供了缓存教唆，以便集体通讯内核工程师不错告诉 GPU 跳过将其元素存储在 L2 缓存中。

总的来说，咱们顺服领有专用的集体通讯中枢对于优化通讯计较重迭来说是一个更澄莹的想象。

诚然领有系数这些专用引擎似乎是一个好主意，因为专用资源比通用资源花费更少的功率和面积，但它也可能导致瓶颈。必须提前罗致各样专用资源的比例，这可能导致这种均衡对于各样不同的责任负载来说是空虚的。某些资源老是未得到充分运用，而其他资源老是会为止性能。在某些方面，在责任负载仍在演变时过早地专门化您的架构可能是一个冒险的决定。

行状器架构

Trainium2 和 Trainium2-Ultra 行状器的构建块即是咱们所说的 Trainium2“物理行状器”。每个 Trainium2 物理行状器齐有一个专有的架构，占用 18 个机架单元 (RU)，由一个 2 机架单元 (2U) CPU 机头托盘构成，该托盘集聚到八个 2U 计较托盘。在行状器的背面，系数计较托盘齐使用肖似于 GB200 NVL36 的无源铜背板集聚在一齐形成一个 4×4 2D 环面，不同之处在于，对于 GB200 NVL36，背板将每个 GPU 集聚到多个 NVSwitches，而在 Trainium2 上，莫得使用交换机，系数集聚齐仅仅两个加速器之间的点对点集聚。

每个 2U 计较托盘齐有两个 Trainium 芯片，莫得 CPU。这与 GB200 NVL72 架构不同，在 GB200 NVL72 架构中，每个计较托盘齐在消失托盘中同期配备 CPU 和 GPU。每个 Trainium2 计较托盘经常也被称为 JBOG，即“一堆 GPU”，因为每个计较托盘齐莫得任何 CPU，无法沉静运行。

因此，每台 Trainium2 行状器可容纳 16 个 Trainium2 芯片。两台 16 芯片 Trainium2 行状器可装入一个机架。对于 Trainium2-Ultra SKU，每台行状器由四台物理行状器构成，每台物理行状器有 16 个芯片，因此统共可容纳 64 个芯片，占用整整两个机架。咱们将更详确地形貌机架布局，并鄙人方提供立面图。

每个计较托盘通过行状器正面的外部 PCIe 5.0 x16 DAC 无源铜缆集聚到 CPU 托盘。

CPU 托盘（Tray）

CPU 托盘里面有 PCIe 交换机，用于将计较托盘与腹地 NVMe 磁盘集聚起来，这么 Trainium2 就不错使用 GPUDirect-Storage 拜谒存储，而无需通过 CPU。

每台行状器共有 16 个腹地 NVMe，它们齐不错通过 Trainium2 芯片平直拜谒。此外，还有集聚到主 CPU0 的表率 80Gbit/s 弹性块存储链路和用于 AWS 前端量聚的主 100Gbit/s Nitro 卡，称为弹性相聚适配器 (ENA)。

每个 CPU 托盘包含两个 Intel Xeon Sapphire Rapids CPU 和 32 个用于 DDR5 RAM 的 DIMM 插槽，最多可容纳 2 TB 的 CPU 内存，并使用肖似的机架级 48V DC 母线配电系统。

计较托盘

正如著述滥觞提到的，有两个 SKU：

Trainium2 （Trn2）

Trainium2-Ultra (Trn2-Ultra)

咱们滥觞来谈谈普通 Trn2 Trainium2 实例的计较托盘。每个 Trn2 计较托盘在 PCB 上有两个 Trainium2 芯片，并有 6 个行状器内推广铜背板集聚器，用于集聚到消失行状器中的其他计较托盘。

此外，对于 Trn2，每个计较托盘最多有 8 个 200G EFAv3 NIC，每个芯片可提供高达 800Gbit/s 的横向推广以太网相聚。从计较托盘集聚到 CPU 托盘的笼子也需要一个重定时器。计较托盘左侧的 Trainium2 芯片将使用集聚到 CPU 托盘的前 8 个通说念，而右侧的 Trainium2 芯片将使用集聚到 CPU 托盘的后 8 个通说念。

Trn2-Ultra SKU 与普通 Trn2 SKU 非常相似，但 Trainium2 芯片上用于横向推广相聚的 PCIe SerDes 通说念较少，Trn2-Ultra 芯片的横向推广带宽仅为 200Gbit/s，而 Trn2 芯片的横向推广带宽为 800Gbit/s。这些通说念用于将 4 台物理行状器（每台 16 个芯片）集聚在一齐，形成一个 64 台的纵向推广相聚。纵向推广相聚将在每个芯片上竣事两个 16 通说念 OSFP-XD 笼（cages），有源电铜缆将插入其中，以将四台每台 16 个芯片的物理行状器集聚在一齐，使 Trn2-Ultra 行状器的横向推广鸿沟达到 64 个芯片。

系统/机架架构

链接究诘系统/机架架构，咱们滥觞究诘普通的 Trn2 SKU。每个机架将研讨两台 Trn2 行状器和四台 12.8T ToR EFAv3 以太网交换机，以提供每芯片高达 800Gbit/s 的横向推广带宽。该机架使用肖似的 48V DC 母线架构，交流到直流的转机发生在机架级，而不是每个机箱单独履行转机。

Trn2-Ultra SKU 由每个推广域 4 个 16 芯片物理行状器构成，因此每个推广域有 64 个芯片，由两个机架构成，设立肖似于 GB200 NVL36x2。为了沿 z 轴形成一个圆环，每个物理行状器使用一组有源铜缆集聚到另外两个物理行状器。

功率预算

咱们估算了每个 SKU 的机架功率，Trn2-Ultra 64-Chip 行状器内的两个机架中的每一个齐需要 24kW 的功率（每个 Trn2-Ultra 64-Chip 行状器 48kW）。

用于放手普通 Trn2 行状器的机架的机架密度为 27kW，但请记取，一个机架中将放手两台 16 芯片 Trn2 行状器。放手 Trn2 行状器的机架的功率密度更高，这是因为需要罕见的 NIC 和使用更多更高基数的 ToR 交换机来复旧每个 Trainium2 芯片的横向推广相聚高达 800Gbit/s，与 Trn2-Ultra 行状器的罕见行状器间 AEC NeuronLinkv3 电缆比较，这会花费更多电力。

Rainier 名目 – 400k Trainium2 集群

咱们以为，最大的 Trainium2 集群部署之一将在印第安纳州。在这里，AWS 面前正在为 Anthropic 部署一个名为“Project Rainier”的包含 40 万个 Trainium2 芯片的集群。

该园区已完成第一期建造，面前共有七栋建筑，每栋建筑的 IT 电力为 65MW，筹谋 455MW。印第安纳州 AWS 园区的第二期将再加多九栋 65MW 建筑，总电力为 1,040MW。咱们以为该园区的 PUE 约为 1.10-1.15，因为园区位于印第安纳州北部。除了 Trainium2 部署外，该数据中心园区还将与 AWS 传统的面向 CPU 的行状器以及 AWS 的 Blackwell 集群部署分享。

咱们以为 AWS 将在包含 16 个机架的闭塞舱中部署 Trainium2 机架，并在闭塞舱的结尾罕见部署 4 个机架，用于相聚交换机、照看交换机和其他行状器。冷空气将通过行状器的前部，而行状器的后部将热空气排入烟囱。

请稳固，400k Trainium2 的原始 flops 比 100k GB200 集群要少。这意味着 Anthropic 将很难与竞争敌手 100k GB200 集群匹敌，因为 Amdahl 定律会并吞它们。在 400k Trainium2 和 EFA 上进行系数缩减将非常贫窭，因此 Anthropic 需要在异步考试方面取得一些相当大的翻新进步。

相聚综合

咱们上头还是提到并究诘了相聚难题的各个部分，但本节将全面解释 Trainium2 架构中使用的系数相聚。

基于 Trainium2 的实例有四种不同类型的相聚：

1.扩大鸿沟：NeuronLinkv3

(A)行状器内 NeuronLinkv3

(B)行状器间 NeuronLinkv3

2.横向推广：Elastic Fabric Adaptor EFAv3

3.前端和存储：弹性相聚适配器 (ENA)、弹性块存储 (EBS)

4.带外照看相聚

NeuronLinkv3 是一个推广相聚，相当于 Nvidia 的 NVLink 互连的 AWS 版块。与 Nvidia 的 NVLink 互连不同，NeuronLinkv3 分为两种类型：行状器内和行状器间。行状器内将每个物理行状器内的 16 个芯片集聚在一齐，而行状器间 NeuronLink 将来自不同物理行状器的芯片集聚在一齐，统共 64 个芯片。每个芯片将具有 640GByte/s 的单向带宽，最多可有 6 个平直邻居。

与 NeuronLinkv3 有限的 64 个芯片全国大小比较，EFAv3 后端/计较结构用于将通讯从数十个机架推广到数千个机架。尽管 EFAv3 不错将更多芯片集聚在一齐，但纰谬是它比 NeuronLink 相聚慢得多。在 Trn2 SKU 上，EFAv3 结构慢 6.4 倍，在 Trn2-Ultra SKU 上，EFAv3 结构慢 25.6 倍。

快速总结一下，前端量聚仅仅一个普通的以太网相聚，用于集聚到互联网、SLURM/Kubernetes 和相聚存储以加载考试数据和模子检查点。每个包含一个 CPU 托盘的 16 个芯片的物理行状器齐有 100Gbit/s 的 ENA 前端量聚和 80Gbit/s 的专用 EBS 块存储相聚。

临了，还有带外照看相聚。它用于再行镜像操作系统，监控节点健康气象，如电扇速率、温度、功耗等。行状器、PDU、交换机、CDU 上的基板照看规定器 (BMC) 经常集聚到此相聚，以监控和规定行状器和各样其他 IT 拓荒。

在接下来的章节中，咱们将深刻探讨并详确解释上述一些相聚主题和结构。

NeuronLinkv3 推广相聚

每个 Trainium2 物理行状器齐有一个铜背板，其中每个芯片使用 JBOG PCB 板上的 PCB 走线集聚到另一个芯片（即左侧 Trainium2 芯片集聚到消失 JBOG 上的右侧 Trainium2 芯片），况兼每个芯片还使用铜背板集聚到其他三个行状器内芯片。NeuronLinkv3 基于 PCIe Gen 5.0，这意味着每条通说念 32Gbit/s（单向）。

每个芯片使用 32 个 PCIe 通说念集聚到其他行状器内芯片，这意味着每个芯片以 128GByte/s（单向）的速率与其每个行状器内邻居通讯。行状器内 NeuronLinkv3 是一个 2x2x2x2 超立方体网格。

2x2x2x2 超立方体 == 4×4 环面的同构诠释注解

意旨的是，2x2x2x2 4D 超立方体与 4×4 2D 环面同构，这意味着每个 Trainium2 物理行状器齐是一个 4×4 2D 环面。

同构的数学诠释注解非常通俗。咱们只需要检查极点和边的总和是否相通，况兼每个极点在每个图中齐有相通数目的橙色、蓝色和绿色邻居。咱们不错鄙人面的可视化中澄莹地看到，两个图齐有 16 个极点、32 条边，况兼两个图中的每个极点齐有一条橙色边、一条绿色边和 2 条蓝色边。由于它们得意系数这些条目，因此它照实是同构的。

行状器间 NeuronLinkv3 推广

在 Trainium2 Trn2-Ultra SKU 中，四台物理行状器集聚在一齐，形成一个“超等行状器”，在推广域内有 64 个芯片。这 64 个芯片集聚在一齐形成一个 4x4x4 3D 环面，其中 Z 轴的点对点带宽仅为 64GByte/s，而 x 轴和 y 轴上的点对点带宽为 128GByte/s，点对点带宽是后者的两倍。每个芯片使用 OSFP-XD 有源电缆集聚到其他物理行状器中的另外两个芯片。这么，芯片就大要形成一个在 Z 轴上具有环绕集聚的链（物理行状器 A -> 物理行状器 B -> 物理行状器 C -> 物理行状器 D -> 物理行状器 A）。

这个 4x4x4 3D 圆环与 TPU 立方体机架想象非常相似，后者亦然 4x4x4 3D 圆环，况兼在系数 3 个轴上齐具有对称的点对点带宽。Trainium2 拓扑和 TPU 拓扑之间的另一个辞别是，TPU 立方体不错通过光学器件在系数六个面上与其他 TPU 立方体集聚，而 Trainium2 不允许这么作念。

由于 PCIe 物理层允许在 NIC 和 NeuronLinkv3 之间再行分派，Amazon 和 Anthropic 可能在推广带宽和推广带宽之间达成了息争。Trainium2 仅有弥漫的 NeuronLinkv3 通说念来创建一个 4x4x2 3D 圆环，其全国大小为 32，况兼在系数 3 个轴上具有对称的点对点带宽，但对于前沿 LLM 考试和推理来说，具有不合称 BW 的全国大小为 64 可能比较小的 32 全国大小要好得多。

无 NeuronLinkv3 PCIe 光学元件

由于每个机架有两个物理行状器，因此物理行状器大要在不到 2 米的距离内形成一个环，因此不错保持在 PCIe AEC 的范围内。

要是 AWS 想象的 Trainium2 架构具有较低的机架功率密度，因此每个机架只可容纳一台行状器，那么他们就必须创建一个四物理行状器环来形成一个 64 芯片超等行状器，然后 AWS 将需要使用 PCIe 光纤，因为最长的集聚当今将跨越四个机架，超出了 PCIe AEC 的范围。

与实质罗致的想象（使用更低廉、更可靠的 AEC）比较，引入 PCIe 光学器件会导致 NeuronLinkv3 行状器间相聚的可靠性下落和成本加多。

Trn2-Ultra-Max-Plus 4x4x16 意见 SKU

由于 NeuronLink 行状器间架构是一个 3D 环面，且沿 z 轴的第5和第6 个邻居上有 AEC 集聚，因此咱们顺服不错沿着热通说念拦阻侧内的物理行状器行将 z 轴从 4 台行状器推广到 16 台行状器。

咱们提议了一个名为“trn-2-ultra-max-plus”的意见 SKU，它将 256 个芯片集聚在一齐，形成一个 4x4x16 3D 圆环，而不是通用版 Trn2-Ultra SKU 中的 64 个芯片。此意见 SKU 中的点对点带宽在 x 和 y 轴上仍为 128GByte/s，在 z 轴上仍为 64GByte/s。

咱们以为，在需要使用 PCIe 光纤集聚单个热通说念拦阻区外的机架排之前，256 个芯片将是 AEC 和无源铜缆不错达到的最大范围。

这种更大鸿沟的全国鸿沟将允许更灵验地考试无法容纳在单个 Trn2-Ultra 64 芯片行状器中的相对较大的模子。这个意见的一个纰谬是，这意味着很多芯片将以点对点拓扑（如圆环）经久地互相集聚。

责任爆炸半径（Job Blast Radius）

当很多芯片以点对点环形拓扑集聚在一齐时，要是环形中只消一个芯片发生故障，系数这个词环形推广域就会变得毫无须处。这会导致婉曲量欠安，就像 TPUv2 pod 和 TPUv3 pod 中看到的那样。对于 Trn2-Ultra，要是 Trn-2-Ultra 行状器中的 64 个芯片中只消一个发生故障，则系数 64 个芯片齐将无法作念出任何有用的责任。在咱们的意见 SKU trn2-ultra-max-plus 中，假定有 256 个芯片集聚在一齐，要是 256 个芯片中只消一个发生故障，则系数256 个芯片齐将被视为故障。

对于 TPUv4，Google 针对这一巨大的功课爆炸半径问题推出了处置决策。该处置决策是在 4x4x4 圆环立方体之间使用可重构光开关，将每个芯片故障的爆炸半径为止在 64 个芯片以内，尽管推广后的 pod 大小为 4k。

正如您鄙人面的 TPU pod 立方体舆图中看到的，不同的用户不错形成一个 3D 圆环，围绕发生故障的立方体（红色立方体）进行路由。诚然这是一个智慧的处置决策，但咱们以为 AWS 并莫得罗致在每个 4x4x4 立方体之间使用 OCS 来竣事巨大的全国鸿沟，因为从软件和硬件的角度来看，OCS 的部署非常复杂，况兼触及使用上流的光学器件。光纤链路需要收发器，由于需求量大，收发器仍然短缺，而且每带宽的成本经常比无源 DAC 铜缆高 10 倍。恰是由于这种非常高的光学成本，您会看到 TPU 立方体内大多使用 DAC 铜缆，这亦然 Nvidia GB200 NVL72 也在其推广相聚中使用铜缆的一个关节原因。

诚然在立方体中，单个芯片故障意味着系数这个词立方体住手行状，但芯片之间的相聚链路故障经常不错绕过。话虽如斯，咱们祈望相聚链路具有较高的可靠性，因此咱们不但愿 Trainium2 NeuronLinkv3 相聚链路发生故障或出现太多波动，因为系数链路齐使用无源铜缆和有源铜缆。

光收发器与无源和有源电铜缆上的相齐集聚的平均故障时期 (MTTF) 存在巨大互异。由于收发器激光器故障，MTTF 可能为 100 万至 1000 万小时，而对于铜缆，典型的 MTTF 约为 1 亿小时，可靠性率高出 10 至 100 倍。

可能比更长的 MTBF 更迫切的是，与光学系统比较，无源/有源铜缆的抖动可能性要低几个数目级。抖动是光学系统中的常见问题，由于激光器和/或模块过热问题，链路会中断一段时期，时期范围从几微秒到几秒不等。抖动会给依赖牢固相聚在芯片之间进行通讯的考试任务带来巨大问题。

由于 TPU 在立方体之间使用光学链路，因此当收发器和/或 OCS 发生故障时，Google 必须创建容错路由。有了容错路由，系数这个词 Pod 就不会因为光学链路中断而住手行状，而是在 LLM 考试责任负载下，Pod 仍能以较小的速率运行。与立方体不作念任何有用责任比较，这种速率的微小下落对于系数这个词物理 TPU 系统的灵验婉曲量来说非常有益。

然则，值得稳固的是，容错 TPU 路由仅有助于处理断开的集聚，但是当出现芯片级故障时，由于 JAX/Pytorch 需要具有 0 个孔的长方体形拓扑，这仍然会导致 64 个芯片的爆炸半径。

咱们以为，实质上 Trainium2 不需要实施容错路由，因为 NeuronLink 故障真的为零，因为链路运行在无源和有源铜缆上，这将使链路可靠性比 TPU 交叉立方体光学系统高出 100 倍。要是出于某种原因，这些无源和有源铜缆链路启动成为紧要空虚开首，那么 Trainium2 Neuronx 集体团队将需要启动实施容错路由。

EFAv3 横向推广以太网相聚

对于 Trainium2，为了在单个互连集群中推广到数万个芯片，AWS 将使用其里面以太网版块，称为 Elastic Fabric Adapter Version (EFAv3)。这将复旧普通 Trn2（16 个芯片）实例每芯片高达 800Gbit/s 的 EFAv3 BW 和 Trn2-Ultra（64 个芯片）每芯片 200Gbit/s 的 EFAv3 带宽。如上所述，Trn2-Ultra SKU 将成为用于最大鸿沟考试和推理责任负载的最常见实例。与 EFAv2 NIC 巨大的 9 微秒数据包延伸比较，EFAv3 NIC 具有较低的 6.5 微秒数据包延伸。数据包延伸是决定集体算法带宽运行速率的主要成分之一。咱们将在行将发表的 NCCL 集体深刻探讨著述中更详确地解释这些意见。

与 Nvidia 参考相聚想象不同，亚马逊面前不会使用铁路优化相聚。这意味着行状器上的系数芯片齐将集聚到消失机架中的消失平直交换机，称为“机架顶部”（ToR）相聚想象。

这意味着，与大多数流量只消一跳的轨说念优化拓扑比较，平均而言，每个流量在系数这个词相聚中需要更多的跳数。在 ToR 想象中，必须有非常好的自得当路由，以幸免不同并发流量之间的旅途突破。诸如在队伍对 ID 中而不是表率元组中进行哈希处理之类的时间将允许加多相聚中的熵，以为止流量突破的数目。

下图走漏了咱们模拟的这种无禁锢机架顶部结构的热图，其中浅蓝色暗示由于拥塞导致带宽减少，深蓝色暗示接近满线速率。如您所见，使用 ToR 拓扑不错达到线速率，但由于系数 8 个流齐进入一个交换机，因此仍然存在相当大的拥塞，婉曲量变得愈加不牢固，拥塞严重导致这些流的带宽减少。

尽管 ToR 相聚的性能较差，但咱们以为 AWS 罗致使用基于 ToR 的相聚而不是轨说念优化想象的原因是 ToR 更低廉，况兼可靠性更高，因为 ToR 架构不错使用从 AI 芯片到第一个平直交换机的铜线。比较之下，在轨说念优化的情况下，每个 AI 芯片可能集聚到远方的机架，因此需要在很多链路上使用光学器件。AWS 里面尝试过轨说念优化，但由于上头究诘的可靠性影响以及出于快速部署的愿望，他们罗致相持使用 ToR 架构。减少光学器件的使用还不错减少因 AI 飞扬带来的需求空前增长而导致的全球光学收发器接续短缺所变成的问题。

ToR 的另一个刚正是，经常在 AI 集群中，横向推广 NIC 到其第一个平直交换机只消一个单点故障。在 ToR 架构中，将横向推广 NIC 集聚到 ToR 交换机的无源 DAC 铜缆的 MTTF 培植了 100 倍，抖动减少了 100 倍。AWS 将在 NIC 和 ToR 交换机之间使用 400G QSFP-DD 到两个 200G QSFP56 无源铜缆。

对于机架内 ToR 交换机，咱们以为大多数交换机将是基于 Marvell Teralynx 6.4T 和 12.8T 交换机芯片的白盒交换机。对于 ToR 交换机，AWS 在 Broadcom 和 Marvell 商用交换机 ASIC 之间进行多供应商采购。

对于 Leaf 和 Spine 交换机，AWS 将使用基于 Broadcom Tomahawk4 硅片的 1U 25.6T 白盒交换机。

AWS 不使用多个交换机来构成基于机箱的模块化交换机，因为这种成立的爆炸半径很大。要是机箱发生故障，则机箱集聚的系数线卡和链路齐会发生故障。这可能触及数百个 Trainium2 芯片。由于潜在的爆炸半径很大，大多数超大鸿沟企业经常对这些物理机箱模块化交换机不感趣味。

相背，AWS 更可爱使用诬捏模块化交换机，通过无源和有源铜缆将单个机架中的多个披萨盒 1U 交换机集聚起来，形成一个模块化交换机。这么，就不会出现单一故障影响大范围的爆炸半径。

EBS+ENA+OOB

每个 Trainium2 物理行状器齐将有一个专用的 80Gbit/s 集聚，集聚到称为“弹性块存储”（EBS）的 AWS 托管块存储，并有一个称为“弹性相聚适配器”（ENA）的 100Gbit/s 前端集聚。这些将提供对正常相聚流量的快速拜谒，举例容器拉取、SLURM/Kubernetes 照看流量。这两个相聚齐使用 AWS 的里面 Nitro DPU 卡，可将 VPC 和安全功能卸载到硬件上，从而开释 CPU 资源。最迫切的是，这些 Nitro 卡不错准确测量通过 NIC 的流量，从而准确地为该流量计费！请稳固，物理张量和 AI 流量不会在这些相聚上运行，而只会在咱们之前究诘过的 EFAv3 横向推广和 NeuronLinkv3 纵向推广相聚上运行。

相齐集聚器和电缆成本

对于 Trainium2 的两个 SKU，TE 将是背板的惟一供应商，每台行状器将包括 48 个集聚器和 1,536 条铜缆。与 GB200 NVL72 中 NVLink 电缆从 GPU 集聚到 NVLink 交换机不同，对于 Trainium2，电缆是每个芯片之间的点对点。Trn2-Ultra SKU 还将配备 AEC 电缆，Astera Labs 将提供这些电缆。咱们以为，每块芯片的相齐集聚器和电缆总成本将达到近 1,000 好意思元。

对于 Trn2 SKU，诚然莫得行状器间的 NeuronLinkv3 AEC 电缆，但 EFAv3 带宽加多到每芯片 800Gbit/s 将卓绝这一从简，从而将相齐集聚器和电缆的总成本加多到每芯片约 1.2 万好意思元。

软件

想象定制 AI 芯片最具挑战性的方面之一是软件，其中 ML 编译器和集成到现存 ML 科学家责任经由中的坚强用户体验齐至关迫切。

以前，在 Trainium 上，使用 AI 芯片的惟一门路是通过 Pytorch XLA Lazy Tensor ML 编译器，这是一个非常倒霉的 API，会导致宽广空虚，况兼短少可移植性。咱们以为 AWS Trainium 团队还是通过提供使用肖似 Triton 的tile编程谈话编写内核的平直拜谒权限进行了改良。

此外，Trainium 软件团队当今为 JAX 提供了测试版复旧，JAX 是一个更面向 XLA ML 编译器和环形拓扑静态编译 AI 芯片的 ML 框架。Trainium 和 Trainium2 与 TPU 非常相似，因为它们齐是具有环形拓扑的巨型拖沓阵列芯片，因此 Trainium 软件当今通过 XLA 复旧 JAX 将是一个更合适的软件堆栈。

XLA

Pytorch XLA 的梦思是使用惰性张量来追踪编译计较图和系数 Pytorch 操作，仅在遭遇需要具体化的图部分时才在 Trainium 拓荒上运行图。这对于通俗模子很灵验，但惰性张量的问题在于，只消向 ML 模子架构添加任何类型的复杂性，就会出现宽广空虚。使用宽广规定流语句会破损惰性张量，尤其是在宽广使用数据依赖规定流语句的情况下。此外，对于具罕有十万个计较操作的巨型计较图（举例大型 LLM 模子），由于 Python 的速率经常较慢，惰性追踪图也会导致高开销。

Pytorch XLA 的贪图是，一朝 Pytorch<>XLA 将计较图追踪到 StableHLO，XLA 便大要履行图优化，举例删除形成象征的子图并进行垂直 + 水平会通。然后，XLA 会将其裁减为依赖于硬件的图，以进行矢量化和活水线优化，从而竣事高性能内核。

Pytorch 2.0 引入了一种使用称为“TorchDynamo”的 Python 字节码解释拿获计较图的新风物。

TorchDynamo 大要将计较图拿获到称为“Aten IR”的 IR 中，其入网算 IR 中的每个节点齐是一个 Aten 操作。然后，编译器后端不错将此 Aten IR 图行为输入，并将其裁减到特定于域的里面 IR，举例 XLA 的 StableHLO 或电感器 IR。咱们顺服这个 API 更恰当 Trainium XLA，况兼通过 Dyanmo API 对 Pytorch 的复旧会更好。

对于 TorchDyanmo XLA 的不应许音书是，面前，TorchDynamo 将考试图理解为三个图（前向、后向、优化器法子），而不是一个圆善的图，这导致与 LazyTensor XLA 比较性能欠安。好音书是，Meta 团队正在勉力在 TorchDynamo 中拿获圆善图，这项责任将推广到 TorchDynamo XLA。

对于 Trainium 来说，坏音书是，Pytorch<>XLA 代码旅途经常未被 Meta 里面闲居使用，因此主要由 AWS 和 Google Pytorch 团队可贵，这是 Pytorch<>XLA 代码旅途存在宽广空虚的主要原因。此外，与 Jax 团队比较，Google 和 AWS Pytorch 团队齐是二流团队，因为 Deepmind 和 Anthropic 等实验室的责任负载通过 Jax 堆栈运行。

通常，Meta 不会对 Pytorch AMD 代码旅途进行闲居的里面现场测试。咱们将在行将发表的著述“MI300X 与 H100 与 H200 的考试比较”中更多地究诘 AMD 的性能。AWS 团队应该与 Meta 合营，对 Trainium2 进行里面现场测试，以得当 Meta 的里面分娩考试责任负载，让 Pytorch Trainium2 软件体验愈加出色！

传统上，在使用 Google 的里面 XLA ML 编译器时，JAX 仅适用于 Google 的 TPU 芯片，但咱们以为 JAX 的编程模子也非常恰当 Trainium2，因为 TPU 和 Trainium2 齐是具有巨型脉动阵列和使用 3D 圆环的静态编译图的 AI 芯片。这意味着 Trainium 不错直不雅地集聚到 XLA Compiler 插件和 PRJT 运行时。JAX 的逻辑长方体网格和轴非常恰当 Trainium 拓扑。AWS 最近文告了其 JAX <> Trainium 集成的公开测试版，这是一个非常令东说念主愉快的标的。

NKI 内核谈话

神经元内核谈话 (NKI：The Neuron Kernel Language) - 发音为“Nicky” - 是 Trainium 领域专用谈话，用于编写肖似于 NVIDIA 的 CUDA 和 OpenAI 的 Triton 谈话的内核。与 Nvidia 的 CUDA 谈话不同，NKI 基于肖似于 OpenAI 的 Triton 编程谈话的分块编程。NKI 将使专科设施员大要在 Trainium2 芯片上竣事接近光速 (SOL) 的性能。

除了 AWS 我方的民众文档和内核示例外，为了传播 NKI 内核谈话的常识和教师，亚马逊还与斯坦福大学合营，为学生打发功课，重心是编写推行全国的内核，举例会通卷积 + max_pool。咱们可爱他们前进的标的，因为为了与 CUDA 生态系统竞争，AWS 必须遴选生态系统和开源风物围绕 NKI 内核谈话进行教师。

散布式调试和分析器具

亚马逊还提供了非常细粒度的内核级和散布式系统级调试和分析器具，肖似于 Nvidia 在其 GPU 生态系统上提供的器具。咱们以为这是正确的标的，因为它使众人级最终用户大要稳固到其考试和推理责任负载的瓶颈并加以处置。

与 Nvidia Nsight Compute 肖似，您不错稽查内核级别的分析并稽查 Tensor Engine 的行为以及 SRAM 寄存器压力等。

Nvidia 生态系管辖有 Nvidia Nsight Systems 和 Pytorch Perfetto 分析器等效器具，而 Trainium2 生态系管辖有 Neuron 散布式事件追踪。这将允许 ML 工程师调试散布式性能问题，并稽查通讯与计较的重迭进程。在某些方面，这比开箱即用的 Pytorch 分析器更好，因为它会自动合并系数品级的系数追踪，而不是让最终用户手动编写剧原来合并可能包含空虚的品级。

此外，亚马逊还公布了大部分 ISA 和精准的周期时期，这使得开发东说念主员在调试和分析时赢得比在基于 Nvidia 的 GPU 上开发更积极的体验。这意味着从硬件中索求确凿的光速要比 Nvidia GPU 容易得多，因为 Nvidia GPU 会特别向用户隐敝细节。

集体交流库（Collective Communication Library）

对于集体通讯，AWS 创建了一个 Trainium 专用库，称为NeuronX Collective Communication Library。该库肖似于 Nvidia 的 NCCL，因为它提供了一组集体通讯算法和集体，举例 All Reduce、All Gather、Reduce Scatter 等，它们是专为 Trainium2 2D/3D 环面拓扑构建的。最终用户不错平直通过 C++ 接口（或通过 Python 绑定）拜谒此库，XLA ML 编译器也不错在复杂化阶段自动插入集体调用。Trainium2 还不错使用一种名为“GPUDirect RDMA”的时间平直与其 EFAv3 NIC 通讯，而无需通过 CPU。此外，借助 AWS 的里面 libfabric，他们大要绕过内核操作系统，从而进一步减少数据包延伸。

Trainium NeuronX 集体通讯库面前不复旧 All to All 集体。All to All 集体对于搀和众人 (MoE) 模子中使用的众人并行性极为迫切。此问题于 2022 年 10 月洞开，应许约 3 周前以来一直处于非行为状态，因为 AWS 但愿复旧DataBrick 的 DBRX MoE 模子和其他行将推出的模子。Databricks 最近文告与 AWS 建立“合营伙伴关系”，使用 Trainium2 芯片进行考试和推理。咱们顺服 Databricks 将成为 GenAI 推理的第二大外部 Trainium2 客户。

超等行状器对之间的异步检查点

为了加速检查点速率，Trainium NeuronX 软件包提供了对跨主机检查点冗余的复旧。Trainium2 芯片不会在检查点相对较慢的 S3 或 AWS Managed Lustre 存储时处于闲置状态，而是会快速检查点到我方行状器的 CPU RAM 和/或腹地 NVMe 存储，之后 Trainium2 芯片将链接处理其责任负载。

这种风物的问题在于，要是行状器硬崩溃，则莫得检查点的冗余或其他备份，从而导致很多 Trainium 小时的计较时期亏空。

为了处置这个问题，Trainium2 行状器将大要以成对的相貌冉冉地将检查点复制到其他行状器。此外，为了幸免 Trainium2 行状器将其检查点复制到相邻行状器时导致考试速率降速，AWS 宣称大要在莫得来自考试责任负载的通讯时安排检查点复制流量。

责任负载编排

Trainium2 的最终用户将接受肖似 SLURM 的风物或肖似 Kubernetes 的风物来和洽他们的责任负载。

就肖似 SLURM 的风物而言，AWS 提供两种托管行状，AWS ParallelCluster 和 AWS Batch。AWS ParallelCluster 基本上仅仅托管的 SLURM。SLURM 仅仅一个近似 Linux/bash 原语的编排器。很多来自学术实验室环境的 ML 工程师/科学家齐非常可爱 SLURM，因为它使用 bash 和 Linux 原语。在运行交互式功课和进行代码开发时，这是一种令东说念主陈赞的体验。这与 Kubernetes 容器风物形成了显然对比，后者并不确凿复旧开箱即用的交互式功课。

在肖似托管 Kubernetes 的责任负载编排风物方面，AWS 提供弹性 Kubernetes 行状 (EKS) 和弹性容器行状 (ECS)。Kubernetes 并非为批处理功课（举例需要成组休养的考试责任负载）而构建的现成行状。Kubernetes 非常恰当面向行状的责任负载（举例推理），其中每个副本大多数时候仅仅一台行状器，因此不错仅仅一个 Kubernetes pod。

让 Kubernetes 适用于考试的批处理/群组休养需要宽广责任，让 Kubernetes 适用于交互式责任负载也需要宽广责任。因此，咱们实质上只看到 OpenAI 和字节卓绝等最大的 AI 实验室接受这种风物，因为它们有武艺让里面集群工程师构建器具来让 Kubernetes 适用于群组考试和交互式责任负载。

临了，Kubernetes 经常更恰当用于代理考试/推理以及为 AI 代理启动和关闭容器。咱们但愿这些容器大要得到强化，以注意 AGI 泄漏。

自动被迫和主动健康检查

当在单个责任负载中推广到数万个 AI 芯片刻，可靠性是确保考试责任负载得胜完成的迫切方面。这即是 Nvidia 发布 DCGM 器具的原因。借助 DCGM 会诊，最终用户不错通过数字和全芯片自检检测集群中发生的 80% 的静默数据损坏 (SDC)。咱们看到的另一个意旨且相关的功能是Neocloud 巨头Nebius ，它在其操作职责任负载休养设施中竣事了一项功能，他们不错通过 CRON 功课安排主动健康检查。

这些主动健康检查不错检查 Nvidia GPU 的 NVLink 相聚并运行一系列其他检查。此外，他们大要通过禁用 NVLink 仅使用一个节点自行测试其 InfiniBand 结构，而不是像其他用户在测试使用多个节点的 IB 时所作念的那样。

对于 Trainium2 软件堆栈，也有肖似的器具。不错使用 NCCOM 腹地测试来自我测试 NeuronLink 和/或横向推广 EFAv3 集聚。此外，AWS 还提供了一个测试，该测试运行一个微型考试责任负载并确保该责任负载的实质输出与黄金责任负载相匹配。

https://semianalysis.com/2024/12/03/amazons-ai-self-sufficiency-trainium2-architecture-networking/

半导体杰作公众号保举

专注半导体领域更多原创内容

关怀全球半导体产业动向与趋势

*免责声明：本文由作家原创。著述内容系作家个东说念主不雅点，半导体行业不雅察转载仅为了传达一种不同的不雅点，不代表半导体行业不雅察对该不雅点赞同或复旧，要是有任何异议，接待辩论半导体行业不雅察。

今天是《半导体行业不雅察》为您分享的第3967期内容，接待关怀。

『半导体第一垂直媒体』

及时 专科 原创 深度

公众号ID：icbank

可爱咱们的内容就点“在看”分享给小伙伴哦