模型訓練剖析

為了理解可以應用於提高模型訓練速度和記憶體利用效率的效能最佳化技術，熟悉 GPU 在訓練期間如何被利用以及計算強度如何根據所執行操作而變化是很有幫助的。

讓我們首先透過一個 GPU 利用率和模型訓練執行的激勵示例來開始探索。為了演示，我們需要安裝一些庫。

pip install transformers datasets accelerate nvidia-ml-py3

nvidia-ml-py3 庫允許我們從 Python 內部監視模型的記憶體使用情況。您可能熟悉終端中的 nvidia-smi 命令 - 該庫允許直接在 Python 中訪問相同的資訊。

然後，我們建立一些虛擬資料：100 到 30000 之間的隨機 token ID 和分類器的二進位制標籤。總共，我們得到 512 個序列，每個序列長度為 512，並將它們儲存在 PyTorch 格式的 Dataset 中。

>>> import numpy as np
>>> from datasets import Dataset


>>> seq_len, dataset_size = 512, 512
>>> dummy_data = {
...     "input_ids": np.random.randint(100, 30000, (dataset_size, seq_len)),
...     "labels": np.random.randint(0, 2, (dataset_size)),
... }
>>> ds = Dataset.from_dict(dummy_data)
>>> ds.set_format("pt")

為了使用 Trainer 列印 GPU 利用率和訓練執行的統計摘要，我們定義了兩個輔助函式。

>>> from pynvml import *


>>> def print_gpu_utilization():
...     nvmlInit()
...     handle = nvmlDeviceGetHandleByIndex(0)
...     info = nvmlDeviceGetMemoryInfo(handle)
...     print(f"GPU memory occupied: {info.used//1024**2} MB.")


>>> def print_summary(result):
...     print(f"Time: {result.metrics['train_runtime']:.2f}")
...     print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
...     print_gpu_utilization()

讓我們驗證我們是否從空閒的 GPU 記憶體開始。

>>> print_gpu_utilization()
GPU memory occupied: 0 MB.

看起來不錯：正如我們預期的那樣，在載入任何模型之前，GPU 記憶體沒有被佔用。如果您的機器上不是這種情況，請確保停止所有正在使用 GPU 記憶體的程序。但是，並非所有空閒的 GPU 記憶體都可以被使用者使用。當模型載入到 GPU 時，核心也會被載入，這可能會佔用 1-2GB 的記憶體。為了檢視它有多少，我們將一個微小的張量載入到 GPU 中，這也會觸發核心的載入。

>>> import torch


>>> torch.ones((1, 1)).to("cuda")
>>> print_gpu_utilization()
GPU memory occupied: 1343 MB.

我們看到僅核心就佔用了 1.3GB 的 GPU 記憶體。現在讓我們看看模型使用了多少空間。

載入模型

首先，我們載入 google-bert/bert-large-uncased 模型。我們將模型權重直接載入到 GPU，以便我們可以檢查僅權重佔用了多少空間。

>>> from transformers import AutoModelForSequenceClassification


>>> model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-large-uncased").to("cuda")
>>> print_gpu_utilization()
GPU memory occupied: 2631 MB.

我們可以看到，僅模型權重就佔用了 1.3 GB 的 GPU 記憶體。具體數值取決於您使用的特定 GPU。請注意，在較新的 GPU 上，模型有時會佔用更多空間，因為權重以最佳化的方式載入，從而加快了模型的使用速度。現在我們也可以快速檢查是否獲得了與 nvidia-smi CLI 相同的結果。

nvidia-smi

Tue Jan 11 08:58:05 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.91.03    Driver Version: 460.91.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  On   | 00000000:00:04.0 Off |                    0 |
| N/A   37C    P0    39W / 300W |   2631MiB / 16160MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      3721      C   ...nvs/codeparrot/bin/python     2629MiB |
+-----------------------------------------------------------------------------+

我們得到了與之前相同的數字，您還可以看到我們使用的是具有 16GB 記憶體的 V100 GPU。所以現在我們可以開始訓練模型，看看 GPU 記憶體消耗如何變化。首先，我們設定一些標準訓練引數。

default_args = {
    "output_dir": "tmp",
    "eval_strategy": "steps",
    "num_train_epochs": 1,
    "log_level": "error",
    "report_to": "none",
}

如果您計劃執行多個實驗，為了在實驗之間正確清除記憶體，請在實驗之間重新啟動 Python 核心。

普通訓練時的記憶體利用率

讓我們使用 Trainer 並以批次大小為 4 訓練模型，而不使用任何 GPU 效能最佳化技術。

