想象你要搬家。所有家当要先装进一辆货车运到新家，装车的人必须反复盘算：什么时候车厢里同时堆的箱子最多？如果那一瞬间超了载重，整个搬家计划就得重来——要么少带东西，要么把大件拆开分几趟。 AI 训练就是一场持续数小时甚至数天的“搬家”。显卡的内存是车厢，每一步计算都会产生一些暂存的“纸箱”，用完后又被扔掉。纸箱不会同时都在车上——有的已经送进“新家”（释放），有的还等着下一趟计算（存活）。“车厢最重的那一刻”，就是计算机要应付的压力巅峰。 这篇文章用“搬家装车”的思路，带你搞懂三件事：怎么靠搬家清单提前算出最危险的时刻（静态估算）、怎么在地磅上实测验证（profiling 工具）、以及怎么拆箱换包装化解风险（优化策略）。看懂峰值，你就看懂了显存规划的一半。

岗位关联

掌握峰值内存分析，你可以在训练前预判显存瓶颈，在 OOM 发生时快速定位元凶，把显存规划从“碰运气”变成“算得清”——这是大模型时代 AI 工程师的基本功。

目录
  1. 可视化演示
  2. 内存基本概念
  3. 峰值内存定义
  4. 拓扑序估算
  5. Profiling工具
  6. 峰值来源分析
  7. 优化策略选择
  8. 规划与验证
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示计算图张量生命周期与峰值内存估算、激活值检查点优化及内存复用策略，交互切换查看峰值定位与减负方法。

本节配有交互动效（进入视口后自动加载，约 3KB）。

生命周期：张量时间区间与峰值 张量从创建到消费的存活区间 共存张量叠加形成峰值区间 M_peak=max_load(t) 拓扑序估算峰值显存

内存基本概念

搬家之前，你得先把家当分成两堆：证件、首饰这类要一路贴身带的，和装进纸箱、随取随用的杂物。AI 程序的显存占用也天然分两拨：一拨从程序启动到结束都赖在显存里不走，另一拨随计算任务来了又走。

前者叫静态内存，装的是模型参数、嵌入表、常量缓存；后者叫动态内存，装的是激活值、中间结果和梯度。想算清楚“货车最重的那一刻”，先得给每一箱标好重量——也就是算出一个张量到底占多少显存。

张量占用的字节数只由两个因素决定：元素总数和单元素字节数。元素总数等于形状各维度的乘积，单元素字节数由数据类型决定——FP32 每元素 4 字节，FP16 和 BF16 只有 2 字节。

这就好比纸箱的重量等于体积乘纸板密度。先给每只纸箱称好重，后面的装车模拟才有依据。静态内存像一个永远占着车厢角落的大衣柜，动态内存才是不断上下车的那些箱子——两者的称重规则，用同一杆秤就能称。

S=N×b,N=C×H×WS = N \times b, \qquad N = C \times H \times W
一个张量的显存占用 = 元素总数 × 每元素字节数，元素总数是各维度大小的乘积。
  1. 元素总数 NN 等于张量各维度大小的乘积，形状为 (C,H,W) 的张量共有 C×H×W 个元素。
    N=C×H×WN = C \times H \times W
  2. 单元素字节数 bb 由数据类型决定，FP32 取 4 字节，FP16/BF16 取 2 字节。
    b=dtype_sizeb = \text{dtype\_size}
  3. 张量总字节数 S元素总数乘以单元素字节数，得到该张量在显存中的占用字节数。
    S=N×bS = N \times b

峰值内存定义

装车时你会发现，车厢里的纸箱数量忽多忽少：这箱刚搬上车，那箱已经被搬进新家。最危险的时刻，不是今天总共搬了多少箱，而是某一瞬间车厢里同时堆着多少箱——再多一箱，货车就可能压垮。

计算图执行也一样。每一步算子执行都会做两件事：把不再需要的张量释放掉（纸箱下车），把新算出的张量分配出来（纸箱上车）。“存活”张量就是那些已上车还没下车的箱子，它们的字节数之和就是这一瞬间的显存占用。

把每一步的瞬时占用画成一条曲线，曲线的最高点就是峰值内存。它决定了你的卡够不够用，也决定了 batch size 和模型规模的天花板。

