transformers GPU 训练内存解剖：从权重到激活的显存构成与运算密度分析

本文以 Hugging Face Transformers 官方文档《训练过程解剖》（docs/source/ar/model_memory_anatomy.md）为主体，完整复现"用 pynvml 监控 GPU 显存 → 加载 BERT-large → 基线训练 → 拆解显存构成"的实验链路，并结合仓库中 TrainingArguments 与 Trainer 的源码实现，解释训练时每参数 18 字节显存开销（混合精度 + AdamW）从何而来，帮助读者建立"哪里吃显存、哪里可优化"的工程判断力。

一、为什么需要解剖训练时的 GPU 显存

用较小的批量大小训练模型，显存往往就能被填满，但"模型本身只占一部分显存"这一事实常让人困惑：把 BERT-large 加载到 GPU 上只占约 1.3GB，基线训练却要占满 16GB 的 V100。要回答"训练究竟多占了什么显存"，必须先理解两件事：训练过程中 GPU 上执行了哪些运算，以及这些运算各自对显存的需求。官方文档给出的方法论就是：先实测各阶段显存占用，再逐项拆解占用来源。

二、实验环境搭建：用 pynvml 从 Python 内部监控显存

文档第一步是安装监控与训练所需的库：

pip install transformers datasets accelerate nvidia-ml-py3

其中 nvidia-ml-py3（即 pynvml）允许在 Python 进程内直接读取 NVIDIA 设备的显存信息，等价于在终端运行 nvidia-smi。文档定义了核心训练所需的两个辅助函数：

>>> 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, 1, (dataset_size)),
... }
>>> ds = Dataset.from_dict(dummy_data)
>>> ds.set_format("pt")

这里构造了 512 条长度 512 的序列，token ID 取自 100～30000 的随机整数，标签为二分类。set_format("pt") 使数据集直接以 PyTorch 张量格式喂给 Trainer，与 Trainer.train 的标准数据接口兼容。

>>> 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()

print_gpu_utilization 读取 0 号设备的 used 显存；print_summary 则从 trainer.train() 返回的 result.metrics 中取出 train_runtime 与 train_samples_per_second 两项指标，这两个键是 Trainer 训练结束回调中写回的标准指标名。

文档同时给出一个重要实验卫生准则：计划跑多个实验时，应在实验之间重启 Python 内核，以确保显存被正确释放。

三、基线测量：空载显存与 CUDA 上下文开销

实验首先确认起点是"干净的"：

>>> print_gpu_utilization()
GPU memory occupied: 0 MB.

如果本地不是 0MB，需先关闭所有占用该 GPU 的进程。接着文档揭示了一个容易被忽略的开销——CUDA 上下文（context）。哪怕只把一个 1×1 的小矩阵搬到 GPU 上：

>>> import torch

>>> torch.ones((1, 1)).to("cuda")
>>> print_gpu_utilization()
GPU memory occupied: 1343 MB.

可以看到，仅初始化 CUDA 上下文就消耗约 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.

2631MB 减去上下文开销 1343MB，约 1.3GB 即 BERT-large（约 3.35 亿参数）fp32 权重的占用。文档提醒：具体数值随 GPU 型号而变，且较新架构上，权重可能因加载时做了针对加速的内存布局优化而占更多空间。

可用命令行 nvidia-smi 交叉验证同一进程的真实占用。文档记录到的 V100（16GB）输出节选：

|   0  Tesla V100-SXM2...  On   | 00000000:00:04.0 Off |                    0 |
| N/A   37C    P0    39W / 300W |   2631MiB / 16160MiB |      0%      Default |

Python 侧 pynvml 读到的 2631MB 与 nvidia-smi 完全一致，验证了监控方法的有效性。

五、基线训练：小批量为何填满 16GB 显存

设置一组标准训练参数后，用 [Trainer] 以批量大小 4、不使用任何 GPU 优化技术进行训练：

default_args = {
    "output_dir": "tmp",
    "eval_strategy": "steps",
    "num_train_epochs": 1,
    "log_level": "error",
    "report_to": "none",
}

对照源码可以核对这些参数的含义：training_args.py 中 per_device_train_batch_size 的默认值是 8（field(default=8, ...)），文档显式传 4 正是为了得到可复现的小批量基线；per_device_train_batch_size * num_devices * gradient_accumulation_steps 共同决定有效批量。

>>> 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.

关键观察：批量大小仅 4 的小训练就占用了 14.9GB，接近 V100 16GB 的满容量，远超 1.3GB 的模型权重。而更大的批量通常带来更快收敛或更好的最终效果——我们希望按模型需求而非 GPU 限制来调批量。要"多占的这部分显存"到底去了哪，需要解剖训练过程的运算与显存构成。

六、解剖模型运算：按计算密度分三层

官方文档将 Transformer 的运算按算术强度（arithmetic intensity）分为三类，这一分层是分析性能瓶颈的起点：

1. 矩阵乘法（matmul）：线性层与多头注意力的 QKV/输出投影、前馈层，全部是对批量做矩阵-矩阵乘法。这是 Transformer 训练中计算密度最高的部分，最适合 GPU 的大规模并行。
2. 统计归约运算（reduction）：Softmax 与 LayerNorm 等。它们先跨一个维度做归约算出统计量（如均值、方差），再"map"回原张量，计算密度低于 matmul，且需要多次访问显存。
3. 逐元素运算（element-wise）：偏置、Dropout、激活函数、残差连接等，是计算密度最低的一类。

