# 在 AWS Trainium 上进行训练

[AWS Trainium](https://aws.amazon.com/ai/machine-learning/trainium/) 通过 `torch-neuronx` 包原生运行 PyTorch，该包会注册 `neuron` 设备。TRL 训练器无需修改代码即可在其上运行。Neuron 会为遇到的每个输入形状编译一个图，每次编译大约需要一分钟，因此本页的大部分内容都涉及如何保持形状静态。有关安装和 Transformers 方面的信息，请参阅 [Transformers AWS Trainium 指南](https://huggingface.co/docs/transformers/main/en/aws_trainium)。

## 安装

按照 [AWS Neuron 设置指南](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/pytorch/index.html) 安装 `torch-neuronx`，然后在同一环境中安装 TRL。

```bash
pip install trl
```

## 监督微调

[SFTTrainer](/docs/trl/v1.18.0/zh/sft_trainer#trl.SFTTrainer) 快速入门可直接运行：

```python
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](https://github.com/huggingface/trl/issues/7660)），因此请将其关闭：

```python
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，它保持解码形状固定：

```python
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 秒。

> [!NOTE]
> 目前，静态缓存需要一个 bf16 模型：如果使用默认的 fp32 模型，键和值最终会采用不同的数据类型（[#7664](https://github.com/huggingface/trl/issues/7664)）。

## 数据并行

在包含一个 Trainium2 芯片的 `trn2.3xlarge` 实例上，默认的 [逻辑 NeuronCore 配置](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/logical-neuroncore-config.html)（`LNC=2`）会将芯片的 8 个物理 NeuronCore 配对为 4 个逻辑 NeuronCore。

要在所有四个 NeuronCore 上进行训练，请使用 [torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html) 为每个 NeuronCore 启动一个进程：

```bash
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 运行时无法在这些核心之间设置集合通信。