3、大模型显存占用分析：模型参数、梯度、优化器状态、激活值等各部分的显存占用计算与实测方法

大模型部署与微调过程中，显存（GPU Memory）是最核心的稀缺资源之一。许多工程师在实践中最常遇到的困惑是：明明模型参数量不大，为什么一跑训练或长序列推理就 OOM（Out of Memory）？这往往是因为对显存占用的构成缺乏系统性认知。显存占用并非只有模型权重这一项，而是由模型参数、梯度、优化器状态、激活值、临时缓冲区以及 CUDA 上下文等多个部分共同组成。本章将深入拆解每一部分的显存计算公式、实测方法，并结合实际案例帮助读者建立完整的显存预算（Memory Budget）意识。

大模型显存占用全景拆解图 训练态（Training） ① 模型参数 (W) 大小 = 参数量 × 字节数 例：7B FP16 → 14 GB ② 梯度 (G) 大小 = 与参数同形状 例：7B FP16 → 14 GB ③ 优化器状态 (Adam) FP32 主权重 + 一阶 + 二阶动量 例：7B → 84 GB (12字节/参数) ④ 激活值 (Activation) 与 batch_size、seq_len 强相关 例：7B seq=2048 bs=1 → ~2-4 GB 推理态（Inference） ⑤ KV Cache 2 × L × H × seq_len × bs × 字节 例：7B 32层 4096维 seq=2048 → ~2 GB ⑥ 临时缓冲区 (Workspace) CUDA context、cuDNN/cuBLAS workspace 通常 0.5 ~ 2 GB 固定开销 ⑦ 推理权重 可量化至 INT8/INT4 压缩 例：7B INT8 → 7 GB, INT4 → 3.5 GB ⑧ 推理峰值激活 单层激活 × 层数（可重计算优化） 通常远小于训练态 总显存 = 参数 + 梯度 + 优化器状态 + 激活值 + 临时开销 | 训练态约为推理态的 4~12 倍

上图清晰地展示了训练态与推理态在显存占用上的本质差异。训练态需要同时保存参数、梯度、优化器状态和激活值，而推理态则主要关注权重和 KV Cache。接下来，我们将逐一深入分析每个部分的计算逻辑与实测手段。

3.1 模型参数与梯度的显存计算

核心原理：参数与梯度的线性关系

模型参数（Weights）的显存占用是最直观的部分，其计算公式为：显存 = 参数量 × 每个参数占用的字节数。以 FP32（4字节）和 FP16/BF16（2字节）为例，一个 7B（70亿参数）的模型在 FP16 下占用约 14GB，在 FP32 下占用约 28GB。梯度（Gradients）在训练过程中与参数具有完全相同的形状和精度，因此在反向传播时，梯度张量会额外占用与参数等量的显存。这意味着，仅参数和梯度两项，7B 模型在 FP16 训练时就需要 28GB。

关键公式： 训练态基础显存 = 2 × 参数量 × 字节数（参数 + 梯度）。这是任何训练任务都无法绕开的最小显存开销。

判断方法：通过模型配置快速估算

在实际操作中，我们通常通过模型的配置文件（如 config.json）中的 hidden_size、num_hidden_layers、vocab_size 等字段计算参数量。对于 Transformer 架构，参数量约等于 12 × hidden_size² × num_layers（仅计算注意力与 FFN 部分，忽略 embedding）。例如 LLaMA-7B 的 hidden_size 为 4096，层数为 32，计算得 12 × 4096² × 32 ≈ 6.4B，加上 embedding 后接近 7B。

案例拆解：从零计算一个 13B 模型的参数显存

假设我们使用一个 13B 模型进行全参数微调（Full Fine-tuning），精度为 BF16。首先确认参数量为 13B，则参数占用 13B × 2 字节 = 26GB，梯度同样为 26GB。仅这两项就需要 52GB 显存。如果使用一张 80GB 的 A100，剩余可用显存为 28GB，这部分需要容纳优化器状态和激活值。若优化器使用 Adam（需 12 字节/参数），则优化器状态需要 13B × 12 = 156GB，远超单卡容量。此时必须采用 ZeRO Stage 2 或 3 进行优化器状态切分，或者改用 LoRA 等参数高效微调方法。

常见误区： 很多初学者认为“模型能加载进显存就能训练”。实际上，加载模型仅需参数显存，而训练需要参数 + 梯度 + 优化器状态，通常为参数显存的 4~6 倍。务必在训练前用 torch.cuda.max_memory_allocated() 实测峰值。

3.2 优化器状态的显存开销与 ZeRO 优化

核心原理：Adam 优化器的“隐形”大块头

