在訓練神經網路模型時，我們更偏愛大規模的資料集和複雜的網路結構。更大規模的資料集和更復雜的網路結構可以讓我們的模型表徵能力更強，但同時也對計算的時間和空間提出了挑戰。
序言
如果我們只用單機單GPU來跑，則會出現一卡有難（執行時間長、視訊記憶體不足等），十卡圍觀的狀態。很多神經網路庫都提供了分散式訓練的API，但是如果不瞭解內在機理，我們仍然很難得到滿意的效率和效果。
首先需要明確的是神經網路模型的訓練過程中，哪些部分可以並行，哪些部分不能並行。
對於一個 Batch 中的第 i 個樣例xi 來說，透過 forward 求得 lossi 再透過 backward 求得梯度 Gi 的過程是相互獨立的，不同樣例可以同時進行計算。但是，當所有樣例的 backward 完成後，我們需要求再使用
更新訓練引數，這個過程依賴所有樣例的計算結果 Gi ，不能並行。
Parameter Server
模型的訓練可以分成 5 步：
1. 在訓練前將資料和初始化引數載入進 GPU0 中，如果無法一次載入進來也可以分片載入；
2. 引數伺服器 GPU0 將一個 Batch 的資料切成 3 份，交給工作伺服器 GPU1、GPU2、GPU3（Map）；
3. 引數伺服器 GPU0 將模型（引數）複製 3 份，交給工作伺服器 GPU1、GPU2、GPU3（Map）；
4. 工作伺服器 GPU1、GPU2、GPU3 利用資料和模型求得 loss 和梯度；
5. 將梯度求平均，在引數伺服器 GPU0 更新引數（Reducer），並回到第二步（因為 GPU 中已經有資料了，所以不需要再進行第一步）；
隨著幾個工作伺服器 GPU 的增多，我們所需要的訓練時間會越來越短。在 Pytorch 中，我們只需要這樣呼叫 API，就可以實現 PS 並行：
import torchimport torch.nn as nn# Define dataset...train_dataset = ...train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)# 注意這裡模型所在的 GPU 應與 device_ids[0] 和 output_device 一致model = ...model_dp = nn.DataParallel(model.cuda(), device_ids=[0, 1, 2], output_device=0)# 注意要透過 module 獲取模型例項optimizer = optim.SGD(model_dp.module.parameters())for epoch in range(100):for batch_idx, (data, target) in enumerate(train_loader):optimizer.zero_grad()output = model(data)loss = F.nll_loss(output, target)loss.backward()optimizer.step()if batch_idx % args.log_interval == 0:print('Train Epoch: {} [{}/{}]\tLoss: {}'.format(epoch, batch_idx * len(data), len(train_dataset), loss.item()))
左右滑動看全部程式碼
然而，GPU 之間頻寬有限、資料傳輸很耗時。常常跑圖 10 分鐘，打怪 30 秒。Mapper 和 Reducer 都需要佔用大量頻寬。如果是 15M/s 的網路傳輸速度，GPU0 Mapper 完 1G 的資料和模型就要一分多鐘。因此，我們需要一種機制分散 GPU0 的傳輸壓力。
Ring Allreduce
Ring Allreduce 包含 5 個平等的工作 GPU（GPU1、GPU2、GPU3、GPU4、GPU5），第 i 個 GPU 可以從第 i-1 個 GPU 接收資料，可以向第 i+1 個 GPU 傳送資料，構成一個環（所以叫 ring）。
訓練前將資料切片放入 GPU1、GPU2、GPU3、GPU4、GPU5 中，將初始化引數複製到 GPU1、GPU2、GPU3、GPU4、GPU5 中。
GPU 計算模型的梯度 Gi 時，根據模型的層數 j 將 Gi 分成 j 份。我們求也就是在求每一個
。對於一個 5 層的網路，我們可以寫成這樣：
從 i=1 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度 Gii ， GPUi+1 在收到梯度後求和（所以叫 all reduce，每一個 GPU 都是一個 reducer）。
從 i=2 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度 Gii ， GPUi+1 在收到梯度後求和。
從 i=3 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度Gii ，GPUi+1 在收到梯度後求和。
從 i=4 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度Gii ，GPUi+1 在收到梯度後求和。
上述 reduce 過程完成了求和，我們還需要將求和的結果同步到所有 GPU 中。從 i=1 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度 Gii ， 在收到梯度後使用收到的梯度覆蓋當前梯度（更換了一種 reduce 操作）。
從 i=2 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度 Gii ， GPUi+1 在收到梯度後使用收到的梯度覆蓋當前梯度。
從 i=3 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度 Gii ， GPUi+1 在收到梯度後使用收到的梯度覆蓋當前梯度。
從 i=4 開始，GPUi 向 GPUi+1 傳送第 i 層的梯度 Gii ，GPUi+1 在收到梯度後使用收到的梯度覆蓋當前梯度。
最後只需要再除以 N（N 為 GPU 個數）得到的結果就是 了。利用 Ring Allreduce ，我們仍然可以像 Parameter Server 一樣在 2 x（N-1）次傳輸後完成引數的更新過程。
但是，每次傳輸不再是針對中心節點，而是分散在各節點的兩兩之間，大大減小了對頻寬的壓力，實現了頻寬並行。
在 Pytorch 中，我們只需要這樣呼叫 Horovod API 就可以實現 Ring Allreduce：
import torchimport horovod.torch as hvd# 初始化hvd.init()# 設定可用 GPUtorch.cuda.set_device(hvd.local_rank())# Define dataset...train_dataset = ...# 資料切片分給幾個 GPUtrain_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=hvd.size(), rank=hvd.rank())train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)model = ...model.cuda()# 初始化引數複製給幾個 GPUhvd.broadcast_parameters(model.state_dict(), root_rank=0)optimizer = optim.SGD(model.parameters())# 每個 GPU 完成前饋和反饋、all reduce 完成計算平均梯度、更新引數optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())for epoch in range(100):for batch_idx, (data, target) in enumerate(train_loader):optimizer.zero_grad()output = model(data)loss = F.nll_loss(output, target)loss.backward()optimizer.step()if batch_idx % args.log_interval == 0:print('Train Epoch: {} [{}/{}]\tLoss: {}'.format(epoch, batch_idx * len(data), len(train_sampler), loss.item()))
但在實際的應用過程中，我們可能還需要對 Embedding 或是 BatchNormalize 的並行進行定製化處理，否則會對模型的梯度更新有所影響，進而導致 train 和 val 的準確之間有較大的 gap。