用 torchrun 启动 GPT-2 多进程推理：分布式推理基础与 CPU/GPU 混合部署实战

原创已于 2025-05-18 10:43:41 修改 · 619 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#gpt #分布式

于 2025-05-18 10:38:58 首次发布

部署运行你感兴趣的模型镜像

本文介绍如何基于 PyTorch 的 torchrun 工具，搭建一个支持多进程并行的 GPT-2 推理脚本，结合 CPU/GPU 混合使用策略，实现轻量、高效的本地推理并发执行。适合对分布式推理初学者快速上手。

📦 项目结构

本项目位于 dist_infer_demo/ 目录下：

dist_infer_demo/
├── dist_infer_demo.py       # ✅ 本文主程序
├── launch.sh                # 多进程启动脚本
├── gpt2_student_v2/         # 本地训练后的小模型

🚀 分布式推理为什么有用？

在模型部署中，推理通常是瓶颈：

单进程推理速度快但利用率低
多请求并发时常会阻塞
利用多核 CPU / 多卡 GPU 可以提升吞吐

因此，我们使用 PyTorch 内建的 分布式框架（DistributedDataParallel）来实现多进程独立推理。

1️⃣ 多进程启动方式：torchrun

torchrun --nproc_per_node=4 dist_infer_demo.py

参数说明：

参数	含义
`--nproc_per_node`	启动多少个进程（本地核/逻辑线程数）
`dist_infer_demo.py`	每个进程运行的 Python 脚本

这比 mp.spawn() 更易管理，推荐部署时使用。

2️⃣ 核心脚本解读：dist_infer_demo.py

import torch.distributed as dist
import torch.multiprocessing as mp
from transformers import GPT2LMHeadModel, GPT2Tokenizer

def run(rank, world_size):
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # 模型加载
    model = GPT2LMHeadModel.from_pretrained("gpt2_student_v2").eval()
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2_student_v2")

    # 推理输入
    prompt = "你好，今天是"
    inputs = tokenizer(prompt, return_tensors="pt")

    # GPU rank 0 使用 GPU，其他用 CPU
    device = torch.device("cuda" if rank == 0 and torch.cuda.is_available() else "cpu")
    model.to(device)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        logits = model(**inputs).logits
        token_id = int(logits[0, -1].argmax())
        token = tokenizer.decode([token_id])

    print(f"[Rank {rank}] 输出: {token}")

3️⃣ 启动脚本 launch.sh

#!/bin/bash
torchrun --nproc_per_node=4 dist_infer_demo.py

启动 4 个独立进程（模拟 4 并发推理线程）
每个进程会自动分配 rank 值（0~3）
可根据 CPU 核心数灵活调整

4️⃣ 输出示例

[Rank 0] 输出: 很
[Rank 1] 输出: 很
[Rank 2] 输出: 很
[Rank 3] 输出: 很

说明每个进程成功运行模型推理，并输出预测结果。

5️⃣ CPU + GPU 混合部署策略亮点

device = torch.device("cuda" if rank == 0 and torch.cuda.is_available() else "cpu")

rank 0 使用 GPU（通常部署优先）
其余进程使用 CPU（合理利用计算资源）
适用于“1 卡 GPU + 多核 CPU”的轻量推理场景

✅ 小结回顾

项目要素	技术点说明
启动方式	使用 `torchrun` 启动多进程
rank 识别与分配	自动注入 `rank` 与 `world_size` 环境变量
推理并行策略	不依赖梯度计算，每个进程独立执行
多设备调度	rank 0 用 GPU，其余进程用 CPU