当前位置：首页 > news >正文

PyTorch 分布式训练

news 来源：原创 2025/5/14 8:05:38

深入理解 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 distdef 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 用 glooinit_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`