Hugging Face 推出 TorchNeuron，AWS Trainium 原生支持 PyTorch

TorchNeuron 为 PyTorch 在 AWS Trainium 上提供了新的原生后端。Transformers、Diffusers 和 Kernels 原生支持它，使得在 Trainium 上运行模型以及在 HF hub 上分享模型变得前所未有的简单。

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 芯片）上的 HF Mirror 体验是通过optimum-neuron提供的。该库依赖于基于 PyTorch XLA 的 torch-neuronx 包。这在当时是自然的选择，但它意味着要像编译器一样思考：用惰性张量代替 eager 执行，仅支持静态形状，每次图断点都要重新编译。Optimum Neuron 的存在就是为了隐藏这种复杂性，并为此维护自定义代码（建模类、训练器等）。

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

对于 HF Mirror 库来说，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 时强烈建议这样做，以防止频繁的重新编译。

import os

import torch
import torch.distributed as dist
import torch_neuronx  # noqa: F401
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": "AI 的未来是"}],
    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 上训练所使用的 HF Mirror 库与其他任何加速器相同：Transformers、Accelerate、PEFT 和 TRL。你不再需要专用的 trainer，也不需要模型代码的 Neuron 专用分支。标准的 TRL 微调脚本只需少量 Neuron 专用设置即可在 Trainium 实例上运行。

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

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):
    # Neuron 后端目前还无法编译 Dynamo 的自动动态形状（https://github.com/huggingface/trl/issues/7660）
    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:
        # 模型在加载权重时按 TP x FSDP 网格进行分片
        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,
    )
    # UltraChat 还有一个字符串类型的 `prompt` 列，这会让 TRL 将其视为 prompt-completion 数据集。
    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 原生处理。你将一个DistributedConfig与张量并行（TP）和 FSDP 大小一起放入model_init_kwargs，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 组之间分片参数、梯度和优化器状态。

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 个 NeuronCore）上运行了它，这相当于半个trn2.48xlarge实例。

训练会运行至完成，并且trainer.save_model会写出一个标准的 Transformers checkpoint。其中没有任何内容与 Neuron 绑定：你可以原样推送到 HF Mirror Hub，并在任何其他硬件上加载。

Diffusers

推理

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

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

# 分布式设备初始化
dist.init_process_group(backend="neuron") # TP
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", list(range(8))) # TP=8

# Pipeline 初始化
model = "Qwen/Qwen-Image-2.1"
transformer = QwenImage21Transformer2DModel.from_pretrained(
    model,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelConfig(mesh=mesh), # TP
)
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="一只戴着巫师帽的水豚，油画风格", device="cpu")  # 文本编码器保留在 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), # 每个 rank 上使用相同的噪声
).images[0]
if dist.get_rank() == 0:
    image.save("out.png")

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

越来越多的 pipeline 将兼容 TP。如果你喜欢的 pipeline 尚不支持，你也可以按照这份指南为其添加支持。

Kernels

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

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

致谢

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