Pytorch 分布式训练 DistributedDataParallel （实操篇）

首先介绍一下分布式的基础概念：

• group：指进程组，默认为一组
• backwend：指进程使用的通讯后端，Pytorch支持mpi、gloo、nccl、若是使用N卡推荐使用nccl
• world_size：指进程组中的进程数量
  若使用单台机器多GPU，world_size表示使用的GPU数量
  若使用多台机器多GPU，world_size表示使用的机器数量
• rank：指当前进程的序号，用于进程间的通讯，rank=0表示为master。
  若使用单台机器多GPU，rank表示当前正在使用的某个GPU
  若使用多台机器多GPU，rank表示当前正在使用的某台机器
• local_rank：指每台机器上的进程序号
• rank VS local_rank：
  假设目前使用两台机器，每台机器有3块GPU，每个GPU使用一个进程（DDP的最佳使用方式是每一个GPU使用一个进程）
  机器一的rank序号为0,1,2;机器二的rank序号为3,4,5
  机器一的local_rank序号为0,1,2;机器二的local_rank序号为0,1,2

接下来就是实操：

0.导入相关库

我的 pytorch 版本為 1.9.0、torchvision 版本為 0.10.0。

import torch
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler

1.设置local_rank argparse参数

在启动分布式训练时候，需要在命令行使用torch.distributed.launch启动器，该启动器会将当前进程的序号（若每个GPU使用一个进程，也是指GPU序号）通过local_rank参数传递给Python文件。

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()

2.初始化DDP

设置进程通讯的后端为nccl
可参考官方参数设置https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group

dist.init_process_group(backend='nccl')

同步所有进程

dist.barrier()

获取当前进程组内的所有进程数：

world_size = dist.get_world_size()

3.数据并行化

将数据集划分给每一个进程，避免进程间的数据重复

train_sampler = DistributedSampler(train_dataset)
valid_sampler = DistributedSampler(valid_dataset)

• DistributedSampler 分配方式
  假设使用2个GPU且数据集分为0~4份，首先将数据集进行shuffle后分为两份，不足的部分由第一份数据补充，最后再以间隔方式分配数据至不同的GPU中。
  
  接著再將 sampler 傳入 DataLoader
  DataLoader 官方參數設置: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader

train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=256, pin_memory=False, prefetch_factor=2, num_workers=4)
valid_loader = DataLoader(valid_dataset, sampler=valid_sampler, batch_size=256, pin_memory=False, prefetch_factor=2, num_workers=4)

4.模型

假设模型已经load，直接使用.to(device)，将模型分配至CUDA

device = torch.device("cuda", args.local_rank)
model = model.to(devive)

使用SyncBN优化各自进程中数据集的BN计算

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

使用DistributedDataParallel将模型包装起来

 model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

5.set_epoch

在每一次迭代前，通过set_epoch来设置一个random seed，这样每次运行才会将每个epoch的数据都shuffle，并且让每个GPU拿到的数据不同。

train_sampler.set_epoch(epoch)
valid_sampler.set_epoch(epoch)

6.启动分布式训练

这部分有两种写法：torch.distributed.launch和torchrun。

6.1 torch.distributed.launch

参数说明：
nproc_per_node:每台机器有多少进程
nnodes:总共使用多少台机器
node_rank:当前进程属于那台机器
master_addr:master进程的网络地址
master_port:master进程的端口，首先要确认该端口是否被其他进程占用
可参考官方设置：https://pytorch.org/docs/stable/distributed.html#launch-utility

• 运行一台机器多GPU
  CUDA_VISIBLE_DEVICES参数为选择要使用哪一台GPU，0，1表示使用的是第0，1序号的GPU。

$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py

• 运行多台机器多GPU
  假设IP为192.168.1.1 free port为1234

# 第一台机器
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=”192.168.1.1" --master_port=1234 ddp_example.py

# 第二台机器
$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr=”192.168.1.1" --master_port=1234 ddp_example.py

• 运行一台机器多GPU，但是执行两个不同的实验

$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py
# 使用不同 master_port 避免冲突
$ CUDA_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 --master_port 9999 ddp_example.py

6.2 torchrun

PyTorch提供了新的指令替代 torch.distributed.launch，参数的使用方式参考文件： https://pytorch.org/docs/stable/elastic/run.html#launcher-api

• 运行一台机器多GPU

$ CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 --nproc_per_node=2 ddp_example.py

• 运行多台机器多GPU
  参数nnodes为使用机器序号的区间、rdzv_id为唯一的job id、rdzv_endpoint为backwend所执行的endpoint，格式为:，替代原来的master_addr以及master_port。

$ torchrun --nnodes=1:4 --nproc_per_node=$NUM_TRAINERS --rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR ddp_example.py