优化器状态（Optimizer States）是训练显存中的最大变量。以最常用的 AdamW 优化器为例，它为每个参数保存三份 FP32 数据：主权重副本（Master Weights）、一阶动量（Momentum）和二阶动量（Variance）。因此，每个参数需要 4 + 4 + 4 = 12 字节。对于 7B 模型，优化器状态需要 84GB 显存，这远超参数本身。相比之下，SGD 优化器仅需 4 字节/参数（主权重），显存压力小得多。

操作流程：使用 ZeRO 进行显存卸载

DeepSpeed 的 ZeRO（Zero Redundancy Optimizer）通过将优化器状态、梯度、参数分片到不同 GPU 上来消除冗余。具体操作如下：

  1. 在训练脚本中引入 deepspeed.initialize，并配置 ZeRO Stage 2（仅切分优化器状态和梯度）。
  2. 对于 7B 模型，Stage 2 可将优化器显存从 84GB 降至约 84GB / 8（8卡）= 10.5GB。
  3. 若单卡显存仍不足，启用 Stage 3（参数也进行分片），但会增加通信开销。
# deepspeed_config.json 示例
{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    }
  },
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4
}

实战场景：单卡 24GB 训练 7B 模型

假设我们只有一张 RTX 3090（24GB），目标是用 LoRA 微调 7B 模型。LoRA 冻结原参数，仅训练新增的低秩矩阵（如 rank=8）。此时可训练参数量约为 4M（百万级），优化器状态仅需 4M × 12 = 48MB，参数和梯度各占 8MB。显存大头变为冻结的 7B 参数（14GB）和激活值。通过梯度检查点（Gradient Checkpointing）将激活值显存从 4GB 降至 1GB，总显存约 15GB，可顺利训练。

经验总结： 当模型参数量超过 10B 时，全参数训练几乎必须依赖 ZeRO 或多卡并行。对于单卡场景，优先考虑 LoRA/QLoRA，其中 QLoRA 将基础模型量化为 NF4（4-bit），7B 模型仅占 3.5GB，极大释放显存空间。

3.3 激活值显存：动态变化的“隐形杀手”

核心原理：激活值与序列长度的平方关系

激活值（Activations）是前向传播过程中各层输出的中间张量。其显存占用与 batch_size × 序列长度 × 隐藏层维度 × 层数 成正比。更关键的是，在自注意力机制中，注意力矩阵的显存与 序列长度的平方 成正比。例如，当序列长度从 2048 增加到 4096 时，注意力矩阵的显存占用将增加 4 倍。对于 7B 模型，seq_len=2048、batch_size=1 时，激活值约占用 2~4GB；但若 seq_len 提升至 8192，激活值可能飙升至 20GB 以上。

判断方法：使用 PyTorch 的显存分析工具

要精确测量激活值显存，可以使用 PyTorch 的 torch.profiler 或简单的钩子（Hook）记录每一层的输出张量大小。更实用的方法是使用 torch.cuda.max_memory_allocated() 对比开启与关闭梯度检查点时的峰值显存差值。梯度检查点（Gradient Checkpointing）通过在前向传播时丢弃激活值，并在反向传播时重新计算，以时间换空间，可将激活值显存降低约 60%~80%。

案例拆解：长文本训练 OOM 的排障过程

某团队在训练 13B 模型时，设置 seq_len=4096、batch_size=2，A100 80GB 直接 OOM。通过逐项排查：参数 26GB + 梯度 26GB + 优化器状态（ZeRO Stage 2 切分后）约 20GB，合计 72GB，剩余仅 8GB 给激活值。而 13B 模型在 seq_len=4096 时激活值需要约 15GB。解决方案：① 启用梯度检查点，激活值降至 4GB；② 将 batch_size 降为 1；③ 使用 Flash Attention 减少注意力矩阵的显存峰值。最终总显存约 76GB，成功运行。

配置项修改前修改后显存变化
梯度检查点关闭开启激活值 15GB → 4GB
Batch Size21激活值减半
Attention 实现标准Flash Attention峰值降低约 30%
风险提示： 梯度检查点会增加约 20%~30% 的计算时间，并非免费午餐。在推理场景中，激活值显存通常远小于训练态，但 KV Cache 会随序列长度线性增长，长上下文推理时同样需要关注。

3.4 推理态显存：KV Cache 与静态权重

核心原理：KV Cache 的精确计算

在自回归推理中，为了避免重复计算历史 token 的 Key 和 Value，需要将其缓存下来，即 KV Cache。其显存计算公式为：2（K 和 V）× 层数 × 注意力头数 × 头维度 × 序列长度 × 批次大小 × 字节数。以 LLaMA-7B 为例，层数 32，头数 32，头维度 128，则每层 KV 维度为 32 × 128 = 4096。当 seq_len=2048、batch_size=1、FP16 时，KV Cache 占用为 2 × 32 × 4096 × 2048 × 1 × 2 字节 ≈ 2GB。若 batch_size 提升至 8，则直接占用 16GB。

