AWS 公开 TorchNeuron：面向 Trainium 的原生 PyTorch 后端进入公测

AWS 公开 TorchNeuron：面向 Trainium 的原生 PyTorch 后端进入公测

AWS 在 10 月 20 日宣布，面向其 Trainium AI 芯片的开源 PyTorch 后端 TorchNeuron 进入公开测试。AWS 工程师在 PyTorch 博客上发表的较长文章，标注日期为 10 月 19 日，并在 10 月 20 日（北美 PyTorch 大会当天）更新。AWS 的公告说，该代码针对 Amazon EC2 Trainium2 和 Trainium3 实例开源。文章指明许可证为 Apache 2.0，代码托管在 aws-neuron 名下的 GitHub 仓库。本刊 10 月 27 日这一期是在原公告日期一周后报道。两个页面都是 AWS 的自述；PyTorch 文章引述 Deepgram 首席技术官称，在其评估的基准中，Trainium 相对 Deepgram 已有的 GPU 基线，每美元处理的输入音频秒数约高 60%。这是客户自述，不是独立测评或 TorchNeuron 的普适性能保证。

它给开发者带来什么变化

此前，要在 Trainium 上运行 PyTorch 代码，需要使用 AWS 单独的工具链。TorchNeuron 通过 PyTorch 为自定义硬件提供的官方扩展点 PrivateUse1，以及 `torch.accelerator` 接口，把 Trainium 注册为一个普通的 PyTorch 设备。最简单的情形下，开发者用 `.to('neuron')` 迁移模型，训练循环的其余部分保持不变。它支持两种执行模式。即时（eager）模式逐个运行算子，适合调试，文章还加入了一种实验性的“自适应即时”模式，目的是在不改代码的前提下降低启动开销并融合算子。`torch.compile` 模式把模型捕获为计算图并交给 Neuron 编译器，AWS 将其定位为生产用途。文章写明了编译路径：TorchDynamo、AOTAutograd、Torch-MLIR、StableHLO，最后由 Neuron 编译器生成可在芯片上执行的文件。

适用范围及其限制

AWS 列出了覆盖文本、视觉、视频等多类模型的广泛算子支持，BF16、FP32 和 FP16 数据类型，以及标准的分布式接口：DDP、FSDP2、张量并行、流水线并行、DTensor 和 DeviceMesh。节点内的通信走 Trainium 的集合通信核心，节点之间走 Elastic Fabric Adapter。文章坦率地说明了缺口。TorchAO 和 FP8 处于私有测试阶段。非默认流共享同一条硬件队列，因此保持顺序，但并不各自拥有专用引擎。错误会在下一个同步点才暴露，而不是在出错的算子处，TorchNeuron 会带着上下文重新抛出最早的那个错误。自定义计算核用 Neuron Kernel Interface 编写，AWS 还发布了一个此类计算核的开源库，包括注意力和混合专家（MoE）相关的例程。性能分析器的集成把框架事件与硬件活动关联起来。

如何看待这则公告

产品的核心宣称是兼容 PyTorch。AWS 说开发者可以保留 PyTorch 接口，同时用上 Trainium，文中还引用了一位 Meta 工程总监对这一支持表示欢迎的话。自适应即时模式要接近编译图性能的目标，是文章自己陈述的目标。PyTorch 文章确实给出了 Deepgram 上述输入音频成本效率数字，但未提供独立验证或普适性能结论。文中的 10 万次模型运行是以精简但有代表性的配置检查 Hugging Face 模型能否加载、运行、编译以及算子落在设备还是 CPU，并非 10 万次与 GPU 的性能对照。考虑使用 Trainium 的团队，应当用自己的模型分别跑即时和编译两种模式，确认所需算子已被覆盖，并在迁移生产任务之前先阅读测试版的说明。

迁移到 Trainium 时可以检查什么

文章提供了足够的线索，让团队列出一份迁移前的检查清单。首先是设备切换：用 `.to('neuron')` 把模型放到新设备上，再确认训练循环的其余部分不需要修改。其次是执行模式：先用即时模式确认正确性，再用 `torch.compile` 测试性能，并留意文章把“自适应即时”标为实验性。第三是算子覆盖：文章声称覆盖广泛，但每个模型用到的算子不同，缺口只能用自己的模型来发现。第四是分布式方案：确认所用的并行方式属于文章列出的 DDP、FSDP2、张量并行和流水线并行。第五是数据类型：BF16、FP32 和 FP16 已支持，而 TorchAO 与 FP8 仍在私有测试。第六是推理：文章说 Neuron 的推理栈建立在同一套 `torch.compile` 支持之上，vLLM Neuron 是 vLLM 的插件，用于提供大语言模型服务，vLLM Omni Neuron 则用于提供图像和视频生成的扩散模型服务。最后是排错：错误会在下一个同步点暴露，调试时要据此理解报错位置。

对只想了解大意的读者，可以记住三点。第一，它是开源的，目前处于公测阶段，接口和算子覆盖范围仍可能变化。第二，它建立在 PyTorch 的官方扩展点之上，而不是另起一套框架，因此长期维护的风险相对较低，但这是判断，不是文章中的承诺。第三，Deepgram 的成本效率数字只适用于其自述的比较，自己的模型仍须自行测试。

文章还提醒，两个页面都出自 AWS 自己之手，因此相关的兼容性与覆盖范围说法，最好在自己的环境里实际验证一遍。