torch.distributed 通信原语

torch.distributed 是 PyTorch 的分布式通信底层库：管理进程组（process group），提供以 all_reduce 为代表的集合通信原语。日常多卡训练由 DDP 封装，但自定义分布式逻辑（梯度同步、参数广播、指标聚合）必须直接操作这套 API。

核心概念

概念	含义
rank	全局进程编号（0, 1, 2…），rank0 惯例为主进程
world_size	总进程数，单机单卡单进程，N 卡即 world_size=N
local_rank	本机内进程编号，从 0 起，对应可见的 GPU 序号
backend	通信后端，决定底层传输实现

后端选择：

后端	设备	说明
`nccl`	仅 GPU	GPU 训练首选，NVIDIA 专为集合通信优化
`gloo`	CPU/GPU	跨平台，CPU 训练与调试用，GPU 性能差
`mpi`	CPU/GPU	集群 MPI 环境，需源码编译支持

并行策略如何选型（数据/张量/流水线并行）见分布式训练总览，本篇只讲 API 实操。

初始化与清理

dist.init_process_group 建立进程组并完成握手，是所有分布式代码的第一行；结束后调用 destroy_process_group 释放资源。

import os
import torch
import torch.distributed as dist
 
def setup_distributed():
    # torchrun 启动时自动注入这三个环境变量
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
 
    dist.init_process_group(
        backend="nccl",       # GPU 训练用 nccl
        init_method="env://",  # 从环境变量读取主节点地址
    )
 
    # 绑定当前进程到对应 GPU，必须在通信前设置
    torch.cuda.set_device(local_rank)
    print(f"rank:{rank}, local_rank:{local_rank}, world_size:{world_size}")
 
def cleanup():
    dist.destroy_process_group()
 
if __name__ == "__main__":
    setup_distributed()
    # 业务逻辑
    cleanup()

MASTER_ADDR/MASTER_PORT 等多机 rendezvous 变量由 torchrun 自动注入；旧版 python -m torch.distributed.launch 已废弃。

六大通信原语

原语	语义	结果落在谁手里
`all_reduce`	跨进程规约（SUM/MAX/MIN/PRODUCT）	所有 rank
`reduce`	同上，但只写到 dst	仅 dst rank
`all_gather`	收集各 rank 张量	所有 rank（列表）
`broadcast`	把 src 的张量复制给所有人	所有 rank
`scatter`	src 把张量列表按 rank 分发	各 rank 得一片
`barrier`	同步栅栏	—

all_reduce：规约且人人有份

最常用原语，DDP 梯度同步、全局指标聚合都靠它。Ring 等实现算法见AllReduce 集合通信。

tensor = torch.tensor([1.0, 2.0]).cuda()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
# 所有 rank 的 tensor 均变为各 rank 求和结果；原地修改

reduce：规约只留 dst

tensor = torch.tensor([1.0, 2.0]).cuda()
dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)
# 只有 rank0 拿到总和，其余 rank 的 tensor 不变

all_gather：全员收集

world_size = dist.get_world_size()
local_tensor = torch.tensor([1.0]).cuda()
gather_list = [torch.zeros_like(local_tensor) for _ in range(world_size)]
dist.all_gather(gather_list, local_tensor)
# 每个 rank 都拿到完整列表：gather_list[0] 是 rank0 的数据

broadcast：广播

常用于同步权重、模型加载后对齐各进程参数。

tensor = torch.zeros(3).cuda()
if dist.get_rank() == 0:
    tensor = torch.tensor([10.0, 20.0, 30.0]).cuda()
dist.broadcast(tensor, src=0)
# 所有 rank 的 tensor 都变成 [10, 20, 30]

scatter：分片分发

数据并行切分 batch 的底层形态。

world_size = dist.get_world_size()
# 仅 src 进程需要构造完整列表
scatter_list = [torch.full((2,), float(i)).cuda() for i in range(world_size)]
recv = torch.zeros(2).cuda()
dist.scatter(recv, scatter_list, src=0)
# rank_i 收到 scatter_list[i]，其余进程的 scatter_list 被忽略

barrier：栅栏同步

dist.barrier()  # 所有进程都到达后才继续
print("全部进程执行完毕")

典型场景

全局平均 loss——先 all_reduce 求和再除以 world_size：

loss = torch.tensor(2.0).cuda()
dist.all_reduce(loss, op=dist.ReduceOp.SUM)
loss = loss / dist.get_world_size()
if dist.get_rank() == 0:
    print(f"global loss: {loss.item()}")

只在 rank0 打印与保存，配合 barrier 保证其余进程等它写完：

if dist.get_rank() == 0:
    print("只有主进程打印日志、保存 checkpoint")
    torch.save(model.state_dict(), "ckpt.pth")
dist.barrier()  # 等 rank0 落盘后再继续

大模型的分片保存方案见大模型检查点。

与 DDP 的封装关系

DistributedDataParallel 内部封装了本篇全部机制：反向传播结束后自动对各参数梯度做 all_reduce 求平均，用户无需手写梯度同步。自定义需求（如逐层通信、ZeRO 式分片）才下沉到裸 dist API。

from torch.nn.parallel import DistributedDataParallel as DDP
 
model = model.cuda()
model = DDP(model, device_ids=[local_rank])

想少写样板代码，Accelerate 在这套 API 之上再做了一层薄封装。

进程组分组

new_group 创建子进程组，支持局部通信——比如只让前两块卡互相同步：

group = dist.new_group(ranks=[0, 1])
dist.all_reduce(tensor, group=group)  # 仅组内参与，组外进程不进入该调用

踩坑清单

nccl 后端只接受 GPU tensor：传 CPU tensor 直接报错；Tensor 与 Storage 决定了张量的设备归属，通信前确认 .cuda()
所有 rank 的 dist 调用顺序必须一致：if rank==0 分支里出现集合通信而其他 rank 跳过，必然死锁
all_reduce 是原地操作：传入 tensor 被直接改写，需要保留原值先 clone
勿混用 DP 与 DDP：DataParallel 是单进程多线程，根本不走 torch.distributed
多机训练的 MASTER_ADDR/MASTER_PORT 交给 torchrun 设置，不要手写
脚本结束务必调用 dist.destroy_process_group()，否则 NCCL 挂起报错

知识花园

探索