DDP分布式数据并行训练

一、环境准备

查看显卡信息

nvidia-smi

Conda安装Pytorch，创建一个环境并使用，安装对应cuda版本的pytorch。

conda install pytorch torchvision torchaudio pytorch-cuda=<cuda-version> -c pytorch -c nvidia

二、训练方式

分布式数据并行方式（DistributedDataParallel）

PyTorch官方文档-分布式数据并行 Distributed Data Parallel

这是 PyTorch 推荐的分布式训练方式，适用于单机多 GPU 和 多机多 GPU。每个进程分配到一个 GPU，每个 GPU 上运行一个模型副本。所有进程都并行计算并独立反向传播，最终通过 torch.distributed 在不同设备之间同步梯度。

1.分布式数据并行方式（DistributedDataParallel）

1）训练脚本

# ddp_demo.py

import os
import torch
import torchvision
import torch.distributed as dist
import torch.utils.data.distributed
from torchvision import transforms


def main():
    # 从环境变量获取当前进程的本地排名，用于指定使用的GPU
    local_rank = int(os.environ['LOCAL_RANK'])
    device = torch.device(f'cuda:{local_rank}')  # 设置当前进程使用的设备为对应的GPU
    torch.cuda.set_device(device)  # 指定当前GPU设备

    # 初始化分布式进程组，使用NCCL作为后端
    dist.init_process_group(backend='gloo', init_method="env://")

    # 数据预处理：转换为Tensor并进行归一化
    trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
    # 加载MNIST数据集，设置训练集，应用转换，下载数据
    data_set = torchvision.datasets.MNIST("./", train=True, transform=trans, target_transform=None, download=True)
    # 创建分布式采样器，确保每个进程只处理数据集的一部分
    train_sampler = torch.utils.data.distributed.DistributedSampler(data_set)
    # 创建数据加载器，使用分布式采样器
    data_loader_train = torch.utils.data.DataLoader(dataset=data_set, batch_size=256, sampler=train_sampler)

    # 创建ResNet101模型并调整输入层以接受单通道图像
    net = torchvision.models.resnet101(num_classes=10)
    net.conv1 = torch.nn.Conv2d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
    net = net.to(device)  # 将模型移动到当前设备
    # 包装模型为分布式数据并行模型
    net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[device], output_device=device)

    criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数
    opt = torch.optim.Adam(params=net.parameters(), lr=0.001)  # 定义优化器

    print(f"Training on {device}, Rank: {dist.get_rank()}")  # 打印当前使用的CUDA设备和进程号

    # 训练循环
    for epoch in range(10):
        train_sampler.set_epoch(epoch)  # 每个epoch更新采样器
        for step, data in enumerate(data_loader_train):
            images, labels = data  # 获取数据和标签
            images, labels = images.to(device), labels.to(device)  # 移动数据到当前设备
            opt.zero_grad()  # 清空梯度
            outputs = net(images)  # 前向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            opt.step()  # 更新参数
            if step % 10 == 0:
                print(f"Rank {dist.get_rank()}, Device {device}, Loss: {loss.item()}")  # 打印损失和设备信息

    # 仅在rank 0的进程中保存模型参数
    if dist.get_rank() == 0:
        torch.save(net.state_dict(), "my_net.pth")
        print(f"Model saved by Process ID: {os.getpid()}")  # 打印保存模型的进程号

    # 确保销毁进程组
    dist.destroy_process_group()


if __name__ == "__main__":
    main()  # 运行主函数

2）单机多GPU

使用 torchrun 命令启动

--nproc_per_node=2 表示使用两张显卡。

torchrun命令简化启动过程，它可以自动管理多进程运行，处理进程的创建和销毁，无需手动管理。它会自动配置分布式训练所需的环境变量，比如 LOCAL_RANK，使得每个进程能够正确识别自己的 GPU。可以省去很多麻烦。

如果不使用torchrun命令的话，需要更改代码并自己定义进程及管理 LOCAL_RANK，会比较麻烦容易出现bug。

LOCAL_RANK是在分布式训练中用来标识每个进程所使用的 GPU 的变量。在多进程训练中，每个进程通常会绑定到特定的 GPU，以避免多个进程同时使用同一个 GPU，从而引起资源冲突。

torchrun --nproc_per_node=2 ddp_demo.py

训练过程，模型保存在训练脚本同级目录下（仅在 LocalRank 0 的进程中保存模型参数）。

3）多机多GPU

多机多卡训练与单机多卡训练本质上并无区别。无论是同一节点还是不同节点的进程，分布式训练的核心在于通过特定的通信方式汇总梯度，以更新各进程的模型参数。但是与单机多卡不同的是，不同节点之间需要指定通信的IP地址和端口号，因此需额外设置参数。

前提条件

• 网络配置：确保两个节点可以通过 SSH 互相访问，并且能够使用相同的端口（通常是 29500）进行通信。
• PyTorch 安装：确保在两个节点上都安装了 PyTorch，并且具有相同版本。
• CUDA 驱动：确保 CUDA 驱动在两个节点上都正常工作。

使用 torchrun 命令启动训练

假设我有两个节点：192.168.20.43和192.168.20.44。 我们让192.168.20.43作为0节点，只需要在两个节点上分别运行以下命令即可。

# Node 0
torchrun --nproc_per_node=1 --nnodes=2 --node_rank=0 --rdzv_id=123 --master_addr='192.168.20.43' --master_port=29500 ddp_demo.py

# Node 1
torchrun --nproc_per_node=1 --nnodes=2 --node_rank=1 --rdzv_id=123 --master_addr='192.168.20.43' --master_port=29500 ddp_demo.py

改变的是 --node_rank 参数值，分别设置主节点和从节点的排名。多个节点以此类推。 新增 --rdzv_id 参数，rendezvous（会合）ID，用于标识分布式训练会话。确保所有参与的节点使用相同的 ID，以便它们能够找到彼此。

训练图示

192.168.20.43 Node 0 可以看到 192.168.20.44 已经连上来了

192.168.20.43 Node 0 训练过程

192.168.20.44 Node 1 训练过程

错误解决

多GPU并行通信问题，提示在 NCCL 的通信过程中发生了未处理的系统错误，还不知道什么原因导致的，目前换成gloo后端通信能正常跑通训练。

# 初始化分布式进程组，使用GLOO作为后端
dist.init_process_group(backend='gloo', init_method="env://")

4）Rank 和 CUDA 设备编号说明

CUDA 设备编号：每个节点的 GPU 设备是按顺序编号的，比如： 在 192.168.20.43 上，设备是 cuda:0。 在 192.168.20.44 上，设备也是 cuda:0（对应于该节点的第一张 GPU）。

Rank 参数：在分布式训练中，rank 是全局的进程标识符： 在0节点 192.168.20.43 上，rank 为 0。 在1节点 192.168.20.44 上，rank 为 1。