目录
文章目录
- 目录
- 时间线
- Inferentia1
- Trainium1
- Inferentia2
- Trainium2
- Trainium2 Server
- Trainium2 UltraServer
- Trainium2 UltraClusters
- Trainium3
- AWS GPU 实例矩阵与竞品分析
- SuperNode Rack
- Trn2 Server
- Trn2U Server
- ScaleUp 网络
- PCIe Gen5:CPU-Trainium2 ScaleUp
- NeuronLinkv3:Trainium2-Trainium2 ScaleUp
- ScaleOut 网络
- Nitro DPU
- EFAv3 HPN/RDMA 网络
- 专用的光缆中继连接器
- Firefly 光纤插头
- SIDR 路由协议
- 背板布线
- 参考文档
时间线
Inferentia1
Inferentia1 于 2019 年发布,使用 NeuronCore-v1 核心,是一款专为 AI/ML 推理(Inference)而设计的芯片。它旨在处理高吞吐量(一次进行大量预测)和低延迟(快速响应时间),最早应用于 Amazon SageMaker 托管实例中。
计算单元方面:Inferentia1 具有 4 个 NeuronCore-v1(神经元核心),在 FP16/BF16 粒度下提供 16 TeraFLOPS 的性能。包括以下计算单元。
- ScalarEngine(标量引擎)和 VectorEngine(矢量引擎),类似于 NVIDIA GPU 中的 CUDA Core。
- TensorEngine(张量引擎),用于加速矩阵数学,类似于 NVIDIA GPU 中的 TensorCore。
显存方面:集成了 2 块 DDR4 显存。
xPU-xPU 高速互联技术方面:支持 2 个 NeuronLink-v1 连接。
显然的,Inferentia1 的 IC 架构设计思路与 NVIDIA GPU 不同,而是与 Google TPU 一样,前者使用大量较小的张量核心,而后者使用大核心。大型 NeuronCore 核心在 GenAI 工作负载时能够有效减少 Control Unit 方面的开销。因为 AI 加速芯片主要是针对数据并行执行,控制单元较小,执行单元众多。
Trainium1
Trainium1 于 2020 年发布,使用 NeuronCore-v2 核心,是一款用于专注于 AI/ML 训练(Training)的加速芯片,支持 BF16 和 FP32 进行高精度训练,以及支持 TF32、FP16 和 INT8 以提高效率。
计算单元方面:Trainium1 相较于 Inferentia1 减少了 NeuronCore 数量,但在每个 NeuronCore 内的 ScalarEngine、VectorEngine 和 TensorEngine 的数量都增加了一倍。此外,Trainium1 引入了称为 GPSIMDEngine 的通用型计算单元。
- ScalarEngine(标量引擎):管理按元素运算。
- VectorEngine(矢量引擎):处理向量运算和归一化。
- TensorEngine(张量引擎):采用 128×128 脉动阵列,用于矩阵运算,尤其是大量的 GEMM(矩阵乘矩阵)运算,承担了 LLM 训练中 80% 以上的 FLOPS 和功率消耗。
- GPSIMDEngine(通用引擎):具备图灵完备性,支持使用 C/C++ 进行寻址,可以通过 C++ 编写更多的自定义算子,方便开发人员快速实现代码逻辑。
显存方面:集成了 2 块 HBM2e 高速显存。
xPU-xPU 高速互联技术方面:支持 4 个 NeuronLink-v2 连接。
专用的集合通信库芯片(Collective Communication Core):卸载集合通信工作负载,允许通信和计算的重叠运算而不会争抢资源。而在 NVIVIA/AMD 上,集合通信需要 XCCL 等占用 SM/CU 的计算资源。
NeuronCore 里面有 2 个 On-Chip 的 SRAM 高速缓存,SBUF 和 PSUM(PartialSum,求和)。SBUF 用于异步通信的过程中隐藏网络的抖动和延迟,缓存技术使得 128×128 脉动阵列的流水线处于满载状态。同时,TensorEngine 将运算结果输出到 PSUM 中缓存,然后可以在不存入 HBM 的情况下直接通过 NeuronLink 发送到其它 Trn 芯片的 SBUF 中。
Inferentia2
Inferentia2 于 2022 年发布,使用 NeuronCore-v2 核心。
Inferntia2 与 Trainium1 非常相似,但只有更少的 NeuronLink-v2 数量。Inferentia2 被 AWS 广泛应用于 IoT 物联网、语音识别等众多传统 AI 应用及 Gen AI 应用中。
Trainium2
Trainium2 于 2023 年发布,使用 NeuronCore-v3 核心。
Trainium2 是 AWS 专为超过千亿量级参数的 GenAI LLM 大模型训练和推理而打造的 AI 加速芯片。在 re:Invent 上一经发布就引起了行业的轰动,因为 Trainim2 是一个真正意义上的 “训练 + 推理” 芯片,并且从公布的技术参数来看已然有望正面竞争 NVIDIA H100 GPU。
从业务视角来看,Trainium1 面向 NLP(自然语言处理)、计算机视觉、推荐器模型的训练工作,这些模型常用于文本摘要、代码生成、问题解答、图像和视频生成、推荐和欺诈检测等各种应用程序;而目前大模型的参数量级常常到达数千亿甚至数万亿,Trainium2 则是面向 AI 大模型的训练芯片。
与 Trainium1 相比,Trainium2 训练速度提升至 4 倍,大幅降低了模型训练时间,同时能效提升多达 2 倍。用于支撑 Amazon Q、Amazon Nova、Amazon SageMaker、Amazon Bedrock 等 AI/LLM 和 AI/ML 的工作负载。
如上图所示。Trainium2 采用了 Multi-Die Chiplet 架构,并使用台积电的 CoWoS-S/R 先进封装技术,将 2 颗 5nm 的 Compute chips 和 2*2=4 块 24GB HBM 高速显存模块集成在一个紧凑的封装(Package)内。可见,这是一个是典型的 NUMA 架构,且目前没有提及 Multi-Die 间的互联协议,暂且理解为私有协议。
优点是即克服了芯片尺寸的工程极限,又最大限度地缩小了计算和内存之间的距离,实现了高带宽、低延迟。缺点是需要由 Application 来控制 NUMA 亲和。
计算单元方面:每颗 Trainium2 具有 8 个 NeuronCore-v3(神经元核心),提供 FP8 1.3 PFLOPS 性能,和 650 TFLOP/s 的稠密 BF16 性能。
显存方面:集成了 4 块 HBM3e 高速显存,容量 96GB,带宽 2.9TB/s。
xPU-xPU 高速互联技术方面:支持 4 个 NeuronLink-v3 连接。
作为参考,单颗 NVIDIA H100 GPU 的性能参数为 FP8 2 PFLOPS、80GB HBM 高速显存容量、3.35TB/s 内存带宽。即:Trainium2 的算力大约为 H100 的 60%,内存容量和内存带宽与 H100 基本持平。
其实从 NVIDIA 的《NVIDIA H100 Tensor Core GPU Architecture Overview》研究报告可知,在大规模分布式训练中,因受限于 NCCL 通信开销(如 All-Reduce)和内存墙(Memory Wall,显存带宽限制)问题,H100 在 LLM 训练场景中的综合的 MFU(Model FLOPs Utilization,模型浮点运算利用率)大概也就只能到 50% 左右。所以 Trainium2 稍低的算力配置应该是恰当的。
所以 Trainium2 集成了专用的 XCCL 集合通信引擎,通过计算-通信重叠设计(专用通信核心)可以提升 MFU,以此来弥补理论算力的不足。
Trainium2 Server
Trainium2 Server 是一台集成了 16 块 Trainium2 芯片的 SuperNode,可提供 FP8 20.8 PFLOPS(万亿次每秒浮点运算能力)算力和 1.5TB HBM 显存容量,46TB/s 的高速显存带宽。可用于训练数十亿参数的大模型。
- 算力性能:是亚马逊云科技当前最大的 GPU 实例的 1.25 倍。
- 显存容量:是亚马逊云科技当前最大的 GPU 实例的 2.5 倍。
Trainium2 UltraServer
Trainium2 UltraServer 则是一台集成了 64 张 Trainium2 芯片的 SuperNode,可提供 FP8 83.2 PFLOPS 算力,6TB HBM 显存容量,以及 185TB/s 的高速内存带宽。可用于训练数百亿参数的大模型。
- 算力性能:是亚马逊云科技当前最大的 GPU 实例的 5 倍。
- 显存容量:是亚马逊云科技当前最大的 GPU 实例的 10 倍。
Trainium2 UltraServer 代表了当前 AI SuperNode 超级计算的顶级性能,随着 AI 加速芯片数量的增加,单节点性能参数也在线性增加,是 ScaleUP 理念的极致体现。
基于 Trainium2 UltraServer 支撑的 Amazon EC2 Trn2 UltraServer 完全可以用来训练和部署目前最大的 AI 模型,包括语言、多模态和视觉模型。让 AI 工程师能够考虑在单台 64 卡一体机内以更短的时间训练出更加复杂、更加精准的 AI 模型。
Trainium2 UltraClusters
马斯克的 xAI 已经建好了 10 万颗 H100 算力集群,也放出豪言要再买 30 万张 B200;扎克伯格正在用一个超过 10 万 H100 的集群加班加点地训练 Llama4,更不用说微软和 OpenAI 了,10 万量级的 AI 加速芯片已经成为了参与军备竞赛的起步价。
当前(2025 年),亚马逊云科技正在与 Anthropic 合作部署 Rainier 项目,Anthropic 联合创始人兼首席计算官 Tom Brown 宣布下一代 Claude AI 模型将在 Project Rainier 上训练。
Project Rainier 是一个庞大的 AI 超级计算集群,包含约 40 万颗 Trainium2 芯片,预计可提供约 FP8 130 ExaFLOPS 的超强性能,整体算力达 Anthropic 目前用于训练最先进 Claude 模型所需量的 5 倍以上。
该项目完成后,预计将成为全球最大的 AI 计算集群之一,可用于支撑 Anthropic 训练万亿参数量级的超大规模的模型。
自 2023 年与 AWS 建立合作伙伴关系以来,Anthropic 一直在使用 Trainium2 来培训 Claude,该合作伙伴关系始于亚马逊 40 亿美元的投资。Anthropic 本来可以继续使用 NVIDIA GPU(他们也得到了谷歌 20 亿美元的支持),但 Trainium 的成本和可用性似乎起了决定性作用。
但是,据 SemiAnalysis 分析,40 万颗 Trainium2 的原始浮点运算性能仍少于 10 万 GB200 集群。这意味着由于阿姆达尔定律的限制,Anthropic 仍然需要面向竞争对手 10 万 GB200 集群的挑战。并且在 40 万颗 Trainium2 和 EFA 上执行 XCCL 集合通信将非常困难,因此 Anthropic 需要在异步训练方面进行一些重大创新。
Trainium3
2024 年 Re:Invent 大会上,AWS 宣布了下一代 Trainium3 芯片,采用 3nm(纳米)制程,预计将比 Trainium2 提升 2 倍的性能,改善 40% 的能效,将于 2025 年底推出。
估计 Trainium3 Server 单机支持 32 卡,而 UltraServer 则达到 128 卡规模,性能应该能够达到 FP8 166.4 PFLOPS 以上。而在 NeuronLink 层面也会从当前的 HyperCube 拓扑升级为 Switch 架构。
AWS GPU 实例矩阵与竞品分析
当然,Trainium2 UltraServer 本身并不是要去直接取代 NVIDIA GPU。客户的需求总是多样的,所以 AWS 实际上支持推出了包括 H100、H200、B200、L40S 等 NVIDIA GPU 的多种 EC2 实例。
目前,FP16 是训练环节的主要选择,因为它成熟稳定,能满足模型精度需求,并显著提高性能,这仍是 NVIDIA GPU 更具有优势的领域。而 UltraServer 强调的算力是 FP8 的,FP8 正在迅速发展,并逐渐被引入训练环节,特别是在优化计算性能和内存利用率方面。
长久以来,NVIDIA 的 H100 GPU 是 AI 训练的首选,可为从大型语言模型到图像/视频生成等所有功能提供支持。但它们的价格不菲。从理论上讲,Blackwell 的价格是 Hopper 的 40 倍。
所以 AWS 也一直在积极向客户推销 Trainium 芯片,提供计算能力与 NVIDIA 的 H100 GPU 相当的服务器租赁服务,且折扣很大。AWS 将 Trainium 定价为 H100 价格的 25%。
2024 年 AWS 案例研究报告显示:
- EC2 P5 实例具有 8 颗 H100 GPU、640GB vRAM 和 2TB RAM,但成本约为每小时 98.32 美元(按需定价)。
- EC2 Inf1 实例具有 16 颗 Inferentia1,每小时仅需花费约 4.72 美元。就推理请求吞吐量而言,Inf1 实例的运行速度比配备 NVIDIA T4 GPU 的 EC2 G4 实例快 2.3 倍。
- EC2 Inf2 实例每 1000 次推理的成本约为 0.0005 美元,而 EC2 G5 的成本为 0.0008 美元,便宜大约 37.5%。就推理请求吞吐量而言,Inf2 实例比配备 NVIDIA A10 GPU 的 EC2 G5 实例高出约 37%。
- EC2 Trn2 实例相较于前一代 GPU 实例(特别是基于 NVIDIA H200 的 P5e 和 P5en 实例)的性价比要高出 30-40%。非常适合训练和部署具有数十亿参数的大语言模型。
- EC2 Trn2u.48xlarge 实例使用 Trainium2 UltraServer 进行训练比使用 H100 便宜近 40%。Trainium2 UltraServer 远远超过了由 NVIDIA H100 GPU 驱动的 EC2 P5.48xlarge 实例上可用的内存带宽。其计算、内存和网络能力是单一 Trn2 实例的四倍,能够支持训练和部署参数接近数百亿超大规模的模型。
目前 EC2 Trn2 实例已在 AWS 美国东部(俄亥俄)推出,不久后会扩散到其他地区。与此同时,AWS 也提供了更大的 Trn2 UltraServer 配置的预览版。
并且从 Roofline 来看和 Google TPUv6e 和 NV GB200 都很有竞争力。
与 NVIDIA 的 H100 相比,Trainium2 的总拥有成本显示出有利的经济性:
然后按照 4000 卡集群做了一个网络成本分析:
总体成本分析比 H100 好了很多:
性价比分析:
SuperNode Rack
Trn2 Server
每个 Trainium2 SuperNode Rack 占用 18U(RU,机架单元),ScaleUp 纵向扩展了 1 个 CPU HeadTray(2U)和 8 个 Compute Tray(2*8=16U)。
Trn2U Server
Trn2U Server SuperNode 的示意图如下,ScaleUp 纵向扩展了 4 个 Trn2 Server,由两个机架组成。
ScaleUp 网络
Trainium2 的网络基础设施整合了多项技术,包括:
- 基于 NeuronLinkv3 的 ScaleUp 网络
- 基于 EFAv3 的 ScaleOut 网络
- 带外管理网络
- 前端管理网络
- 存储网络
PCIe Gen5:CPU-Trainium2 ScaleUp
从 Rack 图可见,与 NVIDIA NVL72 最大的区别在于,Trn2 Server 将 CPU 和 Trainium2 芯片完全解耦了;而 NVL2 采用的 CPU-GPU SuperChip 集成方案。两者各有优劣:
- 解耦方案:单个 Trainium2 芯片出问题,不会影响到 CPU,故障域更小,增加了灵活性和可维护性。
- 集成方案:CPU-GPU SuperChip 采用 On-Chip NVLink 高速互联,性能更好。
如下图所示,CPU HeadTray 需要通过 PCIe Gen5 AEC 铜缆链接到 8 个 ComputeTray 上。
CPU HeadTray 配置了 8 个 PCIe5 的端口。此外,还配备了 2 块 Nitro DPU,一块 80G 走 EBS 存储,另一块 100G 走前端 ENA/VPC 管理网络。
而每个 Compute Tray 都有 2 颗 Trainium2,但没有 CPU,也称之为 JBOG(Just bundle of GPU)。这意味着 Compute Tray 并不能独立运行,即 Trn2 Server 16 颗 Trainium2 就是最小的 EC2 售卖单元。
最终,单个 Trn2 Server 的 CPU-Trainium2 ScaleUp PCIe 拓扑如下图所示。
NeuronLinkv3:Trainium2-Trainium2 ScaleUp
NeuronLink 是 AWS 的 Trainium2-Trainium2 高速互联私有协议,类似 NVIDIA NVLink 技术和 Google TPU ICI 技术。使 Trainium2 之间可以直接访问彼此的内存。并提供 2TB/s 的内存带宽,延迟仅为 1μs(微秒)。性能高于目前的 NVLink。
如下图所示。另外一个显著区别就是 Trn2 Server 没有引入基于 NeuronLink 的 Switch Tray,所有连接都只是基于 NeuronLinkv3 使用 AEC 有源铜缆完成的 Trainium2 盒子之间的点对点直连;而 NVL72 则引入了独立的 NVLink Switch 设备,所有 GPU 都连接到多个交换机上。
基于物理网络结构的设计,所以 Trn2 Server NeuronLink ScaleUp 网络采用的是 HyperCube 拓扑,如下所示。区别于 NVL72 NVLink Switch ScaleUp 的 Full-Mesh 拓扑。
2x2x2x2 Mesh 实际上和 4x4 的 2D Torus(圆环面)是同构的。
而针对 Trn2U Server 则构成了 4x4x4 的 3D Torus 结构。
ScaleOut 网络
- ScaleUp 网络:构建更强大的 AI 服务器;
- ScaleOut 网络:构建更大规模、更高效率的 AI 服务器集群。
为了基于 Trainium2 UltraServer 实现 ExaFLOPS 级别的 AI 超级计算集群,亚马逊云科技专门设计了 10p10u AI Network 架构(完全区别于 Cloud Network)。
称之为 10p10u,因为它旨在为数千台 Trainium2 UltraServer(13TB/台)提供高达 10PB 的网络容量,并且延迟时间低于 10 微秒,以此来防止出现带宽瓶颈。并且 10p10u 网络具有弹性,可以在几个机架之间部署,也可以将其扩展到跨多个数据中心园区的集群。
为了实现 10p10u,Amazon 集结了一系列的黑科技,包括 Nitro 万卡的端侧网络加速和保序、定制的 EFA 网络接口、专用的光缆中继连接器、Firefly 光纤插头、SIDR 路由管理协议等等。
Nitro DPU
每个 Compute Tray 的前板配备了 8 块 Nitro DPU,都用于接入 TOR 的 EFAv3 网络。
即:为每颗 Trainium2 搭配了 4 个 200G 的 Nitro5 DPU,通过 EFAv3 构建 800Gbps/Trn2 的 ScaleOut 网络。
EFAv3 HPN/RDMA 网络
AWS 的 EFAv3 HPN/RDMA 是一个 FatTree TOR-Leaf-Spine 三层网络架构。
- TOR:12.8T 交换机
- Leaf-Spine:25.6T 交换机
AWS 不使用多个交换机来组成基于机箱的模块化交换机,是因为这种设置的爆炸半径很大。如果机箱发生故障,则机箱连接的所有线卡和链路都会发生故障。这可能涉及数百个 Trainium2 芯片。
专用的光缆中继连接器
为了应对 10p10u 网络的复杂性,AWS 在配线架和光纤连接方面进行了重大改进。他们开发了一种专有的中继连接器,将 16 条独立的光纤电缆组合成一个连接器,所有复杂的组装工作都发生在工厂而不是数据中心机房,大大简化了安装过程,降低了连接错误的风险,并将 AI 机架上的安装时间缩短了 54%。
下图中右侧的接线架使用的是光纤主干电缆,更加整洁,体积也更小。需要管理的连接和线路越少 ,错误就越少,这对于快速构建 AI 基础设施至关重要。
Firefly 光纤插头
AWS 还推出了 Firefly 光纤插头,它可以充当微型信号反射器,允许 AWS 在机架到达数据中心楼层之前测试和验证网络连接。Firefly 光纤插头还可充当保护密封件,防止灰尘颗粒进入光纤连接,确保网络性能的稳定性。
SIDR 路由协议
在十万卡、乃至百万卡的 10p10u AI 网络架构中,其规模异常庞大,路由相对复杂。AWS 开发了 SIDR(Scalable Intent Driven Routing,可扩展意图驱动路由),是一种专为管理这种复杂的 AI 网络架构而设计路由协议。
SIDR 开创性的结合了 “集中规划 + 分散执行” 的特性,是一个混合式的路由协议。该协议使网络能够在不到一秒的时间内响应故障,这比其他方法快十倍。
背板布线
- 褐色:PCIe CPU-Trainium2 连接线,连接 CPU HeadTray 和 8 个 Compute Tray。
- 黄色:EBS 存储连接线,负责连接存储用的,负责存储的解耦。
- 蓝色:ScaleUp 网络连接线,每个 Compute Tray 的 2 个 Trainium2 盒子,共计伸出来 2*2=4 根蓝色的线,16 块 Trainium2 组成 2 个 Hyper-Cube(没有对角线连接的正方体)。
- 紫色:ScaleOut 网络连接线,连接 Compute Tray 和 800G TOR 交换机。如下图所示,具有 64 个端口。
再一个区别就是 Trn2 Server 采用的是风冷方案,而 NVL72 是液冷方案,显然后者维护的成本更好。
参考文档
- https://semianalysis.com/2024/12/03/amazons-ai-self-sufficiency-trainium2-architecture-networking/
- https://mp.weixin.qq.com/s/cARKncahy88DQY9bRcjulQ
- https://mp.weixin.qq.com/s/q3_5hyYBiVRvy3jJEEKx1g