TorchNeuron为PyTorch启用Trainium原生后端

Hugging Face与亚马逊云科技推出TorchNeuron，为PyTorch在AWS Trainium加速器上提供原生后端，帮助研究人员和开发者以接近常规GPU的工作流在云端训练和运行大模型。TorchNeuron将从2026年10月20日起开放公测，用户可通过pip install torch-neuronx安装。该插件将Neuron设备直接接入PyTorch，支持eager执行、torch.compile以及FSDP、张量并行等原生分布式能力，改变过去依赖Optimum Neuron和PyTorch XLA所带来的延迟张量、静态形状和频繁重编译等限制。Hugging Face表示，今后Transformers、Diffusers、Accelerate、PEFT和TRL等库可直接使用Trainium，不再建议通过Optimum Neuron中转。这一变化显著提升可用性，在文本生成推理中，Top 100常用架构的开箱支持率从46%提升至89%。训练方面，开发者可使用标准TRL微调脚本，并保存与Neuron解耦的Transformers模型，直接上传Hugging Face Hub或在其他硬件加载。图像生成场景中，FLUX.1-dev、FLUX.2、Qwen-Image等Diffusers流水线已支持张量并行，Neuron团队也将陆续发布Neuron Kernel Interface内核，用于加速自定义操作和模型推理。该成果降低了AWS Trainium的使用门槛，也强化了Hugging Face软件生态与AWS自研AI芯片的协同。