PyTorch 分布式训练

深入理解 PyTorch 分布式训练：环境变量与进程通信机制

在深度学习模型变得日益庞大之后，单个 GPU 的显存已经无法满足高效训练的需求。此时，“分布式训练（Distributed Training）”技术应运而生，成为加速训练的重要手段。

本文将围绕以下三行典型的 PyTorch 分布式训练代码进行详细解析，并扩展介绍分布式训练的核心概念和实践方法：

local_rank = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
global_rank = int(os.getenv('RANK', -1))
world_size = int(os.getenv('WORLD_SIZE', 1))

一、什么是分布式训练？

分布式训练是指将模型训练过程划分到多个计算设备（通常是多个 GPU，甚至是多台机器）上进行协同处理，目标是加速训练速度和扩展模型容量。

分布式训练可以分为以下几种模式：

• 数据并行（Data Parallelism）：每个 GPU 处理不同的数据子集，同步梯度。
• 模型并行（Model Parallelism）：将模型拆成多个部分，分别部署到不同的 GPU。
• 混合并行（Hybrid Parallelism）：结合模型并行和数据并行。
• 流水线并行（Pipeline Parallelism）：按层切分模型，不同 GPU 处理不同阶段。

二、理解分布式训练的核心概念

1. World Size（全局进程数）

world_size = int(os.getenv('WORLD_SIZE', 1))

• 含义：分布式训练中，所有参与训练的进程总数。通常等于 GPU 总数。
• 作用：用于初始化进程组（torch.distributed.init_process_group()），让每个进程知道集群的规模。

比如你有两台机器，每台 4 块 GPU，那么 world_size = 8。

2. Rank（全局进程编号）

global_rank = int(os.getenv('RANK', -1))

• 含义：标识每个训练进程在所有进程中的唯一编号（从 0 开始）。
• 作用：常用于标记主节点（rank == 0），控制日志输出、模型保存等。

例如：

• rank=0：负责打印日志、保存模型
• rank=1,2,…：只做训练

3. Local Rank（本地进程编号）

local_rank = int(os.getenv('LOCAL_RANK', -1))

• 含义：当前训练进程在本地机器上的 GPU 编号。一般与 CUDA_VISIBLE_DEVICES 配合使用。

• 作用：用于指定该进程应该使用哪块 GPU，如：

  torch.cuda.set_device(local_rank)
  

三、环境变量的设置方式

这些环境变量通常由 分布式启动器 设置。例如使用 torchrun：

torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \
    --master_addr=192.168.1.1 --master_port=12345 train.py

torchrun 会自动为每个进程设置：

• LOCAL_RANK
• RANK
• WORLD_SIZE

也可以手动导出：

export WORLD_SIZE=8
export RANK=3
export LOCAL_RANK=3

四、分布式训练初始化流程（PyTorch 示例）

在 PyTorch 中，典型的初始化流程如下：

import os
import torch
import torch.distributed as dist

def setup_distributed():
    local_rank = int(os.getenv('LOCAL_RANK', -1))
    global_rank = int(os.getenv('RANK', -1))
    world_size = int(os.getenv('WORLD_SIZE', 1))

    torch.cuda.set_device(local_rank)

    dist.init_process_group(
        backend='nccl',  # GPU 用 nccl，CPU 用 gloo
        init_method='env://',
        world_size=world_size,
        rank=global_rank
    )

• init_method='env://'：表示从环境变量中读取初始化信息。
• nccl 是 NVIDIA 的高性能通信库，支持 GPU 间高速通信。

五、分布式训练的代码结构

使用 PyTorch 实现分布式训练的基本框架：

def train():
    setup_distributed()

    model = MyModel().cuda()
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    dataset = MyDataset()
    sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=64)

    for epoch in range(epochs):
        sampler.set_epoch(epoch)
        for batch in dataloader:
            # 正常训练流程

六、Elastic Training（弹性训练）

值得注意的是，示例代码中注释中提到的链接：

# https://pytorch.org/docs/stable/elastic/run.html

这是指 PyTorch 的 弹性分布式训练（Elastic Training），支持在训练过程中动态增加或移除节点，具备高容错性。

• 工具：torch.distributed.elastic
• 启动命令：torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py

该特性对于大规模、长时间训练任务非常重要。

七、总结

这三行代码虽然简单，却是 PyTorch 分布式训练的入口。理解它们，就理解了 PyTorch 在分布式场景下的通信机制和训练框架。

如果你想要进一步深入了解 PyTorch 分布式训练，推荐官方文档：

• PyTorch 分布式训练文档
• torchrun 使用说明