[深度][PyTorch] DDP系列第三篇：實戰與技巧

轉自：https://zhuanlan.zhihu.com/p/250471767

零. 概覽

想要讓你的PyTorch神經網絡在多卡環境上跑得又快又好？那你definitely需要這一篇！

No one knows DDP better than I do!
– – magic_frog（手動狗頭）

本文是DDP系列三篇（基本原理與入門，實現原理與源代碼解析，實戰）中的第三篇。本系列力求深入淺出，簡單易懂，猴子都能看得懂（誤）。

在過去的兩篇文章里，我們已經對DDP的理論、代碼進行了充分、詳細的介紹，相信大家都已經了然在胸。但是，實踐也是很重要的。正所謂理論聯系實踐，如果只掌握理論而不進行實踐，無疑是紙上談兵。

在這篇文章里，我們通過幾個實戰例子，來給大家介紹一下DDP在實際生產中的應用。希望能對大家有所幫助！

在DDP中引入SyncBN
DDP下的Gradient Accumulation的進一步加速
多機多卡環境下的inference加速
保證DDP性能：確保數據的一致性
和DDP有關的小技巧
1. 控制不同進程的執行順序
2. 避免DDP帶來的冗余輸出

請歡快地開始閱讀吧！

**依賴：**pytorch(gpu)>=1.5，python>=3.6

一. 在DDP中引入SyncBN

在這里插入圖片描述

什么是Batch Normalization(BN)? 這里就不多加以介紹。附上BN文章。接下來，讓我們來深入了解下BN在多級多卡環境上的完整實現：SyncBN。

什么是SyncBN？

SyncBN就是Batch Normalization(BN)。其跟一般所說的普通BN的不同在于工程實現方式：SyncBN能夠完美支持多卡訓練，而普通BN在多卡模式下實際上就是單卡模式。

我們知道，BN中有moving mean和moving variance這兩個buffer，這兩個buffer的更新依賴于當前訓練輪次的batch數據的計算結果。但是在普通多卡DP模式下，各個模型只能拿到自己的那部分計算結果，所以在DP模式下的普通BN被設計為只利用主卡上的計算結果來計算moving mean和moving variance，之后再廣播給其他卡。這樣，實際上BN的batch size就只是主卡上的batch size那么大。當模型很大、batch size很小時，這樣的BN無疑會限制模型的性能。

為了解決這個問題，PyTorch新引入了一個叫SyncBN的結構，利用DDP的分布式計算接口來實現真正的多卡BN。

SyncBN的原理

SyncBN的原理很簡單：SyncBN利用分布式通訊接口在各卡間進行通訊，從而能利用所有數據進行BN計算。為了盡可能地減少跨卡傳輸量，**SyncBN做了一個關鍵的優化，即只傳輸各自進程的各自的 小batch mean和 小batch variance，而不是所有數據。**具體流程請見下面：

前向傳播
1. 在各進程上計算各自的 小batch mean和小batch variance
2. 各自的進程對各自的 小batch mean和小batch variance進行all_gather操作，每個進程都得到s的全局量。
3. 注釋：只傳遞mean和variance，而不是整體數據，可以大大減少通訊量，提高速度。
4. 每個進程分別計算總體mean和總體variance，得到一樣的結果
  1. 注釋：在數學上是可行的，有興趣的同學可以自己推導一下。
5. 接下來，延續正常的BN計算。
  1. 注釋：因為從前向傳播的計算數據中得到的batch mean和batch variance在各卡間保持一致，所以，running_mean和running_variance就能保持一致，不需要顯式地同步了！
后向傳播：和正常的一樣

貼一下關鍵代碼，有興趣的同學可以研究下：pytorch源碼

SyncBN與DDP的關系

**一句話總結，當前PyTorch SyncBN只在DDP單進程單卡模式中支持。**SyncBN用到 all_gather這個分布式計算接口，而使用這個接口需要先初始化DDP環境。

復習一下DDP的偽代碼中的準備階段中的DDP初始化階段

d. 創建管理器reducer，給每個parameter注冊梯度平均的hook。
i. 注釋：這一步的具體實現是在C++代碼里面的，即reducer.h文件。
e. （可能）為可能的SyncBN層做準備

這里有三個點需要注意：

