2020-09-09

PyTorch 分布式训练教程

---

DistributedDataParallel内部机制

       DistributedDataParallel通过多进程在多个GPUs间复制模型，每个GPU都由一个进程控制（当然可以让每个进程控制多个GPU，但这显然比每个进程有一个GPU要慢；也可以多个进程在一个GPU上运行）。GPU可以都在同一个节点上，也可以分布在多个节点上。每个进程都执行相同的任务，并且每个进程都与所有其他进程通信。进程或者说GPU之间只传递梯度，这样网络通信就不再是瓶颈。

---

在训练过程中，每个进程从磁盘加载batch数据，并将它们传递到其GPU。每一个GPU都有自己的前向过程，然后梯度在各个GPUs间进行All-Reduce。每一层的梯度不依赖于前一层，所以梯度的All-Reduce和后向过程同时计算，以进一步缓解网络瓶颈。在后向过程的最后，每个节点都得到了平均梯度，这样模型参数保持同步。

这都要求多个进程（可能在多个节点上）同步并通信。Pytorch通过distributed.init_process_group函数来实现这一点。他需要知道进程0位置以便所有进程都可以同步，以及预期的进程总数。每个进程都需要知道进程总数及其在进程中的顺序，以及使用哪个GPU。通常将进程总数称为world_size.Pytorch提供了nn.utils.data.DistributedSampler来为各个进程切分数据，以保证训练数据不重叠。

二、实例讲解

1.普通单卡训练

代码如下（示例）：

import os
from datetime import datetime
import argparse
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

class ConvNet(nn.Module):  
    def __init__(self, num_classes=10):  
        super(ConvNet, self).__init__()  
        self.layer1 = nn.Sequential(  
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),  
            nn.BatchNorm2d(16),  
            nn.ReLU(),  
            nn.MaxPool2d(kernel_size=2, stride=2))  
        self.layer2 = nn.Sequential(  
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),  
            nn.BatchNorm2d(32),  
            nn.ReLU(),  
            nn.MaxPool2d(kernel_size=2, stride=2))  
        self.fc = nn.Linear(7*7*32, num_classes)  
    def forward(self, x):  
        out = self.layer1(x)  
        out = self.layer2(out)  
        out = out.reshape(out.size(0), -1)  
        out = self.fc(out)  
        return out

def main():  
    parser = argparse.ArgumentParser()  
    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N')  
    parser.add_argument('-g', '--gpus', default=1, type=int,  
                        help='number of gpus per node')  
    parser.add_argument('-nr', '--nr', default=0, type=int,  
                        help='ranking within the nodes')  
    parser.add_argument('--epochs', default=2, type=int, metavar='N',  
                        help='number of total epochs to run')  
    args = parser.parse_args()  
    train(0, args) 

训练模型
def train(gpu, args):
   torch.manual_seed(0)
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    batch_size = 100
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data',
                                               train=True,
                                               transform=transforms.ToTensor(),
                                               download=True)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True,
                                               num_workers=0,
                                               pin_memory=True)
    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i + 1) % 100 == 0 and gpu == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
                    epoch + 1,
                    args.epochs,
                    i + 1,
                    total_step,
                    loss.item())
                   )
    if gpu == 0:
        print("Training complete in: " + str(datetime.now() - start))


if __name__ == '__main__':  
    main()

2.分布式训练

使用多进程进行分布式训练，我们需要为每个GPU启动一个进程。每个进程需要知道自己运行在哪个GPU上，以及自身在所有进程中的序号。对于多节点，我们需要在每个节点启动脚本。

首先，我们要配置基本的参数：

def main():  
    parser = argparse.ArgumentParser()  
    parser.add_argument('-n', '--nodes', default=1,  
                        type=int, metavar='N')  
    parser.add_argument('-g', '--gpus', default=1, type=int,  
                        help='number of gpus per node')  
    parser.add_argument('-nr', '--nr', default=0, type=int,  
                        help='ranking within the nodes')  
    parser.add_argument('--epochs', default=2, type=int,  
                        metavar='N',  
                        help='number of total epochs to run')  
    args = parser.parse_args()  
    #########################################################  
    args.world_size = args.gpus * args.nodes                #  
    os.environ['MASTER_ADDR'] = '10.57.23.164'              #  
    os.environ['MASTER_PORT'] = '8888'                      #  
    mp.spawn(train, nprocs=args.gpus, args=(args,))         #  
    #########################################################

