一、显存基础与瓶颈分析

大家好，我是你们的老朋友。今天咱们聊聊显存——这个在深度学习里既烧钱又让人头疼的东西。

说实话，我刚开始做模型训练那会儿，对显存的理解就是「越大越好」。直到有一次，我花大价钱买了块48G的卡，结果跑一个不算大的模型居然OOM了。当时我就懵了。后来才发现，问题根本不在容量，而在带宽和访问模式上。

所以这一章，咱们把显存这玩意儿彻底掰开揉碎。从架构到瓶颈，从理论到工具，一步到位。

1.1 显存架构：GPU怎么「记住」数据

显存，全称是「视频内存」，但在深度学习里，它就是GPU的工作台。你的模型参数、中间激活值、优化器状态，全得放这儿。

现代GPU的显存架构，说白了就是三层结构：

• 全局显存（Global Memory）：最大的那块，几十GB。所有线程都能访问，但速度最慢。延迟大概400-800个时钟周期。
• 共享内存（Shared Memory）：每个线程块独享，几十KB到上百KB。速度极快，延迟只有几个周期。我习惯把它比作「团队内部的小黑板」。
• 寄存器（Registers）：每个线程私有，容量最小（每个线程最多255个32位寄存器）。但速度最快，零延迟。

嗯，这里要注意：很多人以为显存就是全局显存那一块。其实不是。你写的每个变量，都可能在这三层之间来回倒腾。我曾经在优化一个Transformer模型时，发现共享内存用好了，训练速度直接提升了30%。

1.2 带宽与容量：两个不同的瓶颈

很多人把显存带宽和容量混为一谈。我刚开始也犯过这个错。其实它们是两码事：

指标	含义	典型值（A100）	瓶颈场景
容量	能装多少数据	40GB / 80GB	大模型、大Batch Size
带宽	每秒能读写多少数据	1.6 TB/s（HBM2e）	数据搬运频繁、小Batch

你想想看，容量就像你的硬盘大小，带宽就像硬盘的读写速度。硬盘再大，读写慢也是白搭。

我在项目中遇到过这样一个案例：训练一个BERT-Large模型，显存容量明明够用（用了32GB），但训练速度就是上不去。后来用工具一测，发现GPU利用率只有40%多，大部分时间都在等数据从显存搬到计算单元。这就是典型的带宽瓶颈。

1.3 常见瓶颈场景：你大概率遇到过

根据我的经验，显存瓶颈通常出现在以下几个场景：

1. 模型太大，装不下：比如你想在12GB的卡上跑LLaMA-7B。这就像让一个小孩搬一张大桌子——不是不能，但得拆了再装（模型并行、梯度检查点）。
2. Batch Size 太大：激活值占满了显存。我见过有人为了追求「稳定梯度」，把Batch Size设到256，结果显存直接爆了。其实很多时候128就够了。
3. 序列长度太长：Transformer的显存占用和序列长度是平方关系。你把序列长度从512提到1024，显存占用可能翻4倍。
4. 优化器状态太多：Adam优化器需要保存动量和方差，每个参数多占8字节。一个1B参数的模型，光优化器状态就要8GB。
5. 数据加载太慢：CPU来不及喂数据，GPU只能干等。这时候显存利用率很低，但你又不敢减小显存。

为什么会这样？说白了，就是「计算」和「存储」之间的速度差距越来越大。GPU算得越来越快，但显存的读写速度跟不上。这就是所谓的「内存墙」问题。

1.4 性能分析工具：你的「显存CT机」

没有工具，你就是在盲人摸象。我个人的习惯是，拿到一个新模型，先跑一轮性能分析，看看显存到底花在哪儿了。

下面这几个工具，我几乎天天用：

• nvidia-smi：最基础的，看显存占用、温度、功耗。但只能看个大概。
• torch.cuda.memory_summary()：PyTorch自带的，能显示每个张量占了多少显存。我经常用它来排查「显存泄漏」问题。
• Nsight Systems：NVIDIA的官方工具，能看到GPU的完整时间线。哪里在等数据，哪里在计算，一目了然。
• PyTorch Profiler：集成在PyTorch里，能记录每个操作的显存分配和释放。我建议你养成习惯，每次训练都开一下。

下面是一个简单的PyTorch Profiler使用示例：

import torch
from torch.profiler import profile, record_function, ProfilerActivity

# 定义一个简单的模型
model = torch.nn.Linear(1024, 1024).cuda()
input_tensor = torch.randn(64, 1024).cuda()

# 启动性能分析
with profile(
    activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
    record_shapes=True,
    profile_memory=True
) as prof:
    with record_function("model_inference"):
        output = model(input_tensor)
        loss = output.sum()
        loss.backward()

# 打印显存占用最高的操作
print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage", 
    row_limit=10
))

这段代码会告诉你：哪个操作占显存最多，哪个操作最耗时。我上次用它发现，一个不起眼的torch.ones_like()居然占了5%的显存——因为它创建了一个和输入一样大的张量。

1.5 知识体系总览

为了让你更直观地理解这一章的内容，我画了一张图：

这张图把这一章的核心内容串起来了。你从「显存优化」出发，往四个方向展开：架构、带宽与容量、瓶颈场景、分析工具。每个方向都是后续章节的基础。

好了，这一章就到这里。记住一句话：显存优化不是玄学，是工程。下一章咱们会深入具体的优化技术，比如梯度检查点、混合精度训练这些实战技巧。