Hugging Face 发布 TorchNeuron，原生支持 AWS Trainium

TorchNeuron：AWS Trainium 原生 PyTorch 后端发布，简化 HF 模型部署

2026 年 10 月 19 日，Hugging Face 宣布推出 TorchNeuron。这是一个专为 AWS Trainium 设计的 PyTorch 新原生后端。通过该插件，Transformers、Diffusers 和 Kernels 等库实现了对 Trainium 的原生支持，显著降低了在 Trainium 上运行及分享模型的门槛。

核心特性与发布时间

TorchNeuron 作为一个 PyTorch 插件，为 Trainium 提供了新的后端支持，涵盖急切执行模式（Eager Mode）、torch.compile 以及 FSDP 和 TP（张量并行）等本地分布式 API。自 2026 年 10 月 20 日起，TorchNeuron 将作为公开测试版提供，研究人员和开发者可通过 pip install torch-neuronx 进行安装。

架构变革：从 Optimum Neuron 到原生集成

此前，在 AWS Trainium 和 Inferentia 上使用 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 设备。

以下示例展示了如何利用张量并行（TP）生成文本，以充分利用 Neuron 实例性能。鉴于 Neuron 的特性，建议使用静态缓存以防止频繁重新编译：

import os

import torch
import torch.distributed as dist
import torch_neuronx  
from torch.distributed import DeviceMesh
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from transformers.distributed import DistributedConfig
from transformers.generation.configuration_utils import CompileConfig

MODEL_ID = "Qwen/Qwen3.8-27B"
MAX_SEQ_LEN = 1024
PREFILL_CHUNK_SIZE = 512

dist.init_process_group(backend="neuron")
world_size = dist.get_world_size()

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=world_size),
    device_mesh=DeviceMesh("neuron", list(range(world_size))),
).eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": "The future of AI is"}],
    add_generation_prompt=True,
    padding=True,
    tokenizer_kwargs={"pad_to_multiple_of": PREFILL_CHUNK_SIZE},
    return_tensors="pt",
)

generation_config = GenerationConfig(
    max_new_tokens=128,
    do_sample=False,
    use_cache=True,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
    cache_implementation="static",
    max_cache_len=MAX_SEQ_LEN,
    prefill_chunk_size=PREFILL_CHUNK_SIZE,
    compile_config=CompileConfig(backend="neuron", fullgraph=False, dynamic=False),
)

with torch.no_grad():
    output_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        generation_config=generation_config,
    )

if dist.get_rank() == 0:
    prompt_len = inputs["input_ids"].shape[1]
    print(tokenizer.decode(output_ids[0, prompt_len:], skip_special_tokens=True))

训练流程标准化

由于 Neuron 已成为原生 PyTorch 设备，在 Trainium 上的训练现在使用与其他加速器相同的 Hugging Face 库（Transformers, Accelerate, PEFT, TRL）。用户不再需要专门的 Trainer 或特定于 Neuron 的分叉模型代码。标准的 TRL 微调脚本可在 Trainium 实例上运行，仅需少量特定配置。

以下是基于 TRL SFTTrainer的监督微调脚本示例，其结构与 GPU 环境下的用法高度相似：

import os

import torch
from datasets import load_dataset
from transformers.distributed import DistributedConfig

from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser


def main(script_args, training_args, model_args):
    
    torch._dynamo.config.automatic_dynamic_shapes = False

    training_args.model_init_kwargs = {
        "revision": model_args.model_revision,
        "trust_remote_code": training_args.trust_remote_code,
        "attn_implementation": model_args.attn_implementation,
        "dtype": model_args.dtype,
        "device_map": None,
    }
    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
        
        training_args.model_init_kwargs["distributed_config"] = DistributedConfig(
            tp_size=int(os.environ.get("TP_SIZE", "1")),
            fsdp_size=int(os.environ.get("FSDP_SIZE", "1")),
        )

    
    train_dataset = load_dataset(
        script_args.dataset_name,
        name=script_args.dataset_config,
        split=script_args.dataset_train_split,
    )
    
    train_dataset = train_dataset.select_columns(["messages"])

    
    trainer = SFTTrainer(
        model=model_args.model_name_or_path,
        args=training_args,
        train_dataset=train_dataset,
    )

    trainer.train()
    trainer.save_model(training_args.output_dir)
    trainer.end()


