PyTorch 原生支持 Trainium：TorchNeuron 公测

发生了什么

Hugging Face 在 10 月 19 日发文介绍 TorchNeuron：一个让 AWS Trainium 成为 PyTorch 原生设备的插件，支持 eager 模式、torch.compile，以及 DTensor、FSDP2 这类 PyTorch 自己的分布式原语。从 10 月 20 日起它进入公测，pip install torch-neuronx 即可安装。Transformers、Diffusers、Kernels 三个库已经原生支持它，Accelerate、PEFT 和 TRL 也能直接在 Trainium 上运行。

具体内容

过去在 Trainium 和 Inferentia 上用 Hugging Face 的库，要经过 optimum-neuron：它建立在 PyTorch XLA 之上，意味着写代码得像写编译器一样思考——惰性张量而不是即时执行、只能用静态形状、每遇到一次图中断就要重新编译。Optimum Neuron 的存在就是为了藏起这些复杂度，为此维护着一套自己的模型类和训练器。

TorchNeuron 把 neuron 做成了一个普通设备。文章给出的变化：

• 文本生成在 Transformers 里原生可用，前 100 个最常用的文本生成架构里有 89% 开箱即用，Optimum Neuron v0.4.5 时是 46%；代码里传一个 device_map=neuron 就把模型挪到 Neuron 设备上，配合张量并行用满整台实例。生成建议用静态缓存，避免频繁重编译。
• 训练用的还是 Transformers、Accelerate、PEFT 和 TRL，不再需要专门的训练器或 Neuron 版模型代码。文章用一份普通的 TRL SFTTrainer 脚本在 8 颗 Trainium2 芯片（32 个 NeuronCore，半台 trn2.48xlarge）上做了监督微调：张量并行与 FSDP 的大小写进 DistributedConfig，权重加载时就按 TP × FSDP 网格切分，不用再传 --fsdp 或 Accelerate 的 FSDP 配置；--torch_compile true 会自动选中 Neuron 后端。
• 在 Trainium 上要注意三件事：关掉 torch.compile 的自动动态形状，TRL 1.16 及以上会把批次填充到 128 个 token 的倍数以限制输入形状的数量；Transformers 的异步权重加载暂不支持，要设 HF_DEACTIVATE_ASYNC_LOAD=1；保存出来的是标准 Transformers checkpoint，可以直接推到 Hub、在任何硬件上加载。
• Diffusers 里的很多流水线在 eager 与编译模式下都兼容，FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 和 H3 还支持张量并行。

Hugging Face 由此建议 Neuron 用户直接使用这些库，不再经过 Optimum Neuron。

与以往相比

同样的硬件，差别在抽象层：以前是一个为编译器思维打补丁的中间库，现在是 PyTorch 设备列表里多了一个名字。对使用者而言，最直接的变化是写在 GPU 上的 Transformers 与 TRL 脚本可以几乎不改地搬到 Trainium，而不是先学一套 Neuron 专用的模型类。

对读者意味着什么

1. 已经在用 Optimum Neuron 的团队，可以拿一份现有脚本在公测版上跑通一遍，核对支持的架构和 torch.compile 的重编译次数，再决定迁移。
2. 在评估 Trainium 作为训练或推理算力的团队，这次的门槛变化比性能数字更值得关注：文章没有给出吞吐或延迟对比，收益要自己测。
3. 公测意味着还有空白——异步权重加载、动态形状都还不支持，把这些限制写进评估清单再开工。