PyTorch 分布式訓(xùn)練的實現(xiàn)

更新時間：2025年05月15日 10:27:48 作者：handsomeboysk

本文主要介紹了PyTorch 分布式訓(xùn)練的實現(xiàn),包括數(shù)據(jù)并行、模型并行、混合并行和流水線并行等模式,感興趣的可以了解一下

一、什么是分布式訓(xùn)練？

分布式訓(xùn)練是指將模型訓(xùn)練過程劃分到多個計算設(shè)備（通常是多個 GPU，甚至是多臺機器）上進(jìn)行協(xié)同處理，目標(biāo)是加速訓(xùn)練速度和擴展模型容量。

分布式訓(xùn)練可以分為以下幾種模式：

數(shù)據(jù)并行（Data Parallelism）：每個 GPU 處理不同的數(shù)據(jù)子集，同步梯度。
模型并行（Model Parallelism）：將模型拆成多個部分，分別部署到不同的 GPU。
混合并行（Hybrid Parallelism）：結(jié)合模型并行和數(shù)據(jù)并行。
流水線并行（Pipeline Parallelism）：按層切分模型，不同 GPU 處理不同階段。

二、理解分布式訓(xùn)練的核心概念

1. World Size（全局進(jìn)程數(shù)）

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

含義：分布式訓(xùn)練中，所有參與訓(xùn)練的進(jìn)程總數(shù)。通常等于 GPU 總數(shù)。
作用：用于初始化進(jìn)程組（torch.distributed.init_process_group()），讓每個進(jìn)程知道集群的規(guī)模。

比如你有兩臺機器，每臺 4 塊 GPU，那么 world_size = 8。

2. Rank（全局進(jìn)程編號）

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

含義：標(biāo)識每個訓(xùn)練進(jìn)程在所有進(jìn)程中的唯一編號（從 0 開始）。
作用：常用于標(biāo)記主節(jié)點（rank == 0），控制日志輸出、模型保存等。

例如：

rank=0：負(fù)責(zé)打印日志、保存模型
rank=1,2,…：只做訓(xùn)練

3. Local Rank（本地進(jìn)程編號）

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

含義：當(dāng)前訓(xùn)練進(jìn)程在本地機器上的 GPU 編號。一般與 CUDA_VISIBLE_DEVICES 配合使用。
作用：用于指定該進(jìn)程應(yīng)該使用哪塊 GPU，如：
```
torch.cuda.set_device(local_rank)
```

三、環(huán)境變量的設(shè)置方式

這些環(huán)境變量通常由 分布式啟動器 設(shè)置。例如使用 torchrun：

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

torchrun 會自動為每個進(jìn)程設(shè)置：

LOCAL_RANK
RANK
WORLD_SIZE

也可以手動導(dǎo)出：

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

四、分布式訓(xùn)練初始化流程（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://'：表示從環(huán)境變量中讀取初始化信息。
nccl 是 NVIDIA 的高性能通信庫，支持 GPU 間高速通信。

五、分布式訓(xùn)練的代碼結(jié)構(gòu)

使用 PyTorch 實現(xiàn)分布式訓(xùn)練的基本框架：

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:
            # 正常訓(xùn)練流程

六、Elastic Training（彈性訓(xùn)練）

值得注意的是，示例代碼中注釋中提到的鏈接：https://pytorch.org/docs/stable/elastic/run.html

這是指 PyTorch 的 彈性分布式訓(xùn)練（Elastic Training），支持在訓(xùn)練過程中動態(tài)增加或移除節(jié)點，具備高容錯性。

工具：torch.distributed.elastic
啟動命令：torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py

該特性對于大規(guī)模、長時間訓(xùn)練任務(wù)非常重要。

七、總結(jié)

變量名	含義	來源	典型用途
WORLD_SIZE	全局進(jìn)程數(shù)量	torchrun 設(shè)置	初始化進(jìn)程組，全局通信
RANK	當(dāng)前進(jìn)程的全局編號	torchrun 設(shè)置	控制主節(jié)點行為
LOCAL_RANK	當(dāng)前進(jìn)程在本地的 GPU 編號	torchrun 設(shè)置	顯卡綁定：torch.cuda.set_device