# AWS Trainium

[AWS Trainium](https://aws.amazon.com/ai/machine-learning/trainium/) — это инстансы, которые нативно поддерживают PyTorch через пакет `torch-neuronx`, регистрирующий устройство `neuron`. Neuron компилирует вашу модель в один или несколько графов для каждой формы входных данных, поэтому поддержание статических форм входных данных позволяет избежать повторной компиляции.

## Установка

Установите `torch-neuronx`, следуя [инструкции по настройке AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/pytorch/index.html). Затем убедитесь, что PyTorch видит устройства Neuron.

```py
import torch

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

Выводится `True` и количество устройств Neuron, например `True 4`.

## Запуск вывода

`device_map="auto"` не определяет Neuron. Автоматическое размещение зависит от [Accelerate](https://huggingface.co/docs/accelerate/index), который перечисляет доступные устройства, но Accelerate не включает в список устройства Neuron, поэтому модель молча загружается на CPU. Явно передайте `device_map="neuron"` в [from_pretrained()](/docs/transformers/main/ru/main_classes/model#transformers.PreTrainedModel.from_pretrained) или `device="neuron"` в [Pipeline](/docs/transformers/main/ru/main_classes/pipelines#transformers.Pipeline).

Передавайте входные данные для генерации на CPU, а не вызывайте `.to(model.device)`. [generate()](/docs/transformers/main/ru/main_classes/text_generation#transformers.GenerationMixin.generate) перемещает на устройство Neuron только те тензоры, которые потребляет `forward`. Тензоры, растущие на каждом шаге (выходная последовательность, `attention_mask`), остаются на CPU, благодаря чему скомпилированный граф остаётся стабильным. `generate` возвращает результат на том же устройстве, что и входные данные. Подробнее см. в разделе [Keep generation tensors on CPU](./kv_cache#keep-generation-tensors-on-cpu).

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

При использовании статического кэша [generate()](/docs/transformers/main/ru/main_classes/text_generation#transformers.GenerationMixin.generate) автоматически компилирует прямой проход с помощью `torch.compile` на Neuron. Передайте `disable_compile=True` в [generate()](/docs/transformers/main/ru/main_classes/text_generation#transformers.GenerationMixin.generate), чтобы отключить эту функцию.

## Обучение с помощью Trainer

[TrainingArguments](/docs/transformers/main/ru/main_classes/trainer#transformers.TrainingArguments) определяет Neuron и автоматически устанавливает устройство в `neuron`. Все поколения Trainium поддерживают bf16.

```py
from transformers import TrainingArguments

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

### Компиляция модели

Установите `torch_compile=True`, чтобы обернуть модель в `torch.compile`. Это компилирует прямой проход модели, при этом шаг оптимизатора выполняется в обычном режиме. [TrainingArguments](/docs/transformers/main/ru/main_classes/trainer#transformers.TrainingArguments) автоматически устанавливает `torch_compile_backend` в `"neuron"` при обнаружении Neuron.

```py
from transformers import TrainingArguments

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

Каждая новая форма входных данных запускает новую компиляцию. Выровняйте входные данные до фиксированной длины (`padding="max_length"` в токенизаторе) и установите `dataloader_drop_last=True`, чтобы последний батч не имел другого размера.

### Распределённое обучение

Neuron поддерживает распределённый API Transformers. Создайте [DistributedConfig](/docs/transformers/main/ru/expert_parallelism#transformers.DistributedConfig) с нужными размерами параллелизма и передайте его в [from_pretrained()](/docs/transformers/main/ru/main_classes/model#transformers.PreTrainedModel.from_pretrained). Transformers разбивает модель на части при загрузке, а [Trainer](/docs/transformers/main/ru/main_classes/trainer#transformers.Trainer) использует её как есть.

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

Запустите скрипт с помощью `torchrun`.

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

См. [Tensor parallelism](./tensor_parallelism), [Expert parallelism](./expert_parallelism) и [FSDP](./fsdp) для доступных вариантов.

## Ограничения

- `device_map="auto"` загружает модель на CPU вместо Neuron. Явно передайте `device_map="neuron"`.
- Динамические формы вызывают повторную компиляцию. Генерация и обучение замедляются, если формы входных данных меняются от шага к шагу. Используйте статический кэш для генерации и фиксированную длину отступов (padding) для обучения.