當(dāng)一塊GPU不夠用時(shí)，我們就需要使用多卡進(jìn)行并行訓(xùn)練。其中多卡并行可分為數(shù)據(jù)并行和模型并行。具體區(qū)別如下圖所示：

由于模型并行比較少用，這里只對(duì)數(shù)據(jù)并行進(jìn)行記錄。對(duì)于pytorch，有兩種方式可以進(jìn)行數(shù)據(jù)并行：數(shù)據(jù)并行（DataParallel, DP）和分布式數(shù)據(jù)并行（DistributedDataParallel, DDP）。

在多卡訓(xùn)練的實(shí)現(xiàn)上，DP與DDP的思路是相似的：

1、每張卡都復(fù)制一個(gè)有相同參數(shù)的模型副本。

2、每次迭代，每張卡分別輸入不同批次數(shù)據(jù)，分別計(jì)算梯度。

3、DP與DDP的主要不同在于接下來的多卡通信：

DP的多卡交互實(shí)現(xiàn)在一個(gè)進(jìn)程之中，它將一張卡視為主卡，維護(hù)單獨(dú)模型優(yōu)化器。所有卡計(jì)算完梯度后，主卡匯聚其它卡的梯度進(jìn)行平均并用優(yōu)化器更新模型參數(shù)，再將模型參數(shù)更新至其它卡上。

DDP則分別為每張卡創(chuàng)建一個(gè)進(jìn)程，每個(gè)進(jìn)程相應(yīng)的卡上都獨(dú)立維護(hù)模型和優(yōu)化器。在每次每張卡計(jì)算完梯度之后，進(jìn)程之間以NCLL（NVIDIA GPU通信）為通信后端，使各卡獲取其它卡的梯度。各卡對(duì)獲取的梯度進(jìn)行平均，然后執(zhí)行后續(xù)的參數(shù)更新。由于每張卡上的模型與優(yōu)化器參數(shù)在初始化時(shí)就保持一致，而每次迭代的平均梯度也保持一致，那么即使沒有進(jìn)行參數(shù)復(fù)制，所有卡的模型參數(shù)也是保持一致的。

Pytorch官方推薦我們使用DDP。DP經(jīng)過我的實(shí)驗(yàn)，兩塊GPU甚至比一塊還慢。當(dāng)然不同模型可能有不同的結(jié)果。下面分別對(duì)DP和DDP進(jìn)行記錄。

1.DP

Pytorch的DP實(shí)現(xiàn)多GPU訓(xùn)練十分簡單，只需在單GPU的基礎(chǔ)上加一行代碼即可。以下是一個(gè)DEMO的代碼。

import torch
from torch import nn
from torch.optim import Adam
from torch.nn.parallel import DataParallel
class DEMO_model(nn.Module):
  def __init__(self, in_size, out_size):
    super().__init__()
    self.fc = nn.Linear(in_size, out_size)
  def forward(self, inp):
    outp = self.fc(inp)
    print(inp.shape, outp.device)
    return outp
model = DEMO_model(10, 5).to('cuda')
model = DataParallel(model, device_ids=[0, 1]) # 額外加這一行
adam = Adam(model.parameters())
# 進(jìn)行訓(xùn)練
for i in range(1):
  x = torch.rand([128, 10]) # 獲取訓(xùn)練數(shù)據(jù)，無需指定設(shè)備
  y = model(x) # 自動(dòng)均勻劃分?jǐn)?shù)據(jù)批量并分配至各GPU，輸出結(jié)果y會(huì)聚集到GPU0中
  loss = torch.norm(y)
  loss.backward()
  adam.step()

其中model = DataParallel(model, device_ids=[0, 1])這行將模型復(fù)制到0,1號(hào)GPU上。輸入數(shù)據(jù)x無需指定設(shè)備，它將會(huì)被均勻分配至各塊GPU模型，進(jìn)行前向傳播。之后各塊GPU的輸出再合并到GPU0中，得到輸出y。輸出y在GPU0中計(jì)算損失，并進(jìn)行反向傳播計(jì)算梯度、優(yōu)化器更新參數(shù)。

2.DDP

為了對(duì)分布式編程有基本概念，首先使用pytorch內(nèi)部的方法實(shí)現(xiàn)一個(gè)多進(jìn)程程序，再使用DDP模塊實(shí)現(xiàn)模型的分布式訓(xùn)練。

2.1Pytorch分布式基礎(chǔ)

首先使用pytorch內(nèi)部的方法編寫一個(gè)多進(jìn)程程序作為編寫分布式訓(xùn)練的基礎(chǔ)。

import os, torch
import torch.multiprocessing as mp
import torch.distributed as dist
def run(rank, size):
  tensor = torch.tensor([1,2,3,4], device='cuda:'+str(rank)) # ——1—— 
  group = dist.new_group(range(size)) # ——2——
  dist.all_reduce(tensor=tensor, group=group, op=dist.ReduceOp.SUM) # ——3——
  print(str(rank)+ ': ' + str(tensor) + '\n')
