亚马逊让 PyTorch 改一个设置就能上 Trainium：跑起来变容易了，跑得省还要另算

从英伟达换到亚马逊自研芯片 Trainium 的第一道门槛，从改写代码降到了改一个设置；换过去之后能不能省钱，改由调优决定。

我们的看法是：亚马逊 10 月 19 日公布的 TorchNeuron（官方博客称 10 月 20 日在 PyTorch 大会上以公开测试版开源），解决的是「跑起来」这笔账，没有解决「跑得省」这笔账。官方材料里唯一一个每美元产出的数字，来自一家客户在 AWS 团队帮助下数月移植后的自测。对算力账本来说，它短期改变的是试用门槛，不是谁付给谁多少钱。

发生了什么

据 PyTorch 官方博客，TorchNeuron 是 Trainium 的原生 PyTorch 后端，通过 PyTorch 给第三方硬件预留的扩展接口注册成一个叫 neuron 的设备，模型写一句 .to('neuron') 就能在 Trainium2 和 Trainium3 上训练；需要编译时，再指定一次编译后端。eager 模式、torch.compile、分布式训练接口和自写内核接口 NKI 都在支持范围内，许可证是 Apache 2.0（官方文件）。

据 GeekWire 报道，这套软件去年底起私下测试，Annapurna Labs（亚马逊的芯片设计部门）技术软件副总裁 Colin Brace 说，过去要在 Trainium 上跑 PyTorch 代码，得改成调用亚马逊专用的接口；私测期间有 30 多家公司和亚马逊内部团队在用（媒体报道）。About Amazon 的官方文章把它概括为「no rewrite needed」，并称 85% 的 AI 研究者和开发者用 PyTorch 写代码（企业宣称）。

账本上真正的数字只有一个

官方博客引用了语音模型公司 Deepgram 首席技术官的原话：在 AWS 生成式 AI 创新中心协助下，经过「数月」合作把模型移植到 Trainium，在他们评估的基准里，每美元处理的输入音频秒数比现有 GPU 基线多约 60%。GeekWire 补充，Deepgram 现在已在生产环境用 Trainium，这是 Brace 的说法。

这个 60% 需要按口径读：它是客户自己选的基准、自己的 GPU 基线，没有写 GPU 型号和价格，是「企业宣称」；而且它出现在一次数月的、有 AWS 工程团队参与的移植之后，不是改一个设置之后。另一家客户 Cartesia 在同一篇博客里说得更直白：小算子走原生路径就够了，团队把精力放在为反复调用、决定端到端延迟的那几层手写 NKI 内核上。

换句话说，「改一个设置就能跑」和「跑得比 GPU 省」之间，隔着一段工程量。TorchNeuron 把前一段缩到几乎为零，后一段仍然在。

第二笔账为什么还在

AWS Neuron 发布说明（Neuron 2.33.0，10 月 15 日）把限制写得很清楚：

• torch.compile 目前只支持静态形状，必须设 dynamic=False，每遇到一个新的输入形状可能要重新编译一次，官方建议把输入补齐或分桶；
• torch.sparse、torch.func、torch.nested、torch.distributions 不支持，torch.fft 里除了实数到复数的变换都在 CPU 上跑；
• eager 模式下遇到不支持的算子，可以退回 CPU 执行，用一个函数能查到哪些算子退回了；
• 这一版不含 TorchAO 集成和 MXFP8（官方博客说 TorchAO 和 FP8 处在私测，下一版正式提供）。

官方博客还公布了上游测试的分布：8 月一次运行约 15.4 万个测试结果，约 6.9 万个被跳过，其中约 3.9 万是 Neuron 有意跳过，最大一组约 2.2 万个是「还没有原生实现的算子」。博客正文没有用文字写出通过率，10 万个 Hugging Face 模型的扫描结果也只放在图里。

分析推断：这些限制正好落在 GPU 迁移最费工的地方。依据是，大模型推理的输入长度天然不固定，静态形状意味着要么补齐浪费算力、要么多编译几次；退回 CPU 的算子能让代码「跑起来」，却会让单位成本变差。所以「能跑」与「省钱」之间的距离，取决于每个团队的模型有多少落在官方已经调好的路径之外。

为什么这是算力账本上的事

据 GeekWire，Trainium 目前最大的用户是和亚马逊有资本关系的实验室：亚马逊投了数十亿美元的 Anthropic 用它训练和运行 Claude，OpenAI 今年 2 月同意使用 2 吉瓦 Trainium 算力，同一笔交易里亚马逊向 OpenAI 投资 500 亿美元。安迪·贾西在 4 月的股东信里写，包含 Trainium 和 Graviton 的芯片业务年化收入超过 200 亿美元（媒体报道转述股东信）。

分析推断：Trainium 最大的需求目前来自与亚马逊有投资关系的实验室，TorchNeuron 瞄准的是没有这层关系的客户。依据是 GeekWire 对最大用户的描述，以及 About Amazon 的文章把初创公司的工程时间称为「最贵的一项」，并写道开发者能在芯片之间无代价切换时，芯片厂商就得去争取开发者。如果长尾客户真的迁移，Trainium 的收入就不再只靠两三家被投公司；但要让他们迁，亚马逊得拿出第二笔账的答案，也就是公开、可比的每美元产出。

最强反方

最有力的反驳是：对大多数团队来说，迁移最大的成本本来就是改写代码，第二笔账已经被亚马逊替他们算了。依据是官方博客写明，Hugging Face 的 Transformers、Diffusers 已把 Trainium 作为上游支持的设备，vLLM 有 Neuron 插件，亚马逊还开源了覆盖主流稠密和混合专家模型的优化内核库；用标准模型、走标准推理框架的团队，未必需要自己写一行内核。如果亚马逊再把 Trainium 的单价压得足够低，补齐形状造成的浪费也可能被价格吃掉。

我们在这里部分让步：对用主流开源模型、走 vLLM 或 Hugging Face 标准路径的团队，第二笔账确实小很多。我们坚持的是另一半：自研架构、自写 CUDA 内核、输入长度多变的团队，第二笔账仍要自己算，而目前唯一公开的每美元数字，恰好来自一家走了数月定制移植的公司。在亚马逊公布按实例小时的可比价格和第三方实测之前，「改一个设置就省钱」不能当成结论。

我们不知道什么

• Trainium3 与同代 GPU 实例按小时计的可比价格，以及在同一负载下的每美元产出，官方材料没有给。
• 10 万个模型扫描里，多少能在不退回 CPU 的情况下完整跑完，博客只有图，没有文字数字。
• 私测的 30 多家用户里，有多少已经把生产负载迁过去；Anthropic 出现在官方博客的致谢名单里，但它和 OpenAI 是否在生产中用这套软件，没有说明。
• 动态形状支持什么时候到位，发布说明只写了「开发中」。

接下来看什么

一是下一版是否兼容动态形状、FP8 是否如博客所说正式提供；二是除 Deepgram 之外，有没有客户公开带 GPU 型号和价格口径的每美元对比；三是 GeekWire 提到的谷歌 TorchTPU 进展，它的目标同样是让 PyTorch 代码少改甚至不改就能换芯片。