Transformers 文档

模型训练剖析

模型训练剖析

要理解可以应用于提高模型训练效率和内存利用率的性能优化技术,熟悉训练期间GPU的利用方式以及计算强度如何根据执行的操作而变化是有帮助的。

让我们从一个激励性的例子开始,探索GPU的利用率和模型的训练运行。为了演示,我们需要安装一些库:

pip install transformers datasets accelerate nvidia-ml-py3

nvidia-ml-py3 库允许我们从 Python 中监控模型的内存使用情况。你可能对终端中的 nvidia-smi 命令很熟悉 - 这个库允许直接在 Python 中访问相同的信息。

然后,我们创建一些虚拟数据:100到30000之间的随机令牌ID和分类器的二进制标签。 总共,我们得到512个序列,每个序列长度为512,并将它们存储在Dataset中,格式为PyTorch。

>>> 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个主要操作组,按计算强度分组如下。

  1. 张量收缩

    线性层和多头注意力的组件都执行批处理的矩阵-矩阵乘法。这些操作是训练变压器时计算最密集的部分。

  2. 统计归一化

    Softmax 和层归一化比张量收缩计算量小,涉及一个或多个归约操作,其结果通过映射应用。

  3. 元素级运算符

    这些是剩余的操作符:偏置、dropout、激活函数和残差连接。这些是计算强度最低的操作。

这些知识在分析性能瓶颈时可能会有所帮助。

本摘要源自Data Movement Is All You Need: A Case Study on Optimizing Transformers 2020

模型内存的剖析

我们已经看到,训练模型比仅仅将模型放在GPU上使用更多的内存。这是因为在训练过程中有许多组件使用GPU内存。GPU内存上的组件如下:

  1. 模型权重
  2. 优化器状态
  3. 梯度
  4. 保存用于梯度计算的前向激活
  5. 临时缓冲区
  6. 功能特定的内存

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

让我们来看看细节。

模型权重:

  • 4 字节 * 参数数量用于 fp32 训练
  • 6 字节 * 参数数量用于混合精度训练(在内存中维护一个 fp32 模型和一个 fp16 模型)

优化器状态:

  • 8 字节 * 参数数量用于普通 AdamW(保持 2 个状态)
  • 2 字节 * 参数数量,适用于 8 位 AdamW 优化器,如 bitsandbytes
  • 4 字节 * 参数数量,适用于像带有动量的 SGD 这样的优化器(仅维护 1 个状态)

梯度

  • 4 字节 * 参数数量,适用于 fp32 或混合精度训练(梯度始终保持在 fp32)

前向激活

  • 大小取决于许多因素,关键因素包括序列长度、隐藏大小和批量大小。

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

临时内存

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

功能特定内存

然后,您的软件可能会有特殊的内存需求。例如,在使用束搜索生成文本时,软件需要维护输入和输出的多个副本。

forwardbackward 执行速度

对于卷积和线性层,反向传播的计算量是前向传播的2倍,这通常意味着速度会慢约2倍(有时更多,因为反向传播的尺寸往往更不理想)。激活函数通常受带宽限制,通常情况下,激活函数在反向传播时需要读取的数据比前向传播时更多(例如,前向传播时激活函数读取一次,写入一次,反向传播时激活函数读取两次,即前向传播的输出和梯度输出,并写入一次,即梯度输入)。

如你所见,有几个地方可能可以节省GPU内存或加速操作。 既然你已经了解了影响GPU利用率和计算速度的因素,请参考 单GPU高效训练的方法和工具文档页面,以了解 性能优化技术。

< > Update on GitHub