Pytorch中的Distributed Data Parallel與混合精度訓(xùn)練（Apex）

寫在前面：本文內(nèi)容需要在每一個(gè)進(jìn)程設(shè)置相同的隨機(jī)種子，以便所有模型權(quán)重都初始化為相同的值。

1.動(dòng)機(jī)

加速神經(jīng)網(wǎng)絡(luò)訓(xùn)練最簡(jiǎn)單的辦法就是上GPU，如果一塊GPU還是不夠，就多上幾塊。

事實(shí)上，比如BERT和GPT-2這樣的大型語(yǔ)言模型甚至是在上百塊GPU上訓(xùn)練的。

為了實(shí)現(xiàn)多GPU訓(xùn)練，我們必須想一個(gè)辦法在多個(gè)GPU上分發(fā)數(shù)據(jù)和模型，并且協(xié)調(diào)訓(xùn)練過(guò)程。

2.Why Distributed Data Parallel？

Pytorch兼顧了主要神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)的易用性和可控性。而其提供了兩種辦法在多GPU上分割數(shù)據(jù)和模型：即nn.DataParallel 以及 nn.DistributedDataParallel。

nn.DataParallel 使用起來(lái)更加簡(jiǎn)單（通常只要封裝模型然后跑訓(xùn)練代碼就ok了）。但是在每個(gè)訓(xùn)練批次（batch）中，因?yàn)槟Ｐ偷臋?quán)重都是在 一個(gè)進(jìn)程上先算出來(lái) 然后再把他們分發(fā)到每個(gè)GPU上，所以網(wǎng)絡(luò)通信就成為了一個(gè)瓶頸，而GPU使用率也通常很低。

除此之外，nn.DataParallel 需要所有的GPU都在一個(gè)節(jié)點(diǎn)（一臺(tái)機(jī)器）上，且并不支持 Apex 的 混合精度訓(xùn)練.

3.現(xiàn)有文檔的局限性

總的來(lái)說(shuō)，Pytorch的文檔是全面且清晰的，特別是在1.0版本的那些。完全通過(guò)文檔和教程就可以自學(xué)Pytorch，這并不是顯示一個(gè)人有多大佬，而顯然更多地反映了Pytorch的易用性和優(yōu)秀的文檔。

但是好巧不巧的，就是在（Distributed）DataParallel這個(gè)系列的文檔講的就不甚清楚，或者干脆沒(méi)有/不完善/有很多無(wú)關(guān)內(nèi)容。以下是一些例子（抱怨）。

Pytorch提供了一個(gè)使用AWS（亞馬遜網(wǎng)絡(luò)服務(wù)）進(jìn)行分布式訓(xùn)練的教程，這個(gè)教程在教你如何使用AWS方面很出色，但甚至沒(méi)提到 nn.DistributedDataParallel 是干什么用的，這導(dǎo)致相關(guān)的代碼塊很難follow。

而另外一篇Pytorch提供的教程又太細(xì)了，它對(duì)于一個(gè)不是很懂Python中MultiProcessing的人（比如我）來(lái)說(shuō)很難讀懂。因?yàn)樗舜罅康钠v nn.DistributedDataParallel 中的復(fù)制功能（數(shù)據(jù)是怎么復(fù)制的）。然而，他并沒(méi)有在高層邏輯上總結(jié)一下都在扯啥，甚至沒(méi)說(shuō)這個(gè)DistributedDataParallel是咋用的？

這里還有一個(gè)Pytorch關(guān)于入門分布式數(shù)據(jù)并行的(Distributed data parallel)教程。這個(gè)教程展示了如何進(jìn)行一些設(shè)置，但并沒(méi)解釋這些設(shè)置是干啥用的，之后也展示了一些講模型分到各個(gè)GPU上并執(zhí)行一個(gè)優(yōu)化步驟（optimization step）。然而，這篇教程里的代碼是跑不同的(函數(shù)名字都對(duì)不上)，也沒(méi)告訴你怎么跑這個(gè)代碼。和之前的教程一樣，他也沒(méi)給一個(gè)邏輯上分布式訓(xùn)練的工作概括。

