概述

本教程假定你已經(jīng)對于 PyToch 訓(xùn)練一個簡單模型有一定的基礎(chǔ)理解。本教程將展示使用 3 種封裝層級不同的方法調(diào)用 DDP (DistributedDataParallel) 進(jìn)程，在多個 GPU 上訓(xùn)練同一個模型：

• 使用?pytorch.distributed?模塊的原生 PyTorch DDP 模塊

• 使用 ?? Accelerate 對?pytorch.distributed?的輕量封裝，確保程序可以在不修改代碼或者少量修改代碼的情況下在單個 GPU 或 TPU 下正常運行

• 使用 ?? Transformer 的高級 Trainer API ，該 API 抽象封裝了所有代碼模板并且支持不同設(shè)備和分布式場景。

什么是分布式訓(xùn)練，為什么它很重要？

下面是一些非?；A(chǔ)的 PyTorch 訓(xùn)練代碼，它基于 Pytorch 官方在 MNIST 上創(chuàng)建和訓(xùn)練模型的示例。

MNIST 創(chuàng)建和訓(xùn)練模型的示例地址:
https://github.com/pytorch/examples/blob/main/mnist/main.py

我們定義訓(xùn)練設(shè)備 (cuda):

構(gòu)建一些基本的 PyTorch?DataLoaders:

把模型放入 CUDA 設(shè)備:

構(gòu)建 PyTorch?optimizer?(優(yōu)化器)

最終創(chuàng)建一個簡單的訓(xùn)練和評估循環(huán)，訓(xùn)練循環(huán)會使用全部訓(xùn)練數(shù)據(jù)集進(jìn)行訓(xùn)練，評估循環(huán)會計算訓(xùn)練后模型在測試數(shù)據(jù)集上的準(zhǔn)確度：

通常從這里開始，就可以將所有的代碼放入 Python 腳本或在 Jupyter Notebook 上運行它。

然而，只執(zhí)行?python myscript.py?只會使用單個 GPU 運行腳本。如果有多個 GPU 資源可用，您將如何讓這個腳本在兩個 GPU 或多臺機(jī)器上運行，通過分布式訓(xùn)練提高訓(xùn)練速度？這是?torch.distributed?發(fā)揮作用的地方。

PyTorch 分布式數(shù)據(jù)并行

顧名思義，torch.distributed?旨在配置分布式訓(xùn)練。你可以使用它配置多個節(jié)點進(jìn)行訓(xùn)練，例如：多機(jī)器下的單個 GPU，或者單臺機(jī)器下的多個 GPU，或者兩者的任意組合。

為了將上述代碼轉(zhuǎn)換為分布式訓(xùn)練，必須首先定義一些設(shè)置配置，具體細(xì)節(jié)請參閱 DDP 使用教程

DDP 使用教程地址:
https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

首先必須聲明?setup?和?cleanup?函數(shù)。這將創(chuàng)建一個進(jìn)程組，并且所有計算進(jìn)程都可以通過這個進(jìn)程組通信。

?

注意：在本教程的這一部分中，假定這些代碼是在 Python 腳本文件中啟動。稍后將討論使用 ?? Accelerate 的啟動器，就不必聲明?setup?和?cleanup?函數(shù)了

最后一個疑問是，我怎樣把我的數(shù)據(jù)和模型發(fā)送到另一個 GPU 上？

這正是?DistributedDataParallel?模塊發(fā)揮作用的地方， 它將您的模型復(fù)制到每個 GPU 上 ，并且當(dāng)?loss.backward()?被調(diào)用進(jìn)行反向傳播的時候，所有這些模型副本的梯度將被同步地平均/下降 (reduce)。這確保每個設(shè)備在執(zhí)行優(yōu)化器步驟后具有相同的權(quán)重。

下面是我們的訓(xùn)練設(shè)置示例，我們使用了?DistributedDataParallel?重構(gòu)了訓(xùn)練函數(shù)：

?

注意：此處的 rank 是當(dāng)前 GPU 與所有其他可用 GPU 相比的總體 rank，這意味著它們的 rank 為?0 -> n-1

在上述的代碼中需要為每個副本設(shè)備上的模型 (因此在這里是ddp_model的參數(shù)而不是?model?的參數(shù)) 聲明優(yōu)化器，以便正確計算每個副本設(shè)備上的梯度。

最后，要運行腳本，PyTorch 有一個方便的?torchrun?命令行模塊可以提供幫助。只需傳入它應(yīng)該使用的節(jié)點數(shù)以及要運行的腳本即可：

torchrun?--nproc_per_nodes=2?--nnodes=1?example_script.py


上面的代碼可以在在一臺機(jī)器上的兩個 GPU 上運行訓(xùn)練腳本，這是使用 PyTorch 只進(jìn)行分布式訓(xùn)練的情況 (不可以在單機(jī)單卡上運行)。

現(xiàn)在讓我們談?wù)??? Accelerate，一個旨在使并行化更加無縫并有助于一些最佳實踐的庫。

?? Accelerate

