Hugging Face 的 LOGO

Transformers 文档

模型训练剖析

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

模型训练剖析

为了理解可以应用于提高模型训练速度和内存利用效率的性能优化技术，熟悉 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 组主要操作，按计算强度分组如下。

  1. 张量收缩

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

  2. 统计归一化

    Softmax 和层归一化比张量收缩的计算密集度低，涉及一个或多个**归约操作**，其结果通过映射应用。

  3. 逐元素操作符

    这些是其余的操作符：**偏置、Dropout、激活函数和残差连接**。这些是计算密集度最低的操作。

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

此总结源自 数据移动就是你所需要的一切：Transformer 优化案例研究 2020

模型内存剖析

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

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

使用 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 上更新

←PEFT 概述→