峰值曲线就像车厢载重的“心电图”：心跳最高的那一下，才是判断会不会超载的唯一标准。抓不住它，前面搬得再顺利也是白费。

Mpeak=max⁡0≤t≤T∑i∈alive(t)SiM_{peak} = \max_{0 \le t \le T} \sum_{i \in \text{alive}(t)} S_i
峰值内存 = 所有时间步中，存活张量字节数总和的最大值。
  1. 存活张量集合alive(t) 表示第 t 步仍占用显存的张量集合，张量从被分配到释放的区间都属于它。
    i∈alive(t)i \in \text{alive}(t)
  2. 瞬时内存求和把集合内所有张量的字节数 S_i 加起来，就是这一步的瞬时显存占用。
    M(t)=∑i∈alive(t)SiM(t) = \sum_{i \in \text{alive}(t)} S_i
  3. 取最大值遍历全部时间步，最大的 M(t) 就是峰值内存。
    Mpeak=max⁡tM(t)M_{peak} = \max_{t} M(t)

拓扑序估算

搬家有铁一般的顺序：先搬床架才放得下床垫，先装灯泡才能通电试灯。顺序乱了，后面全得返工。计算图也有这样的先后约束——一个算子的输出，必须等前面的算子算完才能出现。

把这种依赖关系展开成一个合法的执行序列，就是拓扑排序。有了这个序列，我们就能像看装车录像一样，从头到尾推演每一步哪些张量下车、哪些上车，算出每个时刻车厢里的总重量。

推演的规则只有两条：执行某算子时，新输出张量上车；某张量最后一次被消费后，立刻下车。内存的递推式写成公式就是这样。

这张“装车时间表”就是静态估算的核心。不需要真卡、不需要真跑，拿一张纸就能预判最危险的时刻——搬家公司最喜欢这种提前量。

M(tk)=M(tk−1)−∑i∈free(tk)Si+∑j∈alloc(tk)Sj,Mpeak=max⁡kM(tk)M(t_k) = M(t_{k-1}) - \sum_{i \in \text{free}(t_k)} S_i + \sum_{j \in \text{alloc}(t_k)} S_j, \qquad M_{peak} = \max_{k} M(t_k)
每步新内存 = 上一步内存 - 释放的张量 + 新分配的张量，最后在所有步骤中取最大值。
  1. 上一步内存M(t_{k-1}) 是第 k 个算子执行前一刻的瞬时内存。
    M(tk−1)M(t_{k-1})
  2. 释放集合free(t_k) 是第 k 步执行完成后被释放的张量，它们的最后一次消费都发生在该步。
    i∈free(tk)i \in \text{free}(t_k)
  3. 分配集合alloc(t_k) 是本算子新产生的输出张量，它们从该步开始存活。
    j∈alloc(tk)j \in \text{alloc}(t_k)
  4. 递推与峰值新内存 = 旧内存 - 释放 + 分配，遍历所有步骤取最大即为峰值。
    Mpeak=max⁡kM(tk)M_{peak} = \max_{k} M(t_k)

Profiling工具

清单估算再准，也怕路上出意外：道路颠簸、箱子移位、临时加货。搬家公司最后都会在车上装地磅，每搬一箱称一笔，事后回放数据看车厢哪一段最重。

Profiling 工具就是训练时的“地磅”。它记录每一次显存的分配和释放事件，关联到具体算子，画出一条真实的显存占用时间线。它与静态估算互补：前者给理论值，后者给实测值。

常用工具三件套：PyTorch Profiler 记录算子级显存事件并输出 trace 时间线；NVIDIA Nsight Systems 提供 kernel 级视图；Nsight Compute 深入单个 kernel 分析资源占用。注意，地磅本身也有重量——profiling 是有开销的。

多花几分钟称重，换来的是“知道哪一刻超载、是谁造成的”。这笔账怎么算都划算——总比到了半路货车抛锚（OOM）再回头强。

