TorchNeuron:在 AWS Trainium 上运行 HF 模型比以往更容易
Hugging Face Blog 发布了这条动态,详情请查看官方发布。

来自 Hugging Face Blog 的官方动态,版本与适用范围以发布方说明为准。
Hugging Face Blog · 查看官方发布中文机器翻译 · 原文附后
正文 · 来源原文
TorchNeuron 是一个 PyTorch 插件,为 Trainium 带来了新的后端,支持 eager 模式、torch.compile 和本机分布式 API(如 FSDP 和 TP)。阅读 Pytorch 博客,了解有关 TorchNeuron 的更多信息。
从 2026 年 10 月 20 日开始,TorchNeuron 将作为公开测试版提供。研究人员和开发人员将能够使用 pip install torch-neuronx 来安装它。
迄今为止,AWS Trainium 和 Inferentia(分别为经济高效的大规模训练和推理而设计的人工智能芯片)上的 Hugging Face 体验已通过 optimization-neuron 提供。该库依赖于 torch-neuronx 包,该包基于 PyTorch XLA。这是当时的自然选择,但这意味着像编译器一样思考:惰性张量而不是急于执行,仅静态形状,每个图形中断都重新编译。 Optimum Neuron 的存在就是为了隐藏这种复杂性,并维护自定义代码(建模类、训练器等)来做到这一点。
TorchNeuron 改变了这一点。它直接集成到 PyTorch 中,并将神经元作为本机设备公开,并具有急切执行、PyTorch 自己的分布式原语(DTensor、FSDP2)和 torch.compile。
对于 Hugging Face 库,Trainium 成为像任何其他设备一样的设备:Transformers、Diffusers、Accelerate、PEFT 和 TRL 直接在其上运行,具有与在任何其他加速器上使用的相同的类和脚本。从现在开始,我们建议 Neuron 用户直接使用这些库,而不是通过 Optimum Neuron。
Transformers 现在原生支持多种模型的文本生成。 TorchNeuron 可以开箱即用地支持前 100 个最常用文本生成模型架构中的 89%,而 Optimum Neuron v0.4.5 的支持率为 46%。您只需在代码中传递 device_map=neuron 参数,您的模型就会被移动到 Neuron 设备。
以下是具有张量并行性的文本生成示例,这将使您能够充分受益于您的 Neuron 实例。请注意,它使用静态缓存进行生成,因为强烈建议 Neuron 使用静态缓存,以防止频繁重新编译。
现在 Neuron 是原生 PyTorch 设备,Trainium 上的训练使用与任何其他加速器相同的 Hugging Face 库:Transformers、Accelerate、PEFT 和 TRL。您不再需要专门的培训师或特定于 Neuron 的模型代码分支。标准 TRL 微调脚本在 Trainium 实例上运行,仅具有一些特定于 Neuron 的设置。
这是一个基于 TRL 的 SFTTrainer 构建的监督微调脚本。如果您之前在 GPU 上使用 TRL 微调过模型,它应该看起来很熟悉。
并行性由 Transformers 本地处理。您将具有张量并行 (TP) 和 FSDP 大小的 DistributedConfig 放入 model_init_kwargs 中,SFTTrainer 将其传递给 from_pretrained,并且模型在加载权重时进行分片。 Trainer 看到模型已经分发,不会再次包装它。您不需要 --fsdp 标志或 Accelerate FSDP 配置。
查看英文原文
TorchNeuron: running HF models on AWS Trainium is easier than ever
Hugging Face Blog 发布了这条动态,详情请查看官方发布。
正文 · 来源原文
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. Read the Pytorch blog to learn more about TorchNeuron.
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.
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.
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.
