一、显存基础与瓶颈分析

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

说实话，我见过太多团队，模型设计得漂漂亮亮，一上GPU训练就崩。为啥？显存爆了。你想想看，一张A100 80G的卡，跑个LLaMA-7B都费劲，更别提更大的模型了。所以，搞懂显存，是高性能训练的第一步。

1.1 GPU显存架构：HBM vs GDDR

先说说显存本身。现在的GPU显存主要有两派：HBM（高带宽内存）和GDDR（图形双倍数据率内存）。

HBM 是高端货，像A100、H100用的就是HBM2e或HBM3。它的特点是带宽极高——A100有2TB/s，H100能到3.35TB/s。但容量有限，单卡最多80G。

GDDR 则是消费级显卡的标配，比如RTX 4090用的GDDR6X。带宽大概1TB/s，但容量可以做得很大，24G、48G都有。

我个人习惯把HBM比作「跑车」，GDDR比作「卡车」。跑车速度快但拉货少，卡车装得多但跑得慢。训练大模型，你既需要速度（带宽），也需要容量。

1.2 显存带宽与容量：你的瓶颈在哪？

我在项目中遇到过不少同学，只看显存容量，不看带宽。结果模型是塞进去了，但训练慢得像蜗牛爬。

举个例子：

GPU型号	显存容量	显存带宽	适用场景
A100 80G	80 GB	2 TB/s	大模型训练
RTX 4090	24 GB	1 TB/s	中小模型/推理
H100	80 GB	3.35 TB/s	超大模型训练

你看，A100和H100容量一样，但H100带宽高了近70%。这意味着同样的模型，H100的数据搬运速度更快，训练吞吐量更高。

那怎么判断自己的瓶颈是带宽还是容量？很简单：

• 如果训练时GPU利用率低（比如<50%），但显存没爆，那多半是带宽瓶颈——数据来不及喂给计算单元。
• 如果训练时直接OOM（Out of Memory），那就是容量瓶颈——显存不够用了。

1.3 训练中的显存占用分布

好，现在咱们看看训练时显存到底被谁吃了。我把它分成四大块：

1. 模型参数：就是权重和偏置。一个FP32的参数占4字节，BF16占2字节。
2. 梯度：反向传播时计算的梯度，大小和参数一样。
3. 优化器状态：比如Adam优化器会保存一阶动量和二阶动量，每个参数额外占8字节（FP32）。
4. 激活值：前向传播时每层的中间结果。这玩意儿是显存大户，尤其是大模型。

我给大家算笔账：一个70亿参数的模型，用FP32训练：

• 参数：7B × 4B = 28 GB
• 梯度：7B × 4B = 28 GB
• 优化器状态（Adam）：7B × 8B = 56 GB
• 激活值：取决于batch size和序列长度，少说几十GB

加起来轻松超过100GB。一张A100才80G，你说怎么训？

1.4 显存瓶颈诊断方法

怎么诊断显存瓶颈？我一般用三板斧：

第一斧：看nvidia-smi

nvidia-smi -l 1  # 每秒刷新一次

看显存使用率（Memory-Usage）和GPU利用率（GPU-Util）。如果显存使用率接近100%，那就是容量瓶颈。如果GPU利用率低但显存使用率不高，可能是带宽或计算瓶颈。

第二斧：用PyTorch Profiler

import torch.profiler as profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True
) as prof:
    # 你的训练代码
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()

print(prof.key_averages().table(sort_by="self_cuda_memory_usage"))

这个能精确看到每个操作占了多少显存。我一般先跑一个小batch，看看哪个层最吃显存。

第三斧：手动估算

嗯，这个方法比较土，但很实用。我习惯在代码里加几行打印：

import torch

def print_memory_usage():
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    print(f"已分配: {allocated:.2f} GB, 已预留: {reserved:.2f} GB")

# 在训练循环中调用
print_memory_usage()

这样你能实时看到显存变化，找到峰值在哪。

知识体系总览

下面这张图是我自己画的，把本章的核心逻辑串起来了：

这张图把本章的核心内容串起来了。从显存架构出发，到关键指标和占用分布，再到诊断方法，最后引出后续的优化策略。你把它当个地图用，学起来会清晰很多。

好了，第一章就到这里。记住，显存优化不是一锤子买卖，而是一个持续诊断、持续优化的过程。下一章咱们聊聊混合精度训练——一个能立竿见影省显存的手段。