分布式 GPU 推理

张量并行将模型分片到多个 GPU 上，并并行化矩阵乘法等计算。它能够将更大的模型尺寸装入内存，并且速度更快，因为每个 GPU 可以处理张量切片。

展开下面的列表以查看哪些模型支持张量并行。打开 GitHub issue 或 pull request 以添加对当前未在下列出的模型的支持。

支持的模型

在 from_pretrained() 中设置 tp_plan="auto" 以启用张量并行推理。

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# initialize distributed environment
rank = int(os.environ["RANK"])
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
torch.distributed.init_process_group("nccl", device_id=device)

# enable tensor parallelism
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    tp_plan="auto",
)

# prepare input tokens
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

# distributed run
outputs = model(inputs)

在 torchrun 上启动上述推理脚本，每个 GPU 4 个进程。

torchrun --nproc-per-node 4 demo.py

您可以从推理速度的显着提升中获益，特别是对于具有大批量大小或长序列的输入。

对于在 Llama 上进行的单次前向传递，序列长度为 512，批量大小不同，您可以期望以下速度提升。

< > 更新在 GitHub 上

Transformers

分布式 GPU 推理