多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且不同批量大小的情况下,预期的加速如下:
![](https://huggingface.co/docs/transformers/v4.47.1/en/huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/Meta-Llama-3-8B-Instruct,%20seqlen%20=%20512,%20python,%20w_%20compile.png)