Hugging Face 推出 TorchNeuron，让 HF 模型在 AWS Trainium 上原生运行

TorchNeuron

TorchNeuron 是一个 PyTorch 插件，为 Trainium 带来新的后端，支持 eager 模式、torch.compile 以及 FSDP 和 TP 等原生分布式 API。阅读 PyTorch 博客以了解更多关于 TorchNeuron 的信息。

从 2026 年 10 月 20 日起，TorchNeuron 将以公开测试版形式提供。研究人员和开发者将能够通过 pip install torch-neuronx 进行安装。

新后端，HF 库中的原生接口

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

TorchNeuron 改变了这一点。它直接集成到 PyTorch 中，并将 neuron 暴露为原生设备，支持即时执行、PyTorch 自有的分布式原语（DTensor、FSDP2）以及 torch.compile。

对于 Hugging Face 库而言，Trainium 变成了和其他设备一样的设备：Transformers、Diffusers、Accelerate、PEFT 和 TRL 可以直接在其上运行，使用的类与脚本与在任何其他加速器上相同。从现在起，我们建议 Neuron 用户直接使用这些库，而不是通过 Optimum Neuron。

Transformers

推理

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 微调脚本只需少量 Neuron 特定设置即可在 Trainium 实例上运行。

以下是一个基于 TRL 的 SFTTrainer 构建的监督微调脚本。如果你之前曾在 GPU 上使用 TRL 微调过模型，它看起来应该很熟悉。

并行性由 Transformers 原生处理。你在 model_init_kwargs 中放入带有张量并行（TP）和 FSDP 大小的 DistributedConfig，SFTTrainer 将其传递给 from_pretrained，模型在加载权重时即被分片。Trainer 看到模型已经分布式，就不会再次包装它。你不需要 --fsdp 标志或 Accelerate FSDP 配置。

脚本中没有任何内容与 Trainium 绑定：相同的代码可在任何硬件上运行。在 Trainium 上运行时，有几点值得了解：

静态形状。脚本禁用了 torch.compile 中的自动动态形状，并使用 wrapped packing 生成最多 2048 个 token 的序列。使用 TRL ≥ 1.16 时，Neuron 上的批次会填充到 128 个 token 的倍数，从而限制输入形状的数量。编译后的图会为具有相同形状的后续步骤重用。

torch.compile 开启。只需传入 --torch_compile true：在 Trainium 实例上，torch.compile 会检测设备并自动选择 Neuron 后端，无需其他配置。

异步权重加载关闭。Transformers 内置的异步权重加载在 Neuron 上不受支持，因此目前我们通过 HF_DEACTIVATE_ASYNC_LOAD=1 将其禁用。

下面的命令使用 torchrun 启动训练，每个 NeuronCore 启动一个进程。TP_SIZE 和 FSDP_SIZE 设置并行布局，它们的乘积就是进程数：此处为 4 x 8 = 32，与 --nproc_per_node 匹配。这 32 个进程形成一个 4 x 8 网格。张量并行将每一层拆分到 TP 组的 4 个核心上，FSDP 则在 8 个 TP 组之间分片参数、梯度和优化器状态。

我们在 8 个 Trainium2 芯片（32 个 NeuronCore）上运行了它，即半个 trn2.48xlarge 实例。

训练运行至完成，trainer.save_model 写入标准的 Transformers 检查点。其中没有任何内容与 Neuron 绑定：你可以原样推送到 Hugging Face Hub，并在任何其他硬件上加载。

Diffusers

推理

Diffusers 库中的许多管道在 eager 和编译模式下与 TorchNeuron 兼容。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 和 H3 也与张量并行兼容，以充分利用你的 Neuron 设备。以下是一个展示如何使用 TP 的示例：

你可以使用以下命令运行这段代码：NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py 。

越来越多的流水线将兼容 TP。如果你喜欢的某个流水线尚不兼容，你也可以按照本指南为其添加支持。

内核

请务必查看 Neuron 团队的内核，因为越来越多的 Neuron 内核将登陆 Hub，以加速你的模型。Neuron 内核可以使用 Neuron 内核接口（NKI）编写，它提供了更底层的控制来优化运算。这样，用户就可以运行自己的自定义运算，并在 Hub 上分享。

例如，你可以使用这个内核来加速你的视觉 Transformer，在 Trainium 上最高可获得 2.4 倍的加速！

致谢

特别感谢 AWS 团队在 TorchNeuron 上的工作，以及他们对我们反馈的积极响应，使 HF 集成尽可能完善。

博客中的更多文章

PyTorch 中的性能分析（第 2 部分）：从 nn.Linear 到融合 MLP

+1

PyTorch 中的性能分析（第 1 部分）：torch.profiler 入门指南