?? Accelerate 是一個庫，旨在無需大幅修改代碼的情況下完成并行化。除此之外，?? Accelerate 附帶的數(shù)據(jù)?pipeline?還可以提高代碼的性能。

?? Accelerate 文檔:
https://hf.co/docs/accelerate

首先，讓我們將剛剛執(zhí)行的所有上述代碼封裝到一個函數(shù)中，以幫助我們直觀地看到差異：

接下來讓我們談?wù)??? Accelerate 如何便利地實現(xiàn)并行化的。上面的代碼有幾個問題：

1. 該代碼有點低效，因為每個設(shè)備都會創(chuàng)建一個?dataloader。

2. 這些代碼只能運行在多 GPU 下，當(dāng)想讓這個代碼運行在單個 GPU 或 TPU 時，還需要額外進(jìn)行一些修改。

?? Accelerate 通過?Accelerator?類解決上述問題。通過它，不論是單節(jié)點還是多節(jié)點，除了三行代碼外，其余代碼幾乎保持不變，如下所示：

?? Accelerate 文檔鏈接：
https://hf.co/docs/accelerate/v0.12.0/en/package_reference/accelerator

借助?Accelerator?對象，您的 PyTorch 訓(xùn)練循環(huán)現(xiàn)在已配置為可以在任何分布式情況運行。使用?Accelerator?改造后的代碼仍然可以通過?torchrun CLI?或通過 ?? Accelerate 自己的 CLI 界面啟動(啟動你的?? Accelerate 腳本)。

通過 CLI 方式啟動的幫助文檔:
https://hf.co/docs/accelerate/v0.12.0/en/basic_tutorials/launch

因此，現(xiàn)在可以盡可能保持 PyTorch 原生代碼不變的前提下，使用 ?? Accelerate 執(zhí)行分布式訓(xùn)練。

早些時候有人提到 ?? Accelerate 還可以使?DataLoaders?更高效。這是通過自定義采樣器實現(xiàn)的，它可以在訓(xùn)練期間自動將部分批次發(fā)送到不同的設(shè)備，從而允許每個設(shè)備只需要儲存數(shù)據(jù)的一部分，而不是一次將數(shù)據(jù)復(fù)制四份存入內(nèi)存，具體取決于配置。因此，內(nèi)存總量中只有原始數(shù)據(jù)集的一個完整副本。該數(shù)據(jù)集會拆分后分配到各個訓(xùn)練節(jié)點上，從而允許在單個實例上訓(xùn)練更大的數(shù)據(jù)集，而不會使內(nèi)存爆炸

使用?notebook_launcher

之前提到您可以直接從 Jupyter Notebook 運行分布式代碼。這來自 ?? Accelerate 的?notebook_launcher?模塊，它可以在 Jupyter Notebook 內(nèi)部的代碼啟動多 GPU 訓(xùn)練。

Notebook Launcher 使用幫助:
https://hf.co/docs/accelerate/v0.12.0/en/basic_tutorials/notebook

使用它就像導(dǎo)入?launcher?一樣簡單：

接著傳遞我們之前聲明的訓(xùn)練函數(shù)、要傳遞的任何參數(shù)以及要使用的進(jìn)程數(shù)（例如 TPU 上的 8 個，或兩個 GPU 上的 2 個）。下面兩個訓(xùn)練函數(shù)都可以運行，但請注意，啟動單次啟動后，實例需要重新啟動才能產(chǎn)生另一個：

或者：

使用 ?? Trainer

終于我們來到了最高級的 API——Hugging Face Trainer.

Hugging Face ?? Transformers API 中?Trainer?文檔地址:
https://hf.co/docs/transformers/main_classes/trainer

它涵蓋了盡可能多的訓(xùn)練類型，同時仍然能夠在分布式系統(tǒng)上進(jìn)行訓(xùn)練，用戶根本不需要做任何事情。

首先我們需要導(dǎo)入 ?? Trainer：

然后我們定義一些?TrainingArguments?來控制所有常用的超參數(shù)。?? Trainer 需要的訓(xùn)練數(shù)據(jù)是字典類型的，因此需要制作自定義整理功能。

最后，我們將訓(xùn)練器子類化并編寫我們自己的?compute_loss.

之后，這段代碼也可以分布式運行，而無需修改任何訓(xùn)練代碼！

Epoch訓(xùn)練損失驗證損失10.8757000.282633

與上面的?notebook_launcher?示例類似，也可以將這個過程封裝成一個訓(xùn)練函數(shù)：

相關(guān)資源

• 要了解有關(guān) PyTorch 分布式數(shù)據(jù)并行性的更多信息，請查看:?
  https://pytorch.org/docs/stable/distributed.html

• 要了解有關(guān) ?? Accelerate 的更多信息，請查看:?
  https://hf.co/docs/accelerate

• 要了解有關(guān) ?? Transformer 的更多信息，請查看:?
  https://hf.co/docs/transformers

原文作者：Zachary Mueller

譯者：innovation64 (李洋)

審校：yaoqi (胡耀淇)

排版：zhongdongy (阿東)