Huấn luyện trên AWS Trainium

AWS Trainium chạy PyTorch một cách bản địa thông qua gói torch-neuronx, trong đó đăng ký thiết bị neuron. Các trainer của TRL chạy trên thiết bị này mà không cần thay đổi mã. Neuron biên dịch một đồ thị cho mỗi hình dạng đầu vào mà nó gặp, và mỗi lần biên dịch mất khoảng một phút, vì vậy phần lớn nội dung của trang này là về việc giữ cho các hình dạng tĩnh. Để cài đặt và phía Transformers, hãy xem hướng dẫn AWS Trainium của Transformers.

Cài đặt

Cài đặt torch-neuronx bằng cách làm theo hướng dẫn thiết lập AWS Neuron, sau đó cài đặt TRL trong cùng môi trường.

pip install trl

Chỉnh huấn có giám sát

Bắt đầu nhanh với SFTTrainer chạy được ngay:

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()

Trên Neuron, pad_to_multiple_of mặc định là 128 (giá trị "auto" của nó), nên các chuỗi được đệm đến bội số của 128. Điều này giới hạn số lượng hình dạng, và do đó số lần biên dịch, xuống max_length / 128. Nếu không có nó, mỗi độ dài chuỗi mới đều kích hoạt một lần biên dịch. Trên một NeuronCore, với per_device_train_batch_size=4, 20 bước của ví dụ trên mất 130 giây với giá trị mặc định và 260 giây với pad_to_multiple_of=None.

Điều tương tự cũng áp dụng cho mọi trainer có tùy chọn pad_to_multiple_of: DPO, KTO, Reward, GRPO, RLOO và Distillation.

Biên dịch mô hình

Đặt torch_compile=True để biên dịch mô hình với backend neuron. Các hình dạng động tự động của Dynamo tạo ra các đồ thị mà backend Neuron chưa thể biên dịch (#7660), vì vậy hãy tắt chúng:

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()

Bước đầu tiên mất khoảng 2 phút để biên dịch. Sau đó, với per_device_train_batch_size=4, Qwen3-0.6B huấn luyện với tốc độ 1,15 giây/bước trên một NeuronCore, so với 1,71 giây/bước khi không biên dịch.

Phương pháp trực tuyến

vLLM không chạy trên Neuron, do đó GRPO và RLOO sẽ sinh bằng Transformers. Với bộ nhớ đệm động mặc định, mỗi token được sinh ra đều có một hình dạng mới, và một bước có thể mất vài phút. Hãy sử dụng bộ nhớ đệm tĩnh, giữ cho các hình dạng giải mã cố định:

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()

Các prompt cũng được đệm (padding) đến pad_to_multiple_of trước khi sinh, nên bộ nhớ đệm tĩnh chỉ chiếm một vài kích thước. Mỗi kích thước mới vẫn mất khoảng 12 phút để biên dịch trên một NeuronCore; sau khi một kích thước được biên dịch, một bước của ví dụ này mất khoảng 25 giây.

Bộ nhớ đệm tĩnh hiện tại cần một mô hình bf16: với mô hình fp32 mặc định, các key và value sẽ kết thúc ở các dtype khác nhau (#7664).

Song song dữ liệu

Trên một instance trn2.3xlarge, chứa một chip Trainium2, cấu hình Logical NeuronCore mặc định (LNC=2) ghép 8 NeuronCore vật lý của chip thành 4 NeuronCore logic.

Để huấn luyện trên cả bốn, hãy sử dụng Accelerate:

accelerate launch --num_processes 4 train.py

Trong cấu hình này, SFT chạy nhanh gấp 3,6 lần trên 4 NeuronCore so với trên 1 NeuronCore.

Để hạn chế một lần chạy trên một tập con của NeuronCores, hãy đặt NEURON_RT_VISIBLE_CORES thành 0,1, 2,3, hoặc 0-3. Các tổ hợp khác như 1,2 không được hỗ trợ: runtime Neuron sẽ không thể thiết lập các thao tác tập thể (collectives) xuyên qua chúng.

Update on GitHub