AWS Trainium

TorchNeuron은 Trainium2와 Trainium3의 네이티브 PyTorch 백엔드입니다. Diffusers는 torch_neuronx 패키지를 통해 이러한 인스턴스에서 실행되며, 이 패키지는 "neuron" 장치를 등록하여 파이프라인이 다른 어떤 장치와 동일한 코드로 실행되도록 합니다.

두 가지 실행 모드가 있습니다:

AWS Neuron 문서를 참고하여 Neuron SDK를 설정하고 torch_neuronx를 설치하세요. 설치 후 import torch를 실행하면 자동으로 로드되고 "neuron" 디바이스가 등록되므로, pipe.to("neuron")만 변경하면 됩니다. TORCH_DEVICE_BACKEND_AUTOLOAD=0로 백엔드 자동 로딩을 비활성화한 경우에만 import torch_neuronx를 추가하세요.

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 대상과 동일한 방식으로 컴파일됩니다. 첫 번째 호출(워밍업)은 컴파일을 수행하기 때문에 느리지만, 동일한 shape을 가진 이후의 호출은 컴파일된 그래프를 재사용합니다.

TorchNeuron은 정적 shape를 요구하므로 dynamic=False를 전달하세요. 새로운 height 또는 width가 해당 shape에 대해 다시 컴파일됩니다. 이미 확인된 shape는 재사용됩니다. 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")

반복 블록을 가진 transformers의 경우, compile_repeated_blocks(...)를 사용하여 해당 블록을 개별적으로 컴파일합니다. 이 블록은 재사용되며 컴파일 시간을 상당히 줄입니다.

Tensor 병렬성

모델 로딩, 샤딩 플랜 및 제한 사항에 대해서는 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를 transformer의 from_pretrained()에 전달하고 나머지 pipeline 구성 요소에는 device를 사용하세요.

NeuronCore마다 프로세스 하나를 시작합니다. 샤딩할 NeuronCore의 수를 --nproc_per_node에 설정합니다.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py