亚马逊Trainium芯片的华丽转身

Trainium芯片已经在人工智能推理领域证明了自己的实力。Anthropic公司已在其生产环境和强化学习训练中成功部署了Trainium2和Trainium3芯片。然而，历史表明，虽然Trainium在推理方面表现出色，但在研究训练领域并非理想选择。推理过程中，模型架构的变化速度相对缓慢，而训练则需要研究人员不断尝试各种模型架构实验和消融研究，需要覆盖广泛API接口的堆栈支持，这恰恰是Trainium过去的短板。

原生PyTorch体验的革命性转变

为了克服这一不足，AWS Neuron团队正积极推动变革，最大的转变是转向一流的原生PyTorch训练体验。过去，开发者需要在训练循环中分散使用xm.step()函数，并与专有的XLA风格执行模型进行艰苦斗争。如今，Trainium已按照PyTorch用户期望的方式运行，支持eager模式、torch.compile和标准API接口。

试用体验与未来展望

在过去的几周里，SemiAnalysis工程团队对Trainium3机架进行了深入测试，整体体验令人印象深刻，我们非常认可其发展方向。当然，目前它仍不完美，我们也曾遇到一些bug。与成熟的CUDA生态系统相比还有差距，但方向是正确的。在智能编程助手的帮助下，该团队正在取得快速进展。

卓越的质量控制系统

与其他加速器公司不同，我们过去常常需要费尽口舌才能获得工程师相对稳定的CI资源。而Trainium则投入高达80,000个物理核心用于自动化质量控制，这一举措令人瞩目。

性能分析工具的突破

本文将详细介绍我们在Trainium3 TorchTitan和TorchNeuron堆栈上的使用经验，以及Trainium的Neuron Explorer分析器。令人惊讶的是，这款分析器的性能甚至超越了Nvidia和AMD的产品。我们制作了一系列演示，展示Trainium分析器如何实现”GPU不可能”的功能。

---

关注微信号：智享开源，及时了解更新信息。

原文链接：https://newsletter.semianalysis.com/p/amazons-trainium-takes-a-strike-at