Hugging Face's logo

Transformers documentation

AWS Trainium

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

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

←Google TPU Intel Gaudi→