Entrenamiento en AWS Trainium

AWS Trainium ejecuta PyTorch de forma nativa a través del paquete torch-neuronx, que registra el dispositivo neuron. Los entrenadores de TRL funcionan en él sin cambios de código. Neuron compila un grafo para cada forma de entrada que detecta, y cada compilación tarda aproximadamente un minuto, por lo que la mayor parte de esta página se centra en mantener las formas estáticas. Para la instalación y el lado de Transformers, consulte la guía de AWS Trainium para Transformers.

Instalación

Instale torch-neuronx siguiendo la guía de configuración de AWS Neuron y, a continuación, instale TRL en el mismo entorno.

pip install trl

Ajuste fino supervisado

La guía de inicio rápido de SFTTrainer se ejecuta tal cual:

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

En Neuron, pad_to_multiple_of tiene un valor predeterminado de 128 (su valor "auto"), por lo que las secuencias se rellenan hasta un múltiplo de 128. Esto limita el número de formas, y por tanto de compilaciones, a max_length / 128. Sin ello, cada nueva longitud de secuencia provoca una compilación. En un NeuronCore, con per_device_train_batch_size=4, 20 pasos del ejemplo anterior tardan 130 s con el valor predeterminado y 260 s con pad_to_multiple_of=None.

Lo mismo aplica a cada entrenador que tenga la opción pad_to_multiple_of: DPO, KTO, Reward, GRPO, RLOO y Distillation.

Compilar el modelo

Establezca torch_compile=True para compilar el modelo con el backend neuron. Las formas dinámicas automáticas de Dynamo generan grafos que el backend Neuron aún no puede compilar (#7660), así que desactívelas:

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

El primer paso compila durante aproximadamente 2 minutos. Después de eso, con per_device_train_batch_size=4, Qwen3-0.6B entrena a 1.15 s/paso en un solo NeuronCore, frente a 1.71 s/paso sin compilación.

Métodos en línea

vLLM no se ejecuta en Neuron, por lo que GRPO y RLOO generan con Transformers. Con la caché dinámica predeterminada, cada token generado tiene una forma nueva y un paso tarda varios minutos. Utilice la caché estática, que mantiene las formas de decodificación fijas:

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

Los prompts también se rellenan hasta pad_to_multiple_of antes de la generación, por lo que la caché estática solo ocupa unos pocos tamaños. Cada nuevo tamaño aún se compila durante unos 12 minutos en un NeuronCore; una vez que un tamaño está compilado, un paso de este ejemplo tarda unos 25 s.

La caché estática requiere un modelo bf16 por ahora: con el modelo fp32 predeterminado, las claves y los valores terminan en diferentes dtypes (#7664).

Paralelismo de datos

En una instancia trn2.3xlarge, que contiene un chip Trainium2, la configuración predeterminada de NeuronCore lógicos (LNC=2) agrupa los 8 NeuronCore físicos del chip en 4 NeuronCore lógicos.

Para entrenar en los cuatro, inicie un proceso por NeuronCore con torchrun:

torchrun --nproc_per_node 4 train.py

En esta configuración, SFT se ejecuta 3,6 veces más rápido en 4 NeuronCores que en 1.

Para restringir una ejecución a un subconjunto de NeuronCores, establezca NEURON_RT_VISIBLE_CORES en 0,1, 2,3 o 0-3. Otras combinaciones, como 1,2, no se admiten: el runtime de Neuron no puede configurar los colectivos entre ellas.