TorchNeuron：在 AWS Trainium 上运行 HF 模型前所未有的简单

来源：https://huggingface.co/blog/hf-torchneuron

TorchNeuron 为 PyTorch 在 AWS Trainium 上提供了一个全新的原生后端。Transformers、Diffusers 和 Kernels 原生支持它，使得在 Trainium 上运行模型以及在 HF Hub 上分享模型从未如此简单。

https://huggingface.co/blog/hf-torchneuron#torchneuronTorchNeuron

TorchNeuron（https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/index.html）是一个 PyTorch 插件，为 Trainium 带来了一个全新的后端，支持 eager 模式、torch.compile 以及 FSDP 和 TP 等原生分布式 API。阅读 PyTorch 博客（https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/）了解更多关于 TorchNeuron 的信息。

从 2026 年 10 月 20 日起，TorchNeuron 将以公开测试版的形式提供。研究人员和开发者可以通过 pip install torch-neuronx 进行安装。

https://huggingface.co/blog/hf-torchneuron#new-backend-native-interface-in-hf-libraries全新后端，HF 库的原生接口

迄今为止，Hugging Face 在 AWS Trainium 和 Inferentia（分别专为大规模高性价比训练和推理设计的 AI 芯片）上的体验一直是通过 optimum-neuron（https://github.com/huggingface/optimum-neuron）来提供的。该库依赖于 torch-neuronx 包，而 torch-neuronx 是基于 PyTorch XLA 构建的。在当时这是自然的选择，但意味着需要像编译器一样思考：使用惰性张量而非 eager 执行、仅支持静态形状、每次图断开都需要重新编译。Optimum Neuron 的存在就是为了隐藏这些复杂性，并维护了自定义代码（模型类、训练器等）来做到这一点。

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

对于 Hugging Face 库来说，Trainium 成为了一个和任何其他设备一样的设备：Transformers、Diffusers、Accelerate、PEFT 和 TRL 可以直接在其上运行，使用与其他加速器相同的类和脚本。从现在起，我们建议 Neuron 用户直接使用这些库，而不是通过 Optimum Neuron。

https://huggingface.co/blog/hf-torchneuron#transformersTransformers

https://huggingface.co/blog/hf-torchneuron#inference推理

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": "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))

https://huggingface.co/blog/hf-torchneuron#training训练

既然 Neuron 是 PyTorch 的原生设备，在 Trainium 上的训练与在任何其他加速器上一样，使用相同的 Hugging Face 库：Transformers、Accelerate、PEFT 和 TRL。你不再需要专门的训练器或针对 Neuron 的模型代码分支。一个标准的 TRL 微调脚本只需几个 Neuron 特有的设置即可在 Trainium 实例上运行。

以下是基于 TRL 的 SFTTrainer 构建的监督微调脚本。如果你之前用 TRL 在 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):
    # 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 原生处理。你将包含张量并行（TP）和 FSDP 大小的 DistributedConfig 放入 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 × 8 = 32，与 --nproc_per_node 匹配。32 个进程形成一个 4 × 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 检查点。其中没有任何内容与 Neuron 绑定：你可以原样将其推送到 Hugging Face Hub，并在任何其他硬件上加载。

https://huggingface.co/blog/hf-torchneuron#diffusersDiffusers

https://huggingface.co/blog/hf-torchneuron#inference-1推理

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="A capybara wearing a wizard hat, oil painting", 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 尚不支持，你也可以按照本指南（https://huggingface.co/docs/diffusers/main/en/training/distributed_inference#writing-a-tensor-parallelism-plan）为其添加支持。

https://huggingface.co/blog/hf-torchneuron#kernelsKernels

请务必关注 Neuron 团队提供的 kernels（https://huggingface.co/aws-neuron/kernels），越来越多的 Neuron kernels 将在 Hub 上可用，为你的模型加速。Neuron kernels 可以使用 Neuron Kernel Interface（https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html）（NKI）编写，它提供了更底层的控制来优化运算。这样，用户可以运行自己的自定义操作并在 Hub 上分享。

例如，你可以使用这个 kernel（https://huggingface.co/kernels/aws-neuron/vit-encoder-neuron-kernels）来加速你的 vision transformer，在 Trainium 上可获得最高 2.4 倍的加速！

https://huggingface.co/blog/hf-torchneuron#acknowledgments致谢

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