AWS Trainium

TorchNeuron é o backend nativo do PyTorch para Trainium2 e Trainium3. O Diffusers roda nessas instâncias por meio do pacote torch_neuronx, que registra um dispositivo "neuron", permitindo que os pipelines sejam executados com o mesmo código de qualquer outro dispositivo.

Dois modos de execução estão disponíveis:

Siga a documentação do AWS Neuron para configurar o SDK do Neuron e instalar o torch_neuronx. Uma vez instalado, import torch o carrega automaticamente e registra o dispositivo "neuron", de modo que pipe.to("neuron") é a única alteração necessária. Adicione import torch_neuronx apenas se você desativou o carregamento automático do backend com TORCH_DEVICE_BACKEND_AUTOLOAD=0.

No modo eager, cada operação é compilada na primeira vez que é executada. No modo compilado, a primeira chamada compila o grafo. Ambos os modos fazem o cache do formato executável Neuron compilado (NEFF) em tmp/neff_cache por padrão, e esse diretório é limpo na reinicialização. Defina TORCH_NEURONX_NEFF_CACHE_DIR para um diretório persistente para manter o cache.

Modo Eager

FLUX.2-klein-4B cabe em um único NeuronCore, então mova o pipeline inteiro para o dispositivo "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")

Se um modelo for grande demais para um único NeuronCore, ou se você tiver vários núcleos e quiser menor latência, distribua os modelos entre os núcleos. Consulte a seção Paralelismo de tensor.

Modo compilado

TorchNeuron registra "neuron" como um nome de backend para torch.compile, de modo que os componentes são compilados como qualquer outro alvo de torch.compile. A primeira chamada (aquecimento) é lenta porque realiza a compilação; chamadas subsequentes com as mesmas formas reutilizam o gráfico compilado.

TorchNeuron requer formas estáticas, então passe dynamic=False. Um novo height ou width compila novamente para essa forma. Formas já vistas são reutilizadas. Alterar num_inference_steps não recompila.

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

Para transformers com blocos repetidos, use compile_repeated_blocks(...) para compilar esses blocos individualmente. Esse bloco é reutilizado e reduz o tempo de compilação consideravelmente.

Paralelismo de tensor

Siga o guia de Paralelismo de tensor para o carregamento de modelos, planos de sharding e limitações. Para adaptar seu exemplo para o Neuron, substitua a configuração do grupo de processos e do dispositivo CUDA pelo seguinte.

No Neuron, inicialize o grupo de processos com backend="neuron" e construa a malha com 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)

Passe este parallel_config para o from_pretrained() do transformer e use device para os componentes restantes do pipeline.

Inicie um processo por NeuronCore. Defina --nproc_per_node para o número de NeuronCores que você deseja dividir.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py