PyTorch 分布式训练的实现_Python

在深度学习模型变得日益庞大之后，单个 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

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

七、总结

变量名	含义	来源	典型用途
world_size	全局进程数量	torchrun 设置	初始化进程组，全局通信
rank	当前进程的全局编号	torchrun 设置	控制主节点行为
local_rank	当前进程在本地的 gpu 编号	torchrun 设置	显卡绑定：torch.cuda.set_device