# 数据并行

[数据并行](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=data_parallelism)会在每个 GPU 上复制相同的模型，并为每个 GPU 分配不同的提示词。当模型本身已能装入单个 GPU 时，它可以提升吞吐量，同时保持单个提示词的延迟不变。如果尚未尝试，请先在单个 GPU 上试用[批量推理](../using-diffusers/batched_inference)。请参阅[分布式推理](../training/distributed_inference)以将其与其他策略进行比较。

您可以使用 [Accelerate](https://huggingface.co/docs/accelerate/index) 或直接使用 [PyTorch Distributed](https://pytorch.org/tutorials/beginner/dist_overview.html) 来设置数据并行。

## Accelerate

Accelerate 会设置进程并为每个进程分配一个 GPU。

使用以下命令安装 Accelerate。

```bash
uv pip install accelerate
```

在名为 `run_distributed.py` 的文件中创建一个 [PartialState](https://huggingface.co/docs/accelerate/v1.15.0/en/package_reference/state#accelerate.PartialState) 以设置分布式环境。将 [DiffusionPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline) 移动到 `PartialState.device` 上，以便为每个进程分配一个 GPU。

```py
import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)
```

在同一个文件中，使用 [split_between_processes](https://huggingface.co/docs/accelerate/v1.15.0/en/package_reference/state#accelerate.PartialState.split_between_processes) 作为上下文管理器，在进程之间拆分提示词。

```py
with distributed_state.split_between_processes(["a dog", "a cat"]) as prompt:
    result = pipeline(prompt).images[0]
    result.save(f"result_{distributed_state.process_index}.png")
```

使用 `accelerate launch` 运行脚本，并将 `--num_processes` 设置为 GPU 的数量。

```bash
accelerate launch --num_processes=2 run_distributed.py
```

> [!TIP]
> 完整示例请参见此[最小多 GPU 推理脚本](https://gist.github.com/sayakpaul/cfaebd221820d7b43fae638b4dfa01ba)。如需了解更多，请参见[使用 Accelerate 进行分布式推理](https://huggingface.co/docs/accelerate/en/usage_guides/distributed_inference#distributed-inference-with-accelerate)指南。

## PyTorch 分布式

您也可以直接使用`torch.distributed`设置数据并行，而无需使用 Accelerate。

调用 [init_process_group](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) 来创建分布式环境。`torchrun` 会通过环境变量设置每个进程的 rank 和进程总数（即 world size），因此您无需自行传入这些参数。

每个进程都需要自己的 GPU 和自己的提示词。`torchrun` 会将 `LOCAL_RANK` 环境变量设置为该进程在其节点上的 GPU 索引，因此请使用它来选择 GPU。使用 `get_rank`（它统计所有节点上的进程）来选择提示词。

```py
import os
import torch
import torch.distributed as dist
from diffusers import DiffusionPipeline

def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")

    pipeline = DiffusionPipeline.from_pretrained(
        "Qwen/Qwen-Image", dtype=torch.bfloat16,
    ).to(device)

    prompts = ["a dog", "a cat"]
    prompt = prompts[rank]

    image = pipeline(prompt).images[0]
    image.save(f"{prompt.replace(' ', '_')}.png")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()
```

将脚本保存为 `run_distributed.py`，并使用 `torchrun` 运行。将 `--nproc-per-node` 设置为 GPU 的数量。`prompts` 列表为每个进程包含一个提示，因此请为使用的每个 GPU 添加一个提示。

```bash
torchrun --nproc-per-node 2 run_distributed.py
```