该分层源自 2020 年论文《Data Movement Is All You Need: A Case Study on Optimizing Transformers》（arXiv:2007.00072），其核心论点是：优化 Transformer 性能的关键在于减少数据搬运，而非增加运算。

七、解剖显存构成：六个组成部分逐项量化

训练时 GPU 显存上共存放六类内容，文档给出了每类的字节成本：

1. 模型权重（Model weights）

• fp32 训练：4 字节 × 参数量；
• 混合精度训练：6 字节 × 参数量——因为内存中同时保留一份 fp32 主权重（用于稳定的权重更新）和一份 fp16 副本（用于前向/反向传播）。

2. 优化器状态（Optimizer states）

• 标准 AdamW：8 字节 × 参数量（分别保存一阶动量与二阶方差两个 fp32 张量）；
• 8 位量化 AdamW（如 bitsandbytes 实现）：2 字节 × 参数量。这一点在仓库中有对应的集成实现，quantizer_bnb_8bit.py 与 quantizer_bnb_4bit.py 等文件即为 BitsAndBytes 8/4 位量化的接入层，说明"用量化优化器压显存"在 Transformers 中是可落地的路径；
• SGD + momentum 一类只保存一个状态的优化器：4 字节 × 参数量。

3. 梯度（Gradients）

• 无论 fp32 还是混合精度训练，梯度始终按 fp32 保存：4 字节 × 参数量。

4. 前向激活（Forward activations）

• 大小取决于序列长度、隐藏维度、批量大小与模型深度，是最难估算的变量。这是为什么"模型放得下，训练却 OOM"：长序列或大批量会把激活撑爆显存。

5. 临时缓存（Temporary caches）

• softmax、矩阵乘法的中间 scratch 张量等。它们在运算结束后即被释放，但峰值时刻可能叠加出额外显存尖峰，直接触发 OOM。文档强调：写代码时要策略性地考虑这些临时变量，必要时在不再需要时显式 del 释放。

6. 功能特定内存（Feature-specific memory）

• 例如 beam search 生成文本时要同时保留多份输入/输出；大词表的 embedding 表本身也可能是显存大户。

每参数字节账本

文档给出的关键结论：用混合精度 + AdamW 训练一个典型模型，每参数需要 18 字节（6 权重 + 8 优化器状态 + 4 梯度），外加激活显存；推理时没有优化器状态与梯度，每参数只需 6 字节（混合精度），外加激活。英文同主题文档（docs/source/en/model_memory_anatomy.md）给出了一个规模化示例：在批量大小 16 下混合精度训练一个 4B 参数模型约需 85GB 显存——权重 6B/param、Adam 状态 8B/param、梯度 4B/param、激活随 batch×seq_len×depth×hidden 浮动，量化 Adam 可用 bitsandbytes 把优化器状态压到 2B/param，与上述账本完全吻合。

组成部分	训练（混合精度）	推理（混合精度）
权重	6 字节/参数	6 字节/参数
AdamW 状态	8 字节/参数	—
梯度	4 字节/参数	—
激活	随批量/序列长度/深度变化	随输入变化
临时张量/功能特定内存	峰值尖峰可致 OOM	beam search 等可致多份副本

八、前向 vs 反向：为什么 backward 大约慢 2 倍

文档最后指出运算速度的不对称性：

• 对卷积与线性层，反向传播的浮点运算量是前向的 2 倍，通常体现为约 2 倍的耗时（有时更多，因为反向涉及的张量尺寸分布对硬件更不友好）；
• 激活类运算多为带宽受限（bandwidth-bound）。例如一个激活函数的前向只需"读一次、写一次"，反向却要读 gradOutput 与前向输出两份、再写 gradInput 一份，读放大使反向更慢。

这两点共同决定了优化空间：显存上可压缩的是优化器状态与激活，速度上可优化的是归约与逐元素运算的访存——对应到仓库中，就是梯度累积、梯度检查点、混合精度与自定义融合内核这几类现成手段。

九、下一步：仓库中可用的显存优化手段

理解了显存解剖之后，文档指引读者结合训练优化专题继续深入。当前仓库英文文档中对应且仍然存在的指南包括：

• 梯度累积（grad_accumulation）：在不增加峰值显存的前提下模拟更大的有效批量；
• 梯度检查点（grad_checkpointing）：以重算换显存，只保留部分中间激活；
• 混合精度训练（mixed_precision_training）：把前向/反向放到低精度，把权重主副本留在 fp32，正是前文"6 字节/参数"的来源；
• 自定义融合内核（kernels）：加速归约与逐元素类运算。

需要说明的是，阿拉伯语文档末尾引用的 perf_train_gpu_one（单卡高效训练）页面在当前 阿拉伯语目录树 中已被注释停用，实际查阅时以上述英文指南与仓库源码（training_args.py、trainer.py）为准。

小结

本文完整继承了官方"训练过程解剖"文档的实验主线：pynvml 显存监控 → 空载/上下文/权重三级基线（0MB → 1343MB → 2631MB）→ 批量 4 基线训练实测 14949MB，并落到两张"账本"上——运算按计算密度分 matmul / 归约 / 逐元素三层，显存按权重（6B）、AdamW 状态（8B）、梯度（4B）、激活、临时缓存与功能特定内存六项拆解，得出"混合精度 + AdamW 训练 = 18 字节/参数、推理 = 6 字节/参数"的核心结论。掌握这套解剖方法后，任何一次训练 OOM 都可以沿着"是权重、状态、梯度还是激活"的路径定位，并选择梯度累积、检查点、低精度或量化优化器这类对症下药的优化。