基于DistributedDataParallel (DDP)的单机多卡数据并行(multiprocessing启动)

这次我们一起学习下分布式的基础内容——基于Distributed Data Parallel (DDP)的单机多卡数据并行。分布式训练是指将模型训练任务分别为多个子任务，在多个计算结点上并行进行训练。一般分为两种场景，其一是大模型分为若干小模型，在多个计算结点跑相同数据以实现加速训练，即模型并行 (Model parallel)；其二是数据集分为若干小集合，在多个计算结点跑相同模型，以加速训练，即数据并行 (Data parallel)。简单讲就是划分谁，谁并行。除此之外，也可以将两种方式进行组合使用。

分布数据并行原理

pytorch框架中的DistributedDataPrallel是基于Ring-All-Reduce通信模式的，Ring-All-Reduce是一种高效的集群通信算法。对于数据并行而言，问题在于如何合并计算结果。举个例子，一共四个计算结点，最简单的莫过于四个计算结点把计算结果都发给某一个主计算结点，主计算结点计算好梯度后，再发给所有计算结点。但是这种方式通信成本较大，单点通信压力大，从而可能影响训练速度，降低系统稳定性。

那么Ring-All-Reduce是如何操作的呢。如果你已经了解过了，可以直接看图1；否则，可以先看我举的例子，再食用图1。举个例子，每个 GPU 上有 4 个数据块需要通信。以下是每个 GPU 上初始的数据：

• GPU 0: [A0, A1, A2, A3]
• GPU 1: [B0, B1, B2, B3]
• GPU 2: [C0, C1, C2, C3]
• GPU 3: [D0, D1, D2, D3]

Ring-All-Reduce 分为两个阶段：

1. 第一阶段：环形传递与累加（Reduce-Scatter）

   • 将数据按块分割，每个设备只传递和累加一个块的数据。

   第1轮传递：

   • GPU 0 将 A1 发送给 GPU 1，接收 GPU 3 的 D0 并累加为 A0+D0。
   • GPU 1 将 B1 发送给 GPU 2，接收 GPU 0 的 A1 并累加为 B1+A1。
   • GPU 2 将 C1 发送给 GPU 3，接收 GPU 1 的 B1 并累加为 C1+B1。
   • GPU 3 将 D1 发送给 GPU 0，接收 GPU 2 的 C1 并累加为 D1+C1。

   这一步完成后，各 GPU 上的数据如下：

   • GPU 0: [A0+D0, A1, A2, A3]
   • GPU 1: [B0, B1+A1, B2, B3]
   • GPU 2: [C0, C1+B1, C2, C3]
   • GPU 3: [D0, D1+C1, D2, D3]

   第2轮传递：

   • GPU 0 将 A2 发送给 GPU 1，接收 GPU 3 的 D1+C1 并累加为 A1+D0+C1。
   • GPU 1 将 B2 发送给 GPU 2，接收 GPU 0 的 A2 并累加为 B2+A2。
   • GPU 2 将 C2 发送给 GPU 3，接收 GPU 1 的 B2 并累加为 C2+B2。
   • GPU 3 将 D2 发送给 GPU 0，接收 GPU 2 的 C2 并累加为 D2+C2。

   以此类推，经过三轮传递，每个 GPU 上最终会保存归约后的一块数据：

   • GPU 0: [A0+D0+C0+B0]
   • GPU 1: [A1+B1+C1+D1]
   • GPU 2: [A2+B2+C2+D2]
   • GPU 3: [A3+B3+C3+D3]

2. 第二阶段：环形广播（All-Gather）

   • 通过环形广播，将每个 GPU 上累加后的结果发送给所有设备。

   经过广播后，每个 GPU 上的数据都相同，最终获得了全局归约的结果：

   • GPU 0, GPU 1, GPU 2, GPU 3: [A0+B0+C0+D0, A1+B1+C1+D1, A2+B2+C2+D2, A3+B3+C3+D3]

Ring-All-Reduce 的高效性

总结下它为什么高效就是数据分组和环形通信

