Buckets:
| # 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/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained)，并使用 `device` 处理流水线中的其余组件。 | |
| 为每个 NeuronCore 启动一个进程。将 `--nproc_per_node` 设置为你希望分片使用的 NeuronCore 数量。 | |
| ```bash | |
| torchrun --nproc_per_node=8 tensor_parallel_neuron.py | |
| ``` | |
Xet Storage Details
- Size:
- 4.4 kB
- Xet hash:
- 531f4bc9664b4a2be608bea83dc78b21808e03843ff94cb78b88e1b70d1d3b70
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.