AWS Trainium

TorchNeuronは、Trainium2およびTrainium3向けのネイティブなPyTorchバックエンドです。Diffusersはtorch_neuronxパッケージを介してこれらのインスタンス上で動作し、このパッケージは"neuron"デバイスを登録します。これにより、パイプラインは他のデバイスと同一のコードで実行できます。

2つの実行モードが利用可能です：

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 に収まらない場合、または複数のコアを所有していてレイテンシを下げたい場合は、代わりにモデルをコア間でシャードしてください。テンソル並列セクションを参照してください。

コンパイルモード

TorchNeuron は "neuron" を torch.compile のバックエンド名として登録するため、コンポーネントは他の torch.compile ターゲットと同様にコンパイルされます。最初の呼び出し（ウォームアップ）はコンパイルを行うため遅くなりますが、同じ形状での後続の呼び出しはコンパイル済みグラフを再利用します。

TorchNeuron は静的な形状を必要とするため、dynamic=Falseを渡してください。新しいheightまたはwidthが指定されると、その形状に対して再コンパイルが行われます。すでに確認済みの形状は再利用されます。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 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ごとに1つのプロセスを起動します。--nproc_per_nodeを、シャードに分割したいNeuronCoreの数に設定します。

torchrun --nproc_per_node=8 tensor_parallel_neuron.py