O=Tprofiled−TbaselineTbaseline×100%O = \frac{T_{profiled} - T_{baseline}}{T_{baseline}} \times 100\%
profiling 开销 = （带分析的时间 - 基线时间） ÷ 基线时间，通常有 5%~30% 的额外耗时。
  1. 基线时间T_baseline 是不开启任何 profiling 时训练循环的执行时间。
    TbaselineT_{baseline}
  2. 带分析时间T_profiled 是开启 profiling 后同样循环的执行时间，包含记录开销。
    TprofiledT_{profiled}
  3. 开销率两者之差除以基线时间并换算成百分比，得到 profiling 的额外开销比例。
    O=Tprofiled−TbaselineTbaseline×100%O = \frac{T_{profiled} - T_{baseline}}{T_{baseline}} \times 100\%

峰值来源分析

回放地磅数据时你会发现，货车最重的那一刻，往往是几件大家具恰好挤在一起：大沙发、大冰箱、大衣柜同时占着车厢。峰值来源分析，就是找出这些“大家具”是谁。

计算图里的“大家具”主要有三类。头号嫌疑是尺寸巨大的激活值张量：Transformer 的中间特征动辄形状为 B×L×H，单张就占几个 GB。

第二个嫌疑是中间结果堆积：有些框架或实现把中间值缓存下来不释放，本该下车的纸箱赖在车上不走。第三个是显存碎片化：小块分配释放把显存切得零零碎碎，空间够但不连续，大张量塞不进去，占用率虚高。

把几只“大家具”从车上请下来各归其位，峰值自然回落。对症下药的前提，是先认出它们是谁——这就是来源分析的价值。

Sact=B×L×H×bS_{act} = B \times L \times H \times b
单个激活张量的显存占用 = 批量大小 × 序列长度 × 隐藏维度 × 单元素字节数。
  1. 批量大小 B一次输入模型的样本数，是激活张量的第一维。
    BB
  2. 序列长度 L文本或序列的 token 数，是激活张量的第二维。
    LL
  3. 隐藏维度 HTransformer 隐藏层宽度，是激活张量的第三维。
    HH
  4. 相乘再乘字节数 b三个维度乘积乘以 b，得到单个激活张量的显存占用。
    Sact=B×L×H×bS_{act} = B \times L \times H \times b

优化策略选择

知道超载原因后，搬家公司有三板斧：大件拆开分批运、腾出空间先搬别的、用轻便的压缩包装替代笨重木箱。AI 训练里对应的招数叫梯度重计算、内存复用、混合精度。

梯度重计算（又称激活检查点）不再保存每一层激活值，只保存少量检查点，反向传播时把中间的激活重新算一遍——用算力换显存。内存复用让生命周期不相交的张量共享同一块显存；混合精度把 FP32 降到 FP16/BF16，字节数直接减半。

以重计算为例，把计算图切成 K 段，只保留每段输入的检查点。激活内存从“全存”变成“存一段加检查点”，峰值大幅下降，代价是训练时间通常增加 20%~30%。

拆开的沙发到新家还得重新组装，多花力气但车不超载。用时间换空间，是显存紧张时最实惠的买卖。

Mact′≈MactK+MckptM'_{act} \approx \frac{M_{act}}{K} + M_{ckpt}
重计算后的激活内存 ≈ 原始激活内存 ÷ 分段数 + 检查点内存，用额外计算换取显存。
  1. 原始激活内存M_act 是不采用重计算时缓存全部激活值所需的显存。
    MactM_{act}
  2. 分段数 K把计算图切成 K 段，只保存每段的输入作为检查点。
    KK
  3. 检查点内存M_ckpt 是全部检查点张量占用的显存总和。
    MckptM_{ckpt}
  4. 重计算后内存任意时刻只有一段激活在内存中，总占用约为 M_act/K 加检查点。
    Mact′≈MactK+MckptM'_{act} \approx \frac{M_{act}}{K} + M_{ckpt}

规划与验证

靠谱的搬家公司出发前会做载重预算：按清单算出最大载重，再加两成余量决定租多大的车。到新家后还要复核家具是否完好、实际载重是否和预估一致。

内存规划也一样：先用静态估算（或上次 profiling 结果）得到峰值，乘上安全系数得到显存预算，据此决定 batch size、序列长度和模型规模能不能塞进一张卡。安全系数通常取 0.2~0.3，用来兜住框架缓存和碎片化等意外。

更重要的是一条纪律：每次修改模型结构或超参后，都要重新做 profiling 回归对比，确保峰值没有异常抬升。预算不是一次性作业，而是伴随训练的长期账本。

