第一篇 pytorch分散式訓練^[1]

在pytorch 1.6.0，torch.distributed中的功能可以分為三個主要的元件：

• Distributed Data-Parallel Training（DDP）是一個廣泛採用的單程式多資料訓練方法。使用DDP，模型會被複制到每個程序，然後每個模型副本會被輸入資料樣本的不同子集。DDP負責梯度通訊以保持模型副本的同步，並將其與梯度計算重疊以加快訓練速度。
• RPC-Based Distributed Training（RPC）是為了支援那些不適用資料並行訓練的通用訓練架構，比如分散式管道並行，引數伺服器正規化，以及DDP和其他訓練方法的結合。它越過機器邊界來幫助管理遠端物件生命週期和擴充套件自動梯度引擎。
• Collective Communication（c10d）庫支援在組內程序之間傳送張量，它同時提供了聚合通訊（collective communication）介面（比如all_reduce和all_gather）和點對點通訊（P2P）介面（比如send和isend）。在1.6.0中DDP和RPC（ProcessGroup Backend）都是基於c10d構建的，前者使用了聚合通訊，而後者使用了P2P通訊。一般地，開發者不需要直接使用這個原生通訊介面，因為上面的DDP和RPC功能可以服務於很多分散式訓練場景。然而，對於有些使用案例這些介面還是有用的。一個例子就是分散式引數均化（distributed parameter averaging），應用希望在反向傳播後不適用DDP來計算梯度，而是計算模型所有引數的均值。這將把通訊從計算中解耦出來，並允許更加細粒度地控制交流的內容。但是另一方面，這也將放棄DDP提供的效能優化。

資料分散式訓練

pytorch為資料分散式訓練提供了多種選擇。隨著應用從簡單到複雜，從原型到產品，常見的開發軌跡可以是：

• 如果資料和模型能放入單個GPU，使用單裝置訓練，此時不用擔心訓練速度；
• 如果伺服器上有多個GPU，並且你在程式碼修改量最小的情況下加速訓練，使用單個機器多GPU DataParallel；
• 如果你想進一步加速訓練並且願意寫一點程式碼來啟動，使用單個機器多個GPU DistributedDataParallel；
• 如果應用程式跨機器邊界擴充套件，使用多機器DistributedDataParallel和啟動指令碼；
• 如果預期有錯誤（比如OOM）或者資源在訓練過程中可以動態連線和分離，使用torchelastic來啟動分散式訓練。

資料分散式訓練也可以和AMP（自動混合精度）協作。

torch.nn.DataParallel

單機器多GPU並行機制，對程式碼的改動很少

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if torch.cuda.device_count()>1:
    model = nn.DataParallel(model)
model.to(device)

但是需要提醒的是：儘管DataParallel使用簡單，它一般無法提供最好的效能。這是因為DataParallel的實現在每個前向傳播中複製了模型，並且它的單程序多執行緒並行自然存在GIL衝突。為了獲取更好的效能，請考慮使用DistributedDataParallel.

torch.nn.parallel.DistributedDataParallel

和DataParallel相比，DistributedDataParallel需要額外一個稱為init_process_group的步驟啟動，DDP使用多程序並行，因此在模型副本之間不存在GIL衝突。另外，模型在DDP構建時廣播（而不是在每次前向傳播），這有助於加速訓練。DDP配備了若干個效能優化技術。

有以下幾個資料可供參考：

• DDP notes：提供初學者樣例和一些設計實現的簡要描述；
• Getting Started with Distributed Data Parallel：解釋一些DDP訓練的常見問題，包括不均衡負載、checkpointing、多裝置模型。注意，DDP可以簡單地和單機器多裝置模型並行結合，在Single-Machine Model Parallel Best Practices會有介紹；
• Launching and configuring distributed data parallel applications會展示如何使用DDP啟動指令碼。
• Pytorch Distributed Trainer with Amazon AWS會示範如何在AWS使用DDP

這裡的多裝置指的是GPU。

接下來進入Getting Started with Distributed Data Parallel。

第二篇 Getting Started with Distributed Data Parallel^[2]

DDP實現了模組級的資料並行，從而在多個機器上執行。使用DDP的應用會生成多個程序，每個程序對應一個DDP例項。DDP使用torch.distributed包中聚合通訊來同步梯度和快取（buffer）。更具體地說，DDP給model.parameters()指定的每個引數註冊一個autograd的hook，在對應梯度在反向傳播中完成計算後這些hook會被銷燬。然後DDP使用這個訊號來觸發程序間的梯度同步。

使用DDP的推薦方法是給每個模型副本生成一個程序，程序可以放置在單個機器或者多個機器上，一個模型副本可以跨越多個GPU，但是程序間不能共享GPU。

基礎使用案例

為了常見DDP模組，首先啟動程序組。

from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
import os

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

考慮到時效性，這裡刪去了原教程中關於32位電腦的程序組啟動方法。

接下來建立玩具模組，用DDP包起來，並喂一些假輸入資料。請注意，由於DDP會從rank 0程序廣播模型狀態到DDP構造器中的所有其他程序，所以你無需擔心不同DDP程序中模型的引數初始值不同。

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

接下來是DDP的使用：

def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    
    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()

後面有空再繼續更。