一、簡要回顧DDP
在上一篇文章中，簡單介紹了Pytorch分散式訓練的一些基礎原理和基本概念。簡要回顧如下：
1，DDP採用Ring-All-Reduce架構，其核心思想為：所有的GPU裝置安排在一個邏輯環中，每個GPU應該有一個左鄰和一個右鄰，裝置從它的左鄰居接收資料，並將資料彙總後傳送給右鄰。透過N輪迭代以後，每個裝置都擁有全域性資料的計算結果。
2，DDP每個GPU對應一個程序，這些程序可以看作是相互獨立的。除非我們自己手動實現，不然各個程序的資料都是不互通的。Pytorch只為我們實現了梯度同步。
3，DDP相關程式碼需要關注三個部分：資料拆分、IO操作、和評估測試。
二、DDP訓練框架的流程
1. 準備DDP環境
在使用DDP訓時，我們首先要初始化一下DDP環境，設定好通訊後端，程序組這些。程式碼很簡單，如下所示：
def setup(rank, world_size):
# 設定主機地址和埠號，這兩個環境變數用於配置程序組通訊的初始化。
# MASTER_ADDR指定了負責協調初始化過程的主機地址，在這裡設定為'localhost'，
# 表示在單機多GPU的設定中所有的程序都將連線到本地機器上。
os.environ['MASTER_ADDR'] = 'localhost'
# MASTER_PORT指定了主機監聽的埠號，用於程序間的通訊。這裡設定為'12355'。
# 注意要選擇一個未被使用的埠號來進行監聽
os.environ['MASTER_PORT'] = '12355'
# 初始化分散式程序組。
# 使用NCCL作為通訊後端，這是NVIDIA GPUs最佳化的通訊庫，適用於高效能GPU之間的通訊。
# rank是當前程序在程序組中的編號，world_size是總程序數（GPU數量），即程序組的大小。
dist.init_process_group("nccl", rank=rank, world_size=world_size)
# 為每個程序設定GPU
torch.cuda.set_device(rank)
2. 準備資料載入器
假設我們已經定義好了dataset，這裡只需要略加修改使用DistributedSampler即可。程式碼如下：
def get_loader(trainset, testset, batch_size, rank, world_size):
train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank)
train_loader = DataLoader(train_set, batch_size=batch_size, sampler=train_sampler)
# 對於測試集來說，可以選擇使用DistributedSampler，也可以選擇不使用，這裡選擇使用
test_sampler = DistributedSampler(test_set, num_replicas=world_size, rank=rank)
test_loader = DataLoader(test_set, batch_size=batch_size, sampler=train_sampler)
# 不使用的程式碼很簡單， 如下所示
# test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False)
rerturn train_loader, test_loader
注：關於testloader要不要使用分散式取樣器，取決於自己的需求。如果測試資料集相對較小，或者不需要頻繁進行測試評估，不使用DistributedSampler可能更簡單，因為每個GPU或程序都會獨立處理完整的資料集，從而簡化了測試流程。然而，對於大型資料集，或當需要在訓練過程中頻繁進行模型評估的情況，使用DistributedSampler可以顯著提高測試的效率，因為它允許每個GPU只處理資料的一個子集，從而減少了單個程序的負載並加快了處理速度。
對於DDP而言，每個程序上都會有一個dataloader，如果使用了DistributedSampler，那麼真的批大小會是batch_size*num_gpus。
有關DistributedSampler的更多細節可以參考：
DDP系列第二篇：實現原理與原始碼解析
3. 準備DDP模型和最佳化器
在定義好模型之後，需要在所有程序中複製模型並用DDP封裝。程式碼如下：
def prepare_model_and_optimizer(model, rank, lr):
# 設定裝置為當前程序的GPU。這裡`rank`代表當前程序的編號，
# `cuda:{rank}` 指定模型應該執行在對應編號的GPU上。
device = torch.device(f"cuda:{rank}")
# 包裝模型以使用分散式資料並行。DDP將在多個程序間同步模型的引數，
# 並且只有指定的`device_ids`中的GPU才會被使用。
model = model.to(device)
model = DDP(model, device_ids=[rank])
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
return model, optimizer
注：在DDP中不同程序之間只會同步梯度，因此為了保證訓練時的引數同步，需要在訓練開始前確保不同程序上模型和最佳化器的初始狀態相同。對於最佳化器而言，當使用PyTorch中內建的最佳化器（如SGD, Adam等）時，只要模型在每個程序中初始化狀態相同，最佳化器在每個程序中建立後的初始狀態也將是相同的。但是，如果是自定義的最佳化器，確保在設計時考慮到跨程序的一致性和同步，特別是當涉及到需要維護跨步驟狀態（如動量、RMS等）時。
確保模型的初始狀態相同有如下兩種方式：
1）引數初始化方法
在DDP中，每個GPU上都會有一個模型。我們可以利用統一的初始化方法，來保證不同GPU上的引數統一性。一個簡單的示例程式碼如下：
def weights_init(m):
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
m.bias.data.fill_(0.01)
torch.manual_seed(42) # 設定隨機種子以確保可重複性
# 設定所有GPU的隨機種子
torch.cuda.manual_seed_all(42)
model = MyModel()
model.apply(weights_init)
注：在初始化時，需要為所有的程序設定好相同的隨機種子，不然weights_init的結果也會不一樣。
2）載入相同的模型權重檔案
另一種方法是在所有程序中載入相同的預訓練權重。這確保了無論在哪個GPU上，模型的起點都是一致的。程式碼如下：
model = MyModel()
model.load_state_dict(torch.load("path_to_weights.pth"))
注：如果你既沒有設定初始化方法，也沒有模型權重。一個可行的方式是手動同步，將rank=0的程序上模型檔案臨時儲存，然後其他程序載入，最後再刪掉臨時檔案。程式碼如下：
def synchronize_model(model, rank, root='temp_model.pth'):
if rank == 0:
# 儲存模型到檔案
torch.save(model.state_dict(), root)
torch.distributed.barrier() # 等待rank=0儲存模型
if rank != 0:
# 載入模型權重
model.load_state_dict(torch.load(root))
torch.distributed.barrier() # 確保所有程序都載入了模型
if rank == 0:
# 刪除臨時檔案
os.remove(root)
模型同步似乎可以省略，在使用torch.nn.parallel.DistributedDataParallel封裝模型時，它會在內部處理所需的同步操作。
4. 開始訓練
訓練時的程式碼，其實和單卡訓練沒有什麼區別。最主要的就是在每個epoch開始的時候，要設定一下sampler的epoch，以保證每個epoch的取樣資料的順序都是不一樣的。程式碼如下：
def train(model, optimizer, criterion, rank, train_loader, num_epochs):
sampler = train_loader.sampler
for epoch in range(num_epochs):
# 在每個epoch開始時更新sampler
sampler.set_epoch(epoch)
model.train()
for batch_idx, (data, targets) in enumerate(dataloader):
data, targets = data.cuda(rank), targets.cuda(rank)
optimizer.zero_grad()
outputs = model(data)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
# 只在rank為0的程序中列印資訊
if rank == 0 and batch_idx % 100 == 0:
print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}")
注：這裡的列印的loss只是rank0上的loss，如果要列印所有卡上的平均loss，則需要使用all_reduce方法。程式碼如下：
# 將損失從所有程序中收集起來並求平均
# 建立一個和loss相同的tensor，用於聚合操作
reduced_loss = torch.tensor([loss.item()]).cuda(rank)
# all_reduce操作預設是求和
dist.all_reduce(reduced_loss)
# 求平均
reduced_loss = reduced_loss / dist.get_world_size()
# 只在rank為0的程序中列印資訊
if rank == 0 and batch_idx % 100 == 0:
print(f"Epoch {epoch}, Batch {batch_idx}, Average Loss: {reduced_loss.item()}")
5. 評估測試
評估的程式碼也和單卡比較類似，唯一的區別就是，如果使用了DistributedSampler，在計算指標時，需要gather每個程序上的preds和gts，然後計算全域性指標。
def evaluate(model, test_loader, rank):
model.eval()
total_preds = []
total_targets = []
with torch.no_grad():
for data, targets in test_loader:
data, targets = data.to(rank), targets.to(rank)
outputs = model(data)
_, preds = torch.max(outputs, 1)
# 收集當前程序的結果
total_preds.append(preds)
total_targets.append(targets)
# 將所有程序的preds和targets轉換為全域性列表
total_preds = torch.cat(total_preds).cpu()
total_targets = torch.cat(total_targets).cpu()
# 使用all_gather將所有程序的資料集中到一個列表中
gathered_preds = [torch.zeros_like(total_preds) for _ in range(dist.get_world_size())]
gathered_targets = [torch.zeros_like(total_targets) for _ in range(dist.get_world_size())]
dist.all_gather(gathered_preds, total_preds)
dist.all_gather(gathered_targets, total_targets)
if rank == 0:
# 只在一個程序中進行計算和輸出
gathered_preds = torch.cat(gathered_preds)
gathered_targets = torch.cat(gathered_targets)
# 計算全域性效能指標
accuracy = (gathered_preds == gathered_targets).float().mean()
print(f'Global Accuracy: {accuracy.item()}')
注：如果test_loader沒有設定DistributedSampler，評估的程式碼可以和單卡程式碼完全一樣，不需要任何修改。
三、完整程式碼
下面以CIFAR100資料集為例，完整展示一下DDP的訓練流程。
import os
import time
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torchvision import datasets, transforms
from torch.nn.parallel import DistributedDataParallel as DDP
# 模型定義
class LeNet(nn.Module):
def __init__(self, num_classes=100):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, num_classes) # CIFAR100 has 100 classes
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = torch.flatten(x, 1)
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def cleanup():
# 銷燬程序組
dist.destroy_process_group()
def get_model():
model = LeNet(100).cuda()
model = DDP(model, device_ids=[torch.cuda.current_device()])
return model
def get_dataloader(train=True):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
rank = dist.get_rank()
# 每個程序建立其獨立的資料目錄，避免I/O衝突
# 這裡使用rank來建立獨立目錄，例如：'./data_0'，'./data_1'等
# 這種方法避免了多個程序同時寫入同一個檔案所導致的衝突
# 注：這是一種簡單的解決方案，但在需要大量磁碟空間的情況下並不高效，因為每個程序都需要儲存資料集的一個完整副本。
dataset = datasets.CIFAR100(root=f'./data_{rank}', train=train, download=True, transform=transform)
sampler = DistributedSampler(dataset, shuffle=train)
loader = DataLoader(dataset, batch_size=64, sampler=sampler)
return loader
def train(model, loader, optimizer, criterion, epoch, rank):
model.train()
# 設定DistributedSampler的epoch
loader.sampler.set_epoch(epoch)
for batch_idx, (data, targets) in enumerate(loader):
data, targets = data.cuda(rank), targets.cuda(rank)
optimizer.zero_grad()
outputs = model(data)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
# 每100個batch計算當前的損失，並在所有程序中進行聚合然後列印
if (batch_idx + 1) % 100 == 0:
# 將當前的loss轉換為tensor，並在所有程序間進行求和
loss_tensor = torch.tensor([loss.item()]).cuda(rank)
dist.all_reduce(loss_tensor)
# 計算所有程序的平均損失
mean_loss = loss_tensor.item() / dist.get_world_size() # 平均損失
# 如果是rank 0，則列印平均損失
if rank == 0:
print(f"Rank {rank}, Epoch {epoch}, Batch {batch_idx + 1}, Mean Loss: {mean_loss}")
def evaluate(model, dataloader, device):
model.eval()
local_preds = []
local_targets = []
with torch.no_grad():
for data, targets in dataloader:
data, targets = data.to(device), targets.to(device)
outputs = model(data)
_, preds = torch.max(outputs, 1)
local_preds.append(preds)
local_targets.append(targets)
# 將本地預測和目標轉換為全域性列表
local_preds = torch.cat(local_preds)
local_targets = torch.cat(local_targets)
# 使用all_gather收集所有程序的預測和目標
world_size = dist.get_world_size()
gathered_preds = [torch.zeros_like(local_preds) for _ in range(world_size)]
gathered_targets = [torch.zeros_like(local_targets) for _ in range(world_size)]
dist.all_gather(gathered_preds, local_preds)
dist.all_gather(gathered_targets, local_targets)
# 只在rank 0進行計算和輸出
if dist.get_rank() == 0:
gathered_preds = torch.cat(gathered_preds)
gathered_targets = torch.cat(gathered_targets)
accuracy = (gathered_preds == gathered_targets).float().mean()
print(f"Global Test Accuracy: {accuracy.item()}")
def main_worker(rank, world_size, num_epochs):
setup(rank, world_size)
model = get_model()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
train_loader = get_dataloader(train=True)
test_loader = get_dataloader(train=False)
start_time = time.time()
for epoch in range(num_epochs): # num of epochs
train(model, train_loader, optimizer, criterion, epoch, rank)
evaluate(model, test_loader, rank)
# 計時結束前同步所有程序，確保所有程序已經完成訓練
dist.barrier()
duration = time.time() - start_time
if rank == 0:
print(f"Training completed in {duration:.2f} seconds")
cleanup()
if __name__ == "__main__":
world_size = 4 # 4塊GPU
num_epochs = 10 # 總共訓練10輪
# 採用mp.spawn啟動
mp.spawn(main_worker, args=(world_size,num_epochs), nprocs=world_size, join=True)
注：
1）關於get_loader函式中資料載入有關部分的問題
dataset = datasets.CIFAR100(root=f'./data_{rank}', train=train, download=True, transform=transform)
上面這段程式碼的最大問題在於，每個程序都會去下載一份資料到該程序對應的目錄，這些目錄之間是物理隔離的。顯然，當要下載的資料集很大時，這種方法並不合適，因為會佔用更多的硬碟資源，並且大量時間會花費在下載資料集上。但是如果不為每個程序設定單獨的目錄，就會造成讀寫衝突，多個程序都去同時讀寫同一個檔案，最終導致資料集載入不成功。
一種更合理的解決方法是，提前下載好檔案，並在建立資料集時設定download為False。程式碼如下：
def download_data(rank):
if rank == 0:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 只在Rank 0 中下載資料集
datasets.CIFAR100(root='./data_cifar100', train=True, download=True, transform=transform)
# 等待rank0下載完成
dist.barrier()
def get_dataloader(train=True):
rank = dist.get_rank()
# 現在只需要下載一次
download_data(rank)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = datasets.CIFAR100(root='./data_cifar100', train=train, download=False, transform=transform) # 設定download為False
sampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=rank, shuffle=train)
loader = DataLoader(dataset, batch_size=64, sampler=sampler, num_workers=4)
return loader
2）關於DDP啟動方式有關的問題
DDP的啟動方式有兩種，一種是torch.distributed.launch，這個工具指令碼設定必要的環境變數，併為每個節點啟動多個程序。通常在從命令列執行指令碼時使用，在新版本的Pytorch 2.+版本中，這種方式是deprecated，不推薦繼續使用。
還有一種就是torch.multiprocessing.spawn： 這個函式在Python指令碼內部程式設計方式啟動多個程序。啟動方式很簡單，分別傳入主入口函式main_worker，然後傳入main_woker的引數，以及GPU數量即可。
四、DP, DDP效能對比
基於上述的程式碼，我還實現了一個DP的程式碼。實驗setting為：
GPU: \(4 \times\) RTX 4090, batch_size=256, optimizer為Adam，學習率為0.001，loss是CE loss。
它們之間的效能對比如下：
| 方式 | 時間 | 準確率 |
|---|---|---|
| DDP | 77秒 | 27.12% |
| DP | 293秒 | 26.76% |
| 單卡訓練 | 248秒 | 26.34% |
沒有調參，網路結構也是個最簡單的LeNet，只訓練了10輪，所以準確率比較低。不過，這個結果還是能說明一些問題的。可以看到DDP耗時最短，DP的時間反而比單卡訓練還長。這主要是因為對於CIFAR100分類，單卡也可以很好地支援訓練，顯示卡並不是效能瓶頸。當使用DP時，模型的所有引數在每次前向和反向傳播後都需要在主GPU上進行聚合，然後再分發到各個GPU。這種多餘的聚合和分發操作會帶來顯著的通訊開銷。並且在DataParallel中，主GPU承擔了額外的資料分發和收集工作，會成為效能瓶頸，導致其他GPU在等待主GPU處理完成時出現閒置狀態。
五、總結
1. How to use DDP all depends on yourself.
在最開始學習DDP的時候，有很多地方是很困惑的。每個部落格的程式碼都有所區別，讓我很是困惑。例如：在testloader到底要不要用DistributedSampler；在計算損失的時候，到底要不要用all_reduce操作來計算mean_loss；在計算指標的時候，到底要不要all_gather。後面瞭解多了之後才發現，到底用不用完全取決自己的需求。
1）mean_loss
由於Pytorch在DDP中會自動同步梯度，因此計算不計算mean_loss對於模型的訓練和引數沒有任何影響。唯一的區別在於列印日誌的時候，是列印全域性的平均損失，還是隻列印某個程序上的損失。如果每張卡上的batch size已經足夠大（例如，設定為128或者更高），列印全域性平均損失和單程序上的損失，一般來說差別不大。
2）測試時DistributedSampler
測試時testloader設不設定DistributedSampler也完全取決於自己的實際需求。如果不設定，那麼就是在每個程序上都會用全部的資料的來進行測試。如果有八塊卡，那麼就相當於在每個卡上都分別測試了一次，一共測試了八次。如果你的測試資料集比較小，比如只有幾百張影像，並且測試的頻率也不高的話，不設定DistributedSampler沒有任何問題，不會有太多的額外開銷。但是如果測試資料集比較大（比如幾萬張影像），並且訓練時每個epoch都要進行測試，那麼最好還是設定一下DistributedSampler，可以有效地減少總體訓練時間。
3）評估時all_gather
至於要不要使用all_gather，則和有沒有使用DistributedSampler相關。如果設定了DistributedSampler，那麼評估時就要使用all_gather來彙總所有程序上的結果，否則列印的只會是某個程序的結果，並不準確。
4）batch_size
此外，testloader在使用DistributedSampler也需要格外注意資料能否被整除。舉個例子，假設我們有8塊卡，每塊卡上的batch_size設定為64，那麼總的batch size就是512。如果我們的訓練資料集只有1000份，為了湊夠完整的兩個batch，DistributedSampler會對資料進行補全（重複部分資料）使得資料總數變為1024份。在這個過程有24份資料被重複評估，這些重複評估的資料可能會對評估結果產生影響。以4分類任務為例，如果類別數量比較均衡，相當於每個類別都有256份資料。在這種情況下，重複評估24份資料，對結果不會有什麼影響。但是如果類別資料並不均衡，有些類別只有十幾份資料，那麼這個重複評估的影響就比較大了。如果正好重複的資料是樣本數量只有十幾份的類別，那麼評估結果將會變得極其不準確！！！在這種情況下，我們需要重寫一個sampler來實現無重複的資料取樣。或者，也可以直接不使用DistributedSampler，在每個程序上都進行一次完整的評估。
5）同步批次歸一化（Synchronized Batch Normalization, SynBN）
之前說過，每個GPU對應一個程序，不同程序的資料一般是不共享的。也就是說，如果模型結構中有BN層，每個GPU上的BN層只能訪問到該GPU上的一部分資料，無法獲得全域性的資料分佈資訊。為此可以使用同步BN層，來利用所有GPU上的資料來計算BN層的mean和variance。程式碼也很簡單，只需要對例項化model之後，轉為同步BN即可。
def get_model():
model = LeNet(100).cuda()
# 轉換所有的BN層為同步BN層
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = DDP(model, device_ids=[torch.cuda.current_device()])
return model
2. 又Out了
以上是藉助Pytorch提供的DDP有關API來搭建自己的分散式訓練程式碼框架的教程，還是有一點小複雜的。現在有很多第三方庫（如HuggingFace的Accelerate、微軟開發的DeepSpeed、Horovod、Pytorch Lightning）對DDP進行了進一步的封裝，使用這些第三方庫可以極大地簡化程式碼。但是，目前我還沒有學習瞭解過這些第三方庫（再次out了，沒有及時學習前沿技術），有機會真應該好好學習一下。尤其是Horovod，它可以跨深度學習框架使用，支援Pytorch、TensorFlow、Keras、MXNet等。Accelerate和DeepSeed也很不錯，做大模型相關基本上都會用到。Pytorch Lightning，顧名思義，讓Pytorch變得更簡單，確實Pytorch Lightning把細節封裝得非常好，程式碼非常簡潔，值得一學。
最後推薦兩個B站上對DDP講解很不錯的幾個影片：
pytorch多GPU並行訓練教程
03 DDP 初步應用（Trainer，torchrun）
知乎上有個帖子也不錯：
DDP系列第一篇：入門教程