def ini_process(rank, size, fn, backend = 'nccl'):  
  os.environ['MASTER_ADDR'] = '127.0.0.1' # ——4——
  os.environ['MASTER_PORT'] = '1234'
  dist.init_process_group(backend, rank=rank, world_size=size) # ——5——
  fn(rank, size) # ——6——
if __name__ == '__main__': # ——7——
  mp.set_start_method('spawn') # ——8—— 
  size = 2 # ——9—— 
  ps = []  
  for rank in range(size):
    p = mp.Process(target=ini_process, args=(rank, size, run)) # ——10—— 
    p.start()  
    ps.append(p)
  for p in ps: # ——11—— 
    p.join()

以上代碼主進(jìn)程創(chuàng)建了兩個(gè)子進(jìn)程，子進(jìn)程之間使用NCCL后端進(jìn)行通信。每個(gè)子進(jìn)程各占用一個(gè)GPU資源，實(shí)現(xiàn)了所有GPU張量求和的功能。細(xì)節(jié)注釋如下：

1、為每個(gè)子進(jìn)程定義相同名稱的張量，并分別分配至不同的GPU，從而能進(jìn)行后續(xù)的GPU間通信。

2、定義一個(gè)通信組，用于后面的all_reduce通信操作。

3、all_reduce操作以及其它通信方式請(qǐng)看下圖：

4、定義編號(hào)(rank)為0的ip和端口地址，讓每個(gè)子進(jìn)程都知道。ip和端口地址可以隨意定義，不沖突即可。如果不設(shè)置，子進(jìn)程在涉及進(jìn)程通信時(shí)會(huì)出錯(cuò)。

5、初始化子進(jìn)程組，定義進(jìn)程間的通信后端（還有GLOO、MPI，只有NCCL支持GPU間通信）、子進(jìn)程rank、子進(jìn)程數(shù)量。只有當(dāng)該函數(shù)在size個(gè)進(jìn)程中被調(diào)用時(shí)，各進(jìn)程才會(huì)繼續(xù)從這里執(zhí)行下去。這個(gè)函數(shù)統(tǒng)一了各子進(jìn)程后續(xù)代碼的開始時(shí)間。

6、執(zhí)行子進(jìn)程代碼。

7、由于創(chuàng)建子進(jìn)程會(huì)執(zhí)行本程序，因此主進(jìn)程的執(zhí)行需要放在__main__里，防止子進(jìn)程執(zhí)行。

8、開始創(chuàng)建子進(jìn)程的方式：spawn、fork。windows默認(rèn)spawn，linux默認(rèn)fork。具體區(qū)別請(qǐng)百度。

9、由于是以NCCL為通信后端的分布式訓(xùn)練，如果不同進(jìn)程中相同名稱的張量在同一GPU上，當(dāng)這個(gè)張量進(jìn)行進(jìn)程間通信時(shí)就會(huì)出錯(cuò)。為了防止出錯(cuò)，限制每張卡獨(dú)占一個(gè)進(jìn)程，每個(gè)進(jìn)程獨(dú)占一張卡。這里有兩張卡，所以最多只能創(chuàng)建兩個(gè)進(jìn)程。

10、創(chuàng)建子進(jìn)程，傳入子進(jìn)程的初始化方法，及子進(jìn)程調(diào)用該方法的參數(shù)。

11、等待子進(jìn)程全部運(yùn)行完畢后再退出主進(jìn)程。 

輸出結(jié)果如下：

正是各進(jìn)程保存在不同GPU上的張量的廣播求和(all_reduce)的結(jié)果。

參考：https://pytorch.org/tutorials/intermediate/dist_tuto.html

2.2Pytorch分布式訓(xùn)練DEMO

我們實(shí)際上可以根據(jù)上面的分布式基礎(chǔ)寫一個(gè)分布式訓(xùn)練，但由于不知道pytorch如何實(shí)現(xiàn)GPU間模型梯度的求和，即官方教程中所謂的ring_reduce（沒找到相關(guān)API），時(shí)間原因，就不再去搜索相關(guān)方法了。這里僅記錄pytorh內(nèi)部的分布式模型訓(xùn)練，即利用DDP模塊實(shí)現(xiàn)。Pytorch版本1.12.1。

import torch,os
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch import nn
def example(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)  # ——1——
    model = nn.Linear(2, 1, False).to(rank) 
    if rank == 0: # ——2——
        model.load_state_dict(torch.load('model_weight')) 
    # model_stat = torch.load('model_weight', {'cuda:0':'cuda:%d'%rank})  #這樣讀取保險(xiǎn)一點(diǎn)
    # model.load_state_dict(model_stat) 
    opt = optim.Adam(model.parameters(), lr=0.0001) # ——3——
    opt_stat = torch.load('opt_weight', {'cuda:0':'cuda:%d'%rank}) # ——4——
    opt.load_state_dict(opt_stat) # ——5——
    ddp_model = DDP(model, device_ids=[rank])# ——6
    inp = torch.tensor([[1.,2]]).to(rank) # ——7——
    labels = torch.tensor([[5.]]).to(rank)
    outp = ddp_model(inp)
    loss = torch.mean((outp - labels)**2)
    opt.zero_grad()
    loss.backward() # ——8——
    opt.step() # ——9
    if rank == 0:# ——10——
        torch.save(model.state_dict(), 'model_weight')
        torch.save(opt.state_dict(), 'opt_weight')
