【分布式訓(xùn)練】基于PyTorch進行多GPU分布式模型訓(xùn)練（補充）

這篇具有很好參考價值的文章主要介紹了【分布式訓(xùn)練】基于PyTorch進行多GPU分布式模型訓(xùn)練（補充）。希望對大家有所幫助。如果存在錯誤或未考慮完全的地方，請大家不吝賜教，您也可以點擊"舉報違法"按鈕提交疑問。

簡介：在PyTorch中使用DistributedDataParallel進行多GPU分布式模型訓(xùn)練。
原文鏈接：https://towardsdatascience.com/distributed-model-training-in-pytorch-using-distributeddataparallel-d3d3864dc2a7

隨著以ChatGPT為代表的大模型的不斷涌現(xiàn)，如何在合理的時間內(nèi)訓(xùn)練大模型逐漸成為一個重要的研究課題。為了解決這個問題，越來越多的從業(yè)者轉(zhuǎn)向分布式訓(xùn)練。分布式訓(xùn)練是使用多個GPU和/或多個機器訓(xùn)練深度學(xué)習(xí)模型的技術(shù)。分布式訓(xùn)練作業(yè)能夠克服單GPU內(nèi)存瓶頸，通過同時利用多個GPU來開發(fā)更大，功能更強大的模型。

本文介紹使用torch.nn.parallel.DistributedDataParallel API在純PyTorch中進行分布式訓(xùn)練的簡介。主要內(nèi)容有：

討論一般的分布式訓(xùn)練方式，尤其是數(shù)據(jù)并行化
介紹torch.dist和DistributedDataParallel的相關(guān)功能，并舉例說明如何使用它們
測試真實的訓(xùn)練腳本，以節(jié)省時間

背景知識

在研究分布式和數(shù)據(jù)并行之前，需要先了解一些關(guān)于分布式訓(xùn)練的背景知識。目前普遍使用的分布式訓(xùn)練基本上有兩種不同形式：數(shù)據(jù)并行化和模型并行化。

數(shù)據(jù)并行化中，模型訓(xùn)練作業(yè)是在數(shù)據(jù)上進行分割的。作業(yè)中的每個GPU接收到自己獨立的數(shù)據(jù)批處理切片。每個GPU使用這些數(shù)據(jù)來獨立計算梯度更新。例如，如果要使用兩個GPU和32的批處理大小，一個GPU將處理前16條記錄的向前和向后傳播，第二個處理后16條記錄的向后和向前傳播。這些梯度更新然后在GPU之間同步，一起平均，最后應(yīng)用到模型(同步步驟在技術(shù)上是可選的，但理論上更快的異步更新策略仍是一個活躍的研究領(lǐng)域)。
模型并行化中，模型訓(xùn)練作業(yè)是在模型上進行分割的。工作中的每個GPU接收模型的一個切片，例如它的層的一個子集。例如，一個GPU負責(zé)它的輸出頭，另一個負責(zé)輸入層，另一個負責(zé)中間的隱藏層。

雖然這兩種技術(shù)各有優(yōu)缺點，但數(shù)據(jù)并行化在這兩種技術(shù)中更容易實現(xiàn)(它不需要了解底層網(wǎng)絡(luò)架構(gòu))，因此通常首先嘗試這種策略(也可以結(jié)合使用這些技術(shù)，例如同時使用模型和數(shù)據(jù)并行化，但這是一個高級主題，不在這里介紹)。

注意：本文是對DistributedDataParallel并行API的介紹，所以不再進一步討論模型并行化的細節(jié)。

數(shù)據(jù)并行工作原理

第一個被廣泛采用的數(shù)據(jù)并行技術(shù)是TensorFlow中的參數(shù)服務(wù)器策略。這個功能實際上早于TensorFlow的第一個版本，早在2012年google內(nèi)部的前身DistBelief中就已經(jīng)實現(xiàn)了。這一策略在下圖中得到了很好的說明：
【分布式訓(xùn)練】基于PyTorch進行多GPU分布式模型訓(xùn)練（補充）,深度學(xué)習(xí)與人工智能,分布式,pytorch,多GPU訓(xùn)練,數(shù)據(jù)并行化,模型并行化
在參數(shù)服務(wù)器策略中，worker和parameter進程的數(shù)量是可變的，每個worker進程在GPU內(nèi)存中維護自己的模型獨立副本。梯度更新計算如下：

