# AWS Trainium

As instâncias [AWS Trainium](https://aws.amazon.com/ai/machine-learning/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](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/pytorch/index.html). Em seguida, verifique se o PyTorch detecta os dispositivos Neuron.

```py
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](https://huggingface.co/docs/accelerate/index) para listar os dispositivos disponíveis, e o Accelerate não lista dispositivos Neuron, então o modelo é carregado silenciosamente na CPU. Passe `device_map="neuron"` explicitamente para [from_pretrained()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained), ou `device="neuron"` para [Pipeline](/docs/transformers/main/pt/main_classes/pipelines#transformers.Pipeline).

Passe os inputs de geração na CPU em vez de chamar `.to(model.device)`. O [generate()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.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](./kv_cache#keep-generation-tensors-on-cpu) para mais detalhes.

```py
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))
```

```txt
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()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.generate) compila o forward pass com `torch.compile` automaticamente no Neuron. Passe `disable_compile=True` para [generate()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.generate) para desativar essa opção.

## Treinar com o Trainer

[TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments) detecta o Neuron e define o dispositivo como `neuron` automaticamente. Todas as gerações do Trainium suportam bf16.

```py
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 do modelo, enquanto o passo do otimizador é executado de forma imediata. O [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments) define `torch_compile_backend` como `"neuron"` automaticamente quando detecta Neuron.

```py
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](/docs/transformers/main/pt/expert_parallelism#transformers.DistributedConfig) com os tamanhos de paralelismo desejados e passe-o para [from_pretrained()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained). O Transformers faz o particionamento do modelo no momento do carregamento, e o [Trainer](/docs/transformers/main/pt/main_classes/trainer#transformers.Trainer) o utiliza diretamente.

```py
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`.

```bash
torchrun --nproc_per_node=4 train.py
```

Consulte [Paralelismo de tensor](./tensor_parallelism), [Paralelismo de especialista](./expert_parallelism) e [FSDP](./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.