TorchNeuron：在 AWS Trainium 上运行 HF 模型从未如此简单

TorchNeuron 是一个 PyTorch 插件，为 Trainium 提供了一个新的后端，支持 eager 模式、torch.compile 以及 FSDP 和 TP 等原生分布式 API。阅读 PyTorch 博客，了解有关 TorchNeuron 的更多信息。TorchNeuron is a PyTorch plugin that brings a new backend for Trainium enabling support for eager mode, torch.compile and native distributed APIs like FSDP and TP. Read the Pytorch blog to learn more about TorchNeuron.

自 2026 年 10 月 20 日起，TorchNeuron 将开放公开测试版。研究人员和开发者可以通过 pip install torch-neuronx 安装。Starting Oct. 20, 2026, TorchNeuron will be available as a public beta. Researchers and developers will be able to install it with pip install torch-neuronx.

HF 库中的全新后端和原生接口 New backend, native interface in HF libraries

迄今为止，Hugging Face 在 AWS Trainium 和 Inferentia 上的使用体验——这两款 AI 芯片分别针对高性价比训练和大规模推理而设计——一直通过 optimum-neuron 提供。该库依赖 torch-neuronx 包，而后者基于 PyTorch XLA。当时这自然是个合适的选择，但也意味着必须像编译器那样思考：使用惰性张量而非即时执行、只支持静态形状，并且每次遇到计算图中断都要重新编译。Optimum Neuron 的存在就是为了隐藏这些复杂性，并为此维护了自定义代码（建模类、训练器等）。To date, the Hugging Face experience on AWS Trainium and Inferentia — AI chips designed respectively for cost-effective training and inference at scale — has been delivered through optimum-neuron. This library relies on the torch-neuronx package, which is based on PyTorch XLA. It was the natural choice at the time, but it meant thinking like a compiler: lazy tensors instead of eager execution, static shapes only, a recompilation for every graph break. Optimum Neuron existed to hide that complexity, and maintained custom code (modeling classes, trainers, etc) to do so.

TorchNeuron 改变了这一点。它直接集成到 PyTorch 中，并将 neuron 作为原生设备提供，支持 eager 执行、PyTorch 自身的分布式原语（DTensor、FSDP2）以及 torch.compile。TorchNeuron changes this. It integrates directly into PyTorch and exposes neuron as a native device, with eager execution, PyTorch's own distributed primitives (DTensor, FSDP2) and torch.compile.

对于 Hugging Face 库而言，Trainium 就像其他设备一样：Transformers、Diffusers、Accelerate、PEFT 和 TRL 都可以直接在其上运行，使用的类和脚本与在其他加速器上完全相同。从现在起，我们建议 Neuron 用户直接使用这些库，而不是通过 Optimum Neuron。For the Hugging Face libraries, Trainium becomes a device like any other: Transformers, Diffusers, Accelerate, PEFT and TRL run on it directly, with the same classes and scripts you would use on any other accelerator. From now on, we recommend that Neuron users use these libraries directly instead of going through Optimum Neuron.

变形金刚 Transformers

推理 Inference

Transformers 现在原生支持多种模型的文本生成。借助 TorchNeuron，使用率最高的前 100 种文本生成模型架构中，有 89% 开箱即用即可获得支持；而 Optimum Neuron v0.4.5 的支持比例为 46%。你只需在代码中传入 device_map=neuron 参数，模型就会被加载到 Neuron 设备上。Text generation is now natively supported in Transformers for a wide range of models. With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box, compared to 46% with Optimum Neuron v0.4.5. You just need to pass the device_map=neuron argument in your code and your model will be moved to the Neuron devices.

以下是一个使用张量并行进行文本生成的示例，可充分发挥 Neuron 实例的性能。请注意，该示例使用静态缓存进行生成，因为强烈建议在 Neuron 上采用这种方式，以避免频繁重新编译。The following is an example of text-generation with tensor parallelism, which will enable to fully benefit from your Neuron instance. Note that it uses a static cache for generation as this is strongly recommended with Neuron to prevent frequent recompilations.

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

培训 Training

现在 Neuron 已成为 PyTorch 原生设备，在 Trainium 上训练时，可以像使用其他加速器一样，使用相同的 Hugging Face 库：Transformers、Accelerate、PEFT 和 TRL。你不再需要专用的训练器，也不需要 Neuron 专属的模型代码分支。标准的 TRL 微调脚本只需设置少量 Neuron 专属参数，就能在 Trainium 实例上运行。Now that Neuron is a native PyTorch device, training on Trainium uses the same Hugging Face libraries as any other accelerator: Transformers, Accelerate, PEFT and TRL. You no longer need a dedicated trainer or a Neuron-specific fork of the model code. A standard TRL fine-tuning script runs on a Trainium instance with only a few Neuron-specific settings.

这是一个基于 TRL 的 SFTTrainer 构建的监督微调脚本。如果你以前用 TRL 在 GPU 上微调过模型，这个脚本应该会让你感到熟悉。Here is a supervised fine-tuning script built on TRL's SFTTrainer. If you have fine-tuned a model with TRL on GPUs before, it should look familiar.

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):
    # The Neuron backend can't compile Dynamo's automatic dynamic shapes yet (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:
        # The model is sharded on the TP x FSDP grid as its weights are loaded
        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")),
        )

    # Prepare dataset
    train_dataset = load_dataset(
        script_args.dataset_name,
        name=script_args.dataset_config,
        split=script_args.dataset_train_split,
    )
    # UltraChat also has a string `prompt` column, which makes TRL treat it as a prompt-completion dataset.
    train_dataset = train_dataset.select_columns(["messages"])

    # Training
    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 配置。Parallelism is handled natively by Transformers. You put a DistributedConfig with the tensor-parallel (TP) and FSDP sizes in model_init_kwargs, SFTTrainer passes it to from_pretrained, and the model is sharded as its weights are loaded. The Trainer sees that the model is already distributed and does not wrap it again. You don't need an --fsdp flag or an Accelerate FSDP config.