在接收到開始信號后，每個worker process為其特定的批處理slice積累梯度。
這些worker以扇出的方式將update發(fā)送到參數(shù)服務(wù)器。
參數(shù)服務(wù)器會一直等待，直到它們擁有所有worker更新，然后對它們負責(zé)的梯度更新參數(shù)空間的那部分梯度求平均。
梯度更新被分散到worker上，然后將它們加起來，應(yīng)用到內(nèi)存中模型權(quán)重的副本上(從而保持worker模型同步)。
一旦每個worker都應(yīng)用了更新，新的一批訓(xùn)練就可以開始了。

雖然很容易實現(xiàn)，但是這個策略有一些主要的限制。最重要的一點是每個附加的參數(shù)服務(wù)器在每個同步步驟中都需要n_workers額外的網(wǎng)絡(luò)調(diào)用——一個O(n)復(fù)雜度代價。計算的總體速度取決于最慢的連接，因此基于大參數(shù)服務(wù)器的模型訓(xùn)練作業(yè)在實踐中效率非常低，將網(wǎng)絡(luò)GPU利用率推到50%或以下。

更現(xiàn)代的分布式訓(xùn)練策略廢除了參數(shù)服務(wù)器，在DistributedDataParallel并行策略中，每個進程都是一個work進程。每個進程仍然在內(nèi)存中維護模型權(quán)重的完整副本，但是批處理slice梯度更新現(xiàn)在是同步的，并且直接在工作進程本身上平均。這是使用從高性能計算領(lǐng)域借來的技術(shù)來實現(xiàn)的：全歸約算法（all-reduce algorithm）
【分布式訓(xùn)練】基于PyTorch進行多GPU分布式模型訓(xùn)練（補充）,深度學(xué)習(xí)與人工智能,分布式,pytorch,多GPU訓(xùn)練,數(shù)據(jù)并行化,模型并行化
該圖顯示了全歸約算法的一種特定實現(xiàn)方式，即循環(huán)全歸約。該算法提供了一種優(yōu)雅的方式來同步一組進程之間的一組變量（在本例中為張量）的狀態(tài)。 向量在直接的worker到worker連接的序列中直接傳遞。這消除了worker與參數(shù)服務(wù)器之間的連接所造成的網(wǎng)絡(luò)瓶頸，從而大大提高了性能。在該方案中，梯度更新計算如下:

每個 worker 維護它自己的模型權(quán)重副本和它自己的數(shù)據(jù)集副本。
在接收到開始信號后，每個工作進程從數(shù)據(jù)集中提取一個分離的批處理，并為該批處理計算一個梯度。
worker使用all-reduce算法來同步他們各自的梯度，本地計算所有節(jié)點上相同的平均梯度。
每個worker都將梯度更新應(yīng)用到它的本地模型副本上。
下一批訓(xùn)練開始。

在2017年百度的一篇論文《Bringing HPC Techniques to Deep Learning》中，這種全精簡策略被提出。它的重要之處在于它基于眾所周知的HPC技術(shù)以及長期存在的開源實現(xiàn)。All-reduce包含在消息傳遞接口(MPI)的標(biāo)準(zhǔn)中，這就是為什么PyTorch不少于三個不同的后端實現(xiàn): Open MPI、NVIDIA NCCL和Facebook Gloo（一般情況下建議使用NVIDIA NCCL）

數(shù)據(jù)分發(fā)流程

1. 流程初始化

必須要知道的一點是將訓(xùn)練腳本修改為使用DistributedDataParallel并行策略并不是簡單的一行更改。具體可以參考另外一篇博客：基于Pytorch的分布式數(shù)據(jù)并行訓(xùn)練

