PyTorch 的分散式訓練功能有效縮短深度學習模型的訓練時間，尤其適用於大型模型和資料集。藉由 Distributed-DataParallel (DDP) 策略，模型訓練得以分配到多個 GPU 上，每個 GPU 負責處理資料的子集，並同步梯度更新，從而提升整體訓練效率。實務上，使用 DistributedSampler 能確保每個 GPU 接收不同的 minibatch，避免模型副本發散。在多 GPU 機器上，CUDA_VISIBLE_DEVICES 環境變數能精確控制參與訓練的 GPU 數量，簡化 GPU 資源管理。程式碼範例展示了單 GPU 和雙 GPU 環境下的訓練流程與輸出差異，並提供 rank 控制輸出，避免重複訊息。對於更進階的多 GPU 訓練需求，Fabric 函式庫提供另一種替代方案，允許更彈性的 GPU 控制與效能最佳化。
使用多個GPU進行分散式訓練
在深度學習中，訓練模型的時間可能非常長，尤其是在單個GPU或機器上。為了加速訓練過程，PyTorch提供了分散式訓練的功能，可以將模型訓練分配到多個GPU和機器上。這種方法可以大大減少訓練時間，特別是在模型開發的實驗階段，需要進行多次訓練迭代來微調模型引數和架構。
PyTorch的Distributed-DataParallel（DDP）策略
PyTorch的DDP是一種基本的分散式訓練策略，透過啟用平行性來實作。它的工作原理是：PyTorch在每個GPU上啟動一個單獨的程式，每個程式都會接收和儲存模型的一份副本。這些副本將在訓練過程中同步更新。
DDP的工作流程
- 模型初始化：在CPU上初始化模型。
- 模型複製：將模型複製到每個GPU上。
- 資料分割：將輸入資料分割成唯一的minibatch，並將其傳遞給每個模型副本。
- 獨立計算：每個GPU上的模型副本獨立計算輸出（logits）。
- 梯度平均：計算梯度並對其進行平均，以更新模型。
DDP的優點
使用DDP可以提高處理資料集的速度，相比於單個GPU。然而，需要注意的是，在裝置之間存在少量的通訊 overhead。
實踐中使用DDP
要使用DDP，需要確保每個GPU接收到不同的、不重疊的minibatch。這可以透過使用DistributedSampler來實作。這樣可以確保模型副本不會發散，並且可以在訓練過程中同步更新。
示例
假設我們有兩個GPU，想要用它們來訓練一個神經網路。每個GPU都會接收到模型的一份副本，然後在每次訓練迭代中，各自接收到不同的minibatch。這些minibatch將被用來計算輸出和梯度，然後梯度將被平均和同步，以更新模型。
分散式訓練與PyTorch
在深度學習中，訓練大型模型需要大量的計算資源。為了加速訓練過程，PyTorch提供了分散式訓練（Distributed Training）的功能，允許我們使用多個GPU進行平行計算。這節將介紹如何使用PyTorch的分散式訓練功能。
分散式訓練的原理
分散式訓練的基本思想是將模型和資料分割到多個GPU上，然後平行地進行前向傳播和反向傳播。每個GPU都有一份模型的副本，資料也被分割到每個GPU上。當所有GPU完成前向傳播和反向傳播後，梯度會被同步到所有GPU上，以便更新模型的引數。
PyTorch的分散式訓練工具
PyTorch提供了幾個工具來支援分散式訓練，包括：
torch.multiprocessing: 這個模組提供了多程式程式設計的功能，可以用來spawn多個程式進行平行計算。DistributedSampler: 這個類別可以用來將資料分割到多個GPU上。DistributedDataParallel: 這個類別可以用來將模型包裝成分散式模型，自動處理梯度的同步和更新。init_process_group和destroy_process_group: 這兩個函式用來初始化和銷毀分散式訓練的程式群。
實作分散式訓練
要實作分散式訓練，需要進行以下步驟：
- 匯入必要的模組和類別。
- 初始化分散式訓練的程式群。
- 將模型包裝成分散式模型。
- 將資料分割到多個GPU上。
- 進行前向傳播和反向傳播。
- 同步梯度和更新模型引數。
以下是實作分散式訓練的示例程式碼：
import torch
import torch.multiprocessing as mp
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group
def ddp_setup(rank, world_size):
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12345"
init_process_group(backend="nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def prepare_dataset():
# 插入資料準備程式碼
def train(model, device, dataloader):
model.train()
for batch in dataloader:
# 進行前向傳播和反向傳播
pass
if __name__ == "__main__":
world_size = 4 # GPU數量
mp.spawn(ddp_setup, args=(world_size,), nprocs=world_size)
model = NeuralNetwork() # 載入模型
device = torch.device("cuda:0") # 設定GPU
model.to(device)
ddp_model = DDP(model, device_ids=[device])
dataloader = prepare_dataset()
train(ddp_model, device, dataloader)
destroy_process_group()
分散式深度學習框架實作
在深度學習任務中，尤其是那些涉及大量資料和計算的任務，單個GPU的計算能力往往不足以滿足需求。因此，分散式深度學習框架的出現為我們提供了一種解決方案，能夠將任務分配到多個GPU上，從而大大提高了計算效率。
資料載入和分散式設定
首先，我們需要載入訓練和測試資料集，並設定分散式環境。這包括建立DataLoader物件，以便批次載入資料，並使用DistributedSampler確保每個GPU接收到的資料是不同的。
train_loader = DataLoader(
dataset=train_ds,
batch_size=2,
shuffle=False,
pin_memory=True,
drop_last=True,
sampler=DistributedSampler(train_ds)
)
模型定義和最佳化器設定
接下來，我們定義神經網路模型和最佳化器。模型需要被轉移到指定的GPU上，並使用DDP（Distributed Data Parallel）包裝器以支援分散式計算。
model = NeuralNetwork(num_inputs=2, num_outputs=2)
model.to(rank)
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
model = DDP(model, device_ids=[rank])
訓練過程
在訓練過程中，我們遍歷每個epoch和批次，對每個批次的資料進行預測、計算損失、反向傳播和最佳化器更新。
for epoch in range(num_epochs):
for features, labels in train_loader:
features, labels = features.to(rank), labels.to(rank)
# insert model prediction and backpropagation code
print(f"[GPU{rank}] Epoch: {epoch+1:03d}/{num_epochs:03d}"
f" | Batchsize {labels.shape[0]:03d}"
f" | Train/Val Loss: {loss:.2f}")
模型評估
在每個epoch結束後，我們評估模型在訓練集和測試集上的準確率。
model.eval()
train_acc = compute_accuracy(model, train_loader, device=rank)
print(f"[GPU{rank}] Training accuracy", train_acc)
test_acc = compute_accuracy(model, test_loader, device=rank)
print(f"[GPU{rank}] Test accuracy", test_acc)
分散式過程終止
最後，當訓練過程完成後，我們需要終止分散式過程，以釋放資源。
destroy_process_group()
內容解密：
DataLoader是PyTorch中用於批次載入資料的工具，它可以自動將資料分批並提供給模型使用。DistributedSampler是用於分散式計算的取樣器，它能夠確保每個GPU接收到的資料是不同的。DDP（Distributed Data Parallel）是PyTorch中的一種分散式計算模式，它可以自動將模型和資料分配到多個GPU上進行計算。torch.optim.SGD是PyTorch中的一種最佳化器，它使用隨機梯度下降法來更新模型引數。model.to(rank)是將模型轉移到指定的GPU上。model.eval()是將模型設定為評估模式，在這種模式下，模型不會更新引數。compute_accuracy是計算模型準確率的函式，它可以根據模型的預測結果和真實標籤計算準確率。
圖表翻譯：
graph LR
A[資料載入] --> B[分散式設定]
B --> C[模型定義和最佳化器設定]
C --> D[訓練過程]
D --> E[模型評估]
E --> F[分散式過程終止]
此圖表描述了分散式深度學習框架的流程，從資料載入、分散式設定、模型定義和最佳化器設定，到訓練過程、模型評估，最後終止分散式過程。
分散式深度學習基礎
在深度學習的世界中，尤其是在大型模型和龐大的資料集上進行訓練時，單一GPU的計算能力往往不足以滿足需求。因此，分散式深度學習應運而生，允許我們利用多個GPU甚至多臺機器來加速模型的訓練過程。
基本概念
- NCCL（NVIDIA Collective Communication Library）：NCCL是一種由NVIDIA開發的函式庫，旨在提供高效的集體通訊操作，特別是在GPU之間的資料傳遞中。它能夠大大提高分散式訓練的效率。
- Rank：在分散式環境中，rank是指每個程式（或GPU）的索引。它用於區分不同的GPU或程式，以便於之間的通訊和資料分配。
- World Size：world size代表了參與分散式訓練的GPU數量。它決定了模型將被拆分成多少部分，並在多少個GPU上進行平行計算。
PyTorch中的分散式訓練
PyTorch提供了強大的支援來實作分散式訓練，主要透過torch.distributed模組實作。以下是一個簡單的例子：
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing as mp
def main(rank, world_size):
# 初始化分散式後端
dist.init_process_group("nccl", rank=rank, world_size=world_size)
# 設定隨機種子以確保可重復性
torch.manual_seed(123)
# 載入模型和資料
model = YourModel()
dataset = YourDataset()
# 封裝模型為DDP模型
model = DDP(model, device_ids=[rank])
# 訓練迴圈
for epoch in range(num_epochs):
# 使用DistributedSampler對資料進行分配
sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=sampler)
for batch in dataloader:
# 進行模型訓練
inputs, labels = batch
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 結束分散式群組
dist.destroy_process_group()
if __name__ == "__main__":
world_size = 4 # 使用4個GPU進行分散式訓練
num_epochs = 3 # 訓練3個epoch
mp.spawn(main, args=(world_size, num_epochs), nprocs=world_size)
分散式訓練環境設定
在進行分散式訓練時，需要將資料集分割成不同的子集，以便每個GPU可以處理不同的子集。這樣可以加速訓練過程，並充分利用GPU的計算資源。
資料集分割
為了實作資料集的分割，需要使用DistributedSampler類別。這個類別可以將資料集分割成不同的子集，每個子集對應一個GPU。以下是設定DistributedSampler的示例：
sampler = DistributedSampler(train_ds)
這裡，train_ds是訓練資料集，DistributedSampler會將其分割成不同的子集。
分散式訓練環境設定
設定分散式訓練環境需要使用ddp_setup函式。這個函式負責設定主節點的地址和埠，初始化過程群組，並設定rank（過程識別器）和world size（總過程數）。以下是ddp_setup函式的示例：
def ddp_setup(rank, world_size):
# 設定主節點的地址和埠
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
# 初始化過程群組
torch.distributed.init_process_group(
backend='nccl',
init_method='env://',
world_size=world_size,
rank=rank
)
# 設定rank和world size
torch.distributed.barrier()
# 指定GPU裝置
device = torch.device(f'cuda:{rank}')
return device
這裡，rank是過程識別器，world_size是總過程數。
分散式訓練過程
分散式訓練過程需要使用mp.spawn函式來啟動多個過程。以下是啟動分散式訓練過程的示例：
if __name__ == '__main__':
# 設定隨機種子
torch.manual_seed(42)
# 啟動多個過程
mp.spawn(
main,
args=(args,),
nprocs=world_size,
join=True
)
這裡，main是主函式，args是命令列引數，world_size是總過程數。
主函式
主函式需要設定分散式訓練環境，載入訓練和測試資料集，設定模型，並進行訓練。以下是主函式的示例：
def main(rank, args):
# 設定分散式訓練環境
device = ddp_setup(rank, args.world_size)
# 載入訓練和測試資料集
train_ds, test_ds = load_data()
# 設定模型
model = MyModel()
# 將模型和資料轉移到目標裝置
model.to(device)
train_ds = train_ds.to(device)
test_ds = test_ds.to(device)
# 包裝模型以啟用梯度同步
model = DDP(model, device_ids=[device])
# 進行訓練
train(model, train_ds, test_ds)
# 評估模型
evaluate(model, test_ds)
# 清理資源
destroy_process_group()
這裡，rank是過程識別器，args是命令列引數，device是GPU裝置。
選擇多GPU機器上的可用GPU
在多GPU機器上，若您希望限制用於訓練的GPU數量，使用CUDA_VISIBLE_DEVICES環境變數是最簡單的方法。假設您的機器有多個GPU，您只想使用其中一個，例如索引為0的GPU。您可以在終端執行以下命令：
CUDA_VISIBLE_DEVICES=0 python some_script.py
或者，如果您的機器有四個GPU，您只想使用第一個和第三個GPU，可以使用：
CUDA_VISIBLE_DEVICES=0,2 python some_script.py
這種設定CUDA_VISIBLE_DEVICES的方法是管理GPU分配的一種簡單有效的方式，不需要修改您的PyTorch指令碼。
287A.9 使用GPU最佳化訓練效能
注意，這應該可以在單GPU和多GPU機器上工作。如果我們在單GPU機器上執行這個程式碼，應該會看到以下輸出：
PyTorch版本：2.2.1+cu117
CUDA可用：True
可用GPU數量：1
[GPU0] Epoch: 001/003 | Batchsize 002 | Train/Val Loss: 0.62
[GPU0] Epoch: 001/003 | Batchsize 002 | Train/Val Loss: 0.32
[GPU0] Epoch: 002/003 | Batchsize 002 | Train/Val Loss: 0.11
[GPU0] Epoch: 002/003 | Batchsize 002 | Train/Val Loss: 0.07
[GPU0] Epoch: 003/003 | Batchsize 002 | Train/Val Loss: 0.02
[GPU0] Epoch: 003/003 | Batchsize 002 | Train/Val Loss: 0.03
[GPU0] 訓練準確率 1.0
[GPU0] 測試準確率 1.0
程式碼輸出看起來與使用單個GPU（第A.9.2節）類別似，這是一個良好的檢查。
現在，如果我們在具有兩個GPU的機器上執行相同的命令和程式碼，應該會看到以下內容：
PyTorch版本：2.2.1+cu117
CUDA可用：True
可用GPU數量：2
[GPU1] Epoch: 001/003 | Batchsize 002 | Train/Val Loss: 0.60
[GPU0] Epoch: 001/003 | Batchsize 002 | Train/Val Loss: 0.59
[GPU0] Epoch: 002/003 | Batchsize 002 | Train/Val Loss: 0.16
[GPU1] Epoch: 002/003 | Batchsize 002 | Train/Val Loss: 0.17
[GPU0] Epoch: 003/003 | Batchsize 002 | Train/Val Loss: 0.05
[GPU1] Epoch: 003/003 | Batchsize 002 | Train/Val Loss: 0.05
[GPU1] 訓練準確率 1.0
[GPU0] 訓練準確率 1.0
[GPU1] 測試準確率 1.0
[GPU0] 測試準確率 1.0
如預期，我們可以看到一些批次是在第一個GPU（GPU0）上處理，而其他批次則是在第二個GPU（GPU1）上處理。然而，我們看到重複的輸出行當列印訓練和測試準確率時。每個過程（換句話說，每個GPU）都獨立地列印測試準確率。如果這讓你感到困擾，你可以使用每個過程的rank來控制你的列印陳述式：
if rank == 0:
print("測試準確率：", accuracy)
這基本上就是DDP如何工作的。如果您對更多細節感興趣，我建議您檢視官方API檔案。
只在第一個過程中列印
摘要
- PyTorch是一個開源函式庫，具有三個核心元件：張量函式庫、自動微分函式和深度學習工具。
- PyTorch的張量函式庫與NumPy等陣列函式庫相似。
- 在PyTorch的背景下，張量是表示標量、向量、矩陣和更高維陣列的陣列樣式資料結構。
- PyTorch張量可以在CPU上執行，但PyTorch張量格式的一個主要優點是其對GPU的支援，可以加速計算。
- PyTorch中的自動微分（autograd）功能允許我們使用反向傳播方便地訓練神經網路，而無需手動匯出梯度。
- PyTorch中的深度學習工具提供了建立自定義深度神經網路的基礎元件。
- PyTorch包含Dataset和DataLoader類別，以設定高效的資料載入管道。
- 在CPU或單個GPU上訓練模型是最容易的。
- 使用DistributedDataParallel是PyTorch中加速訓練的最簡單方法，如果有多個GPU可用。
替代PyTorch API進行多GPU訓練
如果您偏好更直接的方法來在PyTorch中使用多個GPU，您可以考慮新增API，如開源Fabric函式庫。我曾在“加速PyTorch模型訓練：使用混合精確度和完全分片資料平行”中撰寫過相關內容。
隨著深度學習模型日趨複雜，單一GPU已無法滿足訓練需求，分散式訓練成為深度學習領域不可或缺的技術。PyTorch 的 DDP 策略提供了一種有效的多 GPU 訓練方案，其透過將模型和資料分散至各個 GPU，平行計算梯度後再進行平均，從而顯著縮短訓練時間。然而，GPU 間的通訊 overhead 和資料同步仍是潛在的效能瓶頸。
分析 DDP 的實作細節，可以發現DistributedSampler 的運用是確保每個 GPU 接收不重疊資料的關鍵，有效避免模型副本發散。此外，torch.distributed 模組的 init_process_group 和 destroy_process_group 函式則負責初始化和釋放分散式訓練環境，而 NCCL 後端則負責 GPU 間的高效通訊。值得注意的是，透過 CUDA_VISIBLE_DEVICES 環境變數，開發者可以彈性地選擇參與訓練的 GPU，簡化了資源管理。
展望未來，隨著硬體技術的發展和演算法的最佳化，分散式訓練的效率將進一步提升。更精細的資料分割策略、更高效的通訊協定以及更智慧的梯度同步演算法將成為未來的研究熱點。預期未來幾年內，更易用且效能更佳的分散式訓練框架將會出現，進一步降低深度學習的開發門檻。
玄貓認為，熟練掌握 PyTorch 的 DDP 策略以及相關的環境設定技巧，對於提升深度學習模型的訓練效率至關重要。開發者應關注通訊 overhead 的最佳化以及不同分散式訓練策略的比較，才能在實務應用中取得最佳效能。
