亚马逊的 Trainium2 AI 芯片配备 96GB HBM，训练性能提高四倍

本周，亚马逊网络服务（AWS）推出了其新的人工智能（AI）加速器芯片 Trainium2，相比其前身，它显著提升了性能，使 AWS 能够训练具有高达数万亿参数的基础模型（FMs）和大型语言模型（LLMs）。此外，AWS 还为自己设定了一个宏伟目标，即使其客户能够为他们的工作负载访问高达 65 'AI' ExaFLOPS 的性能。

AWS Trainium2 是亚马逊第二代加速器，专为 FMs 和 LLMs 训练而设计。与原始 Trainium 相比，它的训练性能提高了四倍，每瓦性能提高了两倍，内存增加了三倍，总共为 96GB HBM。该芯片由亚马逊的 Annapurna Labs 设计，是一个多瓦片系统级封装，包含两个计算瓦片、四个 HBM 内存堆栈和两个目前用途未公开的芯片。

亚马逊并未披露 Trainium2 的具体性能数据，但表示其 Trn2 实例可以扩展到多达 100，000 个 Trainium2 芯片，以实现 AI 工作负载的低精度计算性能高达 65 ExaFLOPS。这意味着单个 Trainium2 加速器的性能约为 650 TFLOPS。65 EFLOPS 是预计只能在即将上市的高性能 AI 超级计算机，如 Jupiter 上实现的水平。这样的扩展应该会将 3000 亿参数大型语言模型的训练时间从数月缩短到数周。

亚马逊尚未披露 Trainium2 的完整规格，但如果它没有在原始 Trainium 已支持的基础上添加一些功能，我们会感到惊讶。需要提醒的是，该协处理器支持 FP32、TF32、BF16、FP16、UINT8 和可配置的 FP8 数据格式，并提供高达 190 TFLOPS 的 FP16/BF16 计算性能。

或许比单个 AWS Trainium2 加速器的纯性能数据更重要的是，亚马逊有合作伙伴，如 Anthropic，已准备好部署它。

Anthropic 的联合创始人 Tom Brown 表示：「我们正在与 AWS 紧密合作，使用 Trainium 芯片开发我们未来的基础模型。Trainium2 将帮助我们在非常大的规模上构建和训练模型，我们预计它将比第一代 Trainium 芯片在一些关键工作负载上至少快 4 倍。我们与 AWS 的合作将帮助各种规模的组织解锁新的可能性，因为它们使用 Anthropic 的最先进 AI 系统与 AWS 的安全可靠云技术结合。」