7.完整代码

完整的代码参见：https://github.com/chingi071/Pytorch_note/blob/master/ddp_example.py

import os
import argparse
import torch
import torch.nn as nn
from torch.optim import lr_scheduler
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms

def train_epoch(train_loader, optimizer, criterion, lr_scheduler, model, world_size):
        model.train()

        train_running_loss = 0.0
        train_running_acc = 0.0    
        for batch_idx, (data, target) in enumerate(train_loader):
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)

            output = model(data)        
            preds = torch.max(output, 1)[1]

            loss = criterion(output, target)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_running_loss += loss.item()
            train_running_acc += torch.eq(preds, target).sum().item()

        lr_scheduler.step()
        train_loss_value = train_running_loss/ (len(train_dataset) / world_size)
        train_acc_value = train_running_acc/ (len(train_dataset) / world_size)

        return train_loss_value, train_acc_value

def valid_epoch(valid_loader, criterion, model, world_size):
    model.eval()

    valid_running_loss = 0.0
    valid_running_acc = 0.0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(valid_loader):
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)

            outputs = model(data)
            preds = torch.max(outputs, 1)[1]

            loss = criterion(outputs, target)

            valid_running_loss += loss.item() 
            valid_running_acc += torch.eq(preds, target).sum().item()

    valid_loss_value = valid_running_loss/ (len(valid_dataset) / world_size)
    valid_acc_value = valid_running_acc/ (len(valid_dataset) / world_size)

    return valid_loss_value, valid_acc_value

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", default=0, type=int)
    
    args = parser.parse_args()

    dist.init_process_group(backend='nccl')
    dist.barrier()
#     rank = dist.get_rank()
    world_size = dist.get_world_size()


    train_transform = transforms.Compose([transforms.RandomResizedCrop(224),
                                          transforms.RandomHorizontalFlip(),
                                          transforms.ToTensor(),
                                          transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                                               std=[0.229, 0.224, 0.225])])

    valid_transform = transforms.Compose([transforms.Resize(256),
                                          transforms.CenterCrop(224),
                                          transforms.ToTensor(),
                                          transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                                               std=[0.229, 0.224, 0.225])])

    train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)
    valid_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=valid_transform)

    train_sampler = DistributedSampler(train_dataset)
    valid_sampler = DistributedSampler(valid_dataset)

    train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=256,
                              pin_memory=False, prefetch_factor=2, num_workers=4)

    valid_loader = DataLoader(valid_dataset, sampler=valid_sampler, batch_size=256,
                              pin_memory=False, prefetch_factor=2, num_workers=4)

    if torch.cuda.is_available():
        device = torch.device("cuda", args.local_rank)
    else:
        device = torch.device("cpu")

    model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18')
    model.fc = nn.Sequential(nn.Linear(in_features=512, out_features=128), nn.LeakyReLU(),
                             nn.Dropout(0.5), nn.Linear(128, 10))

    model = model.to(device)
    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                      output_device=args.local_rank)

    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.05)
    lr_scheduler_values = lr_scheduler.StepLR(optimizer, step_size = 30, gamma = 0.1)
    criterion = nn.CrossEntropyLoss().to(device)

    num_epochs = 100
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        valid_sampler.set_epoch(epoch)

        train_loss_value, train_acc_value = train_epoch(train_loader, optimizer, criterion, lr_scheduler_values, model, world_size)

        valid_loss_value, valid_acc_value = valid_epoch(valid_loader, criterion, model, world_size)    

        print("Train_local_rank: {} Train_Epoch: {}/{} Training_Loss: {} Training_acc: {:.2f}\
                   ".format(args.local_rank, epoch, num_epochs-1, train_loss_value, train_acc_value))

        print("Valid_local_rank: {} Valid_Epoch: {}/{} Valid_Loss: {} Valid_acc: {:.2f}\
                   ".format(args.local_rank, epoch, num_epochs-1, valid_loss_value, valid_acc_value))

        print('--------------------------------')

    print("finished.")

执行以下命令进行训练

$ CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py

总结

最后，总结PyTorch分布式训练的基本流程

• 设置local_rank参数并使用init_process_group进行初始化
• 使用DsitributedSampler划分数据集
• 将模型分配至CUDA并设置SyncBN及DistributedDataParallel
• 将train、valid Sampler设置set_epoch
• 在sever上执行指令开始训练

参考目录

https://medium.com/ching-i/pytorch-%E5%88%86%E6%95%A3%E5%BC%8F%E8%A8%93%E7%B7%B4-distributeddataparallel-%E5%AF%A6%E4%BD%9C%E7%AF%87-35c762cb7e08