MNIST DDP 分布式数据并行

最新推荐文章于 2025-05-12 05:31:41 发布

the_3rd_bomb

最新推荐文章于 2025-05-12 05:31:41 发布

阅读量623

点赞数 10

文章标签：分布式 MNIST

本文链接：https://blog.csdn.net/the_3rd_bomb/article/details/147875907

版权

Distributed Data Parallel

转自我的个人博客：https://shar-pen.github.io/2025/05/04/torch-distributed-series/3.MNIST_DDP/

The difference between DistributedDataParallel and DataParallel is: DistributedDataParallel uses multiprocessing where a process is created for each GPU, while DataParallel uses multithreading.
By using multiprocessing, each GPU has its dedicated process, this avoids the performance overhead caused by GIL of Python interpreter.

DDP vs DP 的并发模式

DDP使用的是多进程multiprocessing
- 每个 GPU 对应一个独立的 Python 进程。
- 各 GPU/进程之间通过通信（比如 NCCL）同步梯度。
- 进程之间可以并行,每个进程独占一个 GPU，自由度高、效率高。
DP 使用的是多线程（multithreading）
- 一个 Python 主进程控制多个线程，每个线程对应一个 GPU 上的模型副本。
- 所有线程共享同一个 Python 解释器（主进程中的 GIL 环境）。
- 在多线程环境下，同一时刻只能有一个线程执行 Python 字节码

GIL 的性能问题

Python 有个限制叫 GIL（Global Interpreter Lock）：

在 Python 中，进程之间可以并行，线程之间只能并发。
在多线程环境下，同一时刻只能有一个线程执行 Python 字节码。
这意味着虽然多个线程运行在不同 GPU 上，但只要你涉及到 Python 层的逻辑（如 forward 调度、数据调度），就会被 GIL 限制，造成瓶颈。

DDP 的多进程模式就天然绕开了 GIL，每个进程有独立的 Python 解释器和 GIL，不会互相争抢锁。所以执行速度更快、效率更高、更适合大模型和多 GPU 并行。

To use DistributedDataParallel on a host with N GPUs, you should spawn up N processes, ensuring that each process exclusively works on a single GPU from 0 to N-1.

总结下，DDP 用多进程给每个 GPU 配一个独立的进程，这样就不用多个线程去抢 Python 的 GIL，避免了 DataParallel 由于多线程带来的性能开销。

分布式数据并行时，模型（model parameters）/优化器（optimizer states）每张卡都会拷贝一份（replicas），在整个训练过程中 DDP 始终在卡间维持着模型参数和优化器状态的同步一致性；

DDP 将 batch input，通过 DistributedSampler split & 分发到不同的 gpus 上，此时虽然模型/optimizer 相同，但因为数据输入不同，导致 loss 不同，反向传播时计算到的梯度也会不同，如何保证卡间，model/optimizer 的同步一致性，之前 DP 用的 parameter server，而它的问题就在于通信压力都在 server，所以 DDP 对这方面的改进是 ring all-reduce algorithm，将 Server 上的通讯压力均衡转到各个 Worker 上

注意有两个核心概念:

All to one：reduce
one to All：broadcast

方法名	通信结构	通信路径	数据流向	聚合策略	通信瓶颈位置	通信效率	适合场景
Parameter Server	中心化（星型）	所有 Worker ⇄ PS	上传全部梯度 → 聚合 → 下发参数	PS 聚合	PS 带宽和计算压力	❌ 低（集中式瓶颈）	小规模训练，原型实验
Tree All-Reduce	层次化（树型）	节点间按树结构上传/下传	层层上传聚合 → 再层层广播	层次加和 & 广播	上层节点（树根）	✅ 中（ $\log N$ 轮次）	多机多卡，合理拓扑连接
Broadcast + Reduce	两阶段（集中）	所有 → 主节点（reduce） → 所有	所有上传 → 中心聚合 → 广播下发	单节点聚合	主节点	❌ 低	小规模单机多卡
Ring All-Reduce	环形（对称）	相邻节点之间点对点传输	均匀传递/聚合，每轮处理一块数据	分块加和 & 拼接	无集中瓶颈	✅✅ 高（带宽最优）	大规模 GPU 并行，主流方案