留余量、做复盘，搬家和训练才能都不“翻车”。峰值内存分析不是一次测量，而是一套贯穿开发始终的工作流。

Mbudget=Mpeak×(1+α),α∈[0.2,0.3]M_{budget} = M_{peak} \times (1 + \alpha), \qquad \alpha \in [0.2, 0.3]
内存预算 = 峰值估算 × (1 + 安全系数)，安全系数通常取 0.2 到 0.3。
  1. 峰值估算M_peak 来自静态拓扑序推演或上一次 profiling 结果。
    MpeakM_{peak}
  2. 安全系数α 是预留余量比例，通常取 0.2~0.3，覆盖碎片化与框架缓存。
    α∈[0.2,0.3]\alpha \in [0.2, 0.3]
  3. 最终预算预算等于峰值乘上 (1+α)，用来决定 batch size 和模型规模。
    Mbudget=Mpeak×(1+α)M_{budget} = M_{peak} \times (1 + \alpha)

AI 应用场景

大模型训练前的显存预算评估

阶段：训练规划与资源配置

解决的问题：传统做法是“先租卡再试”，batch size 开多大全凭经验，经常出现训练一启动就 OOM，浪费算力和时间。用峰值内存分析，在写训练代码前就能按拓扑序估算各层激活内存，乘上安全系数得到预算。

具体案例：某团队要训练 7B 参数 LLM，先用分析脚本模拟前向/反向的峰值，发现序列长 4096、batch=8 时峰值约 62GB，超过单卡上限；于是把 batch 调到 4 并启用梯度检查点，估算峰值降到 52GB，一次通过，省下多次试错成本。

训练 OOM 故障的快速定位

阶段：训练问题排查

解决的问题：训练中途 OOM，报错只告诉你“显存不足”，不告诉你哪一步、哪个算子。用 profiler 抓完整时间线，对比每一步显存曲线，立刻能看到峰值出现在哪个算子。

具体案例：某 NLP 训练在第 47 步 OOM，用 torch.profiler 抓 trace 后发现峰值来自 attention 的中间 QKV 拼接张量——一个 3.2GB 的临时大激活值没有被及时释放。定位后优化算子执行顺序并做内存复用，峰值下降 18%，训练平稳跑完。

推理服务的 batch 上限调优

阶段：推理部署优化

解决的问题：在线推理服务想在显存允许范围内尽量加大 batch 提吞吐，但盲目调大会导致 OOM 崩掉整个服务。用峰值内存分析算出每个请求的增量内存，设置安全 batch 上限。

具体案例：某文本生成服务用峰值分析脚本绘制“batch 大小 vs 峰值显存”曲线，发现 batch=8 时峰值 11GB，batch=16 时跳到 23GB 接近上限，于是把上限设为 12，吞吐提升 40% 且无 OOM。

Python 演示

python72 行
# -*- coding: utf-8 -*-"""基于拓扑序的峰值内存估算模拟器.用搬家装车思维: 每个算子=一次搬运动作, 张量=纸箱,显存=车厢, 峰值=车厢里纸箱总重量最大的时刻."""  class Tensor:    """张量(纸箱): 记录大小和生产者/消费者算子编号"""     def __init__(self, tid, size, producer, consumers):        self.tid = tid        self.size = size        self.producer = producer        self.consumers = consumers  def build_graph():    """构造一个简化计算图, 返回 (算子列表, 张量字典)"""    ops = [        (0, [], [0, 1]),        (1, [0], [2]),        (2, [1], [3]),        (3, [2, 3], [4]),        (4, [4], [5, 6]),        (5, [5], [7]),        (6, [6, 7], []),    ]    sizes = {0: 4, 1: 6, 2: 8, 3: 5, 4: 9, 5: 7, 6: 3, 7: 5}    tensors = {}    for tid in sizes:        producer = None        consumers = []        for op_id, ins, outs in ops:            if tid in outs:                producer = op_id            if tid in ins:                consumers.append(op_id)        tensors[tid] = Tensor(tid, sizes[tid], producer, consumers)    return ops, tensors  def simulate(ops, tensors):    """按拓扑序模拟显存分配与释放, 返回 (快照列表, 峰值)"""    last_use = {}    for op_id, ins, _ in ops:        for tid in ins:            last_use[tid] = op_id    cur = 0    peak = 0    snapshots = []    for op_id, ins, outs in ops:        alloc = sum(tensors[t].size for t in outs)        peak = max(peak, cur + alloc)      # 执行瞬间: 输入尚未释放        for tid in ins:            if last_use[tid] == op_id:                cur -= tensors[tid].size   # 最后一次消费, 纸箱下车        for tid in outs:            cur += tensors[tid].size       # 新输出上车        snapshots.append(cur)    return snapshots, peak  def main():    ops, tensors = build_graph()    snapshots, peak = simulate(ops, tensors)    print("memory after each op (MB):", snapshots)    print("peak memory (MB):", peak)  if __name__ == "__main__":    main()
目的：模拟一个 7 算子计算图按拓扑序执行时的显存分配与释放，输出每步瞬时内存和峰值内存，演示静态估算的核心逻辑。
  1. last_use[tid] = op_id记录每个张量最后一次被哪个算子消费，决定它的“下车”时机。
  2. peak = max(peak, cur + alloc)算子执行瞬间，输入还没释放而新输出已分配，是瞬态峰值最容易出现的时刻。
  3. cur -= tensors[tid].size最后一次消费完成后，纸箱下车，显存归还。
  4. cur += tensors[tid].size新输出张量上车，显存占用增加。
  5. snapshots.append(cur)记录每一步执行后的内存快照，形成可视化曲线。