這里的為可能的SyncBN層做準備，實際上就是檢測當前是否是DDP單進程單卡模式，如果不是，會直接停止。
這告訴我們，SyncBN需要在DDP環境初始化后初始化，但是要在DDP模型前就準備好。
為什么當前PyTorch SyncBN只支持DDP單進程單卡模式？
- 從SyncBN原理中我們可以看到，其強依賴了all_gather計算，而這個分布式接口當前是不支持單進程多卡或者DP模式的。當然，不排除未來也是有可能支持的。

怎么用SyncBN？

怎么樣才能在我們的代碼引入SyncBN呢？很簡單：

# DDP init
dist.init_process_group(backend='nccl')# 按照原來的方式定義模型，這里的BN都使用普通BN就行了。
model = MyModel()
# 引入SyncBN，這句代碼，會將普通BN替換成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)# 構造DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

又是熟悉的模樣，像DDP一樣，一句代碼就解決了問題。這是怎么做到的呢？

convert_sync_batchnorm的原理：

torch.nn.SyncBatchNorm.convert_sync_batchnorm會搜索model里面的每一個module，如果發現這個module是、或者繼承了torch.nn.modules.batchnorm._BatchNorm類，就把它替換成SyncBN。也就是說，如果你的Normalization層是自己定義的特殊類，沒有繼承過_BatchNorm類，那么convert_sync_batchnorm是不支持的，需要你自己實現一個新的SyncBN！

下面給一下convert_sync_batchnorm的源碼，可以看到convert的過程中，新的SyncBN復制了原來的BN層的所有參數：

    @classmethoddef convert_sync_batchnorm(cls, module, process_group=None):r"""Helper function to convert all :attr:`BatchNorm*D` layers in the model to:class:`torch.nn.SyncBatchNorm` layers."""module_output = moduleif isinstance(module, torch.nn.modules.batchnorm._BatchNorm):module_output = torch.nn.SyncBatchNorm(module.num_features,module.eps, module.momentum,module.affine,module.track_running_stats,process_group)if module.affine:with torch.no_grad():module_output.weight = module.weightmodule_output.bias = module.biasmodule_output.running_mean = module.running_meanmodule_output.running_var = module.running_varmodule_output.num_batches_tracked = module.num_batches_trackedfor name, child in module.named_children():module_output.add_module(name, cls.convert_sync_batchnorm(child, process_group))del modulereturn module_output

二. DDP下的Gradient Accumulation的進一步加速

什么是Gradient Accmulation？

在這里插入圖片描述

Gradient Accumulation，即梯度累加，相信大家都有所了解，是一種增大訓練時batch size的技術，造福了無數硬件條件窘迫的我等窮人。不了解的同學請看這個知乎鏈接。

為什么還能進一步加速？

我們仔細思考一下DDP下的gradient accumulation。

# 單卡模式，即普通情況下的梯度累加
for 每次梯度累加循環optimizer.zero_grad()for _ in range(K):prediction = model(data)loss = loss_fn(prediction, label) / K  # 除以K，模仿loss function中的batchSize方向上的梯度平均，如果本身就沒有的話則不需要。loss.backward()  # 積累梯度，不應用梯度改變optimizer.step()  # 應用梯度改變

我們知道，DDP的gradient all_reduce階段發生在loss_fn(prediction, label).backward()。這意味著，在梯度累加的情況下，假設一次梯度累加循環有K個step，每次梯度累加循環會進行K次 all_reduce！但事實上，每次梯度累加循環只會有一次 optimizer.step()，即只應用一次參數修改，這意味著在每一次梯度累加循環中，我們其實只要進行一次gradient all_reduce即可滿足要求，有K-1次 all_reduce**被浪費了！**而每次 all_reduce的時間成本是很高的！

如何加速

**解決問題的思路在于，對前K-1次step取消其梯度同步。**幸運的是，DDP給我們提供了一個暫時取消梯度同步的context函數 no_sync()（源代碼）。在這個context下，DDP不會進行梯度同步。

所以，我們可以這樣實現加速：

