文章目录
前言
- Data Parallel 并行运算时，各种原因导致单进程多卡的时候只有一张卡在进行运算
- 在使用 Distributed Data Parallel 时，注意 需要在每一个进程设置相同的随机种子，以便所有模型权重都初始化为相同的值。
一、动机
- 加速神经网络训练最简单的办法就是上GPU，如果一块 GPU 还是不够，就多上几块。
- 事实上，比如 BERT 和 GPT-2 这样的大型语言模型甚至是在上百块 GPU 上训练的。
- 为了实现多 GPU 训练，我们必须想一个办法在多个 GPU 上分发数据和模型，并且协调训练过程。
二、Why Distributed Data Parallel？
- Pytorch 兼顾了主要神经网络结构的易用性和可控性。而其提供了两种办法在多 GPU 上分割数据和模型：即
nn.DataParallel以及nn.DistributedDataParallel。- nn.DataParallel 使用起来更加简单（通常只要封装模型然后跑训练代码就ok了）。但是在每个训练批次（batch）中，因为模型的权重都是在 一个进程上先算出来 然后再把他们分发到每个 GPU 上，所以网络通信就成为了一个瓶颈，而 GPU 使用率也通常很低。
- 除此之外，nn.DataParallel 需要所有的 GPU 都在一个节点（一台机器）上，且并不支持 Apex 的 混合精度训练.
三、大图景（The big picture）
- 使用 nn.DistributedDataParallel 进行 Multiprocessing 可以在多个 gpu 之间复制该模型，每个 gpu 由一个进程控制。（如果你想，也可以一个进程控制多个GPU，但这会比控制一个慢得多。也有可能有多个工作进程为每个 GPU 获取数据，但为了简单起见，本文将省略这一点。）这些 GPU 可以位于同一个节点上，也可以分布在多个节点上。每个进程都执行相同的任务，并且每个进程与所有其他进程通信。
- 只有梯度会在进程/GPU之间传播，这样网络通信就不至于成为一个瓶颈了。
- 训练过程中，每个进程从磁盘加载自己的小批（minibatch）数据，并将它们传递给自己的 GPU。每个 GPU 都做它自己的前向计算，然后梯度在 GPU 之间全部约简。每个层的梯度不仅仅依赖于前一层，因此梯度全约简与并行计算反向传播，进一步缓解网络瓶颈。在反向传播结束时，每个节点都有平均的梯度，确保模型权值保持同步（synchronized）。
- 上述的步骤要求需要多个进程，甚至可能是不同结点上的多个进程同步和通信。而 Pytorch 通过它的
distributed.init_process_group函数实现。这个函数需要知道如何找到进程0（process 0），以便所有的进程都可以同步，也知道了一共要同步多少进程。每个独立的进程也要知道总共的进程数，以及自己在所有进程中的阶序（rank）,当然也要知道自己要用哪张 GPU。总进程数称之为 world size。最后，每个进程都需要知道要处理的数据的哪一部分，这样批处理就不会重叠。而 Pytorch 通过nn.utils.data.DistributedSampler来实现这种效果。
四、最小例程与解释
为了展示如何做到这些，这里有一个在 MNIST 上训练的例子，并且之后把它修改为可以在多节点多 GPU 上运行，最终修改的版本还可以支持混合精度运算。
# step1: 首先，我们 import 所有我们需要的库
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 apex.parallel import DistributedDataParallel as DDP
from apex import amp
# step2: 之后，我们训练了一个 MNIST 分类的简单卷积网络
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 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))
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)
if __name__ == '__main__':
main()
执行下面命令，就可以在一个结点上的单个GPU上训练啦~
python src/mnist.py -n 1 -g 1 -nr 0
五、加上 MultiProcessing
- 我们需要一个脚本，用来启动一个进程的每一个GPU。每个进程需要知道使用哪个GPU，以及它在所有正在运行的进程中的阶序（rank）。而且，我们需要在每个节点上运行脚本。
- 现在让我们看看每个函数的变化，这些改变将被单独框出方便查找。
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数.
- args.nr 是当前结点的阶序rank，这个值的取值范围是 0 到 args.nodes - 1.
OK，现在我们一行行看都改了什么
args.world_size = args.gpus * args.nodes
基于结点数以及每个结点的GPU数，我们可以计算 world_size 或者需要运行的总进程数，这和总GPU数相等。os.environ['MASTER_ADDR'] = '10.57.23.164'
告诉 Multiprocessing 模块去哪个 IP 地址找 process 0 以确保初始同步所有进程。os.environ['MASTER_PORT'] = '8888'
同样的，这个是 process 0 所在的端口mp.spawn(train, nprocs=args.gpus, args=(args,))
现在，我们需要生成 args.gpus 个进程, 每个进程都运行 train(i, args), 其中 i 从 0 到 args.gpus - 1。
注意, main() 在每个结点上都运行, 因此总共就有 args.nodes * args.gpus = args.world_size 个进程.
除了第一处和第二处的设置，也可以在终端中运行
export MASTER_ADDR=10.57.23.164
export MASTER_PORT=8888
接下来，需要修改的就是训练函数了，改动的地方依然被框出来啦。
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) #
#############################
...
- Line3： 这里是该进程在所有进程中的全局rank（一个进程对应一个GPU）。这个 rank 在 Line6 会用到
- Line4~6： 初始化进程并加入其他进程。这就叫做
blocking，也就是说只有当所有进程都加入了,单个进程才会运行。这里使用了 nccl 后端，因为 Pytorch 文档说它是跑得最快的。init_method让进程组知道去哪里找到它需要的设置。在这里，它就在寻找名为MASTER_ADDR以及MASTER_PORT的环境变量，这些环境变量在 main 函数中设置过。当然，本来可以把 world_size 设置成一个全局变量，不过本脚本选择把它作为一个关键字参量（和当前进程的全局阶序 global rank 一样）- Line23： 将模型封装为一个 DistributedDataParallel 模型。这将把模型复制到 GPU 上进行处理。
- Line35~39：
nn.utils.data.DistributedSampler确保每个进程拿到的都是不同的训练数据切片。- Line46/Line51： 因为用了 nn.utils.data.DistributedSampler 所以不能用正常的办法做 shuffle。
要在 4 个节点上运行它(每个节点上有 8 个 gpu)，我们需要 4 个终端(每个节点上有一个)。在节点 0 上(由 main 中的第 13 行设置)：
python src/mnist-distributed.py -n 4 -g 8 -nr 0
而在其他的节点上：
python src/mnist-distributed.py -n 4 -g 8 -nr i
- 其中 i∈1,2,3. 换句话说，我们要把这个脚本在每个结点上运行脚本，让脚本运行 args.gpus 个进程以在训练开始之前同步每个进程。
- 注意，脚本中的 batchsize 设置的是每个 GPU 的 batchsize，因此实际的 batchsize 要乘上总共的 GPU 数目（worldsize）。
六、使用Apex进行混合混合精度训练
- 混合精度训练，即组合浮点数 (FP32)和半精度浮点数 (FP16)进行训练，允许我们使用更大的 batchsize，并利用 NVIDIA 张量核进行更快的计算。AWS p3 实例使用了 8 块带张量核的 NVIDIA Tesla V100 GPU。
- 我们只需要修改 train 函数即可，为了简便表示，下面已经从示例中剔除了数据加载代码和反向传播之后的代码，并将它们替换为 … ，不过你可以在这看到完整脚本。
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, optimizer = amp.initialize(model, optimizer, opt_level='O2')
model = DDP(model)
##############################################################
# Data loading code
...
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()
##############################################################
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
##############################################################
optimizer.step()
...
- Line18：
amp.initialize将模型和优化器为了进行后续混合精度训练而进行封装。
注意，在调用 amp.initialize 之前，模型模型必须已经部署在 GPU 上。opt_level 从 O0 （全部使用浮点数）一直到 O3 （全部使用半精度浮点数）。而 O1 和 O2 属于不同的混合精度程度，具体可以参阅APEX的官方文档。注意之前数字前面的是大写字母O。- Line20：
apex.parallel.DistributedDataParallel是一个nn.DistributedDataParallel的替换版本。我们不需要指定GPU，因为 Apex 在一个进程中只允许用一个GPU。且它也假设程序在把模型搬到 GPU 之前已经调用了torch.cuda.set_device(local_rank)(line 10) .- Line37-38： 混合精度训练需要缩放损失函数以阻止梯度出现下溢。不过 Apex 会自动进行这些工作。