>>> from transformers import TrainingArguments, Trainer, logging

>>> logging.set_verbosity_error()


>>> training_args = TrainingArguments(per_device_train_batch_size=4, **default_args)
>>> trainer = Trainer(model=model, args=training_args, train_dataset=ds)
>>> result = trainer.train()
>>> print_summary(result)

Time: 57.82
Samples/second: 8.86
GPU memory occupied: 14949 MB.

我們看到，即使是相對較小的批次大小，也幾乎佔滿了我們 GPU 的全部記憶體。然而，更大的批次大小通常可以帶來更快的模型收斂或更好的最終效能。所以理想情況下，我們希望根據模型的需求而不是 GPU 的限制來調整批次大小。有趣的是，我們使用的記憶體遠超過模型的大小。為了更好地理解這種情況，讓我們看看模型的運算和記憶體需求。

模型操作的剖析

Transformer 架構包含 3 組主要操作，按計算強度分組如下。

張量收縮

線性層和多頭注意力元件都執行批處理的**矩陣-矩陣乘法**。這些操作是訓練 Transformer 中計算最密集的部分。
統計歸一化

Softmax 和層歸一化比張量收縮的計算密集度低，涉及一個或多個**歸約操作**，其結果透過對映應用。
逐元素運算子

這些是其餘的運算子：**偏置、Dropout、啟用函式和殘差連線**。這些是計算密集度最低的操作。

這些知識在分析效能瓶頸時會很有幫助。

此總結源自資料移動就是你所需要的一切：Transformer 最佳化案例研究 2020

模型記憶體剖析

我們已經看到，訓練模型比僅僅將模型放在 GPU 上使用更多的記憶體。這是因為訓練期間有許多元件使用 GPU 記憶體。GPU 記憶體上的元件如下：

模型權重
最佳化器狀態
梯度
為梯度計算儲存的前向啟用
臨時緩衝區
功能特定記憶體

使用 AdamW 進行混合精度訓練的典型模型需要每個模型引數 18 位元組加上啟用記憶體。對於推理，沒有最佳化器狀態和梯度，因此我們可以減去這些。因此，對於混合精度推理，我們最終每個模型引數需要 6 位元組，加上啟用記憶體。

讓我們看看細節。

模型權重

fp32 訓練：4 位元組 * 引數數量
混合精度訓練：6 位元組 * 引數數量（在記憶體中維護一個 fp32 模型和一個 fp16 模型）

最佳化器狀態

普通 AdamW：8 位元組 * 引數數量（維護 2 個狀態）
8 位 AdamW 最佳化器，如 bitsandbytes：2 位元組 * 引數數量
帶有動量的 SGD 等最佳化器：4 位元組 * 引數數量（只維護 1 個狀態）

梯度

無論是 fp32 還是混合精度訓練：4 位元組 * 引數數量（梯度始終以 fp32 形式儲存）

前向啟用

大小取決於許多因素，其中關鍵因素是序列長度、隱藏層大小和批次大小。

這些是前向和後向函式傳入和返回的輸入和輸出，以及為梯度計算儲存的前向啟用。

臨時記憶體

此外，還有各種臨時變數，它們在計算完成後就會釋放，但當下可能會需要額外的記憶體，並可能導致記憶體不足 (OOM)。因此，在編碼時，策略性地考慮這些臨時變數至關重要，有時還需要在不再需要它們時顯式釋放它們。

功能特定記憶體

此外，您的軟體可能有特殊的記憶體需求。例如，當使用束搜尋生成文字時，軟體需要維護輸入和輸出的多個副本。

前向 vs 後向 執行速度

對於卷積層和線性層，反向傳播的浮點運算次數是前向傳播的兩倍，這通常意味著速度會慢兩倍左右（有時會更多，因為反向傳播中的尺寸往往更不規則）。啟用函式通常受頻寬限制，通常情況下，啟用函式在反向傳播中需要讀取的資料比前向傳播中更多（例如，啟用函式前向傳播讀取一次，寫入一次；啟用函式反向傳播讀取兩次，即 gradOutput 和前向傳播的輸出，然後寫入一次，即 gradInput）。

如您所見，我們有幾個潛在的方面可以節省 GPU 記憶體或加快操作速度。既然您已經瞭解了影響 GPU 利用率和計算速度的因素，請參閱單 GPU 高效訓練的方法和工具文件頁面，瞭解效能最佳化技術。

< > 在 GitHub 上更新

Transformers

模型訓練剖析

載入模型

普通訓練時的記憶體利用率

模型操作的剖析

模型記憶體剖析