model = DDP(model)for 每次梯度累加循環optimizer.zero_grad()# 前accumulation_step-1個step，不進行梯度同步，累積梯度。for _ in range(K-1)::with model.no_sync():prediction = model(data)loss = loss_fn(prediction, label) / Kloss.backward()  # 積累梯度，不應用梯度改變# 第K個step，進行梯度同步prediction = model(data)loss = loss_fn(prediction, label) / Kloss.backward()  # 積累梯度，不應用梯度改變optimizer.step()

給一個優雅寫法（同時兼容單卡、DDP模式哦）：

from contextlib import nullcontext
# 如果你的python版本小于3.7，請注釋掉上面一行，使用下面這個：
# from contextlib import suppress as nullcontextif local_rank != -1:model = DDP(model)optimizer.zero_grad()
for i, (data, label) in enumerate(dataloader):# 只在DDP模式下，輪數不是K整數倍的時候使用no_syncmy_context = model.no_sync if local_rank != -1 and i % K != 0 else nullcontextwith my_context():prediction = model(data)loss = loss_fn(prediction, label) / Kloss.backward()  # 積累梯度，不應用梯度改變if i % K == 0:optimizer.step()optimizer.zero_grad()

是不是很漂亮！

三. 多機多卡環境下的inference加速

問題

有一些非常現實的需求，相信大家肯定碰到過：

一般，訓練中每幾個epoch我們會跑一下inference、測試一下模型性能。在DDP多卡訓練環境下，能不能利用多卡來加速inference速度呢？
我有一堆數據要跑一些網絡推理，拿到inference結果。DP下多卡加速比太低，能不能利用DDP多卡來加速呢？

解法

這兩個問題實際是同一個問題。答案肯定是可以的，但是，沒有現成、省力的方法。

測試和訓練的不同在于：

測試的時候不需要進行梯度反向傳播，inference過程中各進程之間不需要通訊。
測試的時候，不同模型的inference結果、性能指標的類型多種多樣，沒有統一的形式。
1. 我們很難定義一個統一的框架，像訓練時model=DDP(model)那樣方便地應用DDP多卡加速。

解決問題的思路很簡單，就是各個進程中各自進行單卡的inference，然后把結果收集到一起。單卡inference很簡單，我們甚至可以直接用DDP包裝前的模型。問題其實只有兩個：

我們要如何把數據split到各個進程中
我們要如何把結果合并到一起

如何把數據split到各個進程中：新的data sampler

大家肯定還記得，在訓練的時候，我們用的 torch.utils.data.distributed.DistributedSampler幫助我們把數據不重復地分到各個進程上去。但是，其分的方法是：每段連續的N個數據，拆成一個一個，分給N個進程，所以每個進程拿到的數據不是連續的。這樣，不利于我們在inference結束的時候將結果合并到一起。

所以，這里我們需要實現一個新的data sampler。它的功能，是能夠連續地劃分數據塊，不重復地分到各個進程上去。直接給代碼：

# 來源：https://github.com/huggingface/transformers/blob/447808c85f0e6d6b0aeeb07214942bf1e578f9d2/src/transformers/trainer_pt_utils.py
class SequentialDistributedSampler(torch.utils.data.sampler.Sampler):"""Distributed Sampler that subsamples indicies sequentially,making it easier to collate all results at the end.Even though we only use this sampler for eval and predict (no training),which means that the model params won't have to be synced (i.e. will not hangfor synchronization even if varied number of forward passes), we still add extrasamples to the sampler to make it evenly divisible (like in `DistributedSampler`)to make it easy to `gather` or `reduce` resulting tensors at the end of the loop."""def __init__(self, dataset, batch_size, rank=None, num_replicas=None):if num_replicas is None:if not torch.distributed.is_available():raise RuntimeError("Requires distributed package to be available")num_replicas = torch.distributed.get_world_size()if rank is None:if not torch.distributed.is_available():raise RuntimeError("Requires distributed package to be available")rank = torch.distributed.get_rank()self.dataset = datasetself.num_replicas = num_replicasself.rank = rankself.batch_size = batch_sizeself.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.batch_size / self.num_replicas)) * self.batch_sizeself.total_size = self.num_samples * self.num_replicasdef __iter__(self):indices = list(range(len(self.dataset)))# add extra samples to make it evenly divisibleindices += [indices[-1]] * (self.total_size - len(indices))# subsampleindices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]return iter(indices)def __len__(self):return self.num_samples

