第一章：显存占用基础

理解显存分配机制、Tensor生命周期、梯度与优化器状态对显存的影响

做算子优化这些年，我见过太多同学一上来就盯着算子实现猛调，结果显存爆了都不知道怎么爆的。说实话，不理解显存是怎么分配的，优化就是瞎忙活。今天咱们就把这块地基打牢。

1.1 显存分配机制：GPU到底怎么分内存？

GPU的显存分配，说白了就是一块大内存，谁用谁申请。但这里有个关键点——显存分配不是即时的。你写了个 torch.zeros(1024, 1024)，PyTorch 并不会立刻去问 GPU 要这 4MB 空间。

我刚开始做优化时就被这个坑过。有一次我监控显存占用，发现程序启动后显存就占了 2GB，可我明明还没创建任何大张量。后来才明白，PyTorch 的 CUDA 缓存分配器会预先申请一大块显存，然后自己管理。

为什么会这样？因为每次向 GPU 驱动申请/释放显存，都要经过 CUDA runtime，这个开销可不小。缓存机制能大幅提升训练速度。

来看个实际例子：

import torch

# 查看当前显存占用
print(f"初始显存占用: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存大小: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

# 创建一个张量
x = torch.randn(1000, 1000, device='cuda')
print(f"创建张量后: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")

# 删除张量
del x
torch.cuda.empty_cache()  # 清空缓存
print(f"删除并清空缓存: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")

运行这段代码你会发现，即使调用了 empty_cache()，缓存可能也不会完全释放。嗯，这里要注意：empty_cache() 只是把缓存池里的显存还给系统，但 PyTorch 内部的一些元数据占用的空间还在。

1.2 Tensor生命周期：张量从生到死

每个张量都有生命周期：创建 → 使用 → 释放。但 GPU 上的张量生命周期管理，比 CPU 复杂得多。

我习惯把张量生命周期分成三个阶段：

• 创建阶段：分配显存，初始化数据。注意，torch.zeros 和 torch.randn 的初始化开销不同。
• 使用阶段：参与计算，可能被多个算子引用。引用计数管理着它的生死。
• 释放阶段：引用计数归零，显存回到缓存池。

你想想看，一个张量什么时候会被释放？不是它离开作用域就立刻释放，而是 Python 的垃圾回收器发现它不再被引用时。这就引出一个常见问题：循环引用导致显存泄漏。

来看一个典型的生命周期管理代码：

def train_step(batch):
    # 创建阶段
    x = batch['image'].cuda()  # 分配显存
    y = batch['label'].cuda()
    
    # 使用阶段
    pred = model(x)
    loss = loss_fn(pred, y)
    
    # 反向传播
    loss.backward()
    
    # 释放阶段 - 函数返回后，x, y, pred, loss 的引用计数减1
    # 如果外部没有保存它们，就会被垃圾回收
    return loss.item()

# 每次调用 train_step，临时张量都会被释放
for batch in dataloader:
    loss = train_step(batch)

这里有个细节：loss.item() 返回的是 Python 浮点数，不是 GPU 张量。所以 loss 张量在函数返回后就可以被回收了。我见过有人直接返回 loss，然后在外面调用 .item()，这样 loss 张量的生命周期就被延长了——虽然影响不大，但养成好习惯总没错。

1.3 梯度与优化器状态：显存占用的隐形大户

很多人只关注模型参数占了多少显存，却忽略了梯度和优化器状态。说实话，这两块加起来往往比模型本身还大。

咱们算笔账：

组件	每个参数占用的显存	说明
模型参数 (fp32)	4 bytes	权重本身
梯度 (fp32)	4 bytes	反向传播时自动分配
优化器状态 (Adam)	8 bytes	一阶动量 + 二阶动量
总计	16 bytes	参数本身的 4 倍

一个 1B 参数的模型，光优化器状态就要占 8GB 显存。我优化过一个 13B 的 LLaMA 模型，用 AdamW 优化器，优化器状态占了 104GB——比模型本身还大。

来看一个实际对比：

import torch
import torch.nn as nn

model = nn.Linear(1024, 1024).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 查看显存占用
param_memory = sum(p.numel() * 4 for p in model.parameters())  # fp32
print(f"模型参数: {param_memory / 1024**2:.2f} MB")

# 前向 + 反向
x = torch.randn(128, 1024, device='cuda')
y = model(x).sum()
y.backward()

# 此时梯度已经分配
grad_memory = sum(p.grad.numel() * 4 for p in model.parameters() if p.grad is not None)
print(f"梯度: {grad_memory / 1024**2:.2f} MB")

# 优化器状态
optimizer.step()
state_memory = 0
for group in optimizer.param_groups:
    for p in group['params']:
        if p in optimizer.state:
            for state_val in optimizer.state[p].values():
                if isinstance(state_val, torch.Tensor):
                    state_memory += state_val.numel() * 4
print(f"优化器状态: {state_memory / 1024**2:.2f} MB")

运行这段代码，你会发现优化器状态占用的显存几乎是参数的两倍。这就是为什么很多大模型训练要用 torch.optim.AdamW 的 foreach 实现——它能减少一些临时缓冲区的分配。

1.4 知识体系总览

说了这么多，咱们用一张图来总结本章的核心逻辑：

这张图把本章的三个核心模块串起来了。你看，显存分配机制决定了底层怎么管理内存，Tensor生命周期决定了每个张量什么时候占、什么时候释放，而梯度与优化器状态则是显存占用的主要来源。三者环环相扣，理解任何一个都不能脱离另外两个。

好了，第一章就到这里。记住一句话：优化显存的第一步，不是想着怎么省，而是搞清楚显存到底用在哪了。后面几章咱们会深入具体的优化技巧，但地基打牢了，后面才能盖高楼。