PyTorch能成为深度学习领域最流行的框架之一，除了动态图机制带来的灵活性，其完善的分布式训练支持也是重要原因。而这套分布式能力的核心，就是分布式通讯包torch.distributed。它不是一个简单的工具函数集合，而是一整套管理多进程通信、协调多卡计算的底层基础设施。理解它的设计思路和正确用法，是从单卡开发迈向多卡训练甚至大规模集群训练的必经之路。
torch.distributed到底是什么
简单来说，torch.distributed是PyTorch提供的分布式通讯模块，它封装了进程间消息传递的底层细节，让多个训练进程能够互相交换梯度、参数和数据。在它出现之前，PyTorch主要依赖torch.nn.DataParallel做多卡训练，但DataParallel基于单进程多线程模型，存在GIL竞争和主卡负载过高的缺陷，现在已经不推荐使用。而torch.distributed配合DistributedDataParallel（DDP）采用的是多进程架构，每个GPU对应一个独立进程，彻底绕开了GIL限制，通信效率也高得多。
这个模块的核心价值可以归纳为三点：第一，提供进程组管理能力，通过init_process_group函数让所有进程互相发现并建立通信信道；第二，提供丰富的集合通信原语，比如梯度聚合常用的all_reduce、参数同步的broadcast、收集各进程数据的all_gather等；第三，支持多种通信后端，包括针对GPU优化的NCCL、面向CPU的Gloo以及高性能计算领域的MPI，可以根据硬件环境灵活选择。
需要注意的一个概念区分是：torch.distributed是通讯层，DDP是构建在它之上的训练封装。DDP内部自动调用all_reduce来同步梯度，而当你需要更精细的控制，比如实现模型并行、自定义同步逻辑或者联邦学习场景时，就需要直接使用torch.distributed的通信原语。
核心API与通信原语详解
使用torch.distributed的第一步是初始化进程组。最常见的写法是通过环境变量传递初始化信息，这样配合torchrun启动器使用最为方便：
import os
import torch
import torch.distributed as dist
# torchrun 会自动注入 RANK、WORLD_SIZE、MASTER_ADDR 等环境变量
dist.init_process_group(backend="nccl") # GPU 训练推荐 nccl 后端
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
# 后续初始化模型并包裹为 DDP
model = MyModel().cuda()
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank]
)初始化过程中有几个关键概念必须弄清楚。rank是全局进程编号，从0开始，world_size是总进程数，local_rank是当前节点内的GPU编号。单机八卡的场景下，world_size为8，每个进程的rank从0到7，local_rank也各自对应一张卡。跨多机时，rank是全局唯一的，而local_rank在每个节点上都会从0开始。
初始化之后，常用的集合通信原语是真正需要花时间理解的部分。all_reduce会对所有进程的同一张量做规约操作（求和、求最大值等）并把结果广播回每个进程，这是梯度同步的基础；broadcast把某个进程的张量发送给所有其他进程，常用于初始化时同步模型参数，保证各卡起点一致；all_gather把每个进程的张量收集拼接，所有进程都能拿到完整结果，常用于分布式推理或评估时汇总预测结果；reduce与all_reduce类似但结果只送到指定进程，适合只在主进程做统计的场景。此外还有barrier，它不做任何数据操作，只是让所有进程在此处等待，直到大家都到达这个同步点，在流程控制中非常实用。
下面这个例子演示了如何用all_reduce统计各进程的损失并求平均，这在分布式验证中很常见：
loss = compute_loss(model, batch) # 各进程各自的损失值 loss_tensor = torch.tensor([loss], device="cuda") # 对所有进程的损失求和，结果广播回每个进程 dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) avg_loss = loss_tensor.item() / dist.get_world_size()
NCCL还是Gloo：后端选择策略
后端决定了通信底层走什么协议。NCCL是NVIDIA推出的集合通信库，针对GPU直连做了深度优化，能充分利用NVLink、InfiniBand等高速互联硬件，GPU训练场景下几乎是无脑之选。Gloo则主要面向CPU场景，或者作为调试环境下的备选方案，它不需要CUDA环境也能运行，对开发调试比较友好。MPI后端多见于传统HPC集群，日常使用相对较少。
一个容易忽略的细节是，NCCL后端只支持GPU张量的通信，如果你试图用NCCL通信一个位于CPU上的张量，会直接报错。反过来，Gloo对CPU张量支持完善，GPU张量支持则有限。因此有一种混合策略：训练通信走NCCL，而某些元数据同步（比如loss数值、控制信号）单独用Gloo进程组处理。可以通过new_group创建指定后端的子进程组来实现：
dist.init_process_group(backend="nccl") # 额外创建一个 gloo 后端的组，用于同步 CPU 上的标量信息 cpu_group = dist.new_group(backend="gloo") flag = torch.tensor([1], dtype=torch.int64) dist.all_reduce(flag, op=dist.ReduceOp.SUM, group=cpu_group)
选择建议可以简化为一句话：GPU训练选NCCL，纯CPU或调试场景选Gloo，遇到CPU和GPU张量都要通信的情况就混用两个进程组。
高频踩坑点与避坑指南
第一个坑是数据切分不当导致各进程训练数据重复。分布式训练中每个进程都要用DistributedSampler来保证数据集被均匀切分且互不重叠，如果忘记使用，每个进程都会训练完整数据集，多卡等于白加。另外调用DistributedSampler时，每个epoch开始前必须调用sampler.set_epoch(epoch)，否则每个epoch的数据顺序完全一样，会削弱随机性对训练效果的影响。
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
for epoch in range(num_epochs):
sampler.set_epoch(epoch) # 不调用则每个 epoch 打乱方式相同
for batch in loader:
train_step(batch)第二个坑是模型保存时的重复写入或多进程冲突。正确做法是只让rank为0的进程保存模型和checkpoint，其他进程跳过，否则多个进程同时写同一个文件会导致文件损坏。同样，日志打印、指标上报等操作也应该做rank判断，避免输出翻倍：
if dist.get_rank() == 0:
torch.save(model.module.state_dict(), "model.pt") # 注意保存 module 而非 DDP 包装对象第三个坑是随机种子不统一引发的隐性bug。虽然DDP会在初始化时同步模型参数，但如果每个进程的数据增强、dropout等随机行为不一致且未被正确协调，评估结果可能不可复现。规范做法是设置相同的基础种子，再结合rank做偏移：seed = base_seed + dist.get_rank()，这样既保证数据增强多样性，又保证可复现性。
第四个坑是集合通信调用不对称导致的死锁。集合通信要求所有进程都参与，如果某个rank在条件分支里少调用了一次all_reduce，其他进程会永远等待，表现为任务卡住且无报错。排查这类问题的关键是检查通信调用是否在所有rank上对称执行。另外barrier的滥用也会拖慢整体速度，因为它强制最慢的进程决定整体进度，仅在确实需要全局同步的边界处使用即可。
掌握torch.distributed并不需要一次性啃下所有API，先把init_process_group、rank与world_size的概念、DDP的基本配合方式以及all_reduce这几个核心点吃透，就足以应对绝大多数多卡训练场景。遇到更复杂的模型并行或流水线并行需求时，再回头深入研究点对点通信和自定义进程组，学习曲线会平缓很多。
PyTorch分布式通讯包torch.distributed分布式训练修改时间：2026-09-02 12:42:50