如何把結果合并到一起: all_gather

通過torch.distributed提供的分布式接口all_gather，我們可以把各個進程的prediction結果集中到一起。

難點就在這里。因為世界上存在著千奇百怪的神經網絡模型，有著千奇百怪的輸出，所以，把數據集中到一起不是一件容易的事情。**但是，如果你的網絡輸出在不同的進程中有著一樣的大小，那么這個問題就好解多了。**下面給一個方法，其要求網絡的prediction結果在各個進程中的大小是一模一樣的：

# 合并結果的函數
# 1. all_gather，將各個進程中的同一份數據合并到一起。
#   和all_reduce不同的是，all_reduce是平均，而這里是合并。
# 2. 要注意的是，函數的最后會裁剪掉后面額外長度的部分，這是之前的SequentialDistributedSampler添加的。
# 3. 這個函數要求，輸入tensor在各個進程中的大小是一模一樣的。
def distributed_concat(tensor, num_total_examples):output_tensors = [tensor.clone() for _ in range(torch.distributed.get_world_size())]torch.distributed.all_gather(output_tensors, tensor)concat = torch.cat(output_tensors, dim=0)# truncate the dummy elements added by SequentialDistributedSamplerreturn concat[:num_total_examples]

完整的流程

結合上面的介紹，我們可以得到下面這樣一個完整的流程。

## 構造測試集
# 假定我們的數據集是這個
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
my_testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
# 使用我們的新sampler
test_sampler = SequentialDistributedSampler(my_testset, batch_size=16)
testloader = torch.utils.data.DataLoader(my_testset, batch_size=16, sampler=test_sampler)# DDP和模型初始化，略。
# ......# 正式訓練和evaluation
for epoch in range(total_epoch_size):# 訓練代碼，略# .......# 開始測試with torch.no_grad():# 1. 得到本進程的predictionpredictions = []labels = []for data, label in testloader:data, label = data.to(local_rank), label.to(local_rank)predictions.append(model(data))labels.append(label)# 進行gatherpredictions = distributed_concat(torch.concat(predictions, dim=0), len(test_sampler.dataset))labels = distributed_concat(torch.concat(labels, dim=0), len(test_sampler.dataset))# 3. 現在我們已經拿到所有數據的predictioin結果，進行evaluate！my_evaluate_func(predictions, labels)

更簡化的解法

如果我們的目的只是得到性能數字，那么，我們甚至可以直接在各個進程中計算各自的性能數字，然后再合并到一起。上面給的解法，是為了更通用的情景。一切根據你的需要來定！
我們可以單向地把predictions、labels集中到 rank=0的進程，只在其進行evaluation并輸出。PyTorch也提供了相應的接口（鏈接，send和recv）。

四. 保證DDP性能：確保數據的一致性

性能期望

從原理上講，當沒有開啟SyncBN時，（或者更嚴格地講，沒有BN層；但一般有的話影響也不大），以下兩種方法訓練出來的模型應該是性能相似的：

進程數為N的DDP訓練
accumulation為N、其他配置完全相同的單卡訓練

如果我們發現性能對不上，那么，往往是DDP中的某些設置出了問題。在DDP系列第二篇中，我們介紹過一個check list，可以根據它檢查下自己的配置。其中，在造成性能對不齊的原因中，最有可能的是數據方面出現了問題。

DDP訓練時，數據的一致性必須被保證：各個進程拿到的數據，要像是accumulation為N、其他配置完全相同的單卡訓練中同個accumulation循環中不同iteration拿到的數據。想象一下，如果各個進程拿到的數據是一樣的，或者分布上有任何相似的地方，那么，這就會造成訓練數據質量的下降，最終導致模型性能下降。

容易錯的點：隨機數種子

為保證實驗的可復現性，一般我們會在代碼在開頭聲明一個固定的隨機數種子，從而使得同一個配置下的實驗，無論啟動多少次，都會拿到同樣的結果。

import random
import numpy as np
import torchdef init_seeds(seed=0, cuda_deterministic=True):random.seed(seed)np.random.seed(seed)torch.manual_seed(seed)# Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.htmlif cuda_deterministic:  # slower, more reproduciblecudnn.deterministic = Truecudnn.benchmark = Falseelse:  # faster, less reproduciblecudnn.deterministic = Falsecudnn.benchmark = Truedef main():# 一般都直接用0作為固定的隨機數種子。init_seeds(0)

