📚 课程目录第4章显存泄漏与碎片诊断实战

4. 显存泄漏与碎片诊断实战

显存问题在推理服务里通常表现为两类截然不同的现象：一类是泄漏——显存占用随请求数单调上升，最终 OOM；另一类是碎片——总空闲显存看起来够用，但申请大块连续显存时失败。两者的诊断路径完全不同，用错工具会浪费大量时间。本节围绕三个工具建立一条可复现的诊断链：torch.cuda.memory_summary 看分配器内部账本，py-spy 定位 Python 侧持有引用的调用栈，Nsight Systems 观察 CUDA 层面的真实分配与释放时序。

一、先分清泄漏与碎片：从 memory_summary 读账本

PyTorch 的 CUDA 缓存分配器把显存分成"已分配（allocated）"和"已缓存（reserved）"两层。泄漏通常表现为 allocated 持续增长且不回落；碎片则表现为 allocated 稳定，但 reserved 远大于 allocated，且 reserved - allocated 的缓存块无法满足新的连续申请。调用方式很简单：

import torch
print(torch.cuda.memory_summary(device=None, abbreviated=False))
# 关键字段：
# Allocated / Reserved / Active / Inactive_split / Non-releasable
# 以及各 size bucket 的 block 数量与碎片统计

阅读时重点看三处：Allocated 曲线是否随迭代单调上升；Reserved 与 Allocated 的差值是否长期维持在高位；Inactive_split 与 Non-releasable 是否异常大。若 Allocated 稳定而 Reserved 高，说明是碎片或缓存未释放；若 Allocated 本身持续增长，则是真正的对象泄漏。为了观察趋势，建议在固定步数打点：

for step in range(N):
    run_inference()
    if step % 50 == 0:
        print(f"step={step} alloc={torch.cuda.memory_allocated()/2**20:.1f}MB "
              f"reserved={torch.cuda.memory_reserved()/2**20:.1f}MB")

下面这张图给出从现象到工具再到结论的判断路径，帮助你在拿到一组显存曲线后快速分流。

二、用 py-spy 定位 Python 侧引用持有

确认是泄漏后，下一步是找到"谁还拿着这个 tensor"。py-spy 的优势是可以在不重启、不修改代码的情况下 attach 到运行中的进程，抓取 Python 调用栈。对推理服务常用两种模式：

# 抓取当前所有线程的调用栈快照
py-spy dump --pid <PID>

# 采样 30 秒，生成火焰图，观察哪些函数持续占用
py-spy record -o profile.svg --pid <PID> --duration 30

判读要点：如果 dump 输出里某个业务函数（例如 cache_result、append_history）反复出现在栈顶附近，且其局部变量或全局容器持有 tensor，就基本锁定泄漏点。常见模式有三类：全局 list 只 append 不清理；functools.lru_cache 缓存了带 tensor 的返回值；异常路径下 try/except 把中间结果存进了类属性。定位后可用 gc.get_referrers() 进一步确认引用链：

import gc
target = some_tensor
for ref in gc.get_referrers(target):
    print(type(ref), ref if not isinstance(ref, dict) else list(ref.keys())[:5])

三、用 Nsight Systems 看 CUDA 分配时序

碎片问题在 Python 层往往看不出根因，需要下沉到 CUDA API 层。Nsight Systems 能记录 cudaMalloc、cudaFree、cudaMallocAsync 以及 kernel 执行的时间线。采集命令：

nsys profile -o infer_trace --trace=cuda,nvtx \
  --capture-range=cudaProfilerApi \
  python serve.py

在时间线上重点观察：是否在稳态推理阶段仍频繁出现 cudaMalloc；每次申请的 block 大小是否忽大忽小（例如 8MB、64MB、200MB 交替）；cudaFree 是否几乎不出现。若稳态阶段仍反复 malloc 且大小不一，说明分配器无法复用已有 block，碎片在累积。此时可结合 PYTORCH_CUDA_ALLOC_CONF 调整：

# 扩大可分割 block 的阈值，减少大块被切碎
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128

expandable_segments:True 让分配器使用可扩展段，缓解外部碎片；max_split_size_mb 限制大块被过度切分。修改后重新采集 Nsight 时间线，对比稳态阶段的 malloc 次数与 block 大小分布，若 malloc 次数显著下降、block 趋于规整，说明碎片得到控制。

四、三工具协同的实操闭环

把三者串成一条可复现的排查链：先用 memory_summary 打点确认是泄漏还是碎片；泄漏走 py-spy dump 找引用持有者，修复后观察 Allocated 是否回落；碎片走 Nsight Systems 看分配时序，调整分配器参数后对比 malloc 频率。一个可直接套用的检查清单如下：

观察项	工具	正常表现	异常信号
Allocated 趋势	memory_summary	稳态波动，不单调上升	随请求数持续增长
Reserved - Allocated	memory_summary	差值稳定且可复用	差值持续扩大
Python 调用栈	py-spy dump	无长期持有 tensor 的容器	全局缓存/闭包累积
cudaMalloc 频率	Nsight Systems	启动后趋于平稳	稳态仍频繁申请
block 大小分布	Nsight Systems	规整、可复用	大小交替、无法复用

排查顺序建议：先确认现象类型（泄漏 vs 碎片），再选工具。对泄漏，py-spy 的调用栈比 Nsight 更直接；对碎片，Nsight 的分配时序比 Python 层日志更有说服力。两者不要混用，否则容易在错误的层面上反复尝试。

注意：torch.cuda.empty_cache() 只能释放缓存分配器中未被引用的 block，对真正的对象泄漏无效。若调用后 Reserved 下降但 Allocated 不降，说明仍有 Python 对象持有显存，应回到 py-spy 继续排查，而不是反复调用 empty_cache。

最后强调一个易被忽略的点：诊断时务必固定输入规模。变长序列、动态 batch 会让显存曲线本身带有噪声，掩盖泄漏或碎片的真实趋势。建议用固定 shape 的合成输入做基线，再逐步引入真实流量，这样 memory_summary 的曲线和 Nsight 的时间线才具备可比性，结论也才站得住。