模型训练剖析¶
要理解可以应用于提高模型训练速度和内存利用率的优化技术，了解 GPU 在训练中的使用情况以及不同操作的计算强度变化是很有帮助的。
我们从一个 GPU 使用情况和模型训练运行的示例开始。为了演示，我们需要安装几个库：
pip install transformers datasets accelerate nvidia-ml-py3
nvidia-ml-py3 库允许我们在 Python 中监控模型的内存使用情况。你可能熟悉终端中的 nvidia-smi 命令——这个库可以直接在 Python 中访问相同的信息。
接下来，我们创建一些虚拟数据：输入 ID 在 100 到 30000 之间的随机值，以及用于分类的二进制标签。总共生成 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 使用情况的摘要统计信息和训练运行情况，我们定义两个辅助函数：
from pynvml import *
def print_gpu_utilization():
nvmlInit()
handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print(f"GPU 内存占用: {info.used//1024**2} MB.")
def print_summary(result):
print(f"时间: {result.metrics['train_runtime']:.2f} 秒")
print(f"每秒样本数: {result.metrics['train_samples_per_second']:.2f}")
print_gpu_utilization()
我们先确认 GPU 内存是空闲的：
print_gpu_utilization()
# 输出:
# GPU 内存占用: 0 MB.
一切正常：在加载任何模型之前，GPU 内存应该是空闲的。如果你的机器不是这样，请确保停止所有占用 GPU 内存的进程。然而，并不是所有空闲的 GPU 内存都可以被用户使用。当模型加载到 GPU 上时，内核也一同加载，这会占用 1-2GB 的内存。我们通过将一个小型张量加载到 GPU 来触发内核的加载。
import torch
torch.ones((1, 1)).to("cuda")
print_gpu_utilization()
# 输出:
# GPU 内存占用: 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 内存占用: 2631 MB.
我们看到模型权重本身占用了 1.3GB 的 GPU 内存。具体数量取决于你使用的 GPU。请注意，在较新的 GPU 上，模型可能会占用更多空间，因为权重是以优化的方式加载的，以加快模型的使用速度。现在，我们可以快速检查一下 nvidia-smi CLI 是否给出相同的结果：
nvidia-smi
我们得到了与之前相同的数字，并且可以看到我们正在使用一个有 16GB 内存的 V100 GPU。现在我们可以开始训练模型，并观察 GPU 内存使用的变化。首先，我们设置一些标准的训练参数：
default_args = {
"output_dir": "tmp",
"eval_strategy": "steps",
"num_train_epochs": 1,
"log_level": "error",
"report_to": "none",
}
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)
我们可以看到即使是相对较小的批大小也几乎填满了 GPU 的全部内存。然而，更大的批大小通常可以加快模型收敛速度或提高最终性能。理想情况下，我们应该根据模型的需求调整批大小，而不是受 GPU 限制。有趣的是，我们使用的内存比模型本身的内存要多得多。为了更好地理解原因，我们来看看模型的操作和内存需求。
模型操作剖析¶
Transformer 架构包括以下三组主要操作，按计算强度分组：
- 张量收缩
- 线性层和多头注意力组件中的所有操作都涉及批处理的 矩阵-矩阵乘法。这些操作是训练变压器最计算密集的部分。
- 统计归一化
- Softmax 和层归一化操作比张量收缩要轻得多，涉及一个或多个 归约操作，其结果通过映射应用。
- 逐元素操作
- 剩下的操作是 偏置、Dropout、激活和残差连接。这些是最不计算密集的操作。
了解这些信息有助于分析性能瓶颈。
模型内存剖析¶
我们已经看到训练模型使用的内存远不止将模型加载到 GPU 上的内存。这是因为训练过程中使用的 GPU 内存有多个组件，包括：
- 模型权重
- 优化器状态
- 梯度
- 用于梯度计算的前向激活
- 临时缓冲区
- 功能特定内存
一个典型的使用混合精度和 AdamW 优化器训练的模型需要每个参数 18 字节，加上激活内存。对于推理，没有优化器状态和梯度，因此可以减去这些部分，最终每个参数需要 6 字节加上激活内存。
我们来看一下这些组件的详细信息：
模型权重：
- fp32 训练：4 字节 × 参数数量
- 混合精度训练：6 字节 × 参数数量（同时维护 fp32 和 fp16 模型）
优化器状态：
- 正常 AdamW：8 字节 × 参数数量（维护两个状态）
- 8 位 AdamW 优化器（如 bitsandbytes）：2 字节 × 参数数量
- 带动量的 SGD 等优化器：4 字节 × 参数数量（仅维护一个状态）
梯度：
- fp32 或混合精度训练：4 字节 × 参数数量（梯度始终保存为 fp32）
前向激活：
- 大小取决于许多因素，关键因素包括序列长度、隐藏大小和批大小。
临时内存：
- 在计算过程中会释放各种临时变量，但在计算过程中这些变量可能会占用额外的内存，导致内存不足。因此，在编程时要战略性地考虑这些临时变量，并在不再需要时显式释放。
功能特定内存：
- 例如，使用束搜索生成文本时，软件需要维护多个输入和输出的副本。
前向 vs 后向执行速度：
- 对于卷积和线性层，后向传播的浮点运算量是前向传播的两倍，通常导致后向传播的速度大约是前向传播的两倍（有时更多，因为后向传播的大小往往更尴尬）。激活通常受带宽限制，通常在后向传播中需要读取更多数据（例如，激活前向传播读取一次，写入一次，激活后向传播读取两次，gradOutput 和前向传播的输出，写入一次，gradInput）。