# AWS Trainium

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

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

| モード | 有効化する方法 | 注記 |
|---|---|---|
| Eager | `pipe.to("neuron")` | 1回ずつディスパッチされる非同期のオペレーション |
| コンパイル | `torch.compile(module, backend="neuron")` | グラフ全体のコンパイル、静的な形状 |

[AWS Neuron ドキュメント](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/)に従って 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"` デバイスに移動します。

```python
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` ターゲットと同様にコンパイルされます。最初の呼び出し（ウォームアップ）はコンパイルを行うため遅くなりますが、同じ形状での後続の呼び出しはコンパイル済みグラフを再利用します。

> [!IMPORTANT]
> TorchNeuron は*静的な形状*を必要とするため、`dynamic=False`を渡してください。新しい`height`または`width`が指定されると、その形状に対して再コンパイルが行われます。すでに確認済みの形状は再利用されます。`num_inference_steps`を変更しても再コンパイルは行われません。

```python
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](../training/distributed_inference#tensor-parallelism)ガイドに従ってください。Neuron用にその例を適応させるには、CUDAのプロセスグループおよびデバイスの設定を、以下のものに置き換えてください。

Neuron では、`backend="neuron"` を指定してプロセスグループを初期化し、`DeviceMesh("neuron", ...)` でメッシュを構築します。

```python
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()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained)に渡し、残りのパイプラインコンポーネントには`device`を使用します。

NeuronCoreごとに1つのプロセスを起動します。`--nproc_per_node`を、シャードに分割したいNeuronCoreの数に設定します。

```bash
torchrun --nproc_per_node=8 tensor_parallel_neuron.py
```