在多个 GPU 上对 HuggingFace 管道进行推理

Question

我正在使用 HuggingFace pipeline 进行推理。以下面的代码为例

from transformers import pipeline

summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0)
    
ARTICLE = """INSERT THE ARTICLE YOU WANT TO SUMMARIZE"""
print(summarizer(ARTICLE, max_length=130, min_length=30, do_sample=False))

参数

device=0

表示推理是在cuda设备上进行的。如果我的机器上有多个可用的 GPU，是否有一种方法以分布式方式执行此推理，以便利用所有 GPU 而不是仅使用一个 GPU，从而减少推理时间？

Answer 1

您可以使用 accelerate 来阅读 使用多个 GPU 进行分布式推理，该库旨在使跨分布式设置轻松训练或运行

inference

。

首先，创建一个 Python 文件并初始化 加速.PartialState创建分布式环境；你的 安装程序会自动检测到，因此您无需显式定义 排名或世界大小。将扩散管道移动到 Distributed_state.device 为每个进程分配一个 GPU。

现在使用 split_ Between_processes 实用程序作为上下文管理器 自动在多个进程之间分配提示。

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 run_distributed.py --num_processes=2