2. PyTorch分布式基础：torch.distributed 包介绍、初始化进程组、后端选择

好，咱们正式开始聊分布式训练。说实话，我第一次接触分布式时，心里也犯嘀咕——不就是多卡跑嘛，能有多复杂？结果一上手，光环境配置就折腾了两天。后来我才明白，分布式训练的核心不在于「怎么跑」，而在于「怎么让多张卡好好说话」。

PyTorch 的 torch.distributed 包，就是负责让这些卡「好好说话」的中间人。今天咱们就把这个包拆开看看，它到底提供了什么，怎么初始化，以及后端该怎么选。

2.1 torch.distributed 包：到底提供了什么？

说白了，torch.distributed 就是一套通信工具箱。它不关心你的模型长什么样，也不管你用什么优化器，它只做一件事：让不同进程间的张量能互相传递。

我习惯把它分成三层来看：

• 通信原语：最底层，比如 all-reduce、broadcast、scatter、gather 这些操作。你不需要手写 TCP 或 MPI 调用，PyTorch 帮你封装好了。
• 进程管理：包括初始化、销毁、获取 rank 和 world size 等。这部分是分布式训练的「骨架」。
• 高级封装：比如 DistributedDataParallel（DDP），它把通信原语和模型训练流程整合在一起，让你几乎不用改代码就能跑多卡。

嗯，这里要注意：torch.distributed 和 torch.nn.DataParallel 是两码事。后者是单进程多线程，前者是多进程。我见过不少新手把两者搞混，结果调试半天发现性能上不去——因为 DataParallel 有 GIL 锁，多卡通信效率很低。

2.2 初始化进程组：init_process_group 详解

所有分布式操作的第一步，就是调用 init_process_group。这个函数干了什么呢？它会在所有参与训练的进程之间建立一条「通信通道」。每个进程都知道自己是谁（rank），也知道总共有多少人（world size）。

来看一个最基础的初始化代码：

import torch
import torch.distributed as dist

def init_distributed():
    dist.init_process_group(
        backend='nccl',          # 后端选择
        init_method='env://',    # 初始化方式
        rank=0,                  # 当前进程编号
        world_size=4             # 总进程数
    )
    print(f"进程 {dist.get_rank()} 已就绪，共 {dist.get_world_size()} 个进程")

这段代码看起来简单，但有几个坑我得提醒你：

init_method 参数有三种常见方式：

我个人习惯用 env://，因为它最灵活。你想想看，在集群上跑任务时，调度器会自动分配节点和端口，你只需要在启动脚本里设置好环境变量就行，代码完全不用改。

2.3 后端选择：NCCL vs Gloo vs MPI

后端选择是分布式训练里最容易让人纠结的地方。我刚开始时也犯过选择困难症，后来发现其实没那么复杂——选对场景，答案就出来了。

先看一张对比表：

说白了，如果你用 GPU 训练，无脑选 NCCL。为什么？因为 NCCL 是 NVIDIA 专门为 GPU 通信设计的，它利用了 NVLink、InfiniBand 这些硬件特性，通信带宽能跑满。我做过测试，同样一个 all-reduce 操作，NCCL 比 Gloo 快了将近 3 倍。

那 Gloo 什么时候用呢？两个场景：

• CPU 训练：NCCL 不支持 CPU，所以只能用 Gloo 或 MPI。
• 调试阶段：Gloo 的错误信息更友好，调试时可以先切到 Gloo，等代码跑通了再换回 NCCL。

至于 MPI，说实话，我很少用。虽然 MPI 理论上性能很高，但安装配置太麻烦，而且不同集群的 MPI 实现（OpenMPI、MPICH）还不兼容。除非你的集群管理员已经装好了 MPI，否则别给自己找麻烦。

2.4 知识体系一览

为了让你更直观地理解这三者的关系，我画了一张图：

这张图其实就概括了本章的核心逻辑：torch.distributed 是顶层接口，它依赖底层的通信后端来实现数据交换。而 init_process_group 就是连接这两层的「桥梁」——你告诉它用哪个后端，它帮你把进程组建好。

好了，这一章的内容就到这里。记住一句话：分布式训练的本质是通信，而通信的核心是选对后端。NCCL 是 GPU 训练的首选，Gloo 适合调试和 CPU 场景，MPI 嘛……除非你特别有经验，否则先放一放。

下一章咱们会深入 DistributedDataParallel 的源码，看看它到底是怎么把模型分布到多卡上的。到时候你会发现，理解了今天的基础，后面就顺理成章了。