TorchNeuron: running HF models on AWS Trainium is easier than ever

本文为原文前 6,000 字符的节选翻译，完整内容请查看原文。

TorchNeuron：在 AWS Trainium 上运行 HF 模型比以往任何时候都更容易

TorchNeuron provides a new native backend for PyTorch on AWS Trainium. Transformers, Diffusers and Kernels natively support it, making it easier than ever to run models on Trainium and to share them on the HF hub. TorchNeuron is a PyTorch plugin that brings a new backend for Trainium enabling support for eager mode, torch.compile and native distributed APIs like FSDP and TP. Starting Oct. 20, 2026, TorchNeuron will be available as a public beta. Researchers and developers will be able to install it with pip install torch-neuronx.

TorchNeuron 为 AWS Trainium 上的 PyTorch 提供了一个新的原生后端。Transformers、Diffusers 和 Kernels 对其提供了原生支持，使得在 Trainium 上运行模型并将其共享到 HF hub 比以往任何时候都更加容易。TorchNeuron 是一个 PyTorch 插件，为 Trainium 带来了新的后端，支持即时（eager）模式、torch.compile 以及诸如 FSDP 和 TP 等原生分布式 API。从 2026 年 10 月 20 日起，TorchNeuron 将作为公开测试版发布。研究人员和开发人员可以通过 pip install torch-neuronx 进行安装。

To date, the Hugging Face experience on AWS Trainium and Inferentia — AI chips designed respectively for cost-effective training and inference at scale — has been delivered through optimum-neuron. This library relies on the torch-neuronx package, which is based on PyTorch XLA. It was the natural choice at the time, but it meant thinking like a compiler: lazy tensors instead of eager execution, static shapes only, a recompilation for every graph break. Optimum Neuron existed to hide that complexity, and maintained custom code (modeling classes, trainers, etc) to do so.

迄今为止，Hugging Face 在 AWS Trainium 和 Inferentia（分别为大规模经济高效的训练和推理而设计的 AI 芯片）上的体验都是通过 optimum-neuron 提供的。该库依赖于基于 PyTorch XLA 的 torch-neuronx 包。在当时，这是自然的选择，但这意味着必须像编译器一样思考：使用惰性张量而非即时执行、仅支持静态形状，并且每次图中断都需要重新编译。Optimum Neuron 的存在是为了隐藏这种复杂性，并为此维护了自定义代码（建模类、训练器等）。

TorchNeuron changes this. It integrates directly into PyTorch and exposes neuron as a native device, with eager execution, PyTorch’s own distributed primitives (DTensor, FSDP2) and torch.compile. For the Hugging Face libraries, Trainium becomes a device like any other: Transformers, Diffusers, Accelerate, PEFT and TRL run on it directly, with the same classes and scripts you would use on any other accelerator. From now on, we recommend that Neuron users use these libraries directly instead of going through Optimum Neuron.

TorchNeuron 改变了这一点。它直接集成到 PyTorch 中，并将 neuron 作为原生设备公开，支持即时执行、PyTorch 自身的分布式原语（DTensor、FSDP2）以及 torch.compile。对于 Hugging Face 库而言，Trainium 变得像其他任何设备一样：Transformers、Diffusers、Accelerate、PEFT 和 TRL 可以直接在其上运行，使用与在其他任何加速器上相同的类和脚本。从现在起，我们建议 Neuron 用户直接使用这些库，而不是通过 Optimum Neuron。

Text generation is now natively supported in Transformers for a wide range of models. With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box, compared to 46% with Optimum Neuron v0.4.5. You just need to pass the device_map=neuron argument in your code and your model will be moved to the Neuron devices.

目前，Transformers 已为广泛的模型原生支持文本生成。借助 TorchNeuron，排名前 100 的最常用文本生成模型架构中有 89% 可以开箱即用，而 Optimum Neuron v0.4.5 的支持率为 46%。您只需在代码中传递 device_map=neuron 参数，模型就会被移动到 Neuron 设备上。

Now that Neuron is a native PyTorch device, training on Trainium uses the same Hugging Face libraries as any other accelerator: Transformers, Accelerate, PEFT and TRL. You no longer need a dedicated trainer or a Neuron-specific fork of the model code. A standard TRL fine-tuning script runs on a Trainium instance with only a few Neuron-specific settings.

既然 Neuron 现在是原生的 PyTorch 设备，在 Trainium 上进行训练所使用的 Hugging Face 库与任何其他加速器相同：Transformers、Accelerate、PEFT 和 TRL。您不再需要专门的训练器或 Neuron 特有的模型代码分支。标准的 TRL 微调脚本只需进行少量 Neuron 特定的设置，即可在 Trainium 实例上运行。