TorchNeuron 是 AWS 为 Trainium 打造的 PyTorch 原生后端，支持 eager mode、torch.compile，以及 FSDP、TP 这类原生分布式 API。从 2026 年 10 月 20 日起公开测试，可用 pip install torch-neuronx 安装。Hugging Face 的 Transformers、Diffusers 和 Kernels 库已原生支持它，意味着在 Trainium 上跑模型和在 GPU 上跑模型的方式几乎一致。

此前，HF 在 Trainium/Inferentia 上的体验靠 optimum-neuron，底层是 torch-neuronx 和 PyTorch XLA。XLA 路线要求开发者像编译器一样思考：惰性张量、静态形状、每次图断裂都要重新编译。optimum-neuron 用大量自定义建模类和训练器隐藏这些复杂性。TorchNeuron 把 neuron 直接暴露为 PyTorch 原生设备，支持 eager 执行、torch.compile 和原生分布式原语（DTensor、FSDP2）。HF 库不再需要专用分支，官方建议 Neuron 用户直接使用 Transformers、Diffusers、Accelerate、PEFT、TRL，而不是经过 Optimum Neuron。

推理方面，Transformers 文本生成原生支持。在 top 100 最常用的文本生成模型架构中，89% 开箱即用，而 Optimum Neuron v0.4.5 只覆盖 46%。只需传 device_map=’neuron’，模型就会被放到 Neuron 设备上。示例展示了张量并行下的文本生成，并建议使用静态缓存，以避免频繁重新编译。

训练方面，标准 TRL 脚本就能跑 Trainium。示例用 SFTTrainer 微调 Qwen 模型，通过 DistributedConfig 设置 TP 和 FSDP 大小，权重加载时自动切分。脚本里没有任何 Trainium 专属代码，GPU 上同样能跑。需要留意几点：脚本关闭了 Dynamo 的自动动态形状，用 wrapped packing 生成最长 2048 token 的序列；TRL ≥ 1.16 时，Neuron 上的批次会填充到 128 token 的倍数，限制形状数量。torch.compile 通过 –torch_compile true 开启，会自动选择 Neuron 后端。Transformers 内建的异步权重加载还不支持 Neuron，需设置 HF_DEACTIVATE_ASYNC_LOAD=1 禁用。示例在 8 个 Trainium2 芯片（32 个 NeuronCores）上运行，即半个 trn2.48xlarge 实例。训练产物是标准 Transformers checkpoint，可推送到 Hub 并在其他硬件加载。

Diffusers 推理方面，多个 pipeline 已兼容 eager 和 compiled 模式，FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1、H3 支持张量并行。示例用 Qwen-Image-2.1 在 TP=8 下生成图像。更多 pipeline 会陆续支持 TP，也可按官方指南自行添加。Kernels 方面，Neuron 团队提供基于 NKI（Neuron Kernel Interface）的算子，可编写自定义操作并分享到 Hub。示例中一个算子能在 Trainium 上为视觉 Transformer 带来最高 x2.4 加速。