一、 寫在前面
隨着深度學習技術的不斷髮展，模型的訓練成本也越來越高。訓練一個高效的通用模型，需要大量的訓練數據和算力。在很多非大模型相關的常規任務上，往往也需要使用多卡來進行並行訓練。在多卡訓練中，最爲常用的就是分佈式數據並行（DistributedDataParallel, DDP）。但是現有的有關DDP的教程和博客比較少，內容也比較分散繁瑣。在大多數情況下，我們只需要學會如何使用即可，不需要特別深入地瞭解原理。爲此，寫下這個系列博客，簡明扼要地介紹一下DDP的使用，拋開繁雜的細節和原理，幫助快速上手使用（All in one blog）。
篇幅較長，分爲上下兩篇：這篇簡要介紹相關背景和理論知識，下篇詳細介紹代碼框架和搭建流程。
二、什麼是分佈式並行訓練
1. 並行訓練
在Pytorch中，有兩種並行訓練方式：
1）模型並行。模型並行通常是指你的模型非常大，大到一塊卡根本放不下，因而需要把模型進行拆分放到不同的卡上。
2）數據並行。數據並行通常用於訓練數據非常龐大的時候，比如有幾百萬張圖像用於訓練模型。此時，如果只用一張卡來進行訓練，那麼訓練時間就會非常的長。並且由於單卡顯存的限制，訓練時的batch size不能設置得過大。但是，對於很多模型的性能而言，由於BN層的使用，都會和batch size的大小正相關。此外，很多基於對比學習的訓練算法，由於其對負樣本的需求，性能也與batch size的大小正相關。因此，我們需要使用多卡訓練，不僅可以訓練加速，並且可以設置更大的batch size來提升性能。
2. 數據並行
在Pytorch中有兩種方式來實現數據並行：
1）數據並行（DataParallel，DP）。DataParallel採用參數服務器架構，其訓練過程是單進程的。在訓練時，會將一塊GPU作爲server，其餘的GPU作爲worker，在每個GPU上都會保留一個模型的副本用於計算。訓練時，首先將數據拆分到不同的GPU上，然後在每個worker上分別進行計算，最終將梯度彙總到server上，在server進行模型參數更新，然後將更新後的模型同步到其他GPU上。這種方式有一個很明顯的弊端，作爲server的GPU其通信開銷和計算成本非常大。它需要和其他所有的GPU進行通信，並且梯度彙總、參數更新等步驟都是由它完成，導致效率比較低。並且，隨着多卡訓練的GPU數量增強，其通信開銷也會線性增長。
不過DataParallel的代碼十分簡潔，僅需在原有單卡訓練的代碼中加上一行即可。
model = nn.DataParallel(model)
如果你的數據集並不大，只有幾千的規模，並且你多卡訓練時的卡也不多，只有4塊左右，那麼DataParallel會是一個不錯的選擇。
關於Parameter Server更詳細的原理介紹，可以參考：
一文讀懂「Parameter Server」的分佈式機器學習訓練原理
2）分佈式數據並行（DistributedDataParallel，DDP）。DDP採用Ring-All-Reduce架構，其訓練過程是多進程的。如果要用DDP來進行訓練，我們通常需要修改三個地方的代碼：數據讀取器dataloader，日誌輸出print，指標評估evaluate。其代碼實現略微複雜，不過我們只需要始終牢記一點即可：每一塊GPU都對應一個進程，除非我們手動實現相應代碼，不然各個進程的數據都是不互通的。Pytorch只爲我們實現了同步梯度和參數更新的代碼，其餘的需要我們自己實現。
三、DDP的基本原理
1. DDP的訓練過程
DDP的訓練過程可以總結爲如下步驟：
1）在訓練開始時，整個數據集被均等分配到每個GPU上。每個GPU獨立地對其分配到的數據進行前向傳播（計算預測輸出）和反向傳播（計算梯度）。
2）同步各個GPU上的梯度，以確保模型更新的一致性，該過程通過Ring-All-Reduce算法實現。
3）一旦所有的GPU上的梯度都同步完成，每個GPU就會使用這些聚合後的梯度來更新其維護的模型副本的參數。因爲每個GPU都使用相同的更新梯度，所以所有的模型副本在任何時間點上都是相同的。
2. Ring-All-Reduce算法
Ring-All-Reduce架構是一個環形架構，所有GPU的位置都是對等的。每個GPU上都會維持一個模型的副本，並且只需要和它相連接的兩個GPU通信。
對於第k個GPU而言，只需要接收來自於第k-1個GPU的數據，並將數據彙總後發送給第k+1個GPU。這個過程在環中持續進行，每個GPU輪流接收、聚合併發送梯度。
經過 N 次的迭代循環後（N是GPU的數量），每個GPU將累積得到所有其他GPU的梯度數據的總和。此時，每個GPU上的梯度數據都是完全同步的。
DDP的通信開銷與GPU的數量無關，因而比DP更爲高效。如果你的訓練數據達到了十萬這個量級，並且需要使用4卡及以上的設備來進行訓練，DDP將會是你的最佳選擇。
關於DDP和Ring-All-Reduce算法的更多實現原理和細節，可以參考：
Bringing HPC Techniques to Deep Learning
Pytorch 分散式訓練 DistributedDataParallel — 概念篇
Technologies behind Distributed Deep Learning: AllReduce
四、如何搭建一個Pytorch DDP代碼框架
1. 與DDP有關的基本概念
在開始使用DDP之前，我們需要了解一些與DDP相關的概念。
| 參數 | 含義 | 查看方式 |
|---|---|---|
| group | 分佈式訓練的進程組，每個group可以進行自己的通信和梯度同步 | Group通常在初始化分佈式環境時創建，並通過torch.distributed.new_group等API創建自定義groups。 |
| world size | 參與當前分佈式訓練任務的總進程數。在單機多GPU的情況下，world size通常等於GPU的數量；在多機情況下，它是所有機器上所有GPU的總和。 | torch.distributed.get_world_size() |
| rank | Rank是指在所有參與分佈式訓練的進程中每個進程的唯一標識符。Rank通常從0開始編號，到world size - 1結束。 | torch.distributed.get_rank() |
| local rank | Local rank是當前進程在其所在節點內的相對編號。例如，在一個有4個GPU的單機中，每個GPU進程的local rank將是0, 1, 2, 3。這個參數常用於確定每個進程應當使用哪個GPU。 | Local rank不由PyTorch的分佈式API直接提供，而通常是在啓動分佈式訓練時由用戶設定的環境變量，或者通過訓練腳本的參數傳入。 |
2. 與DDP有關的一些操作
在DDP中，每個進程的數據是互不影響的（除了採用Ring-All-Reduce同步梯度）。如果我們要彙總或者同步不同進程上的數據，就需要用到一些對應的函數。
1）all_reduce
all_reduce操作會在所有進程中聚合每個進程的數據（如張量），並將結果返回給所有進程。聚合可以是求和、取平均、找最大值等。當你需要獲得所有進程的梯度總和或平均值時，可以使用all_reduce。這在計算全局平均或總和時非常有用，比如全局平均損失。
一個示例代碼如下：
import torch.distributed as dist
tensor_a = torch.tensor([1.0], device=device)
# 所有進程中的tensor_a將會被求和，並且結果會被分配給每個進程中的tensor_a。
dist.all_reduce(tensor_a, op=dist.ReduceOp.SUM)
2）all_gather
all_gather操作用於在每個進程中收集所有進程的數據。它不像all_reduce那樣聚合數據，而是將每個進程的數據保留並彙總成一個列表。當每個進程計算出一個局部結果，並且你需要在每個進程中收集所有結果進行分析或進一步處理時，可以使用all_gather。
一個示例代碼如下：
import torch
import torch.distributed as dist
# 每個進程有一個tensor_a，其值爲當前進程的rank
tensor_a = torch.tensor([rank], device=device) # 假設rank是當前進程的編號
gather_list = [torch.zeros_like(tensor_a) for _ in range(dist.get_world_size())]
# 收集所有進程的tensor_a到每個進程的gather_list
dist.all_gather(gather_list, tensor)
3）broadcast
broadcast操作將一個進程的數據（如張量）發送到所有其他進程中。這通常用於當一個進程生成了某些數據，需要確保其他所有進程都得到相同的數據時。在在開始訓練之前，可以用於同步模型的初始權重或者在所有進程中共享某些全局設置。一個示例代碼如下：
import torch.distributed as dist
tensor_a = torch.tensor([1.0], device=device)
if rank == 0:
tensor_a.fill_(10.0) # 只有rank 0設置tensor_a爲10
dist.broadcast(tensor_a, src=0) # rank 0將tensor_a廣播到所有其他進程
3. 要實現DDP訓練，我們需要解決哪些問題？
1）如何將數據均等拆分到每個GPU
在分佈式訓練中，爲了確保每個GPU都能高效地工作，需要將訓練數據均等地分配到每個GPU上。如果數據分配不均，可能導致某些GPU數據多、某些GPU數據少，從而影響整體的訓練效率。
在PyTorch中，可以使用torch.utils.data.DataLoader結合torch.utils.data.distributed.DistributedSampler。DistributedSampler會自動根據數據集、進程總數（world size）和當前進程編號（rank）來分配數據，確保每個進程獲取到的數據互不重複且均衡分佈。
2）如何在IO操作時避免重複
在使用PyTorch的分佈式數據並行（DDP）進行模型訓練時，由於每個進程都是獨立運行的，IO操作如打印（print）、保存（save）或加載（load）等如果未經特別處理，將會在每個GPU進程上執行。這樣的行爲通常會導致以下問題：重複打印（每個進程都會輸出同樣的打印信息到控制檯，導致輸出信息重複，難以閱讀）、文件寫入衝突（如果多個進程嘗試同時寫入同一個文件，會產生寫入衝突，導致數據損壞或者輸出不正確）、資源浪費（每個進程重複加載相同的數據文件會增加IO負擔，降低效率和浪費資源）。
一個簡單且可行的解決方案是隻在特定進程中進行相關操作，例如，只在rank爲0的進程中執行，如有必要，再同步到其他進程。
3）如何收集每個進程上的數據進行評估
在DDP訓練中，每個GPU進程獨立計算其數據的評估結果（如準確率、損失等），在評估時，可能需要收集和整合這些結果。
通過torch.distributed.all_gather函數，可以將所有進程的評估結果聚集到每個進程中。這樣每個進程都可以獲取到完整的評估數據，進而計算全局的指標。如果只需要全局的彙總數據（如總損失或平均準確率），可以使用torch.distributed.reduce或all_reduce操作直接計算彙總結果，這樣更加高效。
4. 一個最簡單的DDP代碼框架
篇幅太長，見下篇。
五、查資料過程中的一個小驚喜
在查找DDP有關過程中，發現了一些博客和視頻做得很不錯，而且這裏面有一部分是女生做的。博客和視頻的質量都很高，內容安排合理，邏輯表達清晰，參考資料也很全面。我看到的時候，還是很驚豔的，巾幗不讓鬚眉！鏈接如下：
國立中央大學的李馨伊
復旦大學的_Meilinger_