Transformers 文档

多GPU推理

多GPU推理

内置的张量并行(TP)现在可用于某些使用PyTorch的模型。张量并行将模型分片到多个GPU上,支持更大的模型规模,并并行化计算,如矩阵乘法。

要启用张量并行,请将参数 tp_plan="auto" 传递给 from_pretrained()

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

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

# Retrieve tensor parallel model
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    tp_plan="auto",
)

# Prepare input tokens
tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

# Distributed run
outputs = model(inputs)

你可以使用 torchrun 来启动上述脚本,使用多个进程,每个进程映射到一个GPU:

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

PyTorch 张量并行目前支持以下模型:

你可以通过打开GitHub Issue或Pull Request来请求为另一个模型添加张量并行支持。

预期的加速

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

对于在Llama上进行一次前向传递,序列长度为512且不同批量大小的情况下,预期的加速如下:

< > Update on GitHub