AWS Trainium

TorchNeuron — это нативный бэкенд PyTorch для Trainium2 и Trainium3. Diffusers работает на этих инстансах через пакет torch_neuronx, который регистрирует устройство "neuron", благодаря чему пайплайны выполняются с тем же кодом, что и на любом другом устройстве.

Доступны два режима выполнения:

Следуйте документации AWS Neuron, чтобы настроить Neuron SDK и установить torch_neuronx. После установки import torch автоматически загружает его и регистрирует устройство "neuron", поэтому единственным необходимым изменением является pipe.to("neuron"). Добавляйте import torch_neuronx только в том случае, если вы отключили автоматическую загрузку бэкенда с помощью TORCH_DEVICE_BACKEND_AUTOLOAD=0.

В режиме eager каждая операция компилируется при первом запуске. В режиме compiled первый вызов компилирует граф. Оба режима по умолчанию кэшируют скомпилированный исполняемый формат Neuron (NEFF) в tmp/neff_cache, и этот каталог очищается при перезагрузке. Установите TORCH_NEURONX_NEFF_CACHE_DIR в постоянный каталог, чтобы сохранить кэш.

Режим Eager

FLUX.2-klein-4B помещается на один NeuronCore, поэтому переместите весь пайплайн на устройство "neuron".

import torch

from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", dtype=torch.bfloat16)
pipe.to("neuron")

image = pipe(
    prompt="a golden retriever surfing a wave, photorealistic",
    height=1024,
    width=1024,
    num_inference_steps=4,
).images[0]

image.save("output.png")

Если модель слишком велика для одного NeuronCore или у вас несколько ядер и требуется меньшая задержка, распределите модели по ядрам. См. раздел Tensor parallelism.

Скомпилированный режим

TorchNeuron регистрирует "neuron" как имя бэкенда для torch.compile, поэтому компоненты компилируются так же, как и любые другие цели torch.compile. Первый вызов (прогрев) выполняется медленно, так как происходит компиляция; последующие вызовы с теми же формами повторно используют скомпилированный граф.

TorchNeuron требует статических форм, поэтому передайте dynamic=False. Новые значения height или width вызывают повторную компиляцию для соответствующей формы. Уже встреченные формы переиспользуются. Изменение num_inference_steps не приводит к повторной компиляции.

import torch

from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", dtype=torch.bfloat16)
pipe.to("neuron")
pipe.transformer.compile(backend="neuron", fullgraph=True, dynamic=False)

# Warmup — triggers static graph compilation.
pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=4)

# Later calls with the same shapes reuse the compiled graph.
image = pipe(
    prompt="a golden retriever surfing a wave, photorealistic",
    height=1024,
    width=1024,
    num_inference_steps=4,
).images[0]

image.save("output.png")

Для трансформеров с повторяющимися блоками используйте compile_repeated_blocks(...), чтобы компилировать эти блоки по отдельности. Такой блок переиспользуется, что значительно сокращает время компиляции.

Параллелизм тензоров

Следуйте руководству Tensor parallelism для загрузки модели, планов шардирования и ограничений. Чтобы адаптировать его пример для Neuron, замените настройку процесс-группы CUDA и устройств на следующую.

На Neuron инициализируйте группу процессов с backend="neuron" и создайте сетку с помощью DeviceMesh("neuron", ...).

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import DeviceMesh

from diffusers import TensorParallelConfig

dist.init_process_group(backend="neuron")
mesh = DeviceMesh("neuron", list(range(dist.get_world_size())))
device = torch.neuron.current_device()
parallel_config = TensorParallelConfig(mesh=mesh)

Передайте этот parallel_config в метод from_pretrained() трансформера и используйте device для остальных компонентов конвейера.

Запустите один процесс на каждый NeuronCore. Установите значение --nproc_per_node равным количеству NeuronCores, на которые вы хотите распределить модель.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py