常见误区

误区我总觉得只要把模型参数、每一层的输出都加起来，就知道需要多大显存了，峰值内存不就是这个总数吗？

正解峰值内存是“同一时刻”同时存活张量的总和，不是所有张量的累计。就像搬家货车永远装不下全部家当，装车时总是一边卸一边装，最重的那一刻可能只有全部家当的三成。若按累计总数规划，你会多租一辆根本用不上的车。

为何常见：人们习惯用“总量”思维理解存储，忽略了张量有生命周期、会随计算释放这个动态事实。

误区显存不够就把 batch 减半，那 60GB 的峰值不就变成 30GB 了吗？

正解大部分激活值的尺寸与 batch 成正比，但模型参数、优化器状态、梯度张量、嵌入表与 batch 无关。就像一个车队里，司机的体重不会因为乘客减半而减半。batch 减半只对激活相关部分有效，对固定部分毫无影响。

为何常见：大家最容易观察到与 batch 相关的激活张量，忽略了显存里还有其他“常住人口”。

误区打开 profiler 记录一下显存而已，应该对训练速度没有影响吧？

正解profiling 要拦截和记录每一次显存分配、释放和算子事件，本身有开销。开销公式 O=(T_profiled - T_baseline)/T_baseline 通常带来 5%~30% 的额外耗时。就像装地磅称重：每箱过磅都要多花几秒，但换来的是精确记录。

为何常见：profiling 的结果看起来像“免费的额外信息”，人们低估了观测行为自身的代价。

误区碎片化不就是显存不足吗？换一张更大的显卡就好了。

正解碎片化指显存总量足够，但空闲块被切得零零碎碎，放不下一个连续的大张量。就像车厢里还有空间，但被小箱子占满隔档，大沙发塞不进去。换更大的车能缓解，但治本要靠内存复用规划和分配器策略。

为何常见：OOM 报错长得一样，人们容易把所有显存问题都归结为“容量不够”。

误区框架不是会自动复用显存吗？那峰值内存分析是不是就没用了？

正解内存复用只能让“生命周期不相交”的张量共享空间。如果两个大张量同时存活，谁也不让谁，复用也无能为力。峰值分析正是要找出那些“同时存活”的组合，才能决定哪些可以复用、哪些需要重计算或降精度。

为何常见：看到框架宣传内存复用，误以为优化是自动且无条件的，忽略了生命周期约束。

面试题

Q1. 峰值内存指的是训练过程的哪一刻？

单选题 · easy

  • A.训练数据全部读取的时刻
  • B.反向传播全部结束的时刻
  • C.模型参数加载完毕的时刻
  • D.存活张量占用显存最大的时刻← 正确答案

答案：D。峰值内存是某一瞬间所有存活张量占用显存的最大值，对应“车厢最重的那一刻”。其他选项都是单一事件时刻，没有考虑“多个张量同时存活”这个动态概念。

