TorchNeuron：在 AWS Trainium 上
来自 Hugging Face · 期号 2026-10-19
摘要
Hugging Face 宣布通过 TorchNeuron 原生支持 AWS Trainium。 过去要在 Trainium 上跑模型，得用基于 PyTorch XLA 的 Optimum Neuron 库，写代码得像编译器一样思考：惰性张量、静态形状、每次图断裂都要重新编译。TorchNeuron 把 Neuron 变成了 PyTorch 的原生设备，支持 eager 模式、torch.compile 和 FSDP、TP 等原生分布式 API。 Transformers、Diffusers、Accelerate、PEFT、TRL 现在都能直接在 Trainium 上跑，用和 GPU 一样的类和脚本。文本生成方面，89% 的 Top 100 模型架构开箱即用（之前是 46%）。训练脚本也完全通用，一个标准的 TRL SFT 脚本加上几行 Neuron 配置就能跑。 TorchNeuron 将于 2026 年 10 月 20 日开放公开测试，可用 pip install torch-neuronx 安装。
深读
导读 AWS Trainium 芯片的软件生态迎来一次关键升级。Hugging Face 宣布通过 TorchNeuron 原生支持 Trainium，开发者不再需要为 Neuron 芯片维护专用代码分支，Transformers、Diffusers、TRL 等库可以直接使用与 GPU 相同的类和脚本。对于已经在用或考虑用 Trainium 做训练和推理的团队，这意味着迁移成本大幅降低。 关键信息 - TorchNeuron 是一个 PyTorch 插件，将 Neuron 暴露为 PyTorch 的原生设备（device），支持 eager 模式、torch.compile 以及 FSDP、TP 等原生分布式 API。 - 文本生成推理方面，Top 100 最常用模型架构中 89% 开箱即用，而此前基于 Optimum Neuron v0.4.5 的覆盖率是 46%。 - 训练脚本完全通用：一个标准的 TRL SFT 脚本加上几行 Neuron 配置（DistributedConfig、torch.compile）即可在 Trainium 上运行，保存的 checkpoint 是标准 Transformers 格式，可直接推送到 Hub 并在其他硬件上加载。 - 示例训练在 8 个 Trainium2 芯片（32 个 NeuronCore，即半个 trn2.48xlarge 实例）上完成，使用 TP_SIZE=4、FSDP_SIZE=8 的并行布局。 - Diffusers 库中的 FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 和 H3 等管线已支持 tensor parallelism。 - TorchNeuron 将于 2026 年 10 月 20 日开放公开测试版，可通过 pip install torch-neuronx 安装。 逐节讲解 从 Optimum Neuron 到 TorchNeuron：为什么需要换 此前在 Trainium 上运行 Hugging Face 模型依赖 optimum-neuron 库，它基于 torch-neuronx（底层是 PyTorch XLA）。XLA 要求惰性张量、静态形状，每次图断裂都要重新编译。Optimum Neuron 通过定制建模类和训练器来隐藏这些复杂性，但维护成本高、代码与 GPU 不通用。 TorchNeuron 直接集成到 PyTorch 中，将 Neuron 注册为原生设备，支持 eager 执行、PyTorch 自身的分布式原语（DTensor、FSDP2）和 torch.compile。对 Hugging Face 库而言，Trainium 从此与其他加速器无异。 Transformers：推理与训练 推理方面，只需在代码中传入 device_…