• 数据分组：数据块被分组处理，设备可以同时发送和接收数据，减少了等待时间。
• 环形通信：每个设备只与相邻的两个设备通信，避免了网络拥塞，且通信次数是线性的，与设备数量成正比。

图1 Ring-All-Reduce 原理

代码实现

我们以手写数字集为例来进行数据并行的代码实现。首先，导包。

import os  # parameter registration
import time  # record time
import torch
import argparse  # setting parameters
import torch.nn as nn
from model import ConvNet  # loading model
import torch.distributed as dist
import torch.multiprocessing as mp  # multi-process
from torch.cuda.amp import GradScaler  # mixed precision
from torch.utils.data import DataLoader
from torch.utils.data.dataset import TensorDataset
from dl import load_mnist_data, split_data, EarlyStop  # loading data
from torch.utils.data.distributed import DistributedSampler  # distributed data sampler

定义一个简单的卷积网络。

class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1= nn.Sequential(
            nn.Conv2d(1, 32, kernel_size= 3, padding= 1, stride= 2),  # (b, 1, 28, 28)>> (b, 32, 14, 14)
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size= 2, stride= 2)  # (b, 32, 14, 14)>> (b, 32, 7, 7)
        )
        self.layer2= nn.Sequential(
            nn.Conv2d(32, 64, kernel_size= 3, padding= 1, stride= 2),  # (b, 32, 7, 7)>> (b, 64, 4, 4)
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size= 2, stride= 2)  # (b, 64, 4, 4)>> (b, 64, 2, 2)
        )
        self.fc= nn.Sequential(
            nn.Linear(64* 2* 2, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.fc(self.layer2(self.layer1(x)).view(x.shape[0], -1))

注册参数。

def prepare():
    parser= argparse.ArgumentParser()
    parser.add_argument('--gpu', default= '2, 3')
    parser.add_argument('--epochs', type= int, default= 10000, help= 'epoch number')
    parser.add_argument('--lr', type= float, default= 1e-4, help= 'learning rate')
    parser.add_argument('--weight_decay', type= float, default= 5e-5, help= 'weight decay')
    parser.add_argument('--batch_size', type= int, default= 256, help= 'batch size number')
    parser.add_argument('--mnist_file_path', type= str, default= '../data/mnist_test.csv')
    parser.add_argument('--seed', type= int, default= 1206)
    parser.add_argument('--master_addr', type= str, default= 'localhost')
    parser.add_argument('--port', type= str, default= '1206')
    parser.add_argument('--pt_path', type= str, default= '../pt/')
    parser.add_argument('--patience', type= int, default= 5)
    args= parser.parse_args()
    # setting environment variables to enable DDP
    os.environ['MASTER_ADDR']= args.master_addr  # ip address of the master machine, where is others?
    os.environ['MASTER_PORT']= args.port  # port number of the master machine,
    os.environ['CUDA_VISIBLE_DEVICES']= args.gpu  # visible gpu devices
    world_size= torch.cuda.device_count()
    os.environ['WORLD_SIZE']= str(world_size)  # process number
    return args

下面代码是初始化ddp，local_rank是当前进程在其所在计算结点上的GPU编号，例如一个计算结点有4个GPU，那么local_rank的取值就是0-3。它可以用来指定GPU和确定当前进程处理的数据子集。backend=nccl是说分布式通信的后端采用NVIDIA Collective Communication Library，即nccl，它是英伟达提供的高效通信库，专用于多GPU下的分布式训练优化。RANK 是全局进程的唯一标识符，它指示当前进程在整个分布式环境中的位置。它可以用于任务分配和控制通信。那为什么把local_rank给os.environ[rank]呢？因为我们是单机多卡状态，是卡与卡的任务分配和通信。

def init_ddp(local_rank):
    # after this setup, tensors can be moved to GPU via 'a= a.cuda()' rather than 'a= a.to(local_rank)'
    torch.cuda.set_device(local_rank)
    os.environ['RANK']= str(local_rank)
    dist.init_process_group(backend= 'nccl')  # init communication mode

然后我们需要一个generator，用来控制每个计算结点上数据的随机顺序。

def get_ddp_generator(seed= 1206):
    local_rank= dist.get_rank()
    g= torch.Generator()
    g.manual_seed(seed+ local_rank)
    return g

训练代码如下

def train(epoch, net, train_loader, loss_fn, opt, scaler):
    '''
    parameters:
        scaler, Part of the Automatic Mixed Precision (AMP) library. AMP allows you to store tensors with lower precision (
                such as half precision floating-point numbers, i.e. torch. float 16) when training deep learning models, while
                avoiding numerical underflow or overflow issues with higher precision (such as full precision floating-point 
                numbers, i.e. torch. float 32) when calculating gradients.
    '''
    net.train()
    for i, data in enumerate(train_loader):
        imgs= data[0].cuda()
        lbls= data[1].cuda()
        outputs= net(imgs)
        loss= loss_fn(outputs, lbls)
        opt.zero_grad()
        scaler.scale(loss).backward()  # avoiding the problem of overflow caused by numbers that are difficult to represent with float16
        scaler.step(opt)
        # The main purpose of the 'scaler. update()' operation is to adjust the threshold of the scale factor based on
        # the most recent gradient update, in order to ensure more effective management of the gradient range in subse
        # quent iterations and avoid the problem of gradient vanishing or exploding.
        scaler.update()
        print(f'epoch: {epoch+ 1}, step: {i+ 1}, loss: {loss.item()}, device: {dist.get_rank()}', flush= True)

同上，验证代码类似的训练过程。

def valid(net, valid_loader):
    local_rank= dist.get_rank()
    net.eval()
    num= torch.tensor(0.0).cuda()
    correct= torch.tensor(0.0).cuda()
    for imgs, lbls in valid_loader:
        imgs= imgs.cuda()
        lbls= lbls.cuda()
        with torch.no_grad():
            outputs= net(imgs)
            num+= outputs.shape[0]
        correct+= (outputs.argmax(1)== lbls).type(torch.float).sum()
    dist.reduce(num, dst= 0, op= dist.ReduceOp.SUM)  # All reduce on num
    dist.reduce(correct, dst= 0, op= dist.ReduceOp.SUM) # All reduce
    return correct* 1.0/ num

main函数中要注意DataLoader中要传入定义的采样器，train_loader中不需要shuffle，但需要传入generator，从而区别于单计算结点的操作。

def main(local_rank, args):
    # init ddp
    init_ddp(local_rank)
    # load data
    lbls, imgs= load_mnist_data(args.mnist_file_path)
    lbls= lbls.long()
    imgs= imgs.float().view(-1, 1, 28, 28)
    lbl_tr, lbl_val, lbl_te, img_tr, img_val, img_te= split_data(args.seed, lbls, imgs, 0.8, 0.1)
    dataset_tr, dataset_val, dataset_te= TensorDataset(img_tr, lbl_tr), TensorDataset(img_val, lbl_val), TensorDataset(img_te, lbl_te)
    train_sampler, val_sampler, test_sampler= DistributedSampler(dataset_tr), DistributedSampler(dataset_val), DistributedSampler(dataset_te)
    # generator control random behavior (such as data shuffling) to ensure that the order of data on each GPU is different but controllable.
    train_loader= DataLoader(dataset= dataset_tr, batch_size= args.batch_size, shuffle= False, sampler= train_sampler, generator= get_ddp_generator())
    valid_loader= DataLoader(dataset= dataset_val, batch_size= args.batch_size, shuffle= False, sampler= val_sampler)
    test_loader= DataLoader(dataset= dataset_te, batch_size= args.batch_size, shuffle= False, sampler= test_sampler)
    # init net
    net= ConvNet().cuda()
    net= nn.SyncBatchNorm.convert_sync_batchnorm(net)
    net= nn.parallel.DistributedDataParallel(net, device_ids= [local_rank])
    loss_fn= nn.CrossEntropyLoss().cuda()
    opt= torch.optim.SGD(net.parameters(), args.lr)
    scaler= GradScaler()  # mixed precision training
    earlystopping= EarlyStop(args.pt_path, 'cnn.pt', args.patience)
    stop_flag_tensor= torch.tensor([0], dtype= torch.long).cuda()
    # train
    for epoch in range(args.epochs):
        # changing the sampling starting point to ensure that the order of data in each round is different, the model's generalization ability can be improved.        
        train_loader.sampler.set_epoch(epoch)
        train(epoch, net, train_loader, loss_fn, opt, scaler)
        # valid
        acc= valid(net, valid_loader)
        if local_rank== 0:
            earlystopping.check_early_stopping(-acc, net, opt, scaler)
            print(f'epoch: {epoch}, early_stopping.time: {earlystopping.time}, acc: {acc}')
            stop_flag_tensor.fill_(int(earlystopping.flag))
        dist.all_reduce(stop_flag_tensor, op= dist.ReduceOp.SUM)
        if stop_flag_tensor> 0:break
    # destory the process group
    dist.destroy_process_group()

接着，只要用多进程执行main函数就好啦。

if __name__== '__main__':
    args= prepare()
    time_start= time.time()
    mp.spawn(main, args= (args, ), nprocs= torch.cuda.device_count())
    time_end= time.time()
    cost_time= time_end- time_start
    print(f'cost time: {cost_time:.2f} seconds')

日志文件的截图如图2所示，可以看到不同计算结点在一个batch后进行了梯度上的通信。由于模型比较简单，早停patience设置不当，acc可以忽略哈。

图2 运行结果图

generator和set_epoch的使用

可以看到我们dataloader中传入了generator，并且每次epoch开始设置train_loader.sampler.set_epoch(epoch)。generator和set_epoch的使用主要是为了确保数据的分布和顺序在不同设备（如GPU）上保持一致和可控。下面我们分别解释这两个部分的作用和原理，并通过一个具体的例子来帮助理解。

generator的作用

在分布式训练中，通常会将数据划分到多个设备上并行处理。为了确保每个设备上生成的数据是可重复的，并且在每个设备上的训练结果是相对一致的，我们需要设置一个随机数生成器（generator）来控制数据的随机顺序。在代码中，get_ddp_generator()函数使用了一个全局的seed（种子）和local_rank（设备的标识号）来创建并返回一个随机数生成器。这意味着每个设备会有自己独立的generator，但由于使用了全局seed和设备标识，这些generator的生成顺序是可控和可复现的。

set_epoch的作用

在分布式训练中，set_epoch是用来确保数据在每个训练轮次（epoch）中的分布是不同的。特别是在使用DistributedSampler时，set_epoch会改变每一轮次数据划分的起始点，从而确保每一轮次数据是不同的。否则，如果不设置epoch，每一轮次的数据划分可能会是相同的，从而导致模型的训练效果下降。

例子

假设我们有1000张图片的数据集，并且我们使用2个GPU进行分布式训练，每个GPU处理一半的数据。

• 第1轮训练：
  • GPU 0 使用种子1206 + 0，生成随机数控制数据顺序。
  • GPU 1 使用种子1206 + 1，生成不同的随机数，确保数据不同。
  • set_epoch(0)让GPU 0和GPU 1各自从同一位置开始采样，但由于不同的generator，数据顺序不同。
• 第2轮训练：
  • 再次调用set_epoch(1)，这次DistributedSampler改变了采样的起始点，因此GPU 0和GPU 1采样的数据和顺序会与第一轮不同。

具体过程：

• 你可以把数据看作是一副洗好的扑克牌。generator决定了如何洗牌，而set_epoch决定了每一轮次从哪里开始发牌。使用set_epoch就像是在每一轮游戏开始前重新从不同位置切牌一样，这样每次的发牌顺序都不同，有助于模型学习数据的不同组合。

总结

generator确保每个设备上随机行为的可控性。set_epoch确保每个epoch的数据顺序不同，防止模型过拟合特定的数据顺序。

通过这两个机制，分布式训练在不同设备上能保持数据的合理分配和有效的训练。

完整代码

详细代码解释可以查看我的github.

参考资料

[1] Perplexity

[2] 「分布式训练」原理讲解+ 「DDP 代码实现」修改要点

[3] Pytorch 分散式訓練 DistributedDataParallel — 實作篇

[4] DDP-practice/ddp_main.py at main · rickyang1114/DDP-practice (github.com)