03 · 显存与带宽墙

⏱ 30 分钟 ｜ ⭐ 理解这一章，后面十几章都会变简单

🎯 一句话

你的 GPU 不是在算，大部分时间是在等数据。 这一章讲清楚数据存在哪、搬一次要多久， 之后 FlashAttention、量化、KV Cache、算子融合 —— 全都会变成"当然应该这么做"。

🏔️ 一、内存层次：一张必须记住的表

⬇︎ 越往下：容量越大，速度越慢（差好几个数量级）。

🔑 记住这三个比值就够了： - Shared Memory 比 HBM 快约 5 倍（FlashAttention 的全部依据） - HBM 比 PCIe 快约 40 倍（不要随便 .cpu()） - NVLink 比网络快约 20 倍（并行策略选择的依据）

⚖️ 二、算术强度：判断瓶颈的唯一标准

$$\text{算术强度} = \frac{\text{计算量（FLOP）}}{\text{数据搬运量（字节）}}$$

💡 人话：每从显存搬一个字节，能顺带做多少次运算。

📊 常见操作的算术强度（照着这张表判断）

🔑 注意表里第三行：batch=1 的矩阵乘算术强度只有 ~2。 这就是第 15 章的核心 —— 自回归解码每次只生成 1 个 token，就是 batch=1 的矩阵乘， 推理天生是带宽瓶颈，和训练完全相反。

💾 三、训练显存到底花在哪

四大块（以 7B 模型、BF16 混合精度、Adam 为例）：

激活值：唯一和 batch 相关的一块

🔑 这解释了第 1 章那个错误直觉： "OOM 就调小 batch"只对激活有用。 如果你的 OOM 是因为优化器状态装不下，调 batch 到 1 也没用 —— 那时候需要的是 ZeRO 或量化。

🧮 四、动手：估算你的模型要多少显存

def estimate_training_memory(params_B, batch, seq_len, hidden, layers,
                             optimizer="adam", grad_ckpt=False):
    """返回 GB。粗估，用于判断量级和选并行策略。"""
    P = params_B * 1e9
    weights = P * 2 / 1e9                       # BF16 参数
    grads   = P * 2 / 1e9                       # BF16 梯度
    opt     = P * (12 if optimizer == "adam" else 4) / 1e9

    # 激活：不开检查点约 ~34 倍 hidden/层/token（Megatron 的经验系数）
    per_layer = batch * seq_len * hidden * 34 * 2 / 1e9
    acts = per_layer * (layers ** 0.5 if grad_ckpt else layers)
    #                    ↑ 检查点后约降到 √层数 量级

    total = weights + grads + opt + acts
    print(f"参数 {weights:6.1f} GB | 梯度 {grads:6.1f} GB | "
          f"优化器 {opt:6.1f} GB | 激活 {acts:6.1f} GB")
    print(f"合计 ≈ {total:.1f} GB")
    return total

# 7B，batch=4，seq=2048
estimate_training_memory(7, 4, 2048, 4096, 32)
# 参数   14.0 GB | 梯度   14.0 GB | 优化器   84.0 GB | 激活   35.0 GB
# 合计 ≈ 147.0 GB   ⭐ 需要至少 2 张 80G 卡，且必须用 ZeRO

⚠️ 这是粗估，实际还要加：显存碎片（5–10%）、 通信缓冲区、CUDA context（几百 MB）、临时工作区。 留 15% 余量是常见做法。

🚧 五、"带宽墙"这个词的含义

🔑 一个重要的推论： 随着硬件迭代，"少搬数据"的价值只会越来越高。 今天勉强算 compute-bound 的操作，下一代卡上就变成 memory-bound 了。

💡 这也解释了为什么 FP8、KV Cache 量化、算子融合这些技术 越来越受重视 —— 它们都是在直接对抗带宽墙。

🧰 六、由此推出的所有优化（后面章节的地图）

⭐ 看这张表：七个技术，六个是在对抗带宽和显存，只有一个和算力有关。 这就是为什么说这一章是钥匙。

🔗 和站内其他章的关系

✅ 检查点

1. 内存层次里，哪三个比值最该记住？
2. 算术强度的定义是什么？A100 的平衡点是多少？
3. 逐元素加法的算术强度是多少？离平衡点差多远？这推出了什么技术？
4. batch=1 的矩阵乘算术强度是多少？这对推理意味着什么？
5. 训练显存的四大块是什么？7B 模型各占多少？
6. 优化器状态为什么是 12 字节/参数？训练显存的经验公式是什么？
7. 四大块里哪一块和 batch 有关？这修正了什么错误直觉？
8. 什么是"带宽墙"？它随硬件迭代是变好还是变坏？推论是什么？

🛑 可以停在这里

⚡ 走神救援

下一节 👉 04-Roofline与MFU.md ⭐