因此，需要處理的第一個也是最復(fù)雜的新事情是進程初始化。普通的PyTorch訓(xùn)練腳本在單個進程中執(zhí)行其代碼的單一副本。使用數(shù)據(jù)并行模型，情況就更加復(fù)雜了：現(xiàn)在訓(xùn)練腳本的同步副本與訓(xùn)練集群中的GPU數(shù)量一樣多，每個GPU運行在不同的進程中。示例腳本如下：

import os

import torch
import torch.distributed as dist
import torch.multiprocessing as mp


def init_process(rank, size, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)


def train(rank, num_epochs, world_size):
    init_process(rank, world_size)
    print(f"Rank {rank + 1}/{world_size} process initialized.\n")


# rest of the training script goes here!WORLD_SIZE = torch.cuda.device_count()
if __name__ == "__main__":
    mp.spawn(train, args=(NUM_EPOCHS, WORLD_SIZE), nprocs=WORLD_SIZE, join=True)

其中WORLD_SIZE是編排的進程數(shù)量，(全局)rank是當(dāng)前進程在該MPI中的位置。如果這個腳本要在一個有4個GPU的機器上執(zhí)行，WORLD_SIZE應(yīng)該是4(因為torch.cuda.device_count() == 4)，所以是mp.spawn會產(chǎn)生4個不同的進程，它們的rank分別是0、1、2或3。rank為0的進程被賦予一些額外的職責(zé)，因此被稱為主進程。

當(dāng)前進程的rank將作為派生入口點（在本例中為train function）作為其第一個參數(shù)傳遞。在訓(xùn)練時可以執(zhí)行任何工作之前，它需要首先建立與對等點對點的連接。這是dist.init_process_group的工作。在主進程中運行時，此方法在MASTER_ADDR：MASTER_PORT上設(shè)置套接字偵聽器，并開始處理來自其他進程的連接。 一旦所有進程都已連接，此方法將處理建立對等連接，以允許進程進行通信。

請注意，此代碼僅適用于在一臺多GPU機器上進行訓(xùn)練！同一臺機器用于啟動作業(yè)中的每個流程，因此訓(xùn)練只能利用連接到該特定機器的GPU。這使事情變得容易：設(shè)置IPC就像在localhost上找到一個空閑端口一樣容易，該端口對于該計算機上的所有進程都是立即可見的。 跨計算機的IPC更為復(fù)雜，因為它需要配置一個對所有計算機可見的外部IP地址。

2. 流程同步

了解了初始化過程，現(xiàn)在可以開始完成所有worker的train方法的主體。

def train(rank, num_epochs, world_size):
    init_process(rank, world_size)
    for epoch in range(num_epochs):
        print(f"Rank {rank + 1}/{world_size} process initialized.\n")

四個訓(xùn)練過程中的每一個都會運行此函數(shù)直到完成，然后在完成時退出。如果現(xiàn)在（通過python multi_init.py）運行此代碼，將在控制臺上看到類似以下內(nèi)容：

$ python multi_init.py
Rank 4/4 process initialized.
Rank 2/4 process initialized.
Rank 4/4 process initialized.
Rank 2/4 process initialized.
Rank 3/4 process initialized.
Rank 4/4 process initialized.
Rank 2/4 process initialized.
Rank 3/4 process initialized.
Rank 4/4 process initialized.
Rank 2/4 process initialized.
Rank 4/4 process initialized.
Rank 3/4 process initialized.
Rank 2/4 process initialized.
Rank 3/4 process initialized.
Rank 3/4 process initialized.
Rank 1/4 process initialized.
Rank 1/4 process initialized.
Rank 1/4 process initialized.
Rank 1/4 process initialized.
Rank 1/4 process initialized.

這些過程是獨立執(zhí)行的，并且不能保證訓(xùn)練循環(huán)中任一點處于什么狀態(tài)。所以這里需要對初始化過程進行一些仔細的更改。
（1）任何下載數(shù)據(jù)的方法都應(yīng)隔離到主進程中。
否則的話將在所有過程之間復(fù)制數(shù)據(jù)下載過程，從而導(dǎo)致四個過程同時寫入同一文件，這是造成數(shù)據(jù)損壞的原因。

