AWS Trainium

AWS Trainium 实例通过 torch-neuronx 包原生运行 PyTorch，该包会注册 neuron 设备。Neuron 会为它看到的每种输入形状将模型编译为一个或多个图，因此保持输入形状静态可以避免重新编译。

安装

按照 AWS Neuron 设置指南安装 torch-neuronx。然后检查 PyTorch 是否识别 Neuron 设备。

import torch

print(torch.neuron.is_available(), torch.neuron.device_count())

它会打印 True 以及 Neuron 设备的数量，例如 True 4。

运行推理

device_map="auto" 无法检测 Neuron。自动放置依赖于 Accelerate 来列出可用设备，而 Accelerate 不会列出 Neuron 设备，因此模型会静默地加载到 CPU 上。请显式地将 device_map="neuron" 传递给 from_pretrained()，或将 device="neuron" 传递给 Pipeline。

在 CPU 上传入生成输入，而不是调用 .to(model.device)。generate() 只会将 forward 所消费的张量移动到 Neuron 设备。每一步中不断增长的张量（输出序列、attention_mask）会保留在 CPU 上，从而保持编译后的图稳定。generate 会在与输入相同的设备上返回输出。有关更多详细信息，请参阅将生成张量保留在 CPU 上。

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="neuron"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that the French government must provide a certain amount of bread for the citizens of the country. The law is passed by the French parliament. The French parliament is a body of the French government. The French government is a body of the French state. The

使用静态缓存时，generate() 会在 Neuron 上自动使用 torch.compile 编译前向传播。向 generate() 传入 disable_compile=True 即可禁用此行为。

使用 Trainer 进行训练

TrainingArguments 会检测 Neuron 并自动将设备设置为 neuron。所有 Trainium 代际均支持 bf16。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)

编译模型

将 torch_compile=True 设置为 True，以使用 torch.compile 包装模型。这会编译模型的前向传播，而优化器步骤则以急切模式运行。当检测到 Neuron 时，TrainingArguments 会自动将 torch_compile_backend 设置为 "neuron"。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    torch_compile=True,
)

每个新的输入形状都会触发一次新的编译。请将输入填充到固定长度（在分词器中设置 padding="max_length"），并设置 dataloader_drop_last=True，以确保最后一个批次的大小与其他批次一致。

分布式训练

Neuron 支持 Transformers 分布式 API。创建一个具有所需并行规模的 DistributedConfig，并将其传递给 from_pretrained()。Transformers 会在加载时划分模型，Trainer 会直接使用该配置。

import torch
from transformers import AutoModelForCausalLM, DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2),
)

使用 torchrun 启动脚本。

torchrun --nproc_per_node=4 train.py

有关可用选项，请参阅张量并行、专家并行和 FSDP。

局限性

device_map="auto" 会将模型加载到 CPU 而非 Neuron 上。请显式传入 device_map="neuron"。
动态形状会触发重新编译。当输入形状在每一步之间发生变化时，生成和训练都会变慢。生成时请使用静态缓存，训练时请使用固定长度填充。

Update on GitHub