模型训练剖析
要理解可以应用于提高模型训练效率和内存利用率的性能优化技术,熟悉训练期间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个主要操作组,按计算强度分组如下。
张量收缩
线性层和多头注意力的组件都执行批处理的矩阵-矩阵乘法。这些操作是训练变压器时计算最密集的部分。
统计归一化
Softmax 和层归一化比张量收缩计算量小,涉及一个或多个归约操作,其结果通过映射应用。
元素级运算符
这些是剩余的操作符:偏置、dropout、激活函数和残差连接。这些是计算强度最低的操作。
这些知识在分析性能瓶颈时可能会有所帮助。
本摘要源自Data Movement Is All You Need: A Case Study on Optimizing Transformers 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
执行速度
对于卷积和线性层,反向传播的计算量是前向传播的2倍,这通常意味着速度会慢约2倍(有时更多,因为反向传播的尺寸往往更不理想)。激活函数通常受带宽限制,通常情况下,激活函数在反向传播时需要读取的数据比前向传播时更多(例如,前向传播时激活函数读取一次,写入一次,反向传播时激活函数读取两次,即前向传播的输出和梯度输出,并写入一次,即梯度输入)。
如你所见,有几个地方可能可以节省GPU内存或加速操作。 既然你已经了解了影响GPU利用率和计算速度的因素,请参考 单GPU高效训练的方法和工具文档页面,以了解 性能优化技术。
< > Update on GitHub