4、显存瓶颈定位：显存容量与带宽瓶颈、显存碎片化问题、显存泄漏检测、统一内存与显存交换开销

显存，说白了就是 GPU 的「工作台」。你模型再大，算子再复杂，最终都得在显存里铺开才能算。我见过太多团队，模型结构设计得漂漂亮亮，一跑起来直接 OOM，或者吞吐量上不去。这时候别急着怀疑算力，先看看显存是不是被「卡脖子」了。

这一章，咱们就聊聊显存瓶颈的四个核心问题：容量与带宽、碎片化、泄漏检测，以及统一内存的交换开销。每个问题我都踩过坑，咱们一个一个说。

4.1 显存容量与带宽瓶颈

显存容量不够，最直接的表现就是 OOM（Out of Memory）。但很多时候，它不是一下子爆掉的，而是慢慢变慢，甚至卡死。我遇到过一位同事，训练一个 13B 的模型，batch size 设得特别大，结果每次迭代都要等显存回收，吞吐量直接腰斩。

怎么判断是容量瓶颈还是带宽瓶颈？我一般看两个指标：

显存利用率：如果利用率长期接近 100%，说明容量是瓶颈。要么换大显存的卡，要么用梯度检查点（Gradient Checkpointing）或混合精度训练来省显存。
显存带宽利用率：用 nvidia-smi 或 Nsight Systems 看带宽占用。如果带宽跑满了，但计算单元（SM）利用率不高，说明数据搬运成了瓶颈。这时候就该考虑算子融合、减少小张量传输，或者用更高效的数据格式。

核心判断逻辑：

显存占用高 + 计算单元空闲 → 容量瓶颈
带宽跑满 + 计算单元等待 → 带宽瓶颈

举个例子，我曾经优化一个推荐系统模型，发现显存占用只有 60%，但带宽利用率高达 95%。后来定位到是频繁的 torch.cat 操作导致大量小张量搬运。改成预分配大张量后，带宽利用率降到 70%，吞吐量提升了 40%。

我的小技巧：用 nvidia-smi dmon -s pucvmet 实时监控显存带宽和利用率，比看 nvidia-smi 的静态信息有用得多。

4.2 显存碎片化问题

显存碎片化，是个很隐蔽的问题。你明明有足够的显存总量，但就是分配不出一个连续的大块内存。我遇到过最夸张的一次，显存还剩 12GB，但分配一个 4GB 的 tensor 直接报错。为什么？因为显存被切成了一堆小碎片。

碎片化主要来自两个方面：

动态张量分配与释放：训练过程中频繁创建和销毁不同大小的 tensor，就像内存堆一样，时间长了就碎成渣。
框架的缓存机制：PyTorch 的缓存分配器会保留已释放的显存块，以便下次快速分配。但如果分配模式不规律，缓存块之间就会产生大量空洞。

怎么检测碎片化？我一般用两个方法：

查看 torch.cuda.memory_summary()：这个函数会输出详细的显存分配情况，包括大块和小块的数量。如果小块（比如 < 1MB）特别多，说明碎片化严重。
写一个简单的测试脚本：尝试分配一个连续的大张量，如果失败，基本可以确认碎片化。

import torch

# 尝试分配一个连续的大张量
try:
    big_tensor = torch.empty(4 * 1024 * 1024 * 1024 // 4, dtype=torch.float32, device='cuda')
    print("分配成功，显存碎片化不严重")
except RuntimeError as e:
    print(f"分配失败：{e}，可能存在显存碎片化")

避坑指南：我曾经在训练一个 NLP 模型时，因为使用了动态长度的序列，导致显存碎片化严重。后来改用 torch.cuda.empty_cache() 手动清理缓存，但治标不治本。真正有效的办法是：尽量固定张量大小，或者使用 torch.cuda.memory.CUDAPluggableAllocator 替换默认分配器。

4.3 显存泄漏检测

显存泄漏，就是显存只增不减。训练跑着跑着，显存占用越来越高，直到 OOM。我见过最离谱的一次，一个模型训练了 200 步，显存从 8GB 涨到了 80GB，最后直接把服务器搞挂了。

泄漏的常见原因：

忘记 detach 计算图：在循环中累积梯度或张量，导致计算图不断增长。
未释放的中间变量：比如在 torch.no_grad() 块外创建了不需要梯度的张量，但引用没释放。
DataLoader 的 num_workers 问题：某些情况下，子进程的显存不会自动回收。

检测泄漏，我推荐三步走：

用 torch.cuda.memory_allocated() 和 torch.cuda.memory_reserved()：在训练循环的每一步打印这两个值。如果 allocated 持续增长，说明有泄漏。
使用 tracemalloc 或 torch.cuda.memory._record_memory_history()：记录每一步的显存分配堆栈，定位到具体的代码行。
写一个最小复现脚本：把模型和数据简化到最小，逐步排查。

import torch

# 在训练循环中监控显存
for step in range(100):
    # 训练代码...
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    print(f"Step {step}: allocated={allocated:.2f} GB, reserved={reserved:.2f} GB")
    
    # 如果 allocated 持续增长，说明有泄漏
    if step > 10 and allocated > 2 * initial_allocated:
        print("警告：可能存在显存泄漏！")
        break

我的经验：泄漏最常见的地方是 loss.backward() 之后忘记 optimizer.zero_grad()，或者循环中创建了新的 nn.Module 但没有释放。养成好习惯：每个 batch 结束后，显存占用应该回到 baseline。

4.4 统一内存与显存交换开销

统一内存（Unified Memory）是 CUDA 的一个特性，允许 CPU 和 GPU 共享同一块虚拟内存空间。听起来很美，但实际用起来，坑不少。

统一内存的核心问题是：当 GPU 访问的数据不在显存中时，会触发 page fault，从 CPU 内存搬运数据到显存。这个搬运过程是同步的，会阻塞 GPU 计算。我见过一个项目，用了统一内存后，训练速度直接慢了 5 倍。

什么时候该用统一内存？

数据访问模式稀疏且不可预测：比如图神经网络中，节点特征随机访问。
显存容量不足，但 CPU 内存充足：可以自动处理数据交换，但要做好性能下降的心理准备。

什么时候千万别用？

数据访问模式规律且密集：比如卷积神经网络的权重和特征图，手动管理显存比统一内存高效得多。
对延迟敏感的场景：比如推理服务，一次 page fault 可能导致请求超时。

怎么监控统一内存的交换开销？我一般用 nvidia-smi 的 Unified Memory 字段，或者用 Nsight Systems 的 CUDA Unified Memory 事件。如果看到大量的 Page Fault 事件，说明交换开销已经成了瓶颈。

一句话总结：统一内存是「懒人神器」，但性能代价很大。能用显存装下的数据，就别偷懒用统一内存。

知识体系图

下面这张图，帮你理清显存瓶颈定位的四个维度：

显存瓶颈的定位，说白了就是「先看总量，再看分布，最后看交换」。容量不够就省着用，带宽不够就少搬运，碎片化就整理，泄漏就修代码。统一内存嘛，能不用就不用。

嗯，这一章的内容就到这里。显存问题往往是最容易忽略的，但也是最容易优化的。你想想看，很多时候模型跑不动，不是算力不够，而是显存没用好。

📚 课程目录

4、显存瓶颈定位：显存容量与带宽瓶颈、显存碎片化问题、显存泄漏检测、统一内存与显存交换开销

4.1 显存容量与带宽瓶颈

4.2 显存碎片化问题

4.3 显存泄漏检测

4.4 统一内存与显存交换开销

知识体系图