# AWS Trainium

[TorchNeuron](https://github.com/aws-neuron/torch-neuronx), Trainium2 ve Trainium3 için yerel PyTorch arka ucudur. Diffusers, `torch_neuronx` paketi üzerinden bu örneklerde çalışır; bu paket `"neuron"` cihazını kaydeder ve bu sayede pipeline'lar, diğer herhangi bir cihazda olduğuyla aynı kodla çalışır.

İki çalışma modu mevcuttur:

| Mod | Nasıl etkinleştirilir | Notlar |
|---|---|---|
| Eager | `pipe.to("neuron")` | Birer birer dağıtılan işlemler, asenkron |
| Derle | `torch.compile(module, backend="neuron")` | Tüm grafik derlemesi, statik şekiller |

Neuron SDK'yi yapılandırmak ve `torch_neuronx`'i yüklemek için [AWS Neuron belgelerini](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) izleyin. Yüklendikten sonra `import torch` bunu otomatik olarak yükler ve `"neuron"` cihazını kaydeder, bu nedenle tek gereken değişiklik `pipe.to("neuron")`'tür. `TORCH_DEVICE_BACKEND_AUTOLOAD=0` ile arka uç otomatik yüklemeyi devre dışı bıraktıysanız yalnızca `import torch_neuronx` ekleyin.

Eager modunda, her işlem ilk çalıştığında derlenir. Derlenmiş modda, ilk çağrı grafı derler. Her iki mod da derlenmiş Neuron yürütülebilir biçimini (NEFF) varsayılan olarak `tmp/neff_cache` altında önbelleğe alır ve bu dizin yeniden başlatmada temizlenir. Önbelleği korumak için `TORCH_NEURONX_NEFF_CACHE_DIR` değerini kalıcı bir dizine ayarlayın.

## Eager mod

FLUX.2-klein-4B tek bir NeuronCore üzerinde çalışır, bu nedenle tüm pipeline'ı `"neuron"` cihazına taşıyın.

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

Bir model tek bir NeuronCore için çok büyükse veya birden fazla çekirdeğiniz varsa ve daha düşük gecikme süresi istiyorsanız, modelleri çekirdekler arasında bölün (shard) olarak dağıtın. [Tensor parallelism](#tensor-parallelism) bölümüne bakın.

## Derlenmiş mod

TorchNeuron, `"neuron"` adını bir `torch.compile` arka uç adı olarak kaydeder; bu sayede bileşenler herhangi bir `torch.compile` hedefi gibi derlenir. İlk çağrı (ısınma) derleme nedeniyle yavaştır; aynı şekillerle yapılan sonraki çağrılar derlenmiş grafiği yeniden kullanır.

> [!IMPORTANT]
> TorchNeuron *statik şekiller* gerektirir, bu nedenle `dynamic=False` parametresini iletin. Yeni bir `height` veya `width` değeri, o şekil için yeniden derlenir. Daha önce görülmüş şekiller yeniden kullanılır. `num_inference_steps` değerini değiştirmek yeniden derlemeye neden olmaz.

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

Tekrarlayan bloklara sahip transformer'lar için, bu blokları ayrı ayrı derlemek üzere `compile_repeated_blocks(...)` kullanın. Bu blok yeniden kullanılır ve derleme süresini önemli ölçüde azaltır.

## Tensor paralelliği

Model yükleme, bölme planları ve kısıtlamalar için [Tensor parallelism](../training/distributed_inference#tensor-parallelism) kılavuzunu izleyin. Neuron için örneğini uyarlamak üzere CUDA işlem grubu ve cihaz kurulumunu aşağıdakiyle değiştirin.

Neuron üzerinde, süreç grubunu `backend="neuron"` ile başlatın ve mesh'i `DeviceMesh("neuron", ...)` ile oluşturun.

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

Bu `parallel_config` değerini transformer'ın [from_pretrained()](/docs/diffusers/main/tr/api/models/overview#diffusers.ModelMixin.from_pretrained) metoduna iletin ve kalan pipeline bileşenleri için `device` parametresini kullanın.

Her NeuronCore için bir süreç başlatın. `--nproc_per_node` değerini, paylaştırmak istediğiniz NeuronCore sayısına ayarlayın.

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