Parameter Server（PS）和 Broadcast + Reduce 在通信机制上本质相似，区别只在于:

PS 是显式设计了专门的“参数服务器”角色；
Broadcast + Reduce 是“隐式指定”某个节点承担聚合与广播任务。

ring all-reduce:

Reduce-scatter:首先将 gradient 分为 n 块，在第 i 轮 (0<= i < n-1)，每个 gpu j 把第 (i+j) % n 块的数据传给下一个 gpu (j+1 % n)，即每个 gpu 都把自己一个块给下一个做加法，在 n 轮结束后，每个 gpu 上都有一个块是完整的聚合了所有不同 gpu 的 gradient。
All-gather: 将每个 gpu 上的完整聚合后的 gradient 依次传给下一个 gpu，再传递 n-1 次就使所有 gpu 的每块 gradient 都是完整聚合的数据。

虽然传递的数据量还是和 PS 一样，但传输压力平均到每个 gpu 上，不需要单个 worker 承担明显大的压力。

概念/参数名	中文含义	含义解释	示例（2节点 × 每节点4GPU）
`world`	全局进程空间	指整个分布式系统中参与训练的所有进程总和	2 节点 × 4 GPU = 8 个进程
`world_size`	全局进程数	`world` 中的进程总数，参与通信、同步、梯度聚合的总 worker 数	8
`rank`	全局进程编号	当前进程在 `world` 中的唯一编号，范围是 $world_size − 1 ] [0, \text{world\_size} - 1]$	第1节点是 0~3，第2节点是 4~7
`node`	物理节点/机器	实际的服务器或物理机，每个节点运行多个进程，通常对应一台机器	2台服务器（假设每台4 GPU）
`node_rank`	节点编号	当前节点在所有节点中的编号，通常用于标识不同机器	第1台是 0，第2台是 1
`local_rank`	本地GPU编号	当前进程在所在节点上的 GPU 编号，绑定 `cuda(local_rank)`	每台机器上分别为 0~3

简洁点，world 代表所有服务器上的 gpu，rank 代表 world 视角下的 gpu 编号；node 代表某个具体的服务器，node_rank 代表 world 视角下的 node 编号，local_rank 代表 node 视角下的 gpu 编号。

引入 DDP 相关库

import os
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

# 以下是分布式相关的
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler # 分发数据集
from torch.nn.parallel import DistributedDataParallel as DDP # 用 DDP 封装 module 以支持分布式训练
from torch.distributed import init_process_group, destroy_process_group # 初始化和销毁进程组，一个 process 代表一个 gpu 进程

ddp 对原始代码的修改

参数	作用说明
`MASTER_ADDR`	指定主节点（rank=0 所在节点）的 IP 地址或主机名，作为所有进程连接的“服务器”
`MASTER_PORT`	指定主节点上用于通信监听的端口号，所有进程都通过这个端口进行连接与协调

为什么只需要指定主节点的地址和端口？所有进程必须“集合”在一起组成一个通信组（process group）；这个过程需要一个协调者，就像组织会议需要一个人发出会议链接一样；PyTorch DDP 把这个协调角色交给 rank == 0 的进程（主节点）；其它进程只需要“知道去哪找这个协调者”就能完成初始化。

主节点负责协调组网，在 DDP 初始化时，所有节点主动连接主节点，每个节点都会告知主节点自己的地址和端口，主节点收集所有其他进程的网络信息，构建全局通信拓扑，将通信配置信息广播回每个进程，包括每个 rank 要连接哪些 peer，这样每个进程就可以进行后续的双向传输，而不再依赖主节点作为中转。

主节点（rank 0）	工作节点（rank 1,2,…）
在 `MASTER_PORT` 启动一个监听服务（如 TCP server）	主动连接 `MASTER_ADDR:MASTER_PORT`
监听并接受连接，记录加入者信息	与主节点握手，注册自己的 `rank`、地址等
构建通信拓扑，如 Ring 或 NCCL 分组等	一旦接入，就获得组网配置，与其他 worker 点对点通信

ddp 初始化和销毁进程

