使用多个 GPU 进行分布式推理

在分布式设置中，您可以使用 🤗 Accelerate 或 PyTorch Distributed 在多个 GPU 上运行推理，这对于并行生成多个提示很有用。

本指南将向您展示如何使用 🤗 Accelerate 和 PyTorch Distributed 进行分布式推理。

🤗 Accelerate

🤗 Accelerate 是一个旨在简化在分布式设置中进行训练或运行推理的库。它简化了设置分布式环境的过程，使您能够专注于 PyTorch 代码。

首先，创建一个 Python 文件并初始化一个 accelerate.PartialState 来创建分布式环境；您的设置会自动检测，因此您无需显式定义 rank 或 world_size。将 DiffusionPipeline 移动到 distributed_state.device 以将 GPU 分配给每个进程。

现在使用 split_between_processes 实用程序作为上下文管理器，自动在进程数量之间分配提示。

import torch
from accelerate import PartialState
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)
distributed_state = PartialState()
pipeline.to(distributed_state.device)

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")

使用 --num_processes 参数指定要使用的 GPU 数量，并调用 accelerate launch 来运行脚本。

accelerate launch run_distributed.py --num_processes=2

要了解更多信息，请查看使用 🤗 Accelerate 进行分布式推理指南。

PyTorch Distributed

PyTorch 支持 DistributedDataParallel，它可以实现数据并行。

首先，创建一个 Python 文件并导入 torch.distributed 和 torch.multiprocessing 来设置分布式进程组并为每个 GPU 上的推理生成进程。您还应该初始化一个 DiffusionPipeline

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from diffusers import DiffusionPipeline

sd = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True
)

您需要创建一个函数来运行推理；init_process_group 用于使用要使用的后端类型、当前进程的 rank 以及参与的进程数 world_size 创建分布式环境。如果您在 2 个 GPU 上并行运行推理，则 world_size 为 2。

将 DiffusionPipeline 移动到 rank 并使用 get_rank 将 GPU 分配给每个进程，其中每个进程处理不同的提示。

def run_inference(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

    sd.to(rank)

    if torch.distributed.get_rank() == 0:
        prompt = "a dog"
    elif torch.distributed.get_rank() == 1:
        prompt = "a cat"

    image = sd(prompt).images[0]
    image.save(f"./{'_'.join(prompt)}.png")

要运行分布式推理，请调用 mp.spawn 以在 world_size 中定义的 GPU 数量上运行 run_inference 函数。

def main():
    world_size = 2
    mp.spawn(run_inference, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()

完成推理脚本后，使用 --nproc_per_node 参数指定要使用的 GPU 数量，并调用 torchrun 来运行脚本。

torchrun run_distributed.py --nproc_per_node=2

您可以在 DiffusionPipeline 中使用 device_map 将其模型级组件分布到多个设备上。请参阅设备放置指南以了解更多信息。

< > 在 GitHub 上更新

Diffusers

使用多个 GPU 进行分布式推理

🤗 Accelerate

PyTorch Distributed