如何在多GPU上实现HuggingFace Pipeline分布式推理？

可以通过两种主流方式利用多GPU加速推理，根据你的场景选择：

一、数据并行（批量处理多样本，最大化算力利用率） #

适合一次性处理大量待总结的文章，将样本拆分到各个GPU并行处理，显著缩短总推理时间。需要借助accelerate库实现：

1. 先安装依赖：

pip install accelerate

2. 示例代码：

from transformers import pipeline
from accelerate import Accelerator

accelerator = Accelerator()
device = accelerator.device

# 初始化pipeline并绑定当前设备
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=device)

# 准备批量待处理文章
ARTICLES = [
    """文章1的完整内容""",
    """文章2的完整内容""",
    """文章3的完整内容""",
    """文章4的完整内容"""
]

# 拆分数据到各GPU并行推理
with accelerator.split_between_processes(ARTICLES) as local_articles:
    local_results = summarizer(local_articles, max_length=130, min_length=30, do_sample=False)

# 合并所有GPU的推理结果
all_results = accelerator.gather(local_results)

# 仅主进程输出结果
if accelerator.is_main_process:
    for idx, result in enumerate(all_results):
        print(f"文章{idx+1}总结结果：{result['summary_text']}")

二、模型并行（超大型模型单GPU显存不足时使用） #

如果你的模型体积过大，单GPU无法容纳全部参数，可以让transformers自动将模型层分配到多个GPU上：

from transformers import pipeline

# 自动分配模型层到多GPU
summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device_map="auto")

ARTICLE = """待总结的文章内容"""
print(summarizer(ARTICLE, max_length=130, min_length=30, do_sample=False)['summary_text'])

注意事项 #

• 批量处理多样本时，数据并行的提速效果远优于模型并行，优先选择第一种方案
• 确保PyTorch环境已正确识别所有GPU，可通过torch.cuda.device_count()验证GPU数量

内容的提问来源于stack exchange，提问作者andrea