def ddp_setup(rank, world_size):
    """
    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    # rank 0 process
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    # nccl：NVIDIA Collective Communication Library 
    # 分布式情况下的，gpus 间通信
    torch.cuda.set_device(rank)
    init_process_group(backend="nccl", rank=rank, world_size=world_size)

DDP 会在每个 GPU 上运行一个进程，每个进程中都有一套完全相同的 Trainer 副本（包括 model 和 optimizer），各个进程之间通过一个进程池进行通信。

ddp 包装 model

训练函数不需要多大的修改，使用 DistributedDataParallel 包装模型，这样模型才能在各个进程间同步参数。包装后 model 变成了一个 DDP 对象，要访问其参数得这样写 self.model.module.state_dict()

运行过程中单独控制某个进程进行某些操作，比如要想保存 ckpt，由于每张卡里都有完整的模型参数，所以只需要控制一个进程保存即可。需要注意的是：使用 DDP 改写的代码会在每个 GPU 上各自运行，因此需要在程序中获取当前 GPU 的 rank（gpu_id），这样才能对针对性地控制各个 GPU 的行为。

class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_data: DataLoader,
        optimizer: torch.optim.Optimizer,
        gpu_id: int,
        save_every: int, 
    ) -> None:
        self.gpu_id = gpu_id
        self.model = model.to(gpu_id)
        self.train_data = train_data 
        self.optimizer = optimizer
        self.save_every = save_every
        self.model = DDP(model, device_ids=[gpu_id])    # model 要用 DDP 包装一下

    def _run_batch(self, source, targets):
        self.optimizer.zero_grad()
        output = self.model(source)
        loss = F.cross_entropy(output, targets)
        loss.backward()
        self.optimizer.step()

    def _run_epoch(self, epoch):
        b_sz = len(next(iter(self.train_data))[0])
        print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")
        self.train_dataloader.sampler.set_epoch(epoch) # 注意需要在各 epoch 入口调用该 sampler 对象的 set_epoch() 方法，否则每个 epoch 加载的样本顺序都不变
        for source, targets in self.train_data:
            source = source.to(self.gpu_id)
            targets = targets.to(self.gpu_id)
            self._run_batch(source, targets)

    def _save_checkpoint(self, epoch):
        ckp = self.model.state_dict()
        PATH = "checkpoint.pt"
        torch.save(ckp, PATH)
        print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")

    def train(self, max_epochs: int):
        for epoch in range(max_epochs):
            self._run_epoch(epoch)
            if self.gpu_id == 0 and epoch % self.save_every == 0:
                self._save_checkpoint(epoch)

在程序入口初始化进程池；在程序出口销毁进程池

def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int):
    # 初始化进程池
    ddp_setup(rank, world_size)

    # 进行训练
    dataset, model, optimizer = load_train_objs()
    train_data = prepare_dataloader(dataset, batch_size)
    trainer = Trainer(model, train_data, optimizer, rank, save_every)
    trainer.train(total_epochs)
   
    # 销毁进程池
    destroy_process_group()

DistributedSampler

构造 Dataloader 时使用 DistributedSampler 作为 sampler，这个采样器可以自动将数量为 batch_size 的数据分发到各个GPU上，并保证数据不重叠。理解是可以是这样的，但实际是根据 rank 让每个 gpu 能索引到的数据不一样，每个 gpu 上也是有重复的 Dataloader 的，但每个gpu 上 rank 设置不同，Dataloader sample 先根据 shuffle 打乱顺序，再控制不同 rank 能索引到的数据，以实现类似分发的效果。

Rank 0 sees: [4, 7, 3, 0, 6] Rank 1 sees: [1, 5, 9, 8, 2]

def prepare_dataloader(dataset: Dataset, batch_size: int):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        pin_memory=True,
        shuffle=False,                      # 设置了新的 sampler，参数 shuffle 要设置为 False 
        sampler=DistributedSampler(dataset) # 这个 sampler 自动将数据分块后送个各个 GPU，它能避免数据重叠
    )

set_epoch(epoch) 用于设置当前训练 epoch，以确保在分布式训练中每个进程对数据的打乱顺序一致，从而保证每个 rank 分到的数据是互不重叠且可复现的。

当 DistributedSampler 的 shuffle=True 时，它在每个 epoch 会用 torch.Generator().manual_seed(seed) 生成新的随机索引顺序。
但：

如果不调用 set_epoch()，每个进程将使用相同的默认种子；
会导致每个 epoch 每个进程打乱后的样本索引相同 → 重复取样，每个 epoch 的训练数据都一样 → 训练不正确！

你确实可以不手动设置 rank 和 world_size，因为 DistributedSampler 会自动从环境变量中获取它们。
如果你不传入 rank 和 num_replicas，PyTorch 会调用：

torch.distributed.get_world_size() # 获取 world_size
torch.distributed.get_rank() # 获取当前进程 rank

import torch
from torch.utils.data import Dataset, DataLoader, DistributedSampler

# 自定义一个简单的数据集：返回 [0, 1, ..., n-1]
class RangeDataset(Dataset):
    def __init__(self, n):
        self.data = list(range(n))
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]

# 模拟两张卡（进程）下的样本访问情况，并支持 set_epoch
def simulate_distributed_sampler(n=10, world_size=2, num_epochs=2):
    dataset = RangeDataset(n)

    for epoch in range(num_epochs):
        print(f"\nEpoch {epoch}")
        for rank in range(world_size):
            # 设置 shuffle=True 并调用 set_epoch
            sampler = DistributedSampler(
                dataset,
                num_replicas=world_size,
                rank=rank,
                shuffle=True,
            )
            sampler.set_epoch(epoch)  # 关键：确保每轮不同但在所有 rank 一致

            dataloader = DataLoader(dataset, batch_size=1, sampler=sampler)
            data_seen = [batch[0].item() for batch in dataloader]
            print(f"Rank {rank} sees: {data_seen}")

simulate_distributed_sampler(n=10, world_size=2, num_epochs=2)

Epoch 0
Rank 0 sees: [4, 7, 3, 0, 6]
Rank 1 sees: [1, 5, 9, 8, 2]

Epoch 1
Rank 0 sees: [5, 1, 0, 9, 7]
Rank 1 sees: [6, 2, 8, 3, 4]

multiprocessing.spawn 创建多卡进程

使用 torch.multiprocessing.spawn 方法将代码分发到各个 GPU 的进程中执行。在当前机器上启动 nprocs=world_size 个子进程，每个进程执行一次 main() 函数，并由 mp.spawn 自动赋值第一个参数（目的是执行 nprocs 个进程，第一个参数为 0 ~ nprocs-1）。

def start_process(i):
    # Each process is assigned a file to write tracebacks to.  We
    # use the file being non-empty to indicate an exception
    # occurred (vs an expected shutdown).  Note: this previously
    # used a multiprocessing.Queue but that can be prone to
    # deadlocks, so we went with a simpler solution for a one-shot
    # message between processes.
    tf = tempfile.NamedTemporaryFile(
        prefix="pytorch-errorfile-", suffix=".pickle", delete=False
    )
    tf.close()
    os.unlink(tf.name)
    process = mp.Process(
        target=_wrap,
        args=(fn, i, args, tf.name),
        daemon=daemon,
    )
    process.start()
    return i, process, tf.name

if not start_parallel:
    for i in range(nprocs):
        idx, process, tf_name = start_process(i)
        error_files[idx] = tf_name
        processes[idx] = process

可以执行执行以下代码，它展现了 mp 创建进程的效果

import torch.multiprocessing as mp

def run(rank, message):
    print(f"[Rank {rank}] Received message: {message}")

if __name__ == "__main__":
    world_size = 4  # 启动 4 个进程（模拟 4 个GPU）

    mp.spawn(
        fn=run,
        args=("hello world",),   # 注意是 tuple 格式
        nprocs=world_size,
        join=True
    )

效果为:

[Rank 0] Received message: hello world
[Rank 3] Received message: hello world
[Rank 2] Received message: hello world
[Rank 1] Received message: hello world
# 利用 mp.spawn，在整个 distribution group 的 nprocs 个 GPU 上生成进程来执行 fn 方法，并能设置要传入 fn 的参数 args
# 注意不需要传入 fn 的 rank 参数，它由 mp.spawn 自动分配
import multiprocessing as mp
world_size = torch.cuda.device_count()
mp.spawn(
    fn=main, 
    args=(world_size, args.save_every, args.total_epochs, args.batch_size), 
    nprocs=world_size
)

!CUDA_VISIBLE_DEIVES=0,1 python mnist_ddp.py

torchrun

torchrun 是 PyTorch 官方推荐的分布式训练启动工具，它用来自动管理多进程启动、环境变量传递和通信初始化，替代早期的 torch.distributed.launch 工具。

它帮你在每个 GPU 上自动启动一个训练进程；
它设置好 DDP 所需的环境变量（如 RANK, WORLD_SIZE, LOCAL_RANK, MASTER_ADDR 等）；
它会自动将这些参数传递给你的脚本中的 torch.distributed.init_process_group()。

torchrun == python -m torch.distributed.launch --use-env

参数名	类型	说明
`--nproc_per_node`	`int`	每台机器上启动的进程数（默认值为 1）
`--nnodes`	`int`	总节点（机器）数
`--node_rank`	`int`	当前节点编号（范围：`0 ~ nnodes-1`）
`--rdzv_backend`	`str`	rendezvous 后端（默认 `c10d`，一般不改）
`--rdzv_endpoint`	`str`	rendezvous 主地址和端口，格式如 `"localhost:29500"`
`--rdzv_id`	`str`	作业唯一标识，默认 `"default"`
`--rdzv_conf`	`str`	可选的 kv 参数，用逗号分隔，如 `"key1=val1,key2=val2"`
`--max_restarts`	`int`	失败时最多重启次数（默认 3）
`--monitor_interval`	`float`	monitor 进程检查的间隔（秒）
`--run_path`	`str`	若脚本是模块路径形式，比如 `my_module.train`，则用此代替 script
`--tee`	`str`	控制日志输出，可选值为 `"stdout"` 或 `"stderr"`
`--log_dir`	`str`	日志输出目录（默认当前目录）
`--redirects`	`str`	重定向日志，可选：`all`, `none`, `rank`，如 `all:stdout`
`--no_python`	`flag`	若已是 Python 脚本（不用再次 `python` 调用），可加这个 flag

以上的 rendezvous 是每个进程通过 rendezvous 找到主节点，然后加入。之后的通信阶段用 backend, 即 NCCL，在 init_process_group 设置。

最常见的几个参数的用法是

torchrun \
  --nproc_per_node=4 \
  --nnodes=1 \
  --node_rank=0 \
  --rdzv_endpoint=localhost:29500 \
  your_script.py

对比下是否使用 torchrun 时的行为差别

两种 DDP 启动模式的关键区别

对比项	不使用 `torchrun`（手动）	使用 `torchrun`（推荐方式）
启动方式	使用 `mp.spawn(fn, ...)`	使用 `torchrun --nproc_per_node=N`
`rank`, `world_size` 设置方式	手动传入（通过 `spawn` 的参数）	自动由 `torchrun` 设置环境变量
主节点地址 / 端口	你必须手动设置 `MASTER_ADDR/PORT`	`torchrun` 会自动设置这些环境变量
是否需控制进程数量	手动使用 `spawn` 创建	自动由 `torchrun` 创建
是否读取环境变量	❌ 默认不会	✅ 自动从环境变量中读取（如 `RANK`, `LOCAL_RANK`）
脚本能否直接运行（`python train.py`）	❌ 通常不行，需要多进程协调	✅ 支持直接 `torchrun train.py` 运行
是否适用于多机	❌ 手动处理跨节点逻辑	✅ 内建 `--nnodes`, `--node_rank`, 可跨机运行

init_process_group() 的行为

情况	说明
手动传 `rank` 和 `world_size`	常用于 `mp.spawn` 场景（你在代码里传了参数）
不传，内部读取环境变量	如果你用的是 `torchrun`，环境变量如 `RANK`、`WORLD_SIZE` 自动设置了
不传又没用 `torchrun`	❌ 报错：因为 `init_process_group` 找不到必要的通信信息

当你运行：

torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:29500 train.py

它在后台自动设置了以下环境变量（对每个进程）：

RANK=0         # 每个进程唯一编号
WORLD_SIZE=4   # 总进程数
LOCAL_RANK=0   # 当前进程在本节点内的编号
MASTER_ADDR=localhost
MASTER_PORT=29500

而 init_process_group(backend="nccl") 会自动从这些环境变量中解析配置，无需你显式传入。

非 torchrun 完整代码

import os
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms
from time import time
import argparse


# 对 python 多进程的一个 pytorch 包装
import torch.multiprocessing as mp
# 用于收集一些用于汇总的数据
import torch.distributed as dist
# 这个 sampler 可以把采样的数据分散到各个 CPU 上                                      
from torch.utils.data.distributed import DistributedSampler     

# 实现分布式数据并行的核心类        
from torch.nn.parallel import DistributedDataParallel as DDP         

# DDP 在每个 GPU 上运行一个进程，其中都有一套完全相同的 Trainer 副本（包括model和optimizer）
# 各个进程之间通过一个进程池进行通信，这两个方法来初始化和销毁进程池
from torch.distributed import init_process_group, destroy_process_group 



def ddp_setup(rank, world_size):
    """
    setup the distribution process group

    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    # MASTER Node（运行 rank0 进程，多机多卡时的主机）用来协调各个 Node 的所有进程之间的通信
    os.environ["MASTER_ADDR"] = "localhost" # 由于这里是单机实验所以直接写 localhost
    os.environ["MASTER_PORT"] = "12355"     # 任意空闲端口
    init_process_group(
        backend="nccl",                     # Nvidia CUDA CPU 用这个 "nccl"
        rank=rank,                          
        world_size=world_size
    )
    torch.cuda.set_device(rank)
    