但是在DDP訓練中，如果還是像以前一樣，使用0作為隨機數種子，不做修改，就會造成以下后果：

DDP的N個進程都使用同一個隨機數種子
在生成數據時，如果我們使用了一些隨機過程的數據擴充方法，那么，各個進程生成的數據會帶有一定的同態性。
1. 比如說，YOLOv5會使用mosaic數據增強（從數據集中隨機采樣3張圖像與當前的拼在一起，組成一張里面有4張小圖的大圖）。這樣，因為各卡使用了相同的隨機數種子，你會發現，各卡生成的圖像中，除了原本的那張小圖，其他三張小圖都是一模一樣的！
同態性的數據，降低了訓練數據的質量，也就降低了訓練效率！最終得到的模型性能，很有可能是比原來更低的。

所以，我們需要給不同的進程分配不同的、固定的隨機數種子：

def main():rank = torch.distributed.get_rank()# 問題完美解決！init_seeds(1 + rank)

五. 和DDP有關的小技巧

控制不同進程的執行順序

一般情況下，各個進程是各自執行的，速度有快有慢，只有在gradient all-reduce的時候，快的進程才會等一下慢的進程，也就是進行同步。那么，如果我們需要在其他地方進行同步呢？比如說，在加載數據前，如果數據集不存在，我們要下載數據集：

我們只需要在唯一一個進程中開啟一次下載
我們需要讓其他進程等待其下載完成，再去加載數據

怎么解決這個問題呢？torch.distributed提供了一個barrier()的接口，利用它我們可以同步各個DDP中的各個進程！當使用barrier函數時，DDP進程會在函數的位置進行等待，知道所有的進程都跑到了 barrier函數的位置，它們才會再次向下執行。

只在某進程執行，無須同步：

這是最簡單的，只需要一個簡單的判斷，用不到barrier()

if rank == 0:code_only_run_in_rank_0()

簡單的同步:

沒什么好講的，只是一個示范

code_before()
# 在這一步同步
torch.distributed.barrier()
code_after()

在某個進程中執行A操作，其他進程等待其執行完成后再執行B操作：

也簡單。

if rank == 0:do_A()torch.distributed.barrier()
else:torch.distributed.barrier()do_B()

在某個進程中優先執行A操作，其他進程等待其執行完成后再執行A操作：

這個值得深入講一下，因為這個是非常普遍的需求。利用contextlib.contextmanager，我們可以把這個邏輯給優雅地包裝起來！

from contextlib import contextmanager@contextmanager
def torch_distributed_zero_first(rank: int):"""Decorator to make all processes in distributed training wait for each local_master to do something."""if rank not in [-1, 0]:torch.distributed.barrier()# 這里的用法其實就是協程的一種哦。yieldif rank == 0:torch.distributed.barrier()

然后我們就可以這樣騷操作：

with torch_distributed_zero_first(rank):if not check_if_dataset_exist():download_dataset()load_dataset()

優雅地解決了需求！

避免DDP帶來的冗余輸出

問題：

當我們在自己的模型中加入DDP模型時，第一的直觀感受肯定是，終端里的輸出變成了N倍了。這是因為我們現在有N個進程在同時跑整個程序。這不光是對有潔癖的同學造成困擾，其實對所有人都會造成困擾。因為各個進程的速度并不一樣快，在茫茫的輸出海洋中，我們難以debug、把控實驗狀態。

解法：

那么，有什么辦法能避免這個現象呢？下面，筆者給一個可行的方法：**logging模塊+輸出信息等級控制。**即用logging輸出代替所有print輸出，并給不同進程設置不同的輸出等級，只在0號進程保留低等級輸出。舉一個例子：

import logging# 給主要進程（rank=0）設置低輸出等級，給其他進程設置高輸出等級。
logging.basicConfig(level=logging.INFO if rank in [-1, 0] else logging.WARN)
# 普通log，只會打印一次。
logging.info("This is an ordinary log.")
# 危險的warning、error，無論在哪個進程，都會被打印出來，從而方便debug。
logging.error("This is a fatal log!")