其中args.nodes是节点总数，而args.gpus是每个节点的GPU总数（每个节点GPU数是一样的），而args.nr 是当前节点在所有节点的序号。节点总数乘以每个节点的GPU数可以得到world_size，也即进程总数。所有的进程需要知道进程0的IP地址以及端口，这样所有进程可以在开始时同步，一般情况下称进程0是master进程，比如我们会在进程0中打印信息或者保存模型。PyTorch提供了mp.spawn来在一个节点启动该节点所有进程，每个进程运行train(i, args)，其中i从0到args.gpus - 1。

同样，我们要修改训练函数： 

def train(gpu, args):
    ############################################################  
    rank = args.nr * args.gpus + gpu
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=args.world_size,
        rank=rank
    )
    ############################################################  

    torch.manual_seed(0)
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    batch_size = 100
    # define loss function (criterion) and optimizer  
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)

    ###############################################################  
    # Wrap the model  
    model = nn.parallel.DistributedDataParallel(model,
                                                device_ids=[gpu])
    ###############################################################  
    # Data loading code  
    train_dataset = torchvision.datasets.MNIST(
        root='./data',
        train=True,
        transform=transforms.ToTensor(),
        download=True
    )
    ################################################################  
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )
    ################################################################  
    train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=batch_size,
        ##############################  
        shuffle=False,  # 
        ##############################  
        num_workers=0,
        pin_memory=True,
        #############################  
        sampler=train_sampler)  # 
    #############################  
    ... 

这里我们首先计算出当前进程序号：rank = args.nr * args.gpus + gpu，然后就是通过dist.init_process_group初始化分布式环境，其中backend参数指定通信后端，包括mpi, gloo, nccl，这里选择nccl，这是Nvidia提供的官方多卡通信框架，相对比较高效。mpi也是高性能计算常用的通信协议，不过你需要自己安装MPI实现框架，比如OpenMPI。gloo倒是内置通信后端，但是不够高效。init_method指的是如何初始化，以完成刚开始的进程同步；这里我们设置的是env://，指的是环境变量初始化方式，需要在环境变量中配置4个参数：MASTER_PORT，MASTER_ADDR，WORLD_SIZE，RANK，前面两个参数我们已经配置，后面两个参数也可以通过dist.init_process_group函数中world_size和rank参数配置。其它的初始化方式还包括共享文件系统以及TCP，比如init_method='tcp://10.1.1.20:23456'，其实也是要提供master的IP地址和端口。注意这个调用是阻塞的，必须等待所有进程来同步，如果任何一个进程出错，就会失败。

对于模型侧，我们只需要用DistributedDataParallel包装一下原来的model即可，在背后它会支持梯度的All-Reduce操作。对于数据侧，我们nn.utils.data.DistributedSampler来给各个进程切分数据，只需要在dataloader中使用这个sampler就好，值得注意的一点是你要训练循环过程的每个epoch开始时调用train_sampler.set_epoch(epoch)，（主要是为了保证每个epoch的划分是不同的）其它的训练代码都保持不变。

最后就可以执行代码了，比如我们是4节点，每个节点是8卡，那么需要在4个节点分别执行：

python src/mnist-distributed.py -n 4 -g 8 -nr i

要注意的是，此时的有效batch_size其实是batch_size_per_gpu * world_size，对于有BN的模型还可以采用同步BN获取更好的效果：

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

参考

1. Distributed data parallel training in Pytorch https://yangkky.github.io/2019/07/08/distributed-pytorch-tutorial.html (大部分内容来自此处)
2. torch.distributed https://pytorch.org/docs/stable/distributed.html
3. https://my.oschina.net/u/4478364/blog/3212642

---