Treinamento no AWS Trainium

AWS Trainium executa PyTorch nativamente por meio do pacote torch-neuronx, que registra o dispositivo neuron. Os treinadores do TRL rodam nele sem alterações de código. O Neuron compila um grafo para cada forma de entrada que encontra, e cada compilação leva cerca de um minuto, então a maior parte desta página trata de manter as formas estáticas. Para instalação e o lado do Transformers, consulte o guia do Transformers para AWS Trainium.

Instalação

Instale o torch-neuronx seguindo o guia de configuração do AWS Neuron e, em seguida, instale o TRL no mesmo ambiente.

pip install trl

Ajuste fino supervisionado

O início rápido do SFTTrainer funciona como está:

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

No Neuron, pad_to_multiple_of tem valor padrão de 128 (seu valor "auto"), de modo que as sequências são preenchidas até um múltiplo de 128. Isso limita o número de formas, e portanto de compilações, a max_length / 128. Sem isso, cada novo comprimento de sequência dispara uma compilação. Em um NeuronCore, com per_device_train_batch_size=4, 20 passos do exemplo acima levam 130 s com o valor padrão e 260 s com pad_to_multiple_of=None.

O mesmo se aplica a todos os treinadores que possuem a opção pad_to_multiple_of: DPO, KTO, Reward, GRPO, RLOO e Distillation.

Compilando o modelo

Defina torch_compile=True para compilar o modelo com o backend neuron. As formas dinâmicas automáticas do Dynamo geram grafos que o backend Neuron ainda não consegue compilar (#7660), portanto, desative-as:

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

A primeira etapa leva cerca de 2 minutos para compilar. Depois disso, com per_device_train_batch_size=4, o Qwen3-0.6B treina a 1,15 s/etapa em um NeuronCore, contra 1,71 s/etapa sem compilação.

Métodos online

O vLLM não roda no Neuron, portanto o GRPO e o RLOO geram com o Transformers. Com o cache dinâmico padrão, cada token gerado tem um novo shape, e uma etapa leva vários minutos. Use o cache estático, que mantém os shapes de decodificação fixos:

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

Os prompts também são preenchidos até pad_to_multiple_of antes da geração, de modo que o cache estático ocupa apenas alguns tamanhos. Cada novo tamanho ainda leva cerca de 12 minutos para compilar em um NeuronCore; uma vez que um tamanho é compilado, uma etapa deste exemplo leva cerca de 25 s.

O cache estático requer um modelo bf16 por enquanto: com o modelo fp32 padrão, as chaves e os valores acabam em dtypes diferentes (#7664).

Paralelismo de dados

Em uma instância trn2.3xlarge, que contém um chip Trainium2, a configuração padrão de NeuronCore Lógico (LNC=2) agrupa os 8 NeuronCores físicos do chip em 4 NeuronCores lógicos.

Para treinar em todos os quatro, inicie um processo por NeuronCore com torchrun:

torchrun --nproc_per_node 4 train.py

Nessa configuração, o SFT é executado 3,6x mais rápido em 4 NeuronCores do que em 1.

Para restringir uma execução a um subconjunto de NeuronCores, defina NEURON_RT_VISIBLE_CORES como 0,1, 2,3 ou 0-3. Outras combinações, como 1,2, não são suportadas: o runtime do Neuron falha ao configurar coletivas entre elas.