一、显存复用概述

什么是显存复用

显存复用，说白了就是让同一块显存空间在不同时间干不同的事。

我打个比方你就明白了。你租了个办公室，白天办公，晚上睡觉。同一块空间，不同时间段干不同的事——这就是复用。显存复用也是这个道理。

在深度学习训练中，我们会有大量的中间张量（tensor）。比如激活值、梯度、临时缓冲区。这些张量不会同时存在。它们有各自的生命周期。有的在前向传播时创建，反向传播后就释放了。有的只在某个算子内部使用，算完就没了。

显存复用的核心思路就是：把那些生命周期不重叠的张量，安排到同一块显存地址上。

为什么需要显存复用

这个问题我遇到过太多次了。很多同学跑模型时，第一个报错就是 CUDA out of memory。嗯，显存不够用了。

为什么会这样？

我给你算笔账。一个 7B 参数的模型，光参数就要 14GB（FP16）。加上优化器状态（Adam 需要两倍），又是 28GB。再加上激活值、梯度、临时缓冲区……一张 80GB 的 A100 很快就满了。

但问题是，很多显存是浪费的。我见过一个项目，训练过程中峰值显存占用 72GB，但实际同时存活的张量只有 40GB。剩下的 32GB 都是「僵尸张量」——已经用完但还没释放的空间。

显存复用的必要性体现在三个方面：

• 硬件限制：单卡显存有限，买大显存卡太贵。A100 80GB 一张要十几万，H100 更贵。
• 模型越来越大：从 GPT-3 的 175B 到现在的万亿参数模型，显存需求增长远超硬件发展。
• 训练效率：显存不够就得用梯度检查点（checkpointing）或者模型并行，这些都会降低训练速度。

显存复用的核心思想

核心思想其实就四个字：生命周期分析。

每个张量都有三个时间点：创建时间、最后使用时间、销毁时间。如果张量 A 的最后使用时间早于张量 B 的创建时间，那它们就可以共用同一块显存。

我画了一张图，帮你理解这个逻辑：

你看这张图。张量 A 在时间点 1 创建，在时间点 2 最后使用。张量 B 在时间点 3 创建。只要 2 早于 3，A 和 B 就能共用显存。

实际实现时，我们通常用两种策略：

1. 静态复用：在编译阶段分析计算图，提前规划好每个张量的显存地址。适合已知的固定计算图。
2. 动态复用：运行时维护一个显存池，张量用完就回收，新张量从池中分配。适合动态图（比如 PyTorch 的 eager 模式）。

收益分析

显存复用的收益，我习惯从三个维度来看：

维度	收益	典型数据
显存节省	减少峰值显存占用	通常节省 20%-50%
训练速度	允许更大 batch size，减少通信开销	batch size 提升 2-4 倍
硬件利用率	减少显存碎片，提高有效带宽	显存带宽利用率提升 10%-30%

我拿一个实际案例来说。之前优化一个 LLaMA-13B 的微调任务：

• 原始实现：峰值显存 72GB，batch size = 4
• 启用显存复用：峰值显存 44GB，batch size 提到 8
• 再加上激活值检查点：峰值显存 28GB，batch size 提到 16

训练吞吐从原来的 1200 tokens/s 提升到了 3800 tokens/s。三倍多的提升，几乎没花一分钱硬件成本。

嗯，这里要注意一点。显存复用不是万能的。它主要解决的是「中间张量」的显存问题。模型参数、优化器状态这些长期存活的张量，复用空间不大。但即便如此，对于大多数训练任务，中间张量往往占了显存的大头。

我个人习惯在开始任何显存优化之前，先用 torch.cuda.memory_summary() 或者 Nsight Systems 看一下显存分配情况。搞清楚哪些张量占了多少显存、生命周期多长，然后再决定复用策略。这个习惯帮我避过不少坑。

好，显存复用的概念就讲到这里。记住三个关键词：生命周期、地址复用、零开销。后面我们会深入具体的实现技术。