AWS Trainium

TorchNeuron là backend PyTorch native cho Trainium2 và Trainium3. Diffusers chạy trên các instance này thông qua gói torch_neuronx, gói này đăng ký một thiết bị "neuron", do đó các pipeline chạy với cùng mã nguồn như trên bất kỳ thiết bị nào khác.

Có hai chế độ thực thi:

Làm theo tài liệu AWS Neuron để thiết lập Neuron SDK và cài đặt torch_neuronx. Sau khi cài đặt, import torch sẽ tự động nạp và đăng ký thiết bị "neuron", do đó pipe.to("neuron") là thay đổi duy nhất cần thực hiện. Chỉ thêm import torch_neuronx nếu bạn đã tắt tự động nạp backend bằng TORCH_DEVICE_BACKEND_AUTOLOAD=0.

Trong chế độ eager, mỗi thao tác được biên dịch lần đầu tiên khi chạy. Trong chế độ compiled, lần gọi đầu tiên sẽ biên dịch đồ thị. Cả hai chế độ đều lưu bộ nhớ đệm định dạng thực thi Neuron đã biên dịch (NEFF) dưới tmp/neff_cache theo mặc định, và thư mục này sẽ bị xóa khi khởi động lại. Đặt TORCH_NEURONX_NEFF_CACHE_DIR thành một thư mục bền vững để giữ lại bộ nhớ đệm.

Chế độ Eager

FLUX.2-klein-4B vừa đủ trên một NeuronCore, vì vậy hãy di chuyển toàn bộ pipeline sang thiết bị "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")

Nếu một model quá lớn cho một NeuronCore, hoặc bạn có nhiều core và muốn giảm độ trễ, hãy chia model thành các shard trên các core thay vì chạy trên một core. Xem phần Tensor parallelism.

Chế độ biên dịch

TorchNeuron đăng ký "neuron" làm tên backend của torch.compile, do đó các thành phần được biên dịch giống như mọi mục tiêu torch.compile khác. Lệnh gọi đầu tiên (giai đoạn làm nóng) sẽ chậm vì quá trình biên dịch; các lệnh gọi sau với cùng các hình dạng sẽ tái sử dụng đồ thị đã biên dịch.

TorchNeuron yêu cầu hình dạng tĩnh, vì vậy hãy truyền dynamic=False. Một height hoặc width mới sẽ được biên dịch lại cho hình dạng đó. Các hình dạng đã được nhìn thấy sẽ được tái sử dụng. Việc thay đổi num_inference_steps sẽ không biên dịch lại.

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

Đối với các transformer có các khối lặp lại, hãy sử dụng compile_repeated_blocks(...) để biên dịch các khối đó riêng lẻ. Khối này được tái sử dụng và giúp giảm đáng kể thời gian biên dịch.

Song song tensor

Làm theo hướng dẫn Tensor parallelism để tải mô hình, kế hoạch chia nhỏ và các hạn chế. Để điều chỉnh ví dụ của nó cho Neuron, hãy thay thế việc thiết lập nhóm tiến trình và thiết bị CUDA bằng các nội dung sau.

Trên Neuron, khởi tạo nhóm tiến trình với backend="neuron" và xây dựng lưới bằng 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)

Truyền parallel_config này vào from_pretrained() của transformer và sử dụng device cho các thành phần pipeline còn lại.

Khởi chạy một tiến trình cho mỗi NeuronCore. Đặt --nproc_per_node bằng số NeuronCore mà bạn muốn chia nhỏ (shard) trên đó.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py