而官方給的最好的例子，無(wú)疑是ImageNet的訓(xùn)練，然而因?yàn)檫@個(gè)例子要 素 過(guò) 多，導(dǎo)致也看不出來(lái)哪個(gè)部分是用于分布式多GPU訓(xùn)練的。

Apex提供了他們自己的ImageNet的訓(xùn)練例。例子的文檔告訴大家他們的 nn.DistributedDataParallel 是自己重寫的，但是如果連最初的版本都不會(huì)用，更別說(shuō)重寫的了。

而這個(gè)教程很好地描述了在底層，nn.DistributedDataParallel 和 nn.DataParallel 到底有什么不同。然而他并沒(méi)有如何使用 nn.DataParallel 的例程。

4.大綱

本教程實(shí)際上是針對(duì)那些已經(jīng)熟悉在Pytorch中訓(xùn)練神經(jīng)網(wǎng)絡(luò)模型 的人的，本文不會(huì)詳細(xì)介紹這些代碼的任何一部分。

本文將首先概述一下總體情況，然后展示一個(gè)最小的使用GPU訓(xùn)練MNIST數(shù)據(jù)集的例程。之后對(duì)這個(gè)例程進(jìn)行修改，以便在多個(gè)gpu(可能跨多個(gè)節(jié)點(diǎn))上進(jìn)行訓(xùn)練，并逐行解釋這些更改。重要的是，本文還將解釋如何運(yùn)行代碼。

另外，本文還演示了如何使用Apex進(jìn)行簡(jiǎn)單的混合精度分布式訓(xùn)練。

5.大圖景（The big picture）

使用nn.DistributedDataParallel進(jìn)行 Multiprocessing 可以在多個(gè) gpu 之間復(fù)制該模型，每個(gè) gpu 由一個(gè)進(jìn)程控制。（如果你想，也可以一個(gè)進(jìn)程控制多個(gè) GPU，但這會(huì)比控制一個(gè)慢得多。也有可能有多個(gè)工作進(jìn)程為每個(gè) GPU 獲取數(shù)據(jù)，但為了簡(jiǎn)單起見(jiàn)，本文將省略這一點(diǎn)。）這些 GPU 可以位于同一個(gè)節(jié)點(diǎn)上，也可以分布在多個(gè)節(jié)點(diǎn)上。每個(gè)進(jìn)程都執(zhí)行相同的任務(wù)，并且每個(gè)進(jìn)程與所有其他進(jìn)程通信。

只有梯度會(huì)在進(jìn)程 GPU 之間傳播，這樣網(wǎng)絡(luò)通信就不至于成為一個(gè)瓶頸了。

訓(xùn)練過(guò)程中，每個(gè)進(jìn)程從磁盤加載自己的小批（minibatch）數(shù)據(jù)，并將它們傳遞給自己的 GPU。每個(gè) GPU 都做它自己的前向計(jì)算，然后梯度在 GPU 之間全部約簡(jiǎn)。每個(gè)層的梯度不僅僅依賴于前一層，因此梯度全約簡(jiǎn)與并行計(jì)算反向傳播，進(jìn)一步緩解網(wǎng)絡(luò)瓶頸。在反向傳播結(jié)束時(shí)，每個(gè)節(jié)點(diǎn)都有平均的梯度，確保模型權(quán)值保持同步（synchronized）。

上述的步驟要求需要多個(gè)進(jìn)程，甚至可能是不同結(jié)點(diǎn)上的多個(gè)進(jìn)程同步和通信。而Pytorch通過(guò)它的 distributed.init_process_group 函數(shù)實(shí)現(xiàn)。這個(gè)函數(shù)需要知道如何找到進(jìn)程0（process 0），一邊所有的進(jìn)程都可以同步，也知道了一共要同步多少進(jìn)程。每個(gè)獨(dú)立的進(jìn)程也要知道總共的進(jìn)程數(shù)，以及自己在所有進(jìn)程中的階序（rank）,當(dāng)然也要知道自己要用那張 GPU ?？傔M(jìn)程數(shù)稱之為 world size。最后，每個(gè)進(jìn)程都需要知道要處理的數(shù)據(jù)的哪一部分，這樣批處理就不會(huì)重疊。而 Pytorch 通過(guò) nn.utils.data.DistributedSampler 來(lái)實(shí)現(xiàn)這種效果。

