Обучение на AWS Trainium

AWS Trainium работает с PyTorch нативно через пакет torch-neuronx, который регистрирует устройство neuron. Тренировщики TRL запускаются на нём без изменения кода. Neuron компилирует граф для каждой формы входных данных, которую он встречает, и каждая компиляция занимает около минуты, поэтому большая часть этой страницы посвящена сохранению статичности форм. Для установки и информации о стороне Transformers см. руководство по AWS Trainium для Transformers.

Установка

Установите torch-neuronx, следуя инструкции по настройке AWS Neuron, затем установите 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 с со значением по умолчанию и 260 с при pad_to_multiple_of=None.

То же самое относится к каждому тренеру, у которого есть опция 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 обучается со скоростью 1.15 с/шаг на одном NeuronCore, тогда как без компиляции — 1.71 с/шаг.

Онлайн-методы

vLLM не работает на Neuron, поэтому GRPO и RLOO генерируют с помощью Transformers. При использовании динамического кэша по умолчанию каждый сгенерированный токен имеет новую форму, и один шаг занимает несколько минут. Используйте статический кэш, который сохраняет формы декодирования фиксированными:

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 перед генерацией, поэтому статический кэш занимает лишь несколько размеров. Каждый новый размер компилируется примерно 12 минут на одном NeuronCore; после компиляции размера один шаг в этом примере занимает около 25 с.

Статический кэш пока требует модель в bf16: при использовании модели по умолчанию в fp32 ключи и значения оказываются в разных типах данных (#7664).

Параллелизм данных

На инстансе trn2.3xlarge, содержащем один чип Trainium2, конфигурация по умолчанию Logical NeuronCore (LNC=2) объединяет 8 физических NeuronCore чипа в 4 логических NeuronCore.

Для обучения на всех четырёх используйте Accelerate:

accelerate launch --num_processes 4 train.py

В этой конфигурации SFT выполняется в 3,6 раза быстрее на 4 NeuronCore, чем на 1.

Чтобы ограничить выполнение подмножеством NeuronCores, установите NEURON_RT_VISIBLE_CORES в значение 0,1, 2,3 или 0-3. Другие комбинации, например 1,2, не поддерживаются: среда выполнения Neuron не может настроить коллективные операции для них.