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

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

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

一、什么是分布式訓練？

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

分布式訓練可以分為以下幾種模式：

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

二、理解分布式訓練的核心概念

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

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

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

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

2. Rank（全局進程編號）

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

含義：標識每個訓練進程在所有進程中的唯一編號（從 0 開始）。
作用：常用于標記主節(jié)點（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)
```

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

這些環(huán)境變量通常由 分布式啟動器 設置。例如使用 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://'：表示從環(huán)境變量中讀取初始化信息。
nccl 是 NVIDIA 的高性能通信庫，支持 GPU 間高速通信。

五、分布式訓練的代碼結構

使用 PyTorch 實現(xià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:
            # 正常訓練流程

六、Elastic Training（彈性訓練）

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

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

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

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

七、總結

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