深度學習模型和資料集的規模增長速度已經讓 GPU 算力也開始捉襟見肘，如果你的 GPU 連一個樣本都容不下，你要如何訓練大批次模型？透過本文介紹的方法，我們可以在訓練批次甚至單個訓練樣本大於 GPU 記憶體時，在單個或多個 GPU 伺服器上訓練模型。
2018 年的大部分時間我都在試圖訓練神經網路時克服 GPU 極限。無論是在含有 1.5 億個引數的語言模型（如 OpenAI 的大型生成預訓練 Transformer 或最近類似的 BERT 模型）還是饋入 3000 萬個元素輸入的元學習神經網路（如我們在一篇 ICLR 論文《Meta-Learning a Dynamical Language Model》中提到的模型），我都只能在 GPU 上處理很少的訓練樣本。
但在多數情況下，隨機梯度下降演算法需要很大批次才能得出不錯的結果。
如果你的 GPU 只能處理很少的樣本，你要如何訓練大批次模型？
有幾個工具、技巧可以幫助你解決上述問題。在本文中，我將自己用過、學過的東西整理出來供大家參考。
在這篇文章中，我將主要討論 PyTorch 框架。有部分工具尚未包括在 PyTorch（1.0 版本）中，因此我也寫了自定義程式碼。
我們將著重探討以下問題：
在訓練批次甚至單個訓練樣本大於 GPU 記憶體，要如何在單個或多個 GPU 伺服器上訓練模型；
如何儘可能高效地利用多 GPU 機器；
在分散式裝置上使用多個機器的最簡單訓練方法。
在一個或多個 GPU 上訓練大批次模型
你建的模型不錯，在這個簡潔的任務中可能成為新的 SOTA，但每次嘗試在一個批次處理更多樣本時，你都會得到一個 CUDA RuntimeError：記憶體不足。
這位網友指出了你的問題！
但你很確定將批次加倍可以最佳化結果。
你要怎麼做呢？
這個問題有一個簡單的解決方法：梯度累積。
梯度下降最佳化演算法的五個步驟。
與之對等的 PyTorch 程式碼也可以寫成以下五行：
predictions = model(inputs) # Forward pass
loss = loss_function(predictions, labels) # Compute loss function
loss.backward() # Backward pass
optimizer.step() # Optimizer step
predictions = model(inputs) # Forward pass with new parameters
在 loss.backward() 運算期間，為每個引數計算梯度，並將其儲存在與每個引數相關聯的張量——parameter.grad 中。
累積梯度意味著，在呼叫 optimizer.step() 實施一步梯度下降之前，我們會對 parameter.grad 張量中的幾個反向運算的梯度求和。在 PyTorch 中這一點很容易實現，因為梯度張量在不呼叫 model.zero_grad() 或 optimizer.zero_grad() 的情況下不會重置。如果損失在訓練樣本上要取平均，我們還需要除以累積步驟的數量。
以下是使用梯度累積訓練模型的要點。在這個例子中，我們可以用一個大於 GPU 最大容量的 accumulation_steps 批次進行訓練：
model.zero_grad() # Reset gradients tensors for i, (inputs, labels) in enumerate(training_set): predictions = model(inputs) # Forward pass loss = loss_function(predictions, labels) # Compute loss function loss = loss / accumulation_steps # Normalize our loss (if averaged) loss.backward() # Backward pass if (i+1) % accumulation_steps == 0: # Wait for several backward steps optimizer.step() # Now we can do an optimizer step model.zero_grad() # Reset gradients tensors if (i+1) % evaluation_steps == 0: # Evaluate the model when we... evaluate_model() # ...have no gradients accumulated擴充套件到極致
你可以在 GPU 上訓練連一個樣本都無法載入的模型嗎？
如果你的架構沒有太多跳過連線，這就是可能的！解決方案是使用梯度檢查點（gradient-checkpointing）來節省計算資源。
基本思路是沿著模型將梯度在小元件中進行反向傳播，以額外的前饋傳遞為代價，節約儲存完整的反向傳播圖的記憶體。這個方法比較慢，因為我們需要新增額外的計算來減少記憶體要求，但在某些設定中挺有意思，比如在非常長的序列上訓練 RNN 模型（示例參見 https://medium.com/huggingface/from-zero-to-research-an-introduction-to-meta-learning-8e16e677f78a）。
這裡不再贅述，讀者可以檢視以下連結：
TensorFlow：https://github.com/openai/gradient-checkpointing
PyTorch 文件：https://pytorch.org/docs/stable/checkpoint.html
「節約記憶體」（Memory-poor）策略需要 O(1) 的記憶體（但是要求 O(n²) 的計算步）。
充分利用多 GPU 機器
現在我們具體來看如何在多 GPU 上訓練模型。
在多 GPU 伺服器上訓練 PyTorch 模型的首選策略是使用 torch.nn.DataParallel。該容器可以在多個指定裝置上分割輸入，按照批維度（batch dimension）分割，從而實現模組應用的並行化。
DataParallel 非常容易使用，我們只需新增一行來封裝模型：
parallel_model = torch.nn.DataParallel(model) # Encapsulate the model
predictions = parallel_model(inputs) # Forward pass on multi-GPUs
loss = loss_function(predictions, labels) # Compute loss function
loss.backward() # Backward pass
optimizer.step() # Optimizer step
predictions = parallel_model(inputs) # Forward pass with new parameters
但是，DataParallel 有一個問題：GPU 使用不均衡。
在一些設定下，GPU-1 會比其他 GPU 使用率高得多。
這個問題從何而來呢？下圖很好地解釋了 DataParallel 的行為：
使用 torch.nn.DataParallel 的前向和後向傳播。
在前向傳播的第四步（右上），所有平行計算的結果都聚集在 GPU-1 上。這對很多分類問題來說是件好事，但如果你在大批次上訓練語言模型時，這就會成為問題。
我們可以快速計算語言模型輸出的大小：
語言模型輸出中的元素數量。
假設我們的資料集有 4 萬詞彙，每一條序列有 250 個 token、每個 batch 中有 32 條序列，那麼序列中的每一個元素需要 4 個位元組的記憶體空間，模型的輸出大概為 1.2GB。要儲存相關的梯度張量，我們就需要把這個記憶體翻倍，因此我們的模型輸出需要 2.4GB 的記憶體。
這是典型 10GB GPU 記憶體的主要部分，意味著相對於其它 GPU，GPU - 1 會被過度使用，從而限制了並行化的效果。
如果不調整模型和／或最佳化方案，我們就無法輕易減少輸出中的元素數量。但我們可以確保記憶體負載在 GPU 中更均勻地分佈。
多 GPU 機器上的均衡負載
解決辦法是把每部分輸出保留在其 GPU 上，而不是將它們聚集到 GPU-1 上。我們也需要分配損失標準計算，計算損失並進行反向傳播。
幸而，張航開源了一個名為 PyTorch-Encoding 的 PyTorch 包，它包含了這些定製的並行化功能。
我提取並稍稍改動了這個模組，你可以從以下地址下載 gist（parallel.py）來納入並呼叫你的程式碼。它主要包括兩個模組：DataParallelModel 和 DataParallelCriterion，它們的用途如下：
下載地址：https://gist.github.com/thomwolf/7e2407fbd5945f07821adae3d9fd1312
from parallel import DataParallelModel, DataParallelCriterion
parallel_model = DataParallelModel(model) # Encapsulate the model
parallel_loss = DataParallelCriterion(loss_function) # Encapsulate the loss function
predictions = parallel_model(inputs) # Parallel forward pass
# "predictions" is a tuple of n_gpu tensors
loss = parallel_loss(predictions, labels) # Compute loss function in parallel
loss.backward() # Backward pass
optimizer.step() # Optimizer step
predictions = parallel_model(inputs) # Parallel forward pass with new parameters
DataParallelModel 和 torch.nn.DataParallel 的區別在於，前向傳播的輸出（predictions）沒有聚集在 GPU-1 上，而是作為 n_gpu 張量的元組，每個張量分佈在相應的 GPU 上。
DataParallelCriterion 容器封裝了損失函式，並把 n_gpu 張量元組和目標標籤張量作為輸入。它在每個 GPU 上平行計算損失函式，像 DataParallel 分割模型輸入一樣分割目標標籤張量。
下圖說明了 DataParallelModel/DataParallelCriterion 的內部情況：
使用 DataParallelModel 和 DataParallelCriterion。
以下是你可能會遇到的兩個特定案例的解決辦法：
你的模型輸出幾個張量：你可能想分解它們：output_1, output_2 = zip(*predictions)
有時候你並不想使用並行損失函式：收集 CPU 上的所有張量：gathered_predictions = parallel.gather(predictions)
分散式訓練：在多臺機器上訓練
在更大的批次上訓練時，我們要如何控制多個伺服器的算力呢？
最簡單的選擇是使用 PyTorch 的 DistributedDataParallel，它幾乎可以說是以上討論的 DataParallel 的直接替代元件。
但要注意：儘管程式碼看起來很相似，但在分散式設定中訓練模型要改變工作流程，因為你必須在每個節點上啟動一個獨立的 Python 訓練指令碼。正如我們將看到的，一旦啟動，這些訓練指令碼可以透過使用 PyTorch 分散式後端一起同步化。
在實踐中，這意味著每個訓練指令碼將擁有：
它自己的最佳化器，並在每次迭代中執行一個完整的最佳化步驟，不需要進行引數傳播（DataParallel 中的步驟 2）；
一個獨立的 Python 直譯器：這也將避免 GIL-freeze，這是在單個 Python 直譯器上驅動多個並行執行執行緒時會出現的問題。
當多個並行前向呼叫由單個直譯器驅動時，在前向傳播中大量使用 Python 迴圈/呼叫的模型可能會被 Python 直譯器的 GIL 放慢速度。透過這種設定，DistributedDataParallel 甚至在單臺機器設定中也能很方便地替代 DataParallel。
現在我們直接討論程式碼和用途。
DistributedDataParallel 是建立在 torch.distributed 包之上的，這個包可以為同步分散式運算提供低階原語，並能以不同的效能使用多種後端（tcp、gloo、mpi、nccl）。在這篇文章中，我將選擇一種簡單的開箱即用的方式來使用它，但你應該閱讀文件和 Séb Arnold 寫的教程來深入理解這個模組。
文件：https://pytorch.org/docs/stable/distributed.html
教程：https://pytorch.org/tutorials/intermediate/dist_tuto.html
我們將考慮使用具有兩個 4 - GPU 伺服器（節點）的簡單但通用的設定：
主伺服器（伺服器 1）擁有一個可訪問的 IP 地址和一個用於通訊的開放埠。
改寫 Python 訓練指令碼以適應分散式訓練
首先我們需要改寫指令碼，從而令其可以在每臺機器（節點）上獨立執行。我們將實現完全的分散式訓練，並在每個節點的每塊 GPU 上執行一個獨立的程式，因此總共需要 8 個程式。
我們的訓練指令碼有點長，因為需要為同步化初始化分散式後端，封裝模型並準備資料，以在資料的一個子集上來訓練每個程式（每個程式都是獨立的，因此我們需要自行處理）。以下是更新後的程式碼：
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
# Each process runs on 1 GPU device specified by the local_rank argument.
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
# Encapsulate the model on the GPU assigned to the current process
device = torch.device('cuda', arg.local_rank)
model = model.to(device)
distrib_model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
# Restricts data loading to a subset of the dataset exclusive to the current process
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler)
for inputs, labels in dataloader:
predictions = distrib_model(inputs.to(device)) # Forward pass
loss = loss_function(predictions, labels.to(device)) # Compute loss function
loss.backward() # Backward pass
optimizer.step() # Optimizer step
啟動 Python 訓練指令碼的多個例項
我們就快完成了，只需要在每個伺服器上啟動訓練指令碼的一個例項。
為了執行指令碼，我們將使用 PyTorch 的 torch.distributed.launch 工具。它將用來設定環境變數，並用正確的 local_rank 引數呼叫每個指令碼。
第一臺機器是最主要的，它應該對於所有其它機器都是可訪問的，因此擁有一個可訪問的 IP 地址（我們的案例中是 192.168.1.1）以及一個開放埠（在我們的案例中是 1234）。在第一臺機器上，我們使用 torch.distributed.launch 來執行訓練指令碼：
python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script) # Optimizer step
在第二臺機器上，我們類似地啟動指令碼：
python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)
這兩個命令是相同的，除了—node_rank 引數，其在第一臺機器上被設為 0，在第二臺機器上被設為 1（如果再加一臺機器，則設為 2，以此類推…）。