亚马逊宣布推出 TorchNeuron，这是其 Trainium 系列 AI 芯片的开源原生 PyTorch 后端，此举消除了在机器学习工作负载中采用定制芯片的最大实际障碍之一。根据 10 月 20 日北美 PyTorch 会议之前在 PyTorch 博客上发布的官方帖子，开发人员现在可以通过单行更改在 Trainium 上运行现有的 PyTorch 模型代码，而不是围绕专有工具重写工作流程。

该公告立即引起了人们的关注，因为它针对的是分析师长期以来所描述的 Nvidia 最持久的优势：不是原始芯片性能，而是将开发人员与其 GPU 绑定在一起的 CUDA 软件生态系统。路透社称此次发布是亚马逊表示其人工智能芯片将与人工智能领域最流行的开发框架配合使用，而 TechCrunch 则称 TorchNeuron 是该公司迄今为止对 Nvidia 的最大挑战。为了更广泛地了解加速器市场的变化，我们正在进行的 AI 硬件新闻报道 跟踪这些进展。

TorchNeuron 实际改变了什么

TorchNeuron 的核心承诺是简单性。在 PyTorch 博客上发布的示例中，只需使用“.to('neuron')”将模型移动到新设备即可在 Trainium 上运行现有的训练循环。根据该帖子，针对 PyTorch 推荐的与设备无关的 API“torch.accelerator”编写的代码无需修改即可运行。对于已编译的工作负载，开发人员使用 torch.compiler.set_default_backend("neuron") 设置一次默认后端，然后调用 model.compile() ，无需进一步更改。

在底层，TorchNeuron 通过 PrivateUse1（PyTorch 的自定义硬件后端的官方扩展点）注册 Trainium，这意味着该芯片的行为就像标准 PyTorch 设备，而不是固定在框架上的附件。 AWS 将集成描述为跨越多种执行模式：

  • 急切执行用于研究和调试，具有实验性的“自适应急切模式”，可减少启动开销
  • torch.compile 用于提高生产性能，将模型捕获到 Neuron 编译器优化的图表中，并原生支持 FlexAttention
  • 神经元内核接口 (NKI) 适合需要定制低级内核的研究人员

根据公告，后端可在 Trainium2 和较新的 Trainium3 芯片上运行。

直接打击 Nvidia 的软件护城河

英伟达对人工智能计算的控制力从来不仅仅依赖于更快的硬件。 CUDA 是该公司已有两年历史的软件平台，是大多数人工智能研究代码编写、调整和部署的地方，而每一个挑战者芯片——从 AMD 的 Instinct 系列到谷歌的 TPU 和亚马逊的 Trainium——都必须应对要求开发人员离开它的成本。 GeekWire、Benzinga 和 GuruFocus 对 TorchNeuron 发布的报道一致将该声明视为对护城河的攻击。

亚马逊并不是为了赢得开发者的青睐而试图将开发者从 PyTorch 中拉出来。相反，它将 Trainium 引入 PyTorch，因此该框架的庞大用户群可以在不改变习惯的情况下采用该芯片。 PyTorch 博客文章包含了维护该框架的 Meta 的认可：“对 AWS Trainium 的本机支持扩展了这种选择，使研究人员和工程师能够使用他们已经知道的急切执行和分布式 API 进行实验和迭代，”Meta 的 PyTorch 工程总监 Jana van Greunen 在帖子中说道。

这一框架对于 AWS 的云业务很重要。 Trainium 实例已作为 Nvidia GPU 的低成本替代品进行销售，用于训练和推理，并且这些芯片已经吸引了大批 AI 客户。让芯片感觉是 PyTorch 原生的直接解决了阻碍更广泛采用的反对意见：移植工作。

技术范围，从笔记本到集群规模训练

PyTorch 博客上列出的功能表明亚马逊的目标远远超出了业余爱好者的实验。 TorchNeuron 涵盖了文本、多模态和视觉语言模型以及图像和视频生成、状态空间模型和世界模型中看到的算子集，具有 BF16、FP32 和 FP16 精度； TorchAO 和 FP8 支持在下一个版本正式发布之前处于内测阶段。

分布式训练是将单 GPU 代码转变为集群规模运行的学科，通过标准 PyTorch API 支持 - DDP、FSDP2、SimpleFSDP、张量和管道并行性、ZeRO 式分片，以及用于多维进程组的 DTensor 和 DeviceMesh。帖子称，流和事件将 PyTorch 的语义映射到 Trainium 的硬件队列上。

对于推理，Amazon 的 Neuron 堆栈建立在相同的基础上：vLLM Neuron 服务于大型语言模型，而 vLLM Omni Neuron 则服务于图像和视频生成的扩散模型。 Neuron 内核接口允许研究人员为 eager 和编译执行编写自定义运算符，并且 AWS 开源了 NKI-lib，这是一组经过 Trainium 优化的内核，涵盖主流 LLM 和多模式场景的 FlashAttention 和 MoE 块式计算。

还包括性能工具。 Neuron Explorer 扩展了标准的“torch.profiler”，以跟踪从框架到硬件的执行情况，将每个事件链接回其源代码行，并且可以导出 Parquet 格式的指标，代理可以使用这些指标来进行自动瓶颈分析。

可用性和下一步发展

根据博客文章，TorchNeuron 将于 10 月 20 日在北美 PyTorch 会议上开源发布，兼容性扩展到 Hugging Face Transformers v5。 AWS 将 Trainium 描述为一个从头开始设计的芯片组合，用于云规模的人工智能训练和推理。

战略背景是每个主要云提供商同时是 Nvidia 加速器的最大买家和潜在竞争对手的市场。微软、谷歌和 Meta 都开发定制芯片，而包括华为在内的中国公司正在构建自己的软件堆栈以摆脱 CUDA 的引力。通过与 PyTorch 内部的开发人员会面，而不是要求他们迁移，亚马逊选择了参与这场竞赛的破坏性最小的路径，而且如果开发人员的体验持续下去，最有可能将真正的训练工作负载从 Nvidia 硬件上转移出去。

TorchNeuron 是否真正削弱 CUDA 的优势将取决于规模性能和现实世界模型的长尾，而发布材料并未回答这些问题。但阻止大多数 PyTorch 开发人员尝试 Trainium（重写）的障碍现在已降低到一行代码。

---

保持人工智能领先地位

获取最新的人工智能新闻、分析和突破——尽在一个地方。

阅读更多人工智能新闻 →