AWS Trainium

TorchNeuron 是 Trainium2 和 Trainium3 的原生 PyTorch 后端。Diffusers 通过 torch_neuronx 包在这些实例上运行，该包注册了一个 "neuron" 设备，因此流水线可以使用与其他设备相同的代码来运行。

有两种可用的执行模式：

模式	如何激活	备注
Eager	`pipe.to("neuron")`	一次分派一个操作，异步
编译	`torch.compile(module, backend="neuron")`	整图编译，静态形状

请按照 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，因此将整个 pipeline 移至 "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 来说过大，或者您拥有多个核心并希望降低延迟，请将模型分片到各个核心上。请参阅 Tensor parallelism 部分。

编译模式

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()，并使用 device 处理其余的 pipeline 组件。

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

torchrun --nproc_per_node=8 tensor_parallel_neuron.py

Update on GitHub