Hugging Face推出TorchNeuron后端，Trainium原生支持Transformers与Diffusers
Hugging Face称TorchNeuron把Neuron作为原生PyTorch设备暴露，Transformers、Diffusers、Accelerate、PEFT、TRL可直接在AWS Trainium上运行；公共测试版自 2026 年 10 月 20 日起提供，可用pip install torch-neuronx安装。
AI解读：此前Hugging Face在Trainium和Inferentia上的体验依赖optimum-neuron（基于PyTorch XLA）：只能使用惰性张量和静态形状，每次图断点都会触发重新编译，库中维护了大量自定义模型代码和训练器。TorchNeuron改变了这一层，Hugging Face表示今后建议Neuron用户直接使用Transformers等库，而不再经过Optimum Neuron。
文本生成方面，Hugging Face称TorchNeuron开箱支持前 100 个最常用文本生成模型架构中的 89%，而Optimum Neuron v0.4.5为 46%。使用时传入device_map=neuron即可把模型放到Neuron设备上。
训练侧，Hugging Face报告在一个由 8 颗Trainium2芯片（32 个NeuronCore）组成的半台trn2.48xlarge实例上，用TRL的SFTTrainer跑通了有监督微调，产出的checkpoint未绑定Neuron，可上传Hub并在其他硬件加载。
Diffusers的部分管线已兼容TorchNeuron的eager与编译模式，FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1和H3还兼容张量并行；更多管线仍在适配中。
使用中存在若干限制：需要静态形状与静态缓存以避免频繁重编译；Transformers内置的异步权重加载在Neuron上尚不支持，需用HF_DEACTIVATE_ASYNC_LOAD=1 关闭；Neuron后端暂不能编译Dynamo的自动动态形状。
Hugging Face发布TorchNeuron，一个为AWS Trainium提供原生PyTorch后端的插件，支持eager执行、torch.compile，以及FSDP、TP等原生分布式API。按Hugging Face的说法，公共测试版自 2026 年 10 月 20 日起提供，安装命令为pip install torch-neuronx。
从optimum-neuron换到原生设备
Hugging Face介绍，此前在AWS Trainium和Inferentia（分别为大规模训练和推理设计的AI芯片）上的体验通过optimum-neuron交付。该库依赖基于PyTorch XLA的torch-neuronx包，使用时需要按编译器的方式思考：惰性张量而非eager执行、仅支持静态形状、每次图断点都要重新编译。Optimum Neuron的职责就是隐藏这些复杂度，并为此维护自定义的建模类和训练器等代码。
TorchNeuron直接集成进PyTorch，把neuron暴露为一个原生设备，具备eager执行、PyTorch自带的分布式原语（DTensor、FSDP2）和torch.compile。对Hugging Face各库而言，Trainium成为与其他加速器一样的设备，Transformers、Diffusers、Accelerate、PEFT和TRL可用相同的类和脚本直接运行。Hugging Face据此建议Neuron用户今后直接使用这些库，而不再经由Optimum Neuron。
推理：89% 对 46% 的架构覆盖
Hugging Face称，Transformers的文本生成现已原生支持多种模型，TorchNeuron开箱支持前 100 个最常用文本生成模型架构中的 89%，Optimum Neuron v0.4.5为 46%。使用时传入device_map=neuron，模型即被移至Neuron设备。
示例代码演示了张量并行的文本生成，并采用静态缓存，Hugging Face表示这在Neuron上被强烈推荐，可避免频繁重编译。
训练：半台trn2.48xlarge跑通SFT
Hugging Face表示，Neuron成为原生PyTorch设备后，Trainium上的训练使用与其他加速器相同的库：Transformers、Accelerate、PEFT和TRL，不再需要专用训练器或Neuron专用分支。标准的TRL微调脚本只需少量Neuron专属设置即可在Trainium实例上运行。
文中给出的SFT脚本通过DistributedConfig指定张量并行（TP）和FSDP规模，在加载权重时完成分片，Trainer识别到模型已分布式后不再重复包装，不需要 --fsdp标志或Accelerate的FSDP配置。Hugging Face称脚本中没有任何内容绑定Trainium，同一份代码可在任意硬件上运行。
Hugging Face报告的实际运行使用 8 颗Trainium2芯片（32 个NeuronCore），即半台trn2.48xlarge实例，训练跑完，trainer.save_model写出标准Transformers checkpoint；Hugging Face称该checkpoint未绑定Neuron，可直接推送到Hub并在其他硬件加载。
- 静态形状：脚本关闭torch.compile的自动动态形状，采用wrapped packing生成最长 2048 token的序列；TRL ≥ 1.16 时Neuron上的批次会补齐到 128 token的倍数，减少输入形状数量，相同形状的编译图在后续步骤复用。
- torch.compile：传入 --torch_compile true即可，在Trainium实例上会自动检测设备并选用Neuron后端。
- 异步权重加载：Transformers内置的异步权重加载在Neuron上不受支持，目前需用HF_DEACTIVATE_ASYNC_LOAD=1 关闭。
Diffusers与Neuron内核
Hugging Face称Diffusers的许多管线兼容TorchNeuron的eager和编译模式，FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1和H3还兼容张量并行。示例中Qwen-Image-2.1使用TP=8 的DeviceMesh，文本编码器留在CPU上，以NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 运行。Hugging Face表示会有更多管线支持张量并行，未支持的管线可参照其指南添加。
文中还提到Neuron团队的内核将陆续发布到Hub，可用Neuron Kernel Interface（NKI）编写以获得更底层的优化控制，并给出一个据称在Trainium上为视觉Transformer带来最高 2.4 倍加速的内核示例。Hugging Face在致谢中感谢AWS团队对TorchNeuron的工作以及对反馈的响应。