数据并行

数据并行会在每个 GPU 上复制相同的模型，并为每个 GPU 分配不同的提示词。当模型本身已能装入单个 GPU 时，它可以提升吞吐量，同时保持单个提示词的延迟不变。如果尚未尝试，请先在单个 GPU 上运行批量推理。请参阅分布式推理以将其与其他策略进行比较。

您可以使用 Accelerate 或直接使用 PyTorch Distributed 来设置数据并行。

Accelerate

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

使用以下命令安装 Accelerate。

uv pip install accelerate

在名为 run_distributed.py 的文件中创建一个 PartialState 以设置分布式环境。将 DiffusionPipeline 移动到 PartialState.device 上，以便为每个进程分配一个 GPU。

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 作为上下文管理器，在进程之间拆分提示词。

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 的数量。

accelerate launch --num_processes=2 run_distributed.py

完整示例请参见此最小多 GPU 推理脚本。如需了解更多，请参见使用 Accelerate 进行分布式推理指南。

PyTorch 分布式

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

调用 init_process_group 来创建分布式环境。torchrun 会通过环境变量设置每个进程的 rank 以及进程总数（即 world size），因此您无需自行传入这些参数。

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

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 添加一个提示。

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

Update on GitHub

Diffusers

数据并行

Accelerate

PyTorch 分布式