AWS Trainium

AWS Trainium — это инстансы, которые нативно поддерживают PyTorch через пакет torch-neuronx, регистрирующий устройство neuron. Neuron компилирует вашу модель в один или несколько графов для каждой формы входных данных, поэтому поддержание статических форм входных данных позволяет избежать повторной компиляции.

Установка

Установите torch-neuronx, следуя инструкции по настройке AWS Neuron. Затем убедитесь, что PyTorch видит устройства Neuron.

import torch

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

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

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

device_map="auto" не определяет Neuron. Автоматическое размещение зависит от Accelerate, который перечисляет доступные устройства, но Accelerate не включает в список устройства Neuron, поэтому модель молча загружается на CPU. Явно передайте device_map="neuron" в from_pretrained() или device="neuron" в Pipeline.

Передавайте входные данные для генерации на CPU, а не вызывайте .to(model.device). generate() перемещает на устройство Neuron только те тензоры, которые потребляет forward. Тензоры, растущие на каждом шаге (выходная последовательность, attention_mask), остаются на CPU, благодаря чему скомпилированный граф остаётся стабильным. generate возвращает результат на том же устройстве, что и входные данные. Подробнее см. в разделе Keep generation tensors on CPU.

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

При использовании статического кэша generate() автоматически компилирует прямой проход с помощью torch.compile на Neuron. Передайте disable_compile=True в generate(), чтобы отключить эту функцию.

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

TrainingArguments определяет Neuron и автоматически устанавливает устройство в neuron. Все поколения Trainium поддерживают bf16.

from transformers import TrainingArguments

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

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

Установите torch_compile=True, чтобы обернуть модель в torch.compile. Это компилирует прямой проход модели, при этом шаг оптимизатора выполняется в режиме eager. TrainingArguments автоматически устанавливает torch_compile_backend в "neuron" при обнаружении Neuron.

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 с нужными размерами параллелизма и передайте его в from_pretrained(). Transformers разбивает модель на части при загрузке, а Trainer использует её как есть.

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.

torchrun --nproc_per_node=4 train.py

См. Tensor parallelism, Expert parallelism и FSDP для доступных вариантов.

Ограничения

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