操作流程：使用 vLLM 或 TensorRT-LLM 优化 KV Cache

现代推理框架通过 PagedAttention 技术将 KV Cache 分页管理，减少碎片化浪费。操作上，只需在 vLLM 中设置 --max-model-len 和 --gpu-memory-utilization 参数，框架会自动计算可用的 KV Cache 空间。例如设置 --gpu-memory-utilization 0.9 表示允许使用 90% 的显存用于权重和 KV Cache。

# 使用 vLLM 启动推理服务
python -m vllm.entrypoints.openai.api_server \
  --model /path/to/llama-7b \
  --tensor-parallel-size 1 \
  --max-model-len 4096 \
  --gpu-memory-utilization 0.9

实战场景：并发请求下的显存预算

假设我们需要在 24GB 显卡上部署 7B 模型（FP16 权重 14GB），剩余 10GB 用于 KV Cache。若平均请求长度为 1024 token，每个请求的 KV Cache 约为 1GB，则最多支持约 10 个并发请求。若将权重量化至 INT8（7GB），剩余 17GB 可支持约 17 个并发。若进一步使用 INT4（3.5GB），并发数可提升至 20 以上。但需注意量化带来的精度损失。

复盘建议： 推理显存优化的核心是“权重压缩 + KV Cache 管理”。优先使用量化（INT8/INT4）压缩权重，再通过 PagedAttention 或 Continuous Batching 提高 KV Cache 利用率。实测时使用 nvidia-smi 监控显存曲线，确保峰值不超过 95%。

3.5 显存实测方法与工具链

核心原理：实测是验证计算的唯一标准

理论计算只能提供估算值，实际显存占用还受 CUDA 上下文、cuDNN 算法选择、内存碎片等因素影响。因此，实测显存是部署和调优的必要步骤。最常用的工具是 NVIDIA 的 nvidia-smi 和 PyTorch 的 torch.cuda API。

操作流程：三步完成显存实测

  1. 基线测量： 在模型加载前记录 torch.cuda.memory_allocated() 和 torch.cuda.memory_reserved()。
  2. 分段测量： 分别在前向传播后、反向传播后、优化器更新后记录显存峰值，差值即为各阶段新增显存。
  3. 峰值定位： 使用 torch.profiler 导出显存时间线，定位 OOM 发生的具体算子。
import torch
# 记录初始显存
torch.cuda.reset_peak_memory_stats()
initial_mem = torch.cuda.memory_allocated()
# 模拟前向传播
output = model(input_ids)
forward_mem = torch.cuda.memory_allocated()
# 模拟反向传播
loss = output.sum()
loss.backward()
backward_mem = torch.cuda.memory_allocated()
print(f"初始: {initial_mem/1e9:.2f}GB, 前向: {forward_mem/1e9:.2f}GB, 反向: {backward_mem/1e9:.2f}GB")
print(f"峰值: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")

案例拆解：一次完整的显存体检报告

我们对一个 3B 模型进行 LoRA 微调，batch_size=4，seq_len=1024。实测结果如下：初始 CUDA 上下文占用 0.8GB；加载模型权重（FP16）占用 6GB；前向传播后激活值占用 3.2GB；反向传播后梯度占用 0.1GB（仅 LoRA 参数）；优化器更新后额外占用 0.2GB。峰值显存出现在反向传播阶段，总计约 10.3GB。通过分析发现，激活值占比过高，启用梯度检查点后降至 1.1GB，峰值降至 8.2GB，为增大 batch_size 留出空间。

阶段显存占用占比优化手段
CUDA 上下文0.8 GB8%不可优化
模型权重6.0 GB58%量化至 INT8 → 3GB
激活值3.2 GB31%梯度检查点 → 1.1GB
梯度 + 优化器0.3 GB3%LoRA 本身已优化
常见误区： 不要仅依赖 nvidia-smi 查看显存，它显示的是进程占用，而非 PyTorch 实际分配。PyTorch 的缓存分配器会预留显存，导致 nvidia-smi 显示值偏高。务必使用 torch.cuda.memory_allocated() 获取真实张量占用。

总结而言，显存优化是一个系统工程，需要从参数、梯度、优化器、激活值、KV Cache 五个维度分别计算和实测。建议读者在每次训练或部署前，先完成一次显存预算表，再结合本章的实测方法进行验证，从而避免 OOM 的尴尬，并最大化利用宝贵的 GPU 资源。

← 上一章 📖 返回目录 下一章 →