Hugging Face 宣布 TorchNeuron 支持 AWS Trainium

Hugging Face 宣布自 2026 年 10 月 20 日起，TorchNeuron 将以公共测试版形式开放，允许开发者通过 pip install torch-neuronx 安装。该工具作为 PyTorch 插件引入，将 AWS Trainium 和 Inferentia 芯片转变为 PyTorch 的原生设备，支持 eager mode、torch.compile 以及 FSDP 等分布式 API。 这一更新标志着 Hugging Face 在 AWS 硬件上的体验从依赖 optimum-neuron（基于 PyTorch XLA）转向直接集成。旧方案要求像编译器一样思考，仅支持静态形状和惰性张量；而新方案下，Transformers、Diffusers 和 Kernels 库可直接运行在 Trainium 上，无需专用 fork 或特定 trainer。实测数据显示，89% 的前 100 个文本生成模型架构现开箱即用，相比 Optimum Neuron v0.4.5 的 46% 有显著提升。 对于开发者而言，这意味着代码可移植性增强：标准 TRL 微调脚本可在 Trainium 上运行，只需配置 device_map=neuron 即可。不过，目前仍有一些限制，例如不支持异步权重加载（需设置 HF_DEACTIVATE_ASYNC_LOAD=1），且 torch.compile 尚未完全支持自动动态形状，需手动禁用。此外，Diffusers 中的 FLUX.1-dev、Qwen-Image-2.1 等模型已支持 Tensor Parallelism，可通过 NKI 编写自定义内核以进一步提升性能。