GPU如何訓(xùn)練大批量模型？方法在這里

　　擴(kuò)展到極致

　　你可以在 GPU 上訓(xùn)練連一個(gè)樣本都無(wú)法加載的模型嗎?

　　如果你的架構(gòu)沒(méi)有太多跳過(guò)連接，這就是可能的!解決方案是使用梯度檢查點(diǎn)(gradient-checkpointing)來(lái)節(jié)省計(jì)算資源。

　　基本思路是沿著模型將梯度在小組件中進(jìn)行反向傳播，以額外的前饋傳遞為代價(jià)，節(jié)約存儲(chǔ)完整的反向傳播圖的內(nèi)存。這個(gè)方法比較慢，因?yàn)槲覀冃枰砑宇~外的計(jì)算來(lái)減少內(nèi)存要求，但在某些設(shè)置中挺有意思，比如在非常長(zhǎng)的序列上訓(xùn)練 RNN 模型(示例參見(jiàn) 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

　　「節(jié)約內(nèi)存」(Memory-poor)策略需要 O(1) 的內(nèi)存(但是要求 O(n2) 的計(jì)算步)。

　　充分利用多 GPU 機(jī)器

　　現(xiàn)在我們具體來(lái)看如何在多 GPU 上訓(xùn)練模型。

　　在多 GPU 服務(wù)器上訓(xùn)練 PyTorch 模型的首選策略是使用 torch.nn.DataParallel。該容器可以在多個(gè)指定設(shè)備上分割輸入，按照批維度(batch dimension)分割，從而實(shí)現(xiàn)模塊應(yīng)用的并行化。

　　DataParallel 非常容易使用，我們只需添加一行來(lái)封裝模型：

　　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 有一個(gè)問(wèn)題：GPU 使用不均衡。

　　在一些設(shè)置下，GPU-1 會(huì)比其他 GPU 使用率高得多。

　　這個(gè)問(wèn)題從何而來(lái)呢?下圖很好地解釋了 DataParallel 的行為：

　　使用 torch.nn.DataParallel 的前向和后向傳播。

　　在前向傳播的第四步(右上)，所有并行計(jì)算的結(jié)果都聚集在 GPU-1 上。這對(duì)很多分類問(wèn)題來(lái)說(shuō)是件好事，但如果你在大批量上訓(xùn)練語(yǔ)言模型時(shí)，這就會(huì)成為問(wèn)題。

　　我們可以快速計(jì)算語(yǔ)言模型輸出的大?。?/p>

　　語(yǔ)言模型輸出中的元素?cái)?shù)量。

　　假設(shè)我們的數(shù)據(jù)集有 4 萬(wàn)詞匯，每一條序列有 250 個(gè) token、每個(gè) batch 中有 32 條序列，那么序列中的每一個(gè)元素需要 4 個(gè)字節(jié)的內(nèi)存空間，模型的輸出大概為 1.2GB。要儲(chǔ)存相關(guān)的梯度張量，我們就需要把這個(gè)內(nèi)存翻倍，因此我們的模型輸出需要 2.4GB 的內(nèi)存。

　　這是典型 10GB GPU 內(nèi)存的主要部分，意味著相對(duì)于其它 GPU，GPU - 1 會(huì)被過(guò)度使用，從而限制了并行化的效果。

　　如果不調(diào)整模型和/或優(yōu)化方案，我們就無(wú)法輕易減少輸出中的元素?cái)?shù)量。但我們可以確保內(nèi)存負(fù)載在 GPU 中更均勻地分布。

　　多 GPU 機(jī)器上的均衡負(fù)載

　　解決辦法是把每部分輸出保留在其 GPU 上，而不是將它們聚集到 GPU-1 上。我們也需要分配損失標(biāo)準(zhǔn)計(jì)算，計(jì)算損失并進(jìn)行反向傳播。

　　幸而，張航開源了一個(gè)名為 PyTorch-Encoding 的 PyTorch 包，它包含了這些定制的并行化功能。

　　我提取并稍稍改動(dòng)了這個(gè)模塊，你可以從以下地址下載 gist(parallel.py)來(lái)納入并調(diào)用你的代碼。它主要包括兩個(gè)模塊：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 的區(qū)別在于，前向傳播的輸出(predictions)沒(méi)有聚集在 GPU-1 上，而是作為 n_gpu 張量的元組，每個(gè)張量分布在相應(yīng)的 GPU 上。

　　DataParallelCriterion 容器封裝了損失函數(shù)，并把 n_gpu 張量元組和目標(biāo)標(biāo)簽張量作為輸入。它在每個(gè) GPU 上并行計(jì)算損失函數(shù)，像 DataParallel 分割模型輸入一樣分割目標(biāo)標(biāo)簽張量。

　　下圖說(shuō)明了 DataParallelModel/DataParallelCriterion 的內(nèi)部情況：

　　使用 DataParallelModel 和 DataParallelCriterion。

　　以下是你可能會(huì)遇到的兩個(gè)特定案例的解決辦法：

　　你的模型輸出幾個(gè)張量：你可能想分解它們：output_1, output_2 = zip(*predictions)

　　有時(shí)候你并不想使用并行損失函數(shù)：收集 CPU 上的所有張量：gathered_predictions = parallel.gather(predictions)