TorchNeuron: running HF models on AWS Trainium is easier than ever

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

返回文章列表

[-1

]

[0

发布于 2026 年 10 月 19 日

GitHub 上的更新

点赞

12

[

• +6

Régis Pierrard

regisss

关注

Simon Pagezy

pagezyhf

Michael Benayoun

michaelbenayoun

Jingya Huang

Jingya

Alvaro Moran

tengomucho

David Corvoysier

dacorvo

Jeff Boudier

jeffboudier

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

TorchNeuron

TorchNeuron 是一个 PyTorch 插件，为 Trainium 提供了新的后端，支持急切模式、torch.compile 以及原生分布式 API（如 FSDP 和 TP）。阅读 PyTorch 博客以了解更多关于 TorchNeuron 的信息。

自 2026 年 10 月 20 日起，TorchNeuron 将作为公开测试版提供。研究人员和开发者可以通过 pip install torch-neuronx 安装它。

新后端，HF 库中的原生接口

迄今为止，Hugging Face 在 AWS Trainium 和 Inferentia（分别针对大规模训练和推理的高效能 AI 芯片）上的体验是通过 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 设备上。

以下是一个使用张量并行的文本生成示例，这将使你能够充分利用 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 ))

#### 训练

由于 Neuron 现在是原生的 PyTorch 设备，Trainium 上的训练使用与任何其他加速器相同的 Hugging Face 库：Transformers、Accelerate、PEFT 和 TRL。您不再需要专用的训练器或模型代码的 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 将其视为提示-完成数据集

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 后端，无需其他配置。

• 关闭异步权重加载。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 个 NeuronCores）运行，相当于半台 trn2.48xlarge 实例。

训练可完整执行，trainer.save_model 会生成标准 Transformers 检查点。其中内容与 Neuron 无关：你可以直接将其推送到 Hugging Face Hub，并在任何其他硬件上加载使用。

Diffusers

Diffusers 库的许多流水线在急切模式和编译模式下均兼容 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
# 流水线初始化
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 .

越来越多的流水线将兼容TP。如果喜欢的流水线目前不支持，也可以按照此指南添加支持。

内核

请务必查看Neuron团队提供的内核，越来越多的Neuron内核将在Hub上发布，以加速你的模型。Neuron内核可以使用Neuron内核接口（NKI）编写，这为优化操作提供了更低级别的控制。这样用户可以运行自定义操作并在Hub上分享。

例如，你可以使用这个内核在Trainium上将视觉Transformer加速高达2.4倍！

致谢

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

更多博客文章

torch

profile

CUDA

PyTorch性能分析（第二部分）：从nn.Linear到融合MLP

• +1

65

2026年6月11日

PyTorch性能分析（第一部分）：torch.profiler入门指南

183

2026年5月29日

社区

编辑

预览

通过拖拽文本输入框、粘贴或

点击此处

上传图片、音频和视频。

点击此处上传图片

评论

· 注册或登录以发表评论