AWS Trainium

As instâncias AWS Trainium executam PyTorch nativamente por meio do pacote torch-neuronx, que registra o dispositivo neuron. O Neuron compila seu modelo em um ou mais grafos para cada forma de entrada que ele identifica, portanto, manter as formas de entrada estáticas evita recompilações.

Instalação

Instale o torch-neuronx seguindo o guia de configuração do AWS Neuron. Em seguida, verifique se o PyTorch detecta os dispositivos Neuron.

import torch

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

Ele imprime True e o número de dispositivos Neuron, por exemplo True 4.

Executar inferência

device_map="auto" não detecta Neuron. A colocação automática depende do Accelerate para listar os dispositivos disponíveis, e o Accelerate não lista dispositivos Neuron, de modo que o modelo é carregado silenciosamente na CPU. Passe device_map="neuron" explicitamente para from_pretrained(), ou device="neuron" para Pipeline.

Passe os inputs de geração na CPU em vez de chamar .to(model.device). O generate() move apenas os tensors que o forward consome para o dispositivo Neuron. Os tensors que crescem a cada passo (sequência de saída, attention_mask) permanecem na CPU, de modo que o gráfico compilado permanece estável. O generate retorna a saída no mesmo dispositivo dos inputs. Consulte Manter os tensors de geração na CPU para mais detalhes.

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

Com um cache estático, generate() compila o forward pass com torch.compile automaticamente no Neuron. Passe disable_compile=True para generate() para desativar essa opção.

Treinar com o Trainer

TrainingArguments detecta o Neuron e define o dispositivo como neuron automaticamente. Todas as gerações do Trainium suportam bf16.

from transformers import TrainingArguments

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

Compilar o modelo

Defina torch_compile=True para envolver o modelo com torch.compile. Isso compila a passagem de frente (forward pass) do modelo, enquanto o passo do otimizador é executado de forma imediata (eager). O TrainingArguments define torch_compile_backend como "neuron" automaticamente quando detecta Neuron.

from transformers import TrainingArguments

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

Cada novo formato de entrada dispara uma nova compilação. Preencha as entradas até um comprimento fixo (padding="max_length" no tokenizer) e defina dataloader_drop_last=True para que o último lote não tenha um tamanho diferente.

Treinamento distribuído

O Neuron oferece suporte à API distribuída do Transformers. Crie um DistributedConfig com os tamanhos de paralelismo desejados e passe-o para from_pretrained(). O Transformers faz o particionamento do modelo no momento do carregamento, e o Trainer o utiliza diretamente.

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

Execute o script com torchrun.

torchrun --nproc_per_node=4 train.py

Consulte Paralelismo de tensor, Paralelismo de especialista e FSDP para as opções disponíveis.

Limitações

• device_map="auto" carrega o modelo na CPU em vez do Neuron. Passe device_map="neuron" explicitamente.
• Formas dinâmicas acionam a recompilação. A geração e o treinamento ficam mais lentos quando as formas de entrada mudam de uma etapa para outra. Use um cache estático para a geração e preenchimento de comprimento fixo para o treinamento.