class ConvNet(nn.Module):

    def __init__(self):
        super(ConvNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.25)
        )
        self.classifier = nn.Sequential(
            nn.Linear(9216, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        output = F.log_softmax(x, dim=1)
        return output
    

# 自定义Dataset
class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image, label = self.data[idx]
        return image, label


class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_data: DataLoader,
        optimizer: torch.optim.Optimizer,
        gpu_id: int,
        save_every: int,
    ) -> None:
        self.gpu_id = gpu_id
        self.model = model.to(gpu_id)
        self.train_data = train_data
        self.optimizer = optimizer
        self.save_every = save_every                    # 指定保存 ckpt 的周期
        self.model = DDP(model, device_ids=[gpu_id])    # model 要用 DDP 包装一下


    def _run_batch(self, source, targets):
        self.optimizer.zero_grad()
        output = self.model(source)
        loss = F.cross_entropy(output, targets)
        loss.backward()
        self.optimizer.step()

        # 分布式同步 loss
        reduced_loss = loss.detach()
        dist.all_reduce(reduced_loss, op=dist.ReduceOp.SUM)
        reduced_loss /= dist.get_world_size()
        return reduced_loss.item()


    def _run_epoch(self, epoch):
        b_sz = len(next(iter(self.train_data))[0])
        print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")
        self.train_data.sampler.set_epoch(epoch)        # 在各个 epoch 入口调用 DistributedSampler 的 set_epoch 方法是很重要的，这样才能打乱每个 epoch 的样本顺序
        total_loss = 0.0
        num_batches = 0

        for source, targets in self.train_data: 
            source = source.to(self.gpu_id)
            targets = targets.to(self.gpu_id)
            loss = self._run_batch(source, targets)
            total_loss += loss
            num_batches += 1
        
        avg_loss = total_loss / num_batches
        if self.gpu_id == 0:
            print(f"[GPU{self.gpu_id}] Epoch {epoch} | Avg Loss: {avg_loss:.4f}")


    def _save_checkpoint(self, epoch):
        ckp = self.model.module.state_dict()            # 由于多了一层 DDP 包装，通过 .module 获取原始参数 
        PATH = "checkpoint.pt"
        torch.save(ckp, PATH)
        print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")


    def train(self, max_epochs: int):
        for epoch in range(max_epochs):
            self._run_epoch(epoch)
            # 各个 GPU 上都在跑一样的训练进程，这里指定 rank0 进程保存 ckpt 以免重复保存
            if self.gpu_id == 0 and epoch % self.save_every == 0:
                self._save_checkpoint(epoch)


def prepare_dataset():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_data = datasets.MNIST(
        root = './mnist',
        train=True,       # 设置True为训练数据，False为测试数据
        transform = transform,
        # download=True  # 设置True后就自动下载，下载完成后改为False即可
    )
      
    train_set = MyDataset(train_data)
    
    test_data = datasets.MNIST(
        root = './mnist',
        train=False,       # 设置True为训练数据，False为测试数据
        transform = transform,
    )
      
    test_set = MyDataset(test_data)
      
    return train_set, test_set
      

def load_train_objs():
    train_set, test_set = prepare_dataset()  # load your dataset
    model = ConvNet()  # load your model
    
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    return train_set, test_set, model, optimizer


def prepare_dataloader(dataset: Dataset, batch_size: int):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        pin_memory=True,
        shuffle=False,                      # 设置了新的 sampler，参数 shuffle 要设置为 False 
        sampler=DistributedSampler(dataset) # 这个 sampler 自动将数据分块后送个各个 GPU，它能避免数据重叠
    )


