AWS Trainium

Les instances AWS Trainium exécutent PyTorch nativement via le package torch-neuronx, qui enregistre le périphérique neuron. Neuron compile votre modèle en un ou plusieurs graphes pour chaque forme d’entrée qu’il rencontre, il est donc recommandé de conserver des formes d’entrée statiques pour éviter les recompilations.

Installation

Installez torch-neuronx en suivant le guide d’installation AWS Neuron. Vérifiez ensuite que PyTorch détecte les appareils Neuron.

import torch

print(torch.neuron.is_available(), torch.neuron.device_count())

Il affiche True et le nombre de dispositifs Neuron, par exemple True 4.

Exécuter l'inférence

device_map="auto" ne détecte pas Neuron. Le placement automatique s’appuie sur Accelerate pour lister les appareils disponibles, et Accelerate ne liste pas les appareils Neuron, de sorte que le modèle se charge silencieusement sur le CPU. Passez explicitement device_map="neuron" à from_pretrained(), ou device="neuron" à Pipeline.

Transmettez les entrées de génération sur le CPU au lieu d’appeler .to(model.device). generate() ne déplace que les tenseurs consommés par forward vers le périphérique Neuron. Les tenseurs qui augmentent à chaque étape (séquence de sortie, attention_mask) restent sur le CPU, ce qui maintient le graphe compilé stable. generate renvoie la sortie sur le même périphérique que les entrées. Consultez Conserver les tenseurs de génération sur le CPU pour plus de détails.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="neuron"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that the French government must provide a certain amount of bread for the citizens of the country. The law is passed by the French parliament. The French parliament is a body of the French government. The French government is a body of the French state. The

Avec un cache statique, generate() compile automatiquement le passage avant avec torch.compile sur Neuron. Passez disable_compile=True à generate() pour désactiver cette compilation.

Entraîner avec Trainer

TrainingArguments détecte Neuron et définit automatiquement l’appareil sur neuron. Toutes les générations de Trainium prennent en charge bf16.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)

Compiler le modèle

Définissez torch_compile=True pour envelopper le modèle avec torch.compile. Cela compile le passage avant du modèle, tandis que l’étape de l’optimiseur s’exécute de manière immédiate. TrainingArguments définit torch_compile_backend sur "neuron" automatiquement lorsqu’il détecte Neuron.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    torch_compile=True,
)

Chaque nouvelle forme d’entrée déclenche une nouvelle compilation. Remplissez les entrées jusqu’à une longueur fixe (padding="max_length" dans le tokenizer) et définissez dataloader_drop_last=True afin que le dernier lot n’ait pas une taille différente.

Entraînement distribué

Neuron prend en charge l’API distribuée de Transformers. Créez un DistributedConfig avec les tailles de parallélisme souhaitées et passez-le à from_pretrained(). Transformers découpe le modèle au moment du chargement, et Trainer l’utilise tel quel.

import torch
from transformers import AutoModelForCausalLM, DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2),
)

Lancez le script avec torchrun.

torchrun --nproc_per_node=4 train.py

Voir Tensor parallelism, Expert parallelism et FSDP pour les options disponibles.

Limites

• device_map="auto" charge le modèle sur le CPU au lieu de Neuron. Passez explicitement device_map="neuron".
• Les formes dynamiques déclenchent une recompilation. La génération et l’entraînement se ralentissent lorsque les formes des entrées changent d’une étape à l’autre. Utilisez un cache statique pour la génération et un bourrage de longueur fixe pour l’entraînement.