模型训练剖析
为了理解可以应用于提高模型训练速度和内存利用率效率的性能优化技术，熟悉 GPU 在训练过程中的使用方式以及不同操作的计算强度变化非常有帮助。
让我们从探索 GPU 利用率和模型训练运行的激励示例开始。为了演示，我们需要安装一些库。
pip install transformers datasets accelerate nvidia-ml-py3
nvidia-ml-py3 库允许我们从 Python 内部监控模型的内存使用情况。你可能熟悉终端中的 nvidia-smi 命令 - 这个库允许直接在 Python 中访问相同的信息。
然后，我们创建一些虚拟数据：100 到 30000 之间的随机标记 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-2 GB 的内存。为了查看占用量，我们加载一个微小的张量到 GPU 中，这也会触发内核的加载。
>>> import torch
>>> torch.ones((1, 1)).to("cuda")
>>> print_gpu_utilization()
GPU memory occupied: 1343 MB.我们看到，内核本身占用了 1.3 GB 的 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 | +-----------------------------------------------------------------------------+
我们获得了与之前相同的数字，你还可以看到我们使用的是具有 16 GB 内存的 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 的限制来调整批次大小。有趣的是，我们使用的内存远远超过模型的大小。为了更好地理解为什么会出现这种情况，让我们看一下模型的操作和内存需求。
模型操作的剖析
Transformer 架构包括 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。因此，在编码时，务必从策略角度考虑这些临时变量，有时需要在不再需要时明确释放它们。
特定于功能的内存
然后，你的软件可能会有特殊的内存需求。例如，在使用束搜索生成文本时，软件需要维护多个输入和输出副本。
forward 与 backward 执行速度
对于卷积和线性层，反向传播中的浮点运算次数是正向传播的两倍，这通常意味着速度大约慢两倍（有时更多，因为反向传播中的尺寸往往更不方便）。激活通常受带宽限制，并且激活在反向传播中必须读取比正向传播更多的数据（例如，激活正向传播读取一次，写入一次，激活反向传播读取两次，读取 gradOutput 和正向传播的输出，并写入一次，写入 gradInput）。
如你所见，可能有一些地方可以节省 GPU 内存或加快运算速度。现在你已经了解了影响 GPU 利用率和计算速度的因素，请参考 单个 GPU 上高效训练的方法和工具 文档页面，了解性能优化技术。
< > 在 GitHub 上更新