def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int):
    # 初始化进程池, 仅是单个进程 gpu rank 的初始化
    ddp_setup(rank, world_size)

    # 进行训练
    train_set, test_set, model, optimizer = load_train_objs()
    print(f"Train dataset size: {len(train_set)}")
    train_data = prepare_dataloader(train_set, batch_size)
    trainer = Trainer(model, train_data, optimizer, rank, save_every)
    trainer.train(total_epochs)

    # 销毁进程池
    destroy_process_group()


def arg_parser():
    parser = argparse.ArgumentParser(description='MNIST distributed training job')
    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=512, help="Batch size for training")
    parser.add_argument('--save_every', type=int, default=1, help='How often to save a snapshot')
    return parser.parse_args()


r"""
README
执行命令: CUDA_VISIBLE_DEVICES=0,1 python mnist_ddp.py # 用 2 卡训练

注意训练数据是60K条, 训练时输出:
[GPU0] Epoch 0 | Batchsize: 512 | Steps: 59
[GPU1] Epoch 0 | Batchsize: 512 | Steps: 59
512 * 59 = 30208 ~= 30K
排除掉有些 batch_size 不足的情况, 59个 batch 就是 30K, 两个 gpu 平分了数据

"""


if __name__ == "__main__":

    args = arg_parser()
    print(f"Training arguments: {args}")
    
    world_size = torch.cuda.device_count()
    print(f"Using {world_size} GPUs for training")
    
    # 利用 mp.spawn，在整个 distribution group 的 nprocs 个 GPU 上生成进程来执行 fn 方法，并能设置要传入 fn 的参数 args
    # 注意不需要 fn 的 rank 参数，它由 mp.spawn 自动分配
    mp.spawn(
        fn=main, 
        args=(world_size, args.save_every, args.epochs, args.batch_size), 
        nprocs=world_size
    )