if __name__=="__main__":
    os.environ["MASTER_ADDR"] = "localhost"# ——11——
    os.environ["MASTER_PORT"] = "29500"
    world_size = 2
    mp.spawn(example, args=(world_size,), nprocs=world_size, join=True) # ——12——

以上代碼包含模型在多GPU上讀取權(quán)重、進(jìn)行分布式訓(xùn)練、保存權(quán)重等過程。細(xì)節(jié)注釋如下：

1、初始化進(jìn)程組，由于使用GPU通信，后端應(yīng)該寫為NCCL。不過經(jīng)過實(shí)驗(yàn)，即使錯(cuò)寫為gloo，DDP內(nèi)部也會(huì)自動(dòng)使用NCCL作為通信模塊。

2、由于后面使用DDP包裹模型進(jìn)行訓(xùn)練，其內(nèi)部會(huì)自動(dòng)將所有rank的模型權(quán)重同步為rank 0的權(quán)重，因此我們只需在rank 0上讀取模型權(quán)重即可。這是基于Pytorch版本1.12.1，低級(jí)版本似乎沒有這個(gè)特性，需要在不同rank分別導(dǎo)入權(quán)重，則load需要傳入map_location，如下面注釋的兩行代碼所示。

3、這里創(chuàng)建model的優(yōu)化器，而不是創(chuàng)建用ddp包裹后的ddp_model的優(yōu)化器，是為了兼容單GPU訓(xùn)練，讀取優(yōu)化器權(quán)重更方便。

4、將優(yōu)化器權(quán)重讀取至該進(jìn)程占用的GPU。如果沒有map_location參數(shù)，load會(huì)將權(quán)重讀取到原本保存它時(shí)的設(shè)備。

5、優(yōu)化器獲取權(quán)重。經(jīng)過實(shí)驗(yàn)，即使權(quán)重不在優(yōu)化器所在的GPU，權(quán)重也會(huì)遷移過去而不會(huì)報(bào)錯(cuò)。當(dāng)然load直接讀取到相應(yīng)GPU會(huì)減少數(shù)據(jù)傳輸。

6、DDP包裹模型，為模型復(fù)制一個(gè)副本到相應(yīng)GPU中。所有rank的模型副本會(huì)與rank 0保持一致。注意，DDP并不復(fù)制模型優(yōu)化器的副本，因此各進(jìn)程的優(yōu)化器需要我們?cè)诔跏蓟瘯r(shí)保持一致。權(quán)重要么不讀取，要么都讀取。

7、這里開始模型的訓(xùn)練。數(shù)據(jù)需轉(zhuǎn)移到相應(yīng)的GPU設(shè)備。

8、在backward中，所有進(jìn)程的模型計(jì)算梯度后，會(huì)進(jìn)行平均（不是相加）。也就是說，DDP在backward函數(shù)添加了hook，所有進(jìn)程的模型梯度的ring_reduce將在這里執(zhí)行。這個(gè)可以通過給各進(jìn)程模型分別輸入不同的數(shù)據(jù)進(jìn)行驗(yàn)證，backward后這些模型有相同的梯度，且驗(yàn)算的確是所有進(jìn)程梯度的平均。此外，還可以驗(yàn)證backward函數(shù)會(huì)阻斷(block)各進(jìn)程使用梯度，只有當(dāng)所有進(jìn)程都完成backward之后，各進(jìn)程才能讀取和使用梯度。這保證了所有進(jìn)程在梯度上的一致性。

9、各進(jìn)程優(yōu)化器使用梯度更新其模型副本權(quán)重。由于初始化時(shí)各進(jìn)程模型、優(yōu)化器權(quán)重一致，每次反向傳播梯度也保持一致，則所有進(jìn)程的模型在整個(gè)訓(xùn)練過程中都能保持一致。

10、由于所有進(jìn)程權(quán)重保持一致，我們只需通過一個(gè)進(jìn)程保存即可。

11、定義rank 0的IP和端口，使用mp.spawn，只需在主進(jìn)程中定義即可，無需分別在子進(jìn)程中定義。

12、創(chuàng)建子進(jìn)程，傳入：子進(jìn)程調(diào)用的函數(shù)（該函數(shù)第一個(gè)參數(shù)必須是rank）、子進(jìn)程函數(shù)的參數(shù)（除了rank參數(shù)外）、子進(jìn)程數(shù)、是否等待所有子進(jìn)程創(chuàng)建完畢再開始執(zhí)行。

以上就是詳解如何使用Pytorch進(jìn)行多卡訓(xùn)練 的詳細(xì)內(nèi)容，更多關(guān)于Pytorch多卡訓(xùn)練的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

最新評(píng)論