AWS Trainium

TorchNeuron adalah backend PyTorch native untuk Trainium2 dan Trainium3. Diffusers berjalan pada instance ini melalui paket torch_neuronx, yang mendaftarkan perangkat "neuron", sehingga pipeline berjalan dengan kode yang sama seperti pada perangkat lainnya.

Dua mode eksekusi tersedia:

Ikuti dokumentasi AWS Neuron untuk mengatur SDK Neuron dan menginstal torch_neuronx. Setelah terinstal, import torch akan memuatnya secara otomatis dan mendaftarkan perangkat "neuron", sehingga pipe.to("neuron") adalah satu-satunya perubahan yang diperlukan. Tambahkan import torch_neuronx hanya jika Anda menonaktifkan pemuatan otomatis backend dengan TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Dalam mode eager, setiap operasi dikompilasi pada eksekusi pertamanya. Dalam mode compiled, panggilan pertama mengompilasi graf. Kedua mode tersebut meng-cache format eksekutabel Neuron yang dikompilasi (NEFF) di bawah tmp/neff_cache secara default, dan direktori itu dibersihkan saat reboot. Atur TORCH_NEURONX_NEFF_CACHE_DIR ke direktori persisten untuk mempertahankan cache.

Mode Eager

FLUX.2-klein-4B muat di satu NeuronCore, jadi pindahkan seluruh pipeline ke perangkat "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")

Jika model terlalu besar untuk satu NeuronCore, atau Anda memiliki beberapa core dan ingin latensi yang lebih rendah, shard model di antara core. Lihat bagian Tensor parallelism.

Mode terkompilasi

TorchNeuron mendaftarkan "neuron" sebagai nama backend untuk torch.compile, sehingga komponen dikompilasi seperti target torch.compile lainnya. Panggilan pertama (warmup) lambat karena proses kompilasi; panggilan berikutnya dengan bentuk (shapes) yang sama akan menggunakan ulang graf yang telah dikompilasi.

TorchNeuron memerlukan bentuk statis, jadi berikan dynamic=False. Nilai height atau width baru akan dikompilasi ulang untuk bentuk tersebut. Bentuk yang sudah pernah dilihat akan digunakan kembali. Mengubah num_inference_steps tidak memicu kompilasi ulang.

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

Untuk transformer dengan blok berulang, gunakan compile_repeated_blocks(...) untuk mengompilasi blok-blok tersebut secara individual. Blok ini digunakan kembali dan mengurangi waktu kompilasi secara signifikan.

Paralelisme tensor

Ikuti panduan Tensor parallelism untuk pemuatan model, rencana sharding, dan keterbatasannya. Untuk menyesuaikan contohnya untuk Neuron, ganti setup process-group dan perangkat CUDA dengan berikut.

Pada Neuron, inisialisasi grup proses dengan backend="neuron" dan bangun mesh dengan 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)

Lewatkan parallel_config ini ke from_pretrained() transformer dan gunakan device untuk komponen pipeline lainnya.

Luncurkan satu proses per NeuronCore. Atur --nproc_per_node ke jumlah NeuronCore yang ingin Anda shard.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py