脚本中的内容都不依赖 Trainium：相同的代码可以在任何硬件上运行。在 Trainium 上运行时，有几点值得了解：Nothing in the script is tied to Trainium: the same code runs on any hardware. A few things are worth knowing when you run it on Trainium:

• 静态形状。该脚本会在 torch.compile 中禁用自动动态形状，并使用封装式打包生成最长 2048 个 token 的序列。使用 TRL ≥ 1.16 时，Neuron 上的批次会填充至 128 个 token 的倍数，从而限制输入形状的数量。后续步骤中遇到相同形状时，会复用已编译的计算图。
  Static shapes. The script disables automatic dynamic shapes in torch.compile and uses wrapped packing to produce sequences of up to 2048 tokens. With TRL ≥ 1.16, batches on Neuron are padded to a multiple of 128 tokens, limiting the number of input shapes. Compiled graphs are reused for subsequent steps with the same shapes.
• 启用 torch.compile。只需传入 --torch_compile true：在 Trainium 实例上，torch.compile 会检测到设备并自动选择 Neuron 后端，无需进行其他配置。
  torch.compile on. Passing --torch_compile true is all it takes: on a Trainium instance, torch.compile detects the device and picks the Neuron backend automatically, with nothing else to configure.
• 关闭异步权重加载。Transformers 内置的异步权重加载功能不受 Neuron 支持，因此目前我们通过设置 HF_DEACTIVATE_ASYNC_LOAD=1 来禁用它。
  Async weight loading off. The built-in asynchronous weight loading in Transformers is not supported on Neuron, so for now we disable it with 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 组之间分片参数、梯度和优化器状态。The command below launches training with torchrun, which starts one process per NeuronCore. TP_SIZE and FSDP_SIZE set the parallelism layout, and their product is the number of processes: here 4 x 8 = 32, matching --nproc_per_node. The 32 processes form a 4 x 8 grid. Tensor parallelism splits each layer across the 4 cores of a TP group, and FSDP shards the parameters, gradients and optimizer state across the 8 TP groups.

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 实例的一半。We ran it on 8 Trainium2 chips (32 NeuronCores), which is half a trn2.48xlarge instance.

训练会运行至完成，trainer.save_model 会写入一个标准的 Transformers 检查点。其中没有任何内容与 Neuron 绑定：你可以直接将其推送到 Hugging Face Hub，并在任何其他硬件上加载。Training runs to completion, and trainer.save_model writes a standard Transformers checkpoint. Nothing in it is tied to Neuron: you can push it to the Hugging Face Hub as is and load it on any other hardware.

Diffusers Diffusers

推理 Inference

Diffusers 库中的许多 pipeline 都兼容 TorchNeuron 的即时执行模式和编译模式。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 和 H3 也兼容张量并行，可充分利用你的 Neuron 设备。下面的示例展示了如何使用 TP：Many pipelines from the Diffusers library are compatible with TorchNeuron in eager and compiled mode. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 are also compatible with tensor parallelism to make the most of your Neuron devices. The following is an example showing how you can use TP:

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

# Distributed devices initialization
dist.init_process_group(backend="neuron") # TP
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", list(range(8))) # TP=8

# Pipeline initialization
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)

# Image generation
emb, mask, _ = pipe.encode_prompt(prompt="A capybara wearing a wizard hat, oil painting", device="cpu")  # text encoder stays on 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), # same noise on every rank
).images[0]
if dist.get_rank() == 0:
    image.save("out.png")

You can run this piece of code with: NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py.

将会有越来越多的流水线兼容 TP。如果你喜欢的流水线尚不兼容，也可以按照本指南为其添加支持。More and more pipleines will be compatible with TP. If a pipeline you like is not, you can also add support for it following this guide.

内核 Kernels

请务必关注 Neuron 团队的内核，Hub 上将陆续推出越来越多的 Neuron 内核，帮助加速模型。Neuron 内核可以使用 Neuron Kernel Interface (NKI) 编写，它能提供更底层的控制能力，以便优化操作。这样，用户就可以运行自己的自定义操作，并将其分享到 Hub 上。Make sure to check out the kernels from the Neuron team as more and more Neuron kernels will become available on the Hub to speed up your models. Neuron kernels can be written using the Neuron Kernel Interface (NKI), which gives lower control to optimize operations. This way, users can run their own custom operations and share them on the hub.

例如，你可以使用这个内核来加速视觉 Transformer，在 Trainium 上最高可实现 2.4 倍的加速！You can use for example this kernel to speed up your vision transformers with up to a x2.4 speedup on Trainium!

致谢 Acknowledgments

特别感谢 AWS 团队在 TorchNeuron 上所做的工作，以及他们积极响应我们的反馈，尽可能完善 HF 集成。Special thanks to the AWS team for their work on TorchNeuron and for their reactivity to our feedback to make the HF integration as good as possible.