Entraînement sur AWS Trainium

AWS Trainium exécute PyTorch nativement via le package torch-neuronx, qui enregistre le périphérique neuron. Les formateurs TRL s’exécutent dessus sans modification de code. Neuron compile un graphe pour chaque forme d’entrée qu’il rencontre, et chaque compilation prend environ une minute, la majeure partie de cette page porte donc sur le maintien de formes statiques. Pour l’installation et le côté Transformers, consultez le guide Transformers AWS Trainium.

Installation

Installez torch-neuronx en suivant le guide d’installation AWS Neuron, puis installez TRL dans le même environnement.

pip install trl

Ajustement fin supervisé

Le démarrage rapide de SFTTrainer fonctionne tel quel :

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

Sur Neuron, pad_to_multiple_of est défini par défaut à 128 (sa valeur "auto"), de sorte que les séquences sont remplies jusqu’à un multiple de 128. Cela limite le nombre de formes, et donc de compilations, à max_length / 128. Sans cela, chaque nouvelle longueur de séquence déclenche une compilation. Sur un NeuronCore, avec per_device_train_batch_size=4, 20 étapes de l’exemple ci-dessus prennent 130 s avec la valeur par défaut et 260 s avec pad_to_multiple_of=None.

La même règle s’applique à tous les formateurs disposant de l’option pad_to_multiple_of : DPO, KTO, Reward, GRPO, RLOO et Distillation.

Compilation du modèle

Définissez torch_compile=True pour compiler le modèle avec l’arrière-plan neuron. Les formes dynamiques automatiques de Dynamo génèrent des graphes que l’arrière-plan Neuron ne peut pas encore compiler (#7660), il est donc nécessaire de les désactiver :

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

La première étape prend environ 2 minutes pour la compilation. Ensuite, avec per_device_train_batch_size=4, Qwen3-0.6B s’entraîne à 1,15 s/étape sur un seul NeuronCore, contre 1,71 s/étape sans compilation.

Méthodes en ligne

vLLM ne fonctionne pas sur Neuron, GRPO et RLOO génèrent donc avec Transformers. Avec le cache dynamique par défaut, chaque jeton généré a une nouvelle forme, et une étape prend plusieurs minutes. Utilisez le cache statique, qui maintient les formes de décodage fixes :

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

Les prompts sont également rembourrés jusqu'à pad_to_multiple_of avant la génération, de sorte que le cache statique n'occupe que quelques tailles. Chaque nouvelle taille prend environ 12 minutes à compiler sur un seul NeuronCore ; une fois une taille compilée, une étape de cet exemple prend environ 25 s.

Le cache statique nécessite pour l'instant un modèle bf16 : avec le modèle fp32 par défaut, les clés et les valeurs finissent par avoir des types de données différents (#7664).

Parallélisme des données

Sur une instance trn2.3xlarge, qui contient une puce Trainium2, la configuration par défaut des NeuronCores logiques (LNC=2) associe les 8 NeuronCores physiques de la puce en 4 NeuronCores logiques.

Pour entraîner sur les quatre NeuronCore, lancez un processus par NeuronCore avec torchrun :

torchrun --nproc_per_node 4 train.py

Dans cette configuration, SFT s'exécute 3,6 fois plus vite sur 4 NeuronCores que sur 1.

Pour restreindre une exécution à un sous-ensemble de NeuronCores, définissez NEURON_RT_VISIBLE_CORES sur 0,1, 2,3 ou 0-3. D'autres combinaisons, telles que 1,2, ne sont pas prises en charge : le runtime Neuron échoue à configurer les collectives à travers elles.