6. 最小例程與解釋

為了展示如何做到這些，這里有一個(gè)在MNIST上訓(xùn)練的例子，并且之后把它修改為可以在多節(jié)點(diǎn)多GPU上運(yùn)行，最終修改的版本還可以支持混合精度運(yùn)算。

首先，我們import所有我們需要的庫(kù)。

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

之后，我們訓(xùn)練了一個(gè)MNIST分類的簡(jiǎn)單卷積網(wǎng)絡(luò)。

?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

這個(gè) main() 函數(shù)會(huì)接受一些參數(shù)并運(yùn)行訓(xùn)練函數(shù)。

?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)

而這部分則是訓(xùn)練函數(shù)。

?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))

最后，我們要確保 main() 函數(shù)會(huì)被調(diào)用

if?__name__?==?'__main__':
????main()

上述代碼中肯定有一些我們還不需要的額外的東西(例如gpu和節(jié)點(diǎn)的數(shù)量)，但是將整個(gè)框架放置到位是很有幫助的。之后在命令行輸入。

python?src/mnist.py?-n?1?-g?1?-nr?0

就可以在一個(gè)結(jié)點(diǎn)上的單個(gè)GPU上訓(xùn)練啦~

7. 加上MultiProcessing

我們需要一個(gè)腳本，用來(lái)啟動(dòng)一個(gè)進(jìn)程的每一個(gè)GPU。每個(gè)進(jìn)程需要知道使用哪個(gè)GPU，以及它在所有正在運(yùn)行的進(jìn)程中的階序（rank）。而且，我們需要在每個(gè)節(jié)點(diǎn)上運(yùn)行腳本。

現(xiàn)在讓我們康康每個(gè)函數(shù)的變化，這些改變將被單獨(dú)框出方便查找。

?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,))?????????#
????#########################################################

上一節(jié)中一些參數(shù)在這個(gè)地方才需要。

• args.nodes 是我們使用的結(jié)點(diǎn)數(shù)
• args.gpus 是每個(gè)結(jié)點(diǎn)的GPU數(shù).
• args.nr 是當(dāng)前結(jié)點(diǎn)的階序rank，這個(gè)值的取值范圍是 0 到 args.nodes - 1.

OK，現(xiàn)在我們一行行看都改了什么：

• Line 14：基于結(jié)點(diǎn)數(shù)以及每個(gè)結(jié)點(diǎn)的GPU數(shù)，我們可以計(jì)算 world_size 或者需要運(yùn)行的總進(jìn)程數(shù)，這和總GPU數(shù)相等。
• Line 15：告訴Multiprocessing模塊去哪個(gè)IP地址找process 0以確保初始同步所有進(jìn)程。
• Line 16：同樣的，這個(gè)是process 0所在的端口
• Line 17：現(xiàn)在，我們需要生成 args.gpus 個(gè)進(jìn)程, 每個(gè)進(jìn)程都運(yùn)行 train(i, args), 其中 i 從 0 到 args.gpus - 1。注意, main() 在每個(gè)結(jié)點(diǎn)上都運(yùn)行, 因此總共就有 args.nodes * args.gpus = args.world_size 個(gè)進(jìn)程.

除了14，15行的設(shè)置，也可以在終端中運(yùn)行。

export MASTER_ADDR=10.57.23.164 和 export MASTER_PORT=8888

接下來(lái)，需要修改的就是訓(xùn)練函數(shù)了，改動(dòng)的地方依然被框出來(lái)啦。

?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)????#?
????#############################
????...

為了簡(jiǎn)單起見(jiàn)，上面的代碼去掉了簡(jiǎn)單循環(huán)并用 ... 代替，不過(guò)你可以在這里看到完整腳本 。

• Line3：這里是該進(jìn)程在所有進(jìn)程中的全局rank（一個(gè)進(jìn)程對(duì)應(yīng)一個(gè)GPU）。這個(gè)rank在Line6會(huì)用到。

