AWS Trainium

TorchNeuron, Trainium2 ve Trainium3 için yerel PyTorch arka ucudur. Diffusers, torch_neuronx paketi üzerinden bu örneklerde çalışır; bu paket "neuron" cihazını kaydeder ve bu sayede pipeline'lar, diğer herhangi bir cihazda olduğuyla aynı kodla çalışır.

İki çalışma modu mevcuttur:

Neuron SDK'yi yapılandırmak ve torch_neuronx'i yüklemek için AWS Neuron belgelerini izleyin. Yüklendikten sonra import torch bunu otomatik olarak yükler ve "neuron" cihazını kaydeder, bu nedenle tek gereken değişiklik pipe.to("neuron")'tür. TORCH_DEVICE_BACKEND_AUTOLOAD=0 ile arka uç otomatik yüklemeyi devre dışı bıraktıysanız yalnızca import torch_neuronx ekleyin.

Eager modunda, her işlem ilk çalıştığında derlenir. Derlenmiş modda, ilk çağrı grafı derler. Her iki mod da derlenmiş Neuron yürütülebilir biçimini (NEFF) varsayılan olarak tmp/neff_cache altında önbelleğe alır ve bu dizin yeniden başlatmada temizlenir. Önbelleği korumak için TORCH_NEURONX_NEFF_CACHE_DIR değerini kalıcı bir dizine ayarlayın.

Eager mod

FLUX.2-klein-4B tek bir NeuronCore üzerinde çalışır, bu nedenle tüm pipeline'ı "neuron" cihazına taşıyın.

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

Bir model tek bir NeuronCore için çok büyükse veya birden fazla çekirdeğiniz varsa ve daha düşük gecikme süresi istiyorsanız, modelleri çekirdekler arasında bölüştürün. Tensor parallelism bölümüne bakın.

Derlenmiş mod

TorchNeuron, "neuron" adını bir torch.compile arka uç adı olarak kaydeder; bu sayede bileşenler herhangi bir torch.compile hedefi gibi derlenir. İlk çağrı (ısınma) derleme nedeniyle yavaştır; aynı şekillerle yapılan sonraki çağrılar derlenmiş grafiği yeniden kullanır.

TorchNeuron statik şekiller gerektirir, bu nedenle dynamic=False parametresini iletin. Yeni bir height veya width değeri, o şekil için yeniden derlenir. Daha önce görülmüş şekiller yeniden kullanılır. num_inference_steps değerini değiştirmek yeniden derlemeye neden olmaz.

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

Tekrarlayan bloklara sahip transformer'lar için, bu blokları ayrı ayrı derlemek üzere compile_repeated_blocks(...) kullanın. Bu blok yeniden kullanılır ve derleme süresini önemli ölçüde azaltır.

Tensor paralelliği

Model yükleme, parçalama planları ve kısıtlamalar için Tensor parallelism kılavuzunu izleyin. Neuron için örneğini uyarlamak için CUDA işlem grubu ve cihaz kurulumunu aşağıdakiyle değiştirin.

Neuron üzerinde, süreç grubunu backend="neuron" ile başlatın ve mesh'i DeviceMesh("neuron", ...) ile oluşturun.

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)

Bu parallel_config değerini transformer'ın from_pretrained() metoduna iletin ve kalan pipeline bileşenleri için device parametresini kullanın.

Her NeuronCore için bir süreç başlatın. --nproc_per_node değerini, paylaştırmak istediğiniz NeuronCore sayısına ayarlayın.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py