通过多个 CPU 优化推理¶
现在，某些模型在使用 PyTorch 时支持内置的张量并行（TP）。张量并行将模型拆分到多个 GPU 上，从而支持更大规模的模型，并行化诸如矩阵乘法等计算。
要启用张量并行，可以在调用 from_pretrained() 函数时传递参数 tp_plan="auto"：
In [ ]:
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
# 初始化分布式环境
rank = int(os.environ["RANK"]) # 获取当前进程的排名
device = torch.device(f"cuda:{rank}") # 选择对应的 GPU 设备
torch.distributed.init_process_group("nccl", device_id=device) # 初始化分布式进程组
# 获取张量并行模型
model = AutoModelForCausalLM.from_pretrained(
model_id,
tp_plan="auto", # 启用自动张量并行
)
# 准备输入数据
tokenizer = AutoTokenizer.from_pretrained(model_id) # 加载分词器
prompt = "Can I help" # 输入提示
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(device) # 将输入转换为张量并移动到 GPU
# 分布式运行
outputs = model(inputs)
你可以使用 torchrun 来启动上述脚本，每个进程映射到一个 GPU：
In [ ]:
torchrun --nproc-per-node 4 demo.py
