一、簡要回顧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的更多細節可以參考：
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講解很不錯的幾個視頻：
知乎上有個帖子也不錯：