# AWS Trainium

[TorchNeuron](https://github.com/aws-neuron/torch-neuronx) adalah backend PyTorch native untuk Trainium2 dan Trainium3. Diffusers berjalan pada instance ini melalui paket `torch_neuronx`, yang mendaftarkan perangkat `"neuron"`, sehingga pipeline berjalan dengan kode yang sama seperti pada perangkat lainnya.

Dua mode eksekusi tersedia:

| Mode | Cara mengaktifkan | Catatan |
|---|---|---|
| Eager | `pipe.to("neuron")` | Operasi yang dikirim satu per satu, asinkron |
| Compile | `torch.compile(module, backend="neuron")` | Kompilasi grafik utuh, bentuk statis |

Ikuti [dokumentasi AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) untuk mengatur SDK Neuron dan menginstal `torch_neuronx`. Setelah terinstal, `import torch` akan memuatnya secara otomatis dan mendaftarkan perangkat `"neuron"`, sehingga `pipe.to("neuron")` adalah satu-satunya perubahan yang diperlukan. Tambahkan `import torch_neuronx` hanya jika Anda menonaktifkan pemuatan otomatis backend dengan `TORCH_DEVICE_BACKEND_AUTOLOAD=0`.

Dalam mode eager, setiap operasi dikompilasi pada eksekusi pertamanya. Dalam mode compiled, panggilan pertama mengompilasi graf. Kedua mode tersebut meng-cache format eksekutabel Neuron yang dikompilasi (NEFF) di bawah `tmp/neff_cache` secara default, dan direktori itu dibersihkan saat reboot. Atur `TORCH_NEURONX_NEFF_CACHE_DIR` ke direktori persisten untuk mempertahankan cache.

## Mode Eager

FLUX.2-klein-4B muat di satu NeuronCore, jadi pindahkan seluruh pipeline ke perangkat `"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")
```

Jika model terlalu besar untuk satu NeuronCore, atau Anda memiliki beberapa core dan ingin latensi yang lebih rendah, shard model di antara core. Lihat bagian [Tensor parallelism](#tensor-parallelism).

## Mode terkompilasi

TorchNeuron mendaftarkan `"neuron"` sebagai nama backend untuk `torch.compile`, sehingga komponen dikompilasi seperti target `torch.compile` lainnya. Panggilan pertama (warmup) lambat karena proses kompilasi; panggilan berikutnya dengan bentuk (shapes) yang sama akan menggunakan ulang graf yang telah dikompilasi.

> [!IMPORTANT]
> TorchNeuron memerlukan *bentuk statis*, jadi berikan `dynamic=False`. Nilai `height` atau `width` baru akan dikompilasi ulang untuk bentuk tersebut. Bentuk yang sudah pernah dilihat akan digunakan kembali. Mengubah `num_inference_steps` tidak memicu kompilasi ulang.

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

Untuk transformer dengan blok berulang, gunakan `compile_repeated_blocks(...)` untuk mengompilasi blok-blok tersebut secara individual. Blok ini digunakan kembali dan mengurangi waktu kompilasi secara signifikan.

## Paralelisme tensor

Ikuti panduan [Tensor parallelism](../training/distributed_inference#tensor-parallelism) untuk pemuatan model, rencana sharding, dan keterbatasannya. Untuk menyesuaikan contohnya untuk Neuron, ganti setup process-group dan perangkat CUDA dengan berikut.

Pada Neuron, inisialisasi grup proses dengan `backend="neuron"` dan bangun mesh dengan `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)
```

Lewatkan `parallel_config` ini ke [from_pretrained()](/docs/diffusers/main/id/api/models/overview#diffusers.ModelMixin.from_pretrained) transformer dan gunakan `device` untuk komponen pipeline lainnya.

Luncurkan satu proses per NeuronCore. Atur `--nproc_per_node` ke jumlah NeuronCore yang ingin Anda shard.

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