在 AWS Trainium 上进行训练

AWS Trainium 通过 torch-neuronx 包原生运行 PyTorch，该包会注册 neuron 设备。TRL 训练器无需修改代码即可在其上运行。Neuron 会为遇到的每个输入形状编译一个图，每次编译大约需要一分钟，因此本页的大部分内容都涉及如何保持形状静态。有关安装和 Transformers 方面的信息，请参阅 Transformers AWS Trainium 指南。

安装

按照 AWS Neuron 设置指南安装 torch-neuronx，然后在同一环境中安装 TRL。

pip install trl

监督微调

SFTTrainer 快速入门可直接运行：

from datasets import load_dataset
from trl import SFTTrainer

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

在 Neuron 上，pad_to_multiple_of 的默认值为 128（即其 "auto" 值），因此序列会被填充为 128 的倍数。这将形状的数量（进而编译的次数）限制在 max_length / 128 以内。如果没有该设置，每个新的序列长度都会触发一次编译。在单个 NeuronCore 上，当 per_device_train_batch_size=4 时，上述示例的 20 个训练步骤在使用默认设置下耗时 130 秒，而使用 pad_to_multiple_of=None 时耗时 260 秒。

所有具有 pad_to_multiple_of 选项的训练器均适用此规则：DPO、KTO、Reward、GRPO、RLOO 和 Distillation。

编译模型

将 torch_compile=True 设置为 neuron 后端以编译模型。Dynamo 的自动动态形状会生成 Neuron 后端目前尚无法编译的图（#7660），因此请将其关闭：

import torch
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

torch._dynamo.config.automatic_dynamic_shapes = False

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    args=SFTConfig(torch_compile=True),
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

第一步的编译大约需要 2 分钟。之后，在 per_device_train_batch_size=4 的设置下，Qwen3-0.6B 在单个 NeuronCore 上的训练速度为 1.15 秒/步，而未编译时为 1.71 秒/步。

在线方法

vLLM 无法在 Neuron 上运行，因此 GRPO 和 RLOO 使用 Transformers 进行生成。使用默认的 dynamic cache 时，每个生成的 token 都具有新的形状，一个 step 需要几分钟。请使用 static cache，它保持解码形状固定：

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer


# Reward completions close to 200 characters
def reward_len(completions, **kwargs):
    return [-abs(200 - len(completion)) / 100 for completion in completions]


trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    reward_funcs=reward_len,
    args=GRPOConfig(
        per_device_train_batch_size=4,
        num_generations=4,
        max_completion_length=128,
        cache_implementation="static",
        model_init_kwargs={"dtype": "bfloat16"},
    ),
    train_dataset=load_dataset("trl-lib/tldr", split="train"),
)
trainer.train()

提示词在生成前也会被填充至 pad_to_multiple_of，因此静态缓存只需占用少数几种尺寸。每个新尺寸在单个 NeuronCore 上仍需编译约 12 分钟；一旦某个尺寸完成编译，本示例的单步耗时约为 25 秒。

目前，静态缓存需要一个 bf16 模型：如果使用默认的 fp32 模型，键和值最终会采用不同的数据类型（#7664）。

数据并行

在包含一个 Trainium2 芯片的 trn2.3xlarge 实例上，默认的逻辑 NeuronCore 配置（LNC=2）会将芯片的 8 个物理 NeuronCore 配对为 4 个逻辑 NeuronCore。

要在所有四个 NeuronCore 上进行训练，请使用 torchrun 为每个 NeuronCore 启动一个进程：

torchrun --nproc_per_node 4 train.py

在此配置中，SFT 在 4 个 NeuronCore 上的运行速度是在 1 个 NeuronCore 上的 3.6 倍。

要将运行限制在 NeuronCore 的子集上，请将 NEURON_RT_VISIBLE_CORES 设置为 0,1、2,3 或 0-3。其他组合（例如 1,2）不受支持：Neuron 运行时无法在这些核心之间设置集合通信。

Update on GitHub