def train(rank, num_epochs, world_size):
    init_process(rank, world_size)
    for epoch in range(num_epochs):
        if rank == 0:
            downloading_dataset()
            downloading_model_weights()
            dist.barrier()
            print(f"Rank {rank + 1}/{world_size} training process passed data download barrier.\n")

示例中的dist.barrier將阻塞調(diào)用，直到完成主進程（rank == 0）downloading_dataset和downloading_model_weights為止。這樣可以將所有網(wǎng)絡(luò)I / O隔離到一個進程中，并防止其他進程繼續(xù)前進。
（2）數(shù)據(jù)加載器需要使用DistributedSampler。

def get_dataloader(rank, world_size):
    dataset = PascalVOCSegmentationDataset()
    sampler = DistributedSampler(dataset, rank=rank, num_replicas=world_size, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=8, sampler=sampler)

DistributedSampler使用rank和world_size找出如何將整個過程中的數(shù)據(jù)集拆分為不重疊的批次。 worker進程的每個訓(xùn)練步驟都從其本地數(shù)據(jù)集副本中檢索batch_size觀測值。在四個GPU的示例情況下，這意味著有效批大小為8 * 4 = 32。
（3）在正確的設(shè)備中加載張量。
請使用該進程正在管理的設(shè)備的rank來參數(shù)化.cuda()調(diào)用，例如：

batch = batch.cuda(rank)
segmap = segmap.cuda(rank)
model = model.cuda(rank)

（4）必須禁用模型初始化中的任何隨機性。
在整個訓(xùn)練過程中，模型必須啟動并保持同步，這一點非常重要。否則，將獲得不正確的梯度并導(dǎo)致模型無法收斂?？梢允褂靡韵麓a使torch.nn.init.kaiming_normal_之類的隨機初始化方法具有確定性：

torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(0)

（5）任何執(zhí)行文件I / O的方法都應(yīng)與主進程隔離。
由于并發(fā)寫入同一文件而導(dǎo)致的效率低下和潛在的數(shù)據(jù)損壞。同樣，使用簡單的條件邏輯很容易做到這一點：

if rank == 0:
	if not os.path.exists('/spell/checkpoints/'):
		os.mkdir('/spell/checkpoints/')
		torch.save(model.state_dict(), f'/spell/checkpoints/model_{epoch}.pth')

（6）將model包裝在DistributedDataParallel中。

model = DistributedDataParallel(model, device_ids=[rank])

這一步是必須的，也是最后一步，之后模型就可以在分布式數(shù)據(jù)并行模式下訓(xùn)練。

數(shù)據(jù)并行

PyTorch中還有另一種數(shù)據(jù)并行化策略，即torch.nn.DataParallel。該API易于使用。要做的就是包裝模型初始化，如下所示：

model = nn.DataParallel(model)

所有的改動只有一行！為什么不使用它呢？因為在程序的后臺，DataParallel使用多線程而不是多進程來管理其GPU工作器。這極大地簡化了實現(xiàn)：由于工作進程是同一進程的所有不同線程，因此它們都可以訪問相同的共享狀態(tài)，而無需任何其他同步步驟。但是，由于存在全局解釋器鎖GIL，在Python中將多線程用于計算作業(yè)的效果很差。如下一節(jié)中的基準(zhǔn)測試所示，使用DataParallel并行化的模型比使用DistributedDataParallel并行化的模型要慢得多。盡管如此，如果不想花費額外的時間和精力使用多GPU訓(xùn)練，DataParallel是可以考慮的。

基準(zhǔn)測試

【分布式訓(xùn)練】基于PyTorch進行多GPU分布式模型訓(xùn)練（補充）,深度學(xué)習(xí)與人工智能,分布式,pytorch,多GPU訓(xùn)練,數(shù)據(jù)并行化,模型并行化 DistributedDataParallel的效率明顯高于DataParallel，但還遠遠不夠完美。從V100x1切換到V100x4是原始GPU功耗的4倍，但模型訓(xùn)練速度僅為3倍。通過升級到V100x8使計算進一步加倍，只會使訓(xùn)練速度提高約30％,但是DataParallel的效率幾乎達到了DistributedDataParallel的水平。文章來源地址http://www.zghlxwxcb.cn/news/detail-601603.html