多进程 GPT-2 推理通信实战：all_gather_object 实现输出同步与统一收集

本文基于 dist_infer_demo_v2.py，讲解如何在多进程推理场景中使用 PyTorch 的 all_gather_object 实现各 rank 间通信、共享输出、统一收集并排序，适用于部署时构建集中式日志或统一响应结构。

---

📦 项目结构

dist_infer_demo/
├── dist_infer_demo_v2.py     # ✅ 本文主程序
├── launch_v2.sh              # 启动脚本
├── gpt2_student_v2/          # 小模型目录

---

1️⃣ v1 vs v2 差异回顾

特性	v1：dist_infer_demo.py	✅ v2：dist_infer_demo_v2.py
输出方式	每个 rank 单独打印	主进程统一打印
进程间通信	无	使用 all_gather_object
prompt 分发支持	无	✅ 支持多句并行输入

---

2️⃣ 脚本运行方式

bash launch_v2.sh

内容如下：

torchrun --nproc_per_node=4 dist_infer_demo_v2.py

---

3️⃣ 主流程解读

初始化分布式通信

import torch.distributed as dist

dist.init_process_group("gloo")
rank = dist.get_rank()
world_size = dist.get_world_size()

• 自动获取当前进程的编号（rank）
• 总进程数（world_size）

---

自定义每个进程的 prompt

prompts = [
    "你好，今天是",
    "我爱",
    "The sky is",
    "GPT-2 is"
]
prompt = prompts[rank % len(prompts)]

• 每个 rank 处理不同输入
• rank 控制句子分发顺序

---

推理并构造本地输出结构

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

local_result = {
    "rank": rank,
    "prompt": prompt,
    "token": token_text
}

---

使用 all_gather_object 收集所有输出

gathered = [None for _ in range(world_size)]
dist.all_gather_object(gathered, local_result)

• 每个进程都会将自己的 local_result 发送给其他进程
• 最终每个进程都拥有完整的 gathered 输出列表

---

控制打印只在 rank==0 输出

if rank == 0:
    print("📦 全部输出：")
    for item in sorted(gathered, key=lambda x: x["rank"]):
        print(f'[Rank {item["rank"]}] {item["prompt"]} → {item["token"]}')

---

4️⃣ 输出示例

📦 全部输出：
[Rank 0] 你好，今天是 → 周五
[Rank 1] 我爱 → 你
[Rank 2] The sky is → blue
[Rank 3] GPT-2 is → powerful

✅ 所有进程的推理结果统一打印，由 rank 0 控制主输出。

---

🧠 为什么使用 all_gather_object？

功能	描述
多进程数据共享	可传递 Python 任意对象（dict、list、str）
同步统一收集	保证每个 rank 获取完整的输出列表
便于 rank==0 汇总打印	避免输出乱序、打印冲突

---

✅ 关键点汇总

技术点	说明
all_gather_object	同步收集每个进程的输出
多 prompt 分发	每个 rank 使用不同输入，提升效率
主进程打印控制	避免多进程同时打印造成信息混乱
轻量无 GPU 通信依赖	使用 Gloo 后端即可部署（无需 NCCL）

---

📌 总结

• 本文展示了如何使用 all_gather_object 实现多进程 GPT-2 推理结果的同步输出
• dist_infer_demo_v2.py 是构建多 rank 协作式部署服务的关键样例
• 通信简洁、并行灵活、适合 API 日志聚合或控制台集成调试使用

---

🧭 本系列为 GPT-2 分布式推理项目第2篇（共三篇）

• 🧩 第1篇：用 torchrun 启动 GPT-2 多进程推理：分布式推理基础与 CPU/GPU 混合部署实战
• 🚀 第2篇：多进程 GPT-2 推理通信实战：all_gather_object 实现输出同步与统一收集
• 🌐 第3篇：支持生成文本的多进程 GPT-2 分布式推理：generate + KV Cache 实战解析

---

📌 YoanAILab 技术导航页

💡 项目源码 × 实战部署 × 转型经验，一页总览
👉 点击查看完整导航页

📚 包含内容：

• 🧠 GPT-2 项目源码（GitHub）
• ✍️ 优快云 技术专栏合集
• 💼 知乎转型日志
• 📖 公众号 YoanAILab 全文合集