Q2. 一个 FP32 张量的元素总数为 1024×1024，它占用多少 MB？

单选题 · easy

  • A.约 16 MB
  • B.约 8 MB
  • C.约 4 MB← 正确答案
  • D.约 1 MB

答案：C。FP32 每元素 4 字节，元素总数 1024×1024=1048576，乘 4 得 4194304 字节 ≈ 4MB。这就像知道纸箱体积和纸板密度，重量一乘便知。

Q3. 拓扑排序对峰值内存分析的主要作用是？

单选题 · mid

  • A.确定算子执行顺序以估算峰值← 正确答案
  • B.缩小网络中的参数总量
  • C.直接减少模型的参数量
  • D.替换计算图的算子类型

答案：A。拓扑排序把计算图的依赖关系展开成合法执行序列，才能模拟张量“何时上车、何时下车”。它不改变参数数量或算子类型，只是确定“搬运顺序”。

Q4. 梯度重计算（检查点）主要降低哪类显存占用？

单选题 · mid

  • A.中间激活值所占用的显存← 正确答案
  • B.模型参数所占用的显存
  • C.优化器状态所占用的显存
  • D.训练样本所占用的显存

答案：A。梯度重计算将激活值分段保存、反向时重算，直接砍掉的是中间激活的显存。参数和优化器状态依然全量常驻，不受影响——就像只拆大件沙发，不动固定家具。

Q5. 静态估算中张量释放的正确时机是其？

单选题 · mid

  • A.被创建出来的那一时刻
  • B.网络完成所有训练的时刻
  • C.最后一次被消费执行完成时← 正确答案
  • D.下一个算子执行的开始时

答案：C。张量在最后一次被消费的算子执行完成后即可释放，这正是拓扑序模拟中“纸箱下车”的判定标准。过早释放会导致数据丢失，过晚释放会虚增峰值。

Q6. profiling 开销公式中 T_baseline 表示的是？

单选题 · hard

  • A.开启分析器后的执行时间
  • B.显存分配器的调用次数
  • C.模型参数总量的字节数
  • D.未开启分析器的基线时间← 正确答案

答案：D。T_baseline 是未开启 profiling 的基线执行时间，T_profiled 是开启后的时间，两者差值比例即开销率。地磅本身的重量只有和空车对比才能算出来。

Q7. 下列哪一项不是显存峰值升高的常见来源？

单选题 · hard

  • A.序列长度很长的大激活值
  • B.模型参数被复制到了CPU← 正确答案
  • C.未及时释放的中间结果
  • D.显存碎片化导致空间不足

答案：B。大激活值、中间结果堆积、碎片化都是峰值升高的典型来源。参数复制到 CPU 反而释放了显存，不会抬高峰值——就像把大件暂时搬下车，车厢反而更空。

Q8. 下列哪些做法能降低训练时的峰值内存？

多选题 · mid

  • A.使用梯度重计算方法← 正确答案
  • B.随意增大训练的批次大小
  • C.不相交张量内存复用← 正确答案
  • D.采用混合精度来减小存储← 正确答案

答案：ACD。梯度重计算减少激活缓存，内存复用让不相交张量共享空间，混合精度让每个张量字节减半，三者都直接压低峰值。增大 batch 会放大激活张量，反而推高峰值。

Q9. 关于内存预算与安全边界，正确的认识有？

多选题 · hard

  • A.预算只需等于模型参数大小
  • B.预算应以峰值估算乘以余量系数← 正确答案
  • C.安全系数通常选择 0.2 到 0.3← 正确答案
  • D.每次改动后需做回归对比验证← 正确答案

答案：BCD。预算 = 峰值 × (1+α)，α 取 0.2~0.3 覆盖碎片和框架缓存；每次改动都要回归验证峰值。预算只算参数大小会严重低估——货车不可能只装沙发不装别的。

Q10. 张量的生命周期指的是它从？

单选题 · easy

  • A.开始训练到训练结束全程
  • B.被分配到最后一次使用的区间← 正确答案
  • C.数据加载到模型输出的过程
  • D.被保存到被删除的文件区

答案：B。生命周期 = 从张量被分配（纸箱上车）到它最后一次被使用（纸箱下车）的区间。只有落在生命周期内的步骤，该张量才算“存活”并计入峰值。