文章

显存计算详解

#llm#memory#gpu#training

显存计算详解

训练大模型时显存到底用在哪里？四笔账拆开算：模型参数、梯度、优化器状态、激活值。算清楚才能知道需要多少 GPU、选什么并行策略。

一、总公式

训练显存 = 模型参数 + 梯度 + 优化器状态 + 激活值 + 临时缓冲区

其中：
  模型参数    = 参数量 × 每参数字节数
  梯度        = 参数量 × 每参数字节数（和参数等大）
  优化器状态  = 参数量 × 每参数字节数 × 3（Adam: fp32 副本 + m + v）
  激活值      = f(batch, seq_len, hidden_dim, num_layers)

二、逐笔拆开算

2.1 模型参数（Weights）

训练时参数以 FP16 或 BF16 存储（前向和反向用半精度就够了）：

参数量 × 2 bytes (FP16/BF16)

LLaMA-7B:    7B × 2 = 14 GB
LLaMA-13B:  13B × 2 = 26 GB
LLaMA-70B:  70B × 2 = 140 GB   ← H100 80GB 单卡都放不下

注意：这是训练时前向/反向传播用的「工作副本」。FP16 省显存但不省精度——因为累加在 FP32 里做。

2.2 梯度（Gradients）

反向传播算出每个参数的梯度，也是 FP16 存储：

参数量 × 2 bytes

LLaMA-70B: 70B × 2 = 140 GB

梯度和参数量完全一样大——有多少参数，就产生多少梯度。

2.3 优化器状态（Optimizer States）—— 最大的头

Adam 优化器需要为每个参数维护 3 个 FP32 变量：

动量 m:          70B × 4 bytes (FP32) = 280 GB
方差 v:          70B × 4 bytes (FP32) = 280 GB
FP32 参数副本:    70B × 4 bytes (FP32) = 280 GB
────────────────────────────────────────────
合计:                                    840 GB  ← 占训练显存的 72%！

为什么必须存 FP32？ 训练更新参数时，梯度累加需要高精度。FP16 直接更新会导致精度下溢——小梯度直接变成 0。

这就是为什么 DeepSpeed ZeRO-1 只做一件事（分片优化器）就能省 4× 显存，ZeRO-3 分片所有东西能省更多。

2.4 激活值（Activations）

前向传播产生的中间结果，反向传播需要用来算梯度。大小取决于配置：

激活值 ≈ batch_size × seq_len × hidden_dim × num_layers × 系数

以 LLaMA-70B（hidden=8192, layers=80）为例：
  batch=1, seq=4096  →  ≈ 50 GB
  batch=8, seq=4096  →  ≈ 400 GB   ← 显存爆炸

所以大模型训练时 batch size 往往很小（1-2），靠 gradient accumulation 模拟大 batch。

省激活值的技巧：Gradient Checkpointing——不存全部激活值，反向时重新算一遍。用时间换空间。

2.5 临时缓冲区

cuBLAS workspace:   ~1-3 GB (矩阵乘法中间结果)
NCCL buffer:        ~1-2 GB (通信缓冲)
框架开销:           ~2-5 GB (PyTorch 内存管理)

三、实例计算

3.1 LLaMA-70B 单卡训练（做不了，但可以算）

模型参数:      140 GB
梯度:          140 GB
优化器:        840 GB   ← Adam
激活值:         50 GB   (batch=1, seq=4096)
缓冲区:         10 GB
─────────────────────────
总计:         1180 GB

单 H100 80GB:  80 GB
需要多少张:    1180 / 80 ≈ 15 张

所以 LLaMA-70B 训练不可能单卡跑，至少需要 2 个 8 卡 H100 节点 + ZeRO-3 或 FSDP。

3.2 LLaMA-70B 推理（单卡就能跑）

模型参数（FP16）:  140 GB  ← H100 80GB 放不下
模型参数（INT8）:   70 GB  ← H100 80GB 刚好
模型参数（INT4）:   35 GB  ← 绰绰有余

不需要存：梯度、优化器、激活值（用完即丢）

推理显存 ≈ 参数大小，量化后更小。这就是为什么一张 H100 能跑 70B 推理但绝对跑不了训练。

3.3 LLaMA-7B 单卡训练（可以）

模型参数:       14 GB
梯度:           14 GB
优化器:         84 GB   ← 还是大头
激活值:         10 GB
缓冲区:          5 GB
─────────────────────────
总计:          127 GB

单 A100 80GB:   80 GB  ← 差一点，需要 ZeRO-3 offload 或 gradient checkpointing
单 H100 80GB:   80 GB  ← 差一点，同上

所以 7B 模型单卡训练也需要技巧（ZeRO + checkpointing），不能裸跑。

四、为什么显存计算对运维很重要

知道模型的显存需求之后，你才能算：

1. 需要多少 GPU？
   LLaMA-70B 训练 ≈ 1180 GB → 1180 / 80 = 15 张 H100 → 2 个 DGX 节点

2. 需要什么互联？
   如果 TP=8（节点内），需要 NVSwitch 全互联 → 必须买 DGX/HGX，不能买 PCIe

3. 每 GPU 需要多少系统内存？
   CPU offload 时每 GPU 需要额外 64-128 GB DRAM

4. 网络带宽够不够？
   ZeRO-3 每个 step 通信 3× 参数量 → 每 GPU 210 GB → 需要 200 Gbps+ RDMA

这就是为什么 GPU 集群运维知识库的硬件选型和网络设计，最终都要回到模型需求上来。

关联知识

参考资源

学习时间

阶段时间备注
内容创建2026-06-30显存四笔账详解

状态标记

📖 已掌握 — 训练显存四元组、Adam 优化器存储、推理 vs 训练的显存差异 📝 待补充 — KV Cache 显存计算、不同优化器（SGD/AdamW/Lion）的显存对比