# AWS Trainium

[TorchNeuron](https://github.com/aws-neuron/torch-neuronx) 是 Trainium2 和 Trainium3 的原生 PyTorch 后端。Diffusers 通过 `torch_neuronx` 包在这些实例上运行，该包注册了一个 `"neuron"` 设备，因此流水线可以使用与其他设备相同的代码来运行。

有两种可用的执行模式：

| 模式 | 如何激活 | 备注 |
|---|---|---|
| Eager | `pipe.to("neuron")` | 一次分派一个操作，异步 |
| 编译 | `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，因此将整个 pipeline 移至 `"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](./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/zh/api/models/overview#diffusers.ModelMixin.from_pretrained)，并使用 `device` 处理流水线中的其余组件。

为每个 NeuronCore 启动一个进程。将 `--nproc_per_node` 设置为你希望分片使用的 NeuronCore 数量。

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