AWS Trainium

Las instancias de AWS Trainium ejecutan PyTorch de forma nativa a través del paquete torch-neuronx, que registra el dispositivo neuron. Neuron compila tu modelo en uno o más grafos para cada forma de entrada que detecta, por lo que mantener las formas de entrada estáticas evita recompilaciones.

Instalación

Instale torch-neuronx siguiendo la guía de configuración de AWS Neuron. Luego, verifique que PyTorch detecte los dispositivos Neuron.

import torch

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

Imprime True y el número de dispositivos Neuron, por ejemplo True 4.

Ejecutar inferencia

device_map="auto" no detecta Neuron. La colocación automática depende de Accelerate para listar los dispositivos disponibles, y Accelerate no lista los dispositivos Neuron, por lo que el modelo se carga silenciosamente en la CPU. Pasa device_map="neuron" a from_pretrained() de forma explícita, o device="neuron" a Pipeline.

Pase las entradas de generación en la CPU en lugar de llamar a .to(model.device). generate() solo mueve a la Neuron device los tensores que forward consume. Los tensores que crecen en cada paso (secuencia de salida, attention_mask) permanecen en la CPU, por lo que el gráfico compilado se mantiene estable. generate devuelve la salida en el mismo dispositivo que las entradas. Consulte Mantener los tensores de generación en la CPU para más detalles.

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

Con una caché estática, generate() compila la pasada hacia adelante con torch.compile automáticamente en Neuron. Pasa disable_compile=True a generate() para desactivarlo.

Entrenar con Trainer

TrainingArguments detecta Neuron y establece el dispositivo en neuron automáticamente. Todas las generaciones de Trainium admiten bf16.

from transformers import TrainingArguments

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

Compilar el modelo

Establezca torch_compile=True para envolver el modelo con torch.compile. Esto compila la pasada hacia adelante del modelo, mientras que el paso del optimizador se ejecuta de forma inmediata. TrainingArguments establece torch_compile_backend en "neuron" automáticamente cuando detecta Neuron.

from transformers import TrainingArguments

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

Cada nueva forma de entrada provoca una nueva compilación. Rellene las entradas hasta una longitud fija (padding="max_length" en el tokenizador) y establezca dataloader_drop_last=True para que el último lote no tenga un tamaño diferente.

Entrenamiento distribuido

Neuron admite la API distribuida de Transformers. Crea un DistributedConfig con los tamaños de paralelismo que desees y pásalo a from_pretrained(). Transformers fragmenta el modelo en el momento de la carga y Trainer lo utiliza tal cual.

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

Ejecuta el script con torchrun.

torchrun --nproc_per_node=4 train.py

Consulte Tensor parallelism, Expert parallelism y FSDP para conocer las opciones disponibles.

Limitaciones

• device_map="auto" carga el modelo en la CPU en lugar de Neuron. Pasa device_map="neuron" explícitamente.
• Las formas dinámicas provocan recompilación. Tanto la generación como el entrenamiento se ralentizan cuando las formas de entrada cambian de un paso a otro. Utilice una caché estática para la generación y relleno de longitud fija para el entrenamiento.