if __name__ == "__main__":
    parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
    script_args, training_args, model_args = parser.parse_args_and_config()
    main(script_args, training_args, model_args)

Transformers 内部处理并行性。通过在 model_init_kwargs 中设置包含张量并行（TP）和 FSDP 尺寸的 DistributedConfig，SFTTrainer 将其传递给 from_pretrained，模型权重加载时即完成分片。Trainer 识别到模型已分发，无需再次包裹，也不需额外的 --fsdp 标志或 Accelerate FSDP 配置。

该脚本本身不绑定 Trainium，相同代码可跨硬件运行。但在 Trainium 上运行时需注意以下几点：

• 静态形状：脚本在 torch.compile 中禁用自动动态形状，并使用包装器生成最多 2048 个 token 的序列。在 TRL≥1.16 版本中，Neuron 上的批次被填充为 128 个 token 的倍数，以减少输入形状种类，从而复用编译图。
• torch.compile 启用：仅需添加 --torch_compile true。在 Trainium 实例上，torch.compile 会自动检测并选择 Neuron 后端，无需其他配置。
• 异步权重下载禁用：Neuron 暂不支持 Transformers 内置的异步权重加载，目前需通过环境变量 HF_DEACTIVATE_ASYNC_LOAD=1 禁用此功能。

以下命令启动训练，每个神经核启动一个进程。TP_SIZE 和 FSDP_SIZE 定义并行布局，其乘积对应进程数（此处 4 x 8 = 32，匹配 --nproc_per_node）。这 32 个进程形成 4x8 网格：张量并行将每层划分为 4 个核心组成的 TP 组，FSDP 则将参数、梯度和优化器状态划分为 8 个 TP 组。

TP_SIZE=4 FSDP_SIZE=8 HF_DEACTIVATE_ASYNC_LOAD=1 uv run torchrun --nproc_per_node=32 sft_neuron_example.py \
    --model_name_or_path Qwen/Qwen3-8B \
    --dataset_name HuggingFaceH4/ultrachat_200k \
    --dataset_train_split train_sft \
    --output_dir Qwen3-8B-sft-tp4-fsdp8 \
    --dtype bfloat16 \
    --max_length 2048 \
    --packing true \
    --packing_strategy wrapped \
    --learning_rate 2e-5 \
    --warmup_steps 0.03 \
    --num_train_epochs 1 \
    --logging_steps 2 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --save_only_model true \
    --torch_compile true

测试在 8 个 Trainium2 芯片（共 32 个神经核）上进行。训练完成后，trainer.save_model 写入标准 Transformers 检查点。生成的模型文件无神经元特异性标记，可推送至 Hugging Face Hub 并在其他硬件上加载。

Diffusers 支持情况

图像生成推理

Diffusers 库中的许多管道已与 TorchNeuron 兼容。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 和 H3 均支持张量并行。以下示例展示如何使用 TP 运行 Qwen-Image-2.1：

import torch, torch.distributed as dist, torch_neuronx 
from torch.distributed.device_mesh import DeviceMesh 
from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel, TensorParallelConfig


dist.init_process_group(backend="neuron") 
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", list(range(8))) 


model = "Qwen/Qwen-Image-2.1"
transformer = QwenImage21Transformer2DModel.from_pretrained(
    model,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelConfig(mesh=mesh), 
)
pipe = QwenImage21Pipeline.from_pretrained(
    model,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.vae.to(dev)
pipe.transformer.compile_repeated_blocks(backend="neuron", fullgraph=True, dynamic=False)


emb, mask, _ = pipe.encode_prompt(prompt="A capybara wearing a wizard hat, oil painting", device="cpu")  
image = pipe(
    prompt_embeds=emb.to(dev),
    prompt_embeds_mask=mask.to(dev) if mask is not None else None,
    height=1024,
    width=1024,
    num_inference_steps=40,
    generator=torch.Generator("cpu").manual_seed(0), 
).images[0]
if dist.get_rank() == 0:
    image.save("out.png")

上述代码可通过命令 NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py 执行。随着时间推移，更多管道将获得 TP 支持；若所需管道尚未支持，用户可参考指南自行添加。

Kernels 生态

神经元团队提供的内核正在逐步上线 Hub。Neuron Kernels 可使用神经元内核接口（NKI）编写，提供更底层的控制以优化操作。用户可运行自定义操作并在集线器共享。

利用特定内核，视觉转换器在 Trainium 上的速度可提升 2.4 倍。