模型训练解剖

为了理解可以应用哪些性能优化技术来提高模型训练速度和内存利用率，熟悉在训练期间 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")

为了打印 GPU 利用率和使用 Trainer 进行训练运行的摘要统计信息，我们定义了两个辅助函数

>>> 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 并训练模型，而不使用任何 GPU 性能优化技术，批量大小为 4

>>> 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 的限制来调整批量大小。有趣的是，我们使用的内存比模型的大小要多得多。为了更好地理解为什么会这样，让我们看一下模型的运算和内存需求。

模型运算的解剖

Transformers 架构包括 3 个主要的操作组，按计算密集程度分组如下。

张量缩并

线性层和多头注意力的组件都进行批量的矩阵-矩阵乘法。这些操作是训练 transformer 最计算密集的部分。
统计归一化

Softmax 和层归一化不如张量缩并计算密集，并且涉及一个或多个归约操作，其结果然后通过映射应用。
逐元素运算符

这些是剩余的运算符：偏置、dropout、激活和残差连接。这些是计算密集程度最低的操作。

在分析性能瓶颈时，了解这些知识可能很有帮助。

此摘要来源于《数据移动才是你所需要的：Transformer 优化的案例研究 2020》

模型内存的解剖

我们已经看到，训练模型使用的内存比仅仅将模型放在 GPU 上要多得多。这是因为在训练期间有许多组件使用 GPU 内存。GPU 内存上的组件如下

模型权重
优化器状态
梯度
为梯度计算保存的前向激活
临时缓冲区
功能特定内存

使用 AdamW 以混合精度训练的典型模型每个模型参数需要 18 字节加上激活内存。对于推理，没有优化器状态和梯度，因此我们可以减去这些。因此，我们最终得到每个模型参数 6 字节用于混合精度推理，加上激活内存。

让我们看看细节。

模型权重

4 字节 * fp32 训练的参数数量
6 字节 * 混合精度训练的参数数量（在内存中维护一个 fp32 模型和一个 fp16 模型）

优化器状态

8 字节 * 普通 AdamW 的参数数量（维护 2 个状态）
2 字节 * 8 位 AdamW 优化器（如 bitsandbytes）的参数数量
4 字节 * 具有动量的 SGD 等优化器的参数数量（仅维护 1 个状态）

梯度

4 字节 * fp32 或混合精度训练的参数数量（梯度始终保持在 fp32 中）

前向激活

大小取决于许多因素，关键因素是序列长度、隐藏层大小和批量大小。

这里有前向和后向函数传递和返回的输入和输出，以及为梯度计算保存的前向激活。

临时内存

此外，还有各种临时变量，这些变量在计算完成后会被释放，但在计算过程中，这些变量可能需要额外的内存并可能导致 OOM。因此，在编码时，战略性地考虑这些临时变量至关重要，有时要在不再需要它们时立即显式地释放它们。

功能特定内存

那么，您的软件可能具有特殊的内存需求。例如，当使用 beam search 生成文本时，软件需要维护多个输入和输出副本。

`forward` vs `backward` 执行速度

对于卷积和线性层，后向的 FLOPs 是前向的两倍，这通常转化为慢约 2 倍（有时更多，因为后向的大小往往更笨拙）。激活通常受带宽限制，并且激活在后向中读取的数据通常比在前向中读取的数据更多（例如，激活前向读取一次，写入一次，激活后向读取两次，gradOutput 和前向的输出，并写入一次，gradInput）。

正如您所见，我们可能可以在几个地方节省 GPU 内存或加快操作速度。现在您已经了解了哪些因素会影响 GPU 利用率和计算速度，请参阅《在单个 GPU 上进行高效训练的方法和工具》文档页面，了解性能优化技术。

< > 更新在 GitHub 上