• Line4~6：初始化進(jìn)程并加入其他進(jìn)程。這就叫做“blocking”，也就是說(shuō)只有當(dāng)所有進(jìn)程都加入了,單個(gè)進(jìn)程才會(huì)運(yùn)行。這里使用了 nccl 后端，因?yàn)镻ytorch文檔說(shuō)它是跑得最快的。init_method 讓進(jìn)程組知道去哪里找到它需要的設(shè)置。在這里，它就在尋找名為 MASTER_ADDR 以及 MASTER_PORT 的環(huán)境變量，這些環(huán)境變量在 main 函數(shù)中設(shè)置過(guò)。當(dāng)然，本來(lái)可以把world_size 設(shè)置成一個(gè)全局變量，不過(guò)本腳本選擇把它作為一個(gè)關(guān)鍵字參量（和當(dāng)前進(jìn)程的全局階序global rank一樣）

• Line23：將模型封裝為一個(gè) DistributedDataParallel 模型。這將把模型復(fù)制到GPU上進(jìn)行處理。

• Line35~39：nn.utils.data.DistributedSampler 確保每個(gè)進(jìn)程拿到的都是不同的訓(xùn)練數(shù)據(jù)切片。

• Line46/Line51：因?yàn)橛昧?nn.utils.data.DistributedSampler 所以不能用正常的辦法做shuffle。

要在4個(gè)節(jié)點(diǎn)上運(yùn)行它(每個(gè)節(jié)點(diǎn)上有8個(gè)gpu)，我們需要4個(gè)終端(每個(gè)節(jié)點(diǎn)上有一個(gè))。在節(jié)點(diǎn)0上(由 main 中的第13行設(shè)置)：

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

而在其他的節(jié)點(diǎn)上：

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

其中 i∈1,2,3. 換句話說(shuō)，我們要把這個(gè)腳本在每個(gè)結(jié)點(diǎn)上運(yùn)行腳本，讓腳本運(yùn)行 args.gpus 個(gè)進(jìn)程以在訓(xùn)練開(kāi)始之前同步每個(gè)進(jìn)程。

注意，腳本中的batchsize設(shè)置的是每個(gè)GPU的batchsize，因此實(shí)際的batchsize要乘上總共的GPU數(shù)目（worldsize）。

8. 使用Apex進(jìn)行混合混合精度訓(xùn)練

混合精度訓(xùn)練，即組合浮點(diǎn)數(shù) (FP32)和半精度浮點(diǎn)數(shù) (FP16)進(jìn)行訓(xùn)練，允許我們使用更大的batchsize，并利用NVIDIA張量核進(jìn)行更快的計(jì)算。AWS p3實(shí)例使用了8塊帶張量核的NVIDIA Tesla V100 GPU。

我們只需要修改 train 函數(shù)即可，為了簡(jiǎn)便表示，下面已經(jīng)從示例中剔除了數(shù)據(jù)加載代碼和反向傳播之后的代碼，并將它們替換為 ... ，不過(guò)你可以在這看到完整腳本。

????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 將模型和優(yōu)化器為了進(jìn)行后續(xù)混合精度訓(xùn)練而進(jìn)行封裝。注意，在調(diào)用 amp.initialize 之前，模型模型必須已經(jīng)部署在GPU上。opt_level 從 O0 （全部使用浮點(diǎn)數(shù)）一直到 O3 （全部使用半精度浮點(diǎn)數(shù)）。而 O1 和 O2 屬于不同的混合精度程度，具體可以參閱APEX的官方文檔。注意之前數(shù)字前面的是大寫字母O。
• Line20：apex.parallel.DistributedDataParallel 是一個(gè) nn.DistributedDataParallel 的替換版本。我們不需要指定GPU，因?yàn)锳pex在一個(gè)進(jìn)程中只允許用一個(gè)GPU。且它也假設(shè)程序在把模型搬到GPU之前已經(jīng)調(diào)用了 torch.cuda.set_device(local_rank)(line 10)
• Line37-38：混合精度訓(xùn)練需要縮放損失函數(shù)以阻止梯度出現(xiàn)下溢。不過(guò)Apex會(huì)自動(dòng)進(jìn)行這些工作。

這個(gè)腳本和之前的分布式訓(xùn)練腳本的運(yùn)行方式相同。

如果覺(jué)得有用，就請(qǐng)分享到朋友圈吧！