启动代码

CUDA_VISIBLE_DEVICES=0,1 python mnist_ddp.py

torchrun 完整代码

import os
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms
from time import time
import argparse


# 对 python 多进程的一个 pytorch 包装
import torch.multiprocessing as mp
# 用于收集一些用于汇总的数据
import torch.distributed as dist
# 这个 sampler 可以把采样的数据分散到各个 CPU 上                                      
from torch.utils.data.distributed import DistributedSampler     

# 实现分布式数据并行的核心类        
from torch.nn.parallel import DistributedDataParallel as DDP         

# DDP 在每个 GPU 上运行一个进程，其中都有一套完全相同的 Trainer 副本（包括model和optimizer）
# 各个进程之间通过一个进程池进行通信，这两个方法来初始化和销毁进程池
from torch.distributed import init_process_group, destroy_process_group 



def ddp_setup():
    """
    setup the distribution process group

    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    # 用torchrun 后台自动设置的环境变量
    init_process_group(backend="nccl")
    torch.cuda.set_device(int(os.environ['LOCAL_RANK']))
    

class ConvNet(nn.Module):

    def __init__(self):
        super(ConvNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.25)
        )
        self.classifier = nn.Sequential(
            nn.Linear(9216, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        output = F.log_softmax(x, dim=1)
        return output
    

# 自定义Dataset
class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image, label = self.data[idx]
        return image, label


class Trainer:
    def __init__(
        self,
        model: torch.nn.Module,
        train_data: DataLoader,
        optimizer: torch.optim.Optimizer,
        save_every: int,
    ) -> None:
        self.gpu_id = int(os.environ['LOCAL_RANK']) # gpu_id 由 torchrun 自动设置
        self.model = model.to(self.gpu_id)
        self.train_data = train_data
        self.optimizer = optimizer
        self.save_every = save_every                    # 指定保存 ckpt 的周期
        self.model = DDP(model, device_ids=[self.gpu_id])    # model 要用 DDP 包装一下


    def _run_batch(self, source, targets):
        self.optimizer.zero_grad()
        output = self.model(source)
        loss = F.cross_entropy(output, targets)
        loss.backward()
        self.optimizer.step()

        # 分布式同步 loss
        reduced_loss = loss.detach()
        dist.all_reduce(reduced_loss, op=dist.ReduceOp.SUM)
        reduced_loss /= dist.get_world_size()
        return reduced_loss.item()


    def _run_epoch(self, epoch):
        b_sz = len(next(iter(self.train_data))[0])
        print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")
        self.train_data.sampler.set_epoch(epoch)
        total_loss = 0.0
        num_batches = 0

        for source, targets in self.train_data: 
            source = source.to(self.gpu_id)
            targets = targets.to(self.gpu_id)
            loss = self._run_batch(source, targets)
            total_loss += loss
            num_batches += 1
        
        avg_loss = total_loss / num_batches
        if self.gpu_id == 0:
            print(f"[GPU{self.gpu_id}] Epoch {epoch} | Avg Loss: {avg_loss:.4f}")


    def _save_checkpoint(self, epoch):
        ckp = self.model.module.state_dict()            # 由于多了一层 DDP 包装，通过 .module 获取原始参数 
        PATH = "checkpoint.pt"
        torch.save(ckp, PATH)
        print(f"Epoch {epoch} | Training checkpoint saved at {PATH}")


    def train(self, max_epochs: int):
        for epoch in range(max_epochs):
            self._run_epoch(epoch)
            # 各个 GPU 上都在跑一样的训练进程，这里指定 rank0 进程保存 ckpt 以免重复保存
            if self.gpu_id == 0 and epoch % self.save_every == 0:
                self._save_checkpoint(epoch)


def prepare_dataset():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_data = datasets.MNIST(
        root = './mnist',
        train=True,       # 设置True为训练数据，False为测试数据
        transform = transform,
        # download=True  # 设置True后就自动下载，下载完成后改为False即可
    )
      
    train_set = MyDataset(train_data)
    
    test_data = datasets.MNIST(
        root = './mnist',
        train=False,       # 设置True为训练数据，False为测试数据
        transform = transform,
    )
      
    test_set = MyDataset(test_data)
      
    return train_set, test_set
      

def load_train_objs():
    train_set, test_set = prepare_dataset()  # load your dataset
    model = ConvNet()  # load your model
    
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    return train_set, test_set, model, optimizer


def prepare_dataloader(dataset: Dataset, batch_size: int):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        pin_memory=True,
        shuffle=False,                      # 设置了新的 sampler，参数 shuffle 要设置为 False 
        sampler=DistributedSampler(dataset) # 这个 sampler 自动将数据分块后送个各个 GPU，它能避免数据重叠
    )


def main(save_every: int, total_epochs: int, batch_size: int):
    # 初始化进程池, 仅是单个进程 gpu rank 的初始化
    ddp_setup()

    # 进行训练
    train_set, test_set, model, optimizer = load_train_objs()
    print(f"Train dataset size: {len(train_set)}")
    train_data = prepare_dataloader(train_set, batch_size)
    trainer = Trainer(model, train_data, optimizer, save_every)
    trainer.train(total_epochs)

    # 销毁进程池
    destroy_process_group()


def arg_parser():
    parser = argparse.ArgumentParser(description='MNIST distributed training job')
    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=512, help="Batch size for training")
    parser.add_argument('--save_every', type=int, default=1, help='How often to save a snapshot')
    return parser.parse_args()


r"""
README
执行命令: CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 mnist_ddp_torchrun.py
"""


if __name__ == "__main__":

    args = arg_parser()
    print(f"Training arguments: {args}")
    
    world_size = torch.cuda.device_count()
    print(f"Using {world_size} GPUs for training")

    main(
        args.save_every, 
        args.epochs, 
        args.batch_size
    )

启动代码

CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 mnist_ddp_torchrun.py