【PyTorch】torch.distributed()的含义和使用方法-优快云博客

torch.distributed 是 PyTorch 的一个子模块，它提供了支持分布式训练的功能。这意味着它允许开发者将神经网络训练任务分散到多个计算节点上进行。使用分布式训练可以显著加快训练过程，特别是在处理大型数据集和复杂模型时。这个模块支持多种后端，可以在不同的硬件和网络配置上高效运行。

核心组件

torch.distributed 包括以下核心组件：

通信后端：这些后端负责在不同进程或设备间传输数据。常用的后端包括：
- NCCL：针对 NVIDIA GPU 优化的通信库，支持高效的 GPU 间通信。
- Gloo：是一个跨平台的通信库，支持 CPU 和 GPU，适合于内部网络延迟相对较低的情况。
- MPI（消息传递接口）：一种标准的高性能通信协议，用于不同计算节点间的数据传输。
分布式数据并行（Distributed Data Parallel, DDP）：这是一个封装模块，使多个进程可以同时执行模型的前向和反向传播，同时同步梯度。
集合通信操作（Collective Communication Operations）：包括广播（broadcast）、聚集（gather）、散布（scatter）和规约（reduce）操作，这些都是并行计算中常用的操作，用于在多个进程间同步数据。

设置和初始化

在使用 torch.distributed 前，需要进行适当的初始化：

import torch.distributed as dist

def setup(rank, world_size):
    dist.init_process_group(
        backend='nccl',         # 指定后端为 NCCL
        init_method='env://',   # 通过环境变量进行初始化
        world_size=world_size,  # 总的进程数
        rank=rank               # 当前进程的标识
    )

init_process_group：这是初始化分布式环境的关键函数，它设置了后端类型、初始化方法、世界大小（即参与计算的总进程数）和当前进程的排名。
rank：是指当前进程在所有进程中的编号，从0开始。
world_size：是参与当前任务的总进程数。

分布式数据并行的应用

在模型定义后，可以使用 torch.nn.parallel.DistributedDataParallel 来包装模型，从而实现数据的分布式并行处理：

from torch.nn.parallel import DistributedDataParallel as DDP

model = MyModel()
model = model.cuda(rank)  # 将模型移到对应的 GPU
ddp_model = DDP(model, device_ids=[rank])  # 使用 DDP 包装模型