Transformers documentation
GPU 内存使用
GPU 内存使用
在混合精度下以批大小 16 训练一个 4B 参数模型，大约需要 85GB 的 GPU 显存。了解 GPU 显存的占用情况以及训练过程中发生的计算操作，有助于找出可以显存使用量的地方。
┌─────────────────────────── TENSORS ──────────────────────────────────┐ │ │ │ MODEL WEIGHTS ████████████████████████ 6 bytes/param │ │ (mixed precision) ├── fp16 copy (2B) — forward/backward pass │ │ └── fp32 copy (4B) — stable weight updates │ │ │ │ OPTIMIZER STATES ████████████████████████████████ 8 bytes/param │ │ (Adam) ├── fp32 momentum (4B) │ │ └── fp32 variance (4B) │ │ ↳ quantized Adam (bitsandbytes) → 2 bytes/param │ │ │ │ GRADIENTS ████████████████ 4 bytes/param (fp32) │ │ └── computed in backward pass, per parameter │ │ │ │ ACTIVATIONS ████ varies — batch × seq_len × depth × hidden │ │ (forward cache) └── cached for backward; can OOM even if model │ │ fits (long seqs / large batches) │ │ bf16/fp16 if using mixed precision │ │ │ │ TEMPORARY TENSORS ▓ short-lived (softmax, matmul scratch) │ │ └── peak spikes can cause OOM │ │ │ │ OTHER OVERHEAD ▒ beam search caches, large embedding tables │ └──────────────────────────────────────────────────────────────────────┘
GPU 内存中存储着两类内容：张量以及处理它们的算子。
张量
训练需要并生成许多需要存储的张量类型。
模型权重存储在 GPU 上。在混合精度训练中，需要两份权重副本——一份为 fp16，用于前向/反向传播；另一份为 fp32，作为用于稳定权重更新的“主副本”。这相当于每个参数占用 6 字节。
像 Adam 这样的优化器会为每个参数额外存储两个张量，即动量和方差，且两者均为 fp32。这意味着每个参数额外占用 8 字节。
你可以使用不同的优化器，例如 bitsandbytes 中 Adam 的量化版本，将其压缩为每个参数 2 字节。
在反向传播过程中，会为每个参数计算梯度张量。这些梯度以 fp32 格式存储，因此每个参数占用 4 字节。
前向激活值在前向传播过程中计算，并被缓存以供反向传播计算梯度使用。这些激活值的大小会随批大小、序列长度、模型深度和隐藏层大小而变化。这就是为什么即使模型本身能够装入显存，批大小或序列长度仍可能导致 GPU 内存耗尽的原因。
像 softmax 和矩阵乘法这样的操作会创建临时张量，并在每个操作完成后释放它们。如果单个操作的峰值内存占用非常高，可能会产生短暂的内存尖峰，从而导致内存不足。
某些功能会引入额外的开销。例如，束搜索（beam search）需要维护多个输出，而针对大型词表的嵌入表（embedding tables）可能非常大。
运维
主要有三种训练操作类型。
矩阵乘法（matmul）是主要的算子类型，涵盖线性层、QKV 投影、注意力输出投影以及 FFN 层。主要的内存消耗来自注意力分数矩阵，其规模随序列长度的平方增长。这正是长序列成本高昂的原因，也是存在不同 注意力后端 的原因，以避免在内存中实例化完整的矩阵。
归约操作（如 softmax 和 layer norm）会读取整个张量，沿某个维度计算统计量，然后再次读取张量以应用该统计量。这意味着每次操作都需要多次访问内存。
像激活函数和 dropout 这样的逐元素操作会对每个元素独立地应用一个函数，其内存占用与张量大小成正比。
后续步骤
- 请参阅 梯度累积 指南，了解如何在不耗尽 GPU 内存的情况下模拟使用更大有效批大小进行训练。
- 请参阅 梯度检查点 指南，了解如何通过仅存储部分中间激活值来降低内存使用量。
- 请参阅混合精度训练指南，了解如何使用较低精度的数据类型来减少内存占用并加快训练速度。
- 请参阅 Kernels 指南，了解如何使用自定义融合内核加速训练。