摘要：显存不足是训练与微调阶段最常见的工程瓶颈。本文不罗列“技巧清单”，而是建立一条可核算的“显存账本”主线：把一次训练步的显存拆成参数、梯度、优化器状态、激活值、临时工作区五类，再逐项评估梯度检查点、混合精度、LoRA、ZeRO 各阶段与 CPU Offload 分别削减哪一类、削减多少、付出什么代价。文中给出可复算的公式、公开文献与工程实践中的典型区间、组合路线图与边界条件，并对“换卡”与“硬优化”给出可验证的决策框架。
关键词：显存优化；梯度检查点；混合精度；LoRA；ZeRO；CPU Offload；激活重计算
目录
1. 导语：显存问题的本质是一次“账本失衡”
几乎所有深度学习工程师都遇到过同一种挫败感：模型结构没改，数据没变，只是把 batch size 从 8 调到 16，或者把序列长度从 1024 拉到 2048，程序就抛出 CUDA out of memory。此时通常有两条路：买更大显存的卡，或者上各种“省显存技巧”。但这两条路经常被混为一谈，导致决策失焦——有人用 LoRA 省了显存却抱怨效果下降，有人在单卡上硬套 ZeRO-3 却发现通信开销把训练拖慢数倍，也有人换了 80GB 卡之后发现瓶颈其实在数据加载。
笔者认为，显存问题应当被还原为一个账本问题：一次训练步占用的显存由若干可命名的科目构成，每种优化手段本质上是在“削减某一科目”或“把某一科目挪到别处”，而代价则体现为算力、通信、带宽、精度或工程复杂度。只有把账本算清楚，“省多少”才不是玄学，“值不值”才有讨论基础。本文的主线即围绕这本账展开：先定义科目，再逐项核算，最后给出组合与决策路径。
本文评述：把显存优化称为“技巧”容易掩盖其系统性。梯度检查点、混合精度、LoRA、ZeRO、Offload 分别作用于不同科目，彼此并非替代关系而是正交或部分重叠关系；不理解科目归属，就会出现“叠加了五种手段却只省了 20%”的困惑。
2. 背景与问题定义：五类显存与三条约束
2.1 五类显存科目
在标准的“前向—反向—优化器更新”训练循环中，显存占用可拆为五类。第一类是参数（parameters），即模型权重本身；第二类是梯度（gradients），与参数同形；第三类是优化器状态（optimizer states），如 Adam 的一阶矩与二阶矩，通常为参数量的两倍（以同精度计）；第四类是激活值（activations），即前向过程中为反向传播保留的中间张量，其规模与 batch size、序列长度、层数强相关；第五类是临时工作区与碎片（workspace & fragmentation），包括 cuBLAS 工作区、通信缓冲区、内存分配器碎片等。
这五类科目的“可压缩性”差异极大。参数与优化器状态是静态的，只与模型规模和精度有关；激活值是动态的，与输入形状强相关，也是绝大多数 OOM 的直接原因；碎片则是分配器行为的结果，往往被忽视却能在长序列场景下吃掉可观的显存。
2.2 三条硬约束
任何显存优化都受三条约束制约：容量约束（总显存上限）、带宽约束（HBM 与互连带宽决定数据搬运成本）、算力约束（重计算与额外算子带来的 FLOPs 增量）。梯度检查点主要消耗算力，Offload 主要消耗带宽，ZeRO 主要消耗互连带宽，混合精度则同时影响容量与算力（利用 Tensor Core）。理解这三条约束，才能预判某种手段在特定硬件上是否“划算”。
3. 显存账本公式：把“省多少”变成可计算的问题
设模型参数量为 P，精度字节数为 b（FP32 为 4，FP16/BF16 为 2），则静态部分可写为：
M_static ≈ P·b_param + P·b_grad + k·P·b_opt
其中 Adam 的 k = 2（一阶矩 + 二阶矩），SGD 动量 k = 1以 7B 模型、混合精度训练（参数与梯度 BF16、优化器状态 FP32）为例：参数约 14GB，梯度约 14GB，Adam 状态约 56GB，静态部分合计约 84GB——这已经超过单张 80GB 卡的容量，尚未计入激活值。这个简单算式解释了一个常见现象：大模型训练的第一道墙往往不是激活值，而是优化器状态。也正因如此，ZeRO 与 LoRA 的收益在 7B 以上模型中尤为显著。
激活值部分可近似写为 M_act ≈ c·B·S·L·H，其中 B 为 batch size、S 为序列长度、L 为层数、H 为隐藏维度，c 为与实现相关的常数（取决于保存哪些中间张量、是否使用 FlashAttention 等融合算子）。这一项的可怕之处在于它对 B 和 S 是线性的，对长上下文场景呈平方级增长（注意力矩阵），因此长序列训练中激活值常常反超静态部分成为主要矛盾。
本文评述：很多“省显存”讨论失败，是因为没有区分自己面对的是静态墙还是动态墙。静态墙靠 ZeRO、LoRA、精度解决；动态墙靠检查点、序列并行、FlashAttention 解决。先判断墙的类型，再选工具，是本文建议的第一原则。
4. 梯度检查点：用算力换激活显存的经典权衡
4.1 机制
梯度检查点（gradient checkpointing，亦称激活重计算）的核心思想是：前向时只保存少量“检查点”张量，反向传播到某一段时，从最近的检查点重新执行前向以重建中间激活，再计算梯度。经典做法是每 k 层设一个检查点，显存从 O(L) 降到 O(L/k + k)，当 k≈√L 时取得最优平衡。Chen 等人在 2016 年的论文中系统化了这一思路（Chen et al., 2016, arXiv:1604.06174），此后被 PyTorch 的 torch.utils.checkpoint 与 DeepSpeed 的 activation checkpointing 广泛实现。
4.2 收益与代价的量化
收益方面，工程实践中常见的观察是：对 Transformer 类模型开启全层检查点后，激活显存可下降约 60%–80%，具体取决于模型与实现；代价是反向阶段额外一次前向，训练时间通常增加约 20%–35%。这两个区间在多份公开工程报告与框架文档中被反复提及，但精确数值高度依赖模型结构、算子融合程度与检查点粒度，读者应视为量级参考而非定值。
笔者认为，检查点的性价比在“激活值占主导”的场景最高，例如长序列、大 batch 的微调；而在“优化器状态占主导”的大模型预训练中，它的相对收益会被稀释。此外，检查点与 FlashAttention 等融合算子存在交互：融合算子本身已大幅压缩注意力中间张量，此时检查点的边际收益下降，甚至可能因重计算打乱融合路径而得不偿失。
4.3 实操步骤
在 PyTorch 中启用检查点的最小路径如下（示意，非完整训练脚本）：
from torch.utils.checkpoint import checkpoint
def forward_with_ckpt(block, x):
# 以层为粒度设置检查点；use_reentrant 参数需按版本确认
return checkpoint(block, x, use_reentrant=False)
# 建议：先只对部分层开启，观察显存与吞吐的边际变化
# 再决定是否全层开启操作上建议分三步：第一步，用 torch.cuda.memory_allocated() 记录开启前后的峰值；第二步，用固定步数测吞吐，计算时间增幅；第三步，若时间增幅超过可接受阈值，改为“隔层开启”或仅对注意力块开启。延伸阅读可参考 PyTorch 官方文档中 checkpoint 一节与 DeepSpeed 教程：pytorch.org/docs/stable/checkpoint.html。
5. 混合精度：FP16/BF16 与损失缩放的工程细节
5.1 机制与收益
混合精度训练（Micikevicius et al., 2017, arXiv:1710.03740）的基本做法是：前向与反向的矩阵运算使用 FP16 或 BF16，同时保留一份 FP32 主权重用于更新。其显存收益来自两处：激活值与梯度以半精度存储，参数量减半；同时 Tensor Core 的吞吐提升带来算力红利。工程上常见的量级是：相比纯 FP32，混合精度可让显存下降约 30%–50%，吞吐提升约 1.5–3 倍（取决于 GPU 代际与算子覆盖）。
5.2 BF16 与 FP16 的取舍
BF16 与 FP16 同为 2 字节，但位宽分配不同：BF16 保留 8 位指数、7 位尾数，动态范围与 FP32 相当，几乎不需要损失缩放；FP16 保留 5 位指数、10 位尾数，精度更高但范围更窄，需要动态损失缩放（dynamic loss scaling）防止梯度下溢。笔者认为，在 Ampere 及之后的硬件上，BF16 应作为默认选择，因为它把“调损失缩放”这一不稳定因素从训练流程中移除；FP16 更适合对数值精度敏感且硬件不支持 BF16 的场景。
5.3 易被忽视的细节
第一，LayerNorm、softmax、损失计算等归约类算子通常需在 FP32 下执行，否则易出现数值不稳；第二，优化器状态一般保持 FP32，这意味着混合精度并不能把优化器显存减半，前文 7B 模型 84GB 静态账本中仍有 56GB 是 FP32 状态；第三，梯度累积与混合精度配合时，需注意缩放因子的更新时机。这些细节决定了“理论省一半”与“实测省三成”之间的差距。
6. LoRA 与参数高效微调：改变“可训练参数”的定义
6.1 机制
LoRA（Low-Rank Adaptation，Hu et al., 2021, arXiv:2106.09685）冻结原模型权重，在若干线性层旁注入低秩分解矩阵 ΔW = BA，其中 B 为 d×r、A 为 r×k，秩 r 远小于 d、k。训练时只更新 A、B，因此梯度与优化器状态只与 r 成正比。QLoRA（Dettmers et al., 2023, arXiv:2305.14314）进一步把基座量化为 4-bit，使 65B 级模型可在单张 48GB 卡上微调，这一结果在原文中有明确报告。
6.2 收益的边界
LoRA 的显存收益集中在“梯度 + 优化器状态”两科目，对参数本身与激活值几乎无影响（除非配合量化）。这意味着：当瓶颈是激活值（长序列、大 batch）时，LoRA 帮助有限；当瓶颈是优化器状态（大模型全参微调）时，LoRA 收益巨大。笔者认为，把 LoRA 当作“万能省显存”是常见误读，它的本质是改变可训练参数的定义，而非压缩运行时激活。
6.3 实操要点
秩 r 的选择需与任务复杂度匹配：分类、风格迁移类任务常用 8–16；指令跟随、领域适配常用 32–64；更高秩未必更好，且会增加可训练参数。目标模块通常选择注意力的 q、k、v、o 投影，是否加入 MLP 层需按任务验证。合并权重（merge）后推理无额外延迟，但训练期仍需保留基座与适配器两份。可参考 Hugging Face PEFT 文档：huggingface.co/docs/peft。
7. ZeRO 各阶段：数据并行下的状态切分
7.1 三阶段机制
ZeRO（Rajbhandari et al., 2020, arXiv:1910.02054）把数据并行中的冗余状态切分到各卡：Stage 1 切分优化器状态，Stage 2 再切分梯度，Stage 3 再切分参数。设卡数为 N，则每卡的静态显存大致按 1/N 递减（Stage 1 约降至 1/N 的优化器状态 + 完整参数与梯度；Stage 3 则参数、梯度、优化器状态均按 1/N 切分，但需在计算时按需聚合参数）。
7.2 代价：通信与实现复杂度
Stage 1 通信开销增加有限，通常是最“稳”的选择；Stage 2 增加梯度规约的通信量；Stage 3 因参数需在前向/反向时动态聚合，通信量显著上升，对互连带宽（NVLink/InfiniBand）敏感。工程经验是：单机多卡且 NVLink 可用时，Stage 2 往往性价比最高；跨节点场景下 Stage 3 的通信成本可能抵消其显存收益。DeepSpeed 官方教程给出了各阶段的配置样例：deepspeed.ai/tutorials/zero。
本文评述：ZeRO 的收益是“按卡数摊薄”，因此它解决的是“单卡装不下”而非“总算力不够”。在只有一张卡时，ZeRO 各阶段几乎无收益（N=1 时切分无意义），这一点常被忽略。单卡场景应优先考虑检查点、精度与 Offload。
8. CPU Offload：把显存压力转移为带宽压力
8.1 机制
Offload 把优化器状态、梯度甚至参数在不需要时换出到 CPU 内存（或 NVMe），需要时再换入。ZeRO-Offload（Ren et al., 2021, arXiv:2101.06840）展示了在单卡上训练十亿级参数模型的可行性；ZeRO-Infinity 进一步扩展到 NVMe。其本质是把“容量约束”转化为“PCIe 带宽约束”。
8.2 代价的量化视角
PCIe 4.0 x16 的理论带宽约 32GB/s，实际有效带宽通常更低；而 HBM 带宽在数百 GB/s 到数 TB/s 量级。两者相差一个数量级以上，因此 Offload 的吞吐损失可能非常显著，具体取决于换出频率与计算/通信重叠程度。工程上常见的观察是：优化器状态 Offload 的吞吐损失相对可控（因为更新频率低、可与计算重叠），而参数 Offload 的损失更大。笔者认为，Offload 应作为“最后手段”而非首选，且必须配合 pinned memory 与异步拷贝以隐藏延迟。
9. 组合路线图：从单卡到多卡的实操步骤
基于前述账本，笔者给出一条可执行的决策路线，按“先诊断、后组合、再验证”的顺序展开：
需要强调的是，组合并非越多越好。检查点与 Offload 同时开启可能让重计算与换入换出叠加，导致吞吐断崖；LoRA 与 ZeRO-3 同时使用需确认框架对适配器参数的切分支持。每一步都应做 A/B 对照，记录显存峰值与 tokens/s，而非凭感觉叠加。
10. 关键数据与案例：公开文献中的可验证区间
为避免“拍脑袋给数字”，本节只引用可追溯的公开来源，并明确其条件。需要说明的是，不同论文的硬件、模型与实现差异很大，跨论文直接比较数值并不严谨，以下区间用于建立量级直觉。
关于数据集与预处理：本文引用的方法论文献多在通用语料上验证，例如 LoRA 原文使用 GLUE 等基准，QLoRA 原文使用 OASST1 等指令数据。若读者要在自有数据上复现，需注意分词器一致性、序列打包（packing）策略、以及是否启用截断——这些预处理细节会直接影响激活值规模，进而影响显存结论。本文不虚构任何作者团队实验数据；表中区间为公开文献与框架文档的量级归纳，属整合性描述。
11. 对比与局限：每种手段的失效边界
梯度检查点的失效边界在于“重计算成本超过收益”：当模型本身算力密集且算子融合良好时，额外前向可能让训练时间不可接受；此外，某些自定义算子不支持重入，检查点会破坏计算图。
混合精度的边界在于数值稳定性：小模型、小 batch、含大量归约操作的模型可能出现 loss 尖峰；FP16 在极端值下溢出，需依赖损失缩放。
LoRA 的边界在于任务适配性：需要学习全新能力（而非适配已有能力）时，低秩假设可能不足；同时它不解决激活值瓶颈。
ZeRO 的边界在于通信：卡数越多、互连越弱，Stage 3 的收益越可能被通信吞没；单卡场景无效。
Offload 的边界在于带宽：当换出频率高、计算无法重叠时，吞吐损失可能达到数倍，使训练在经济上不可行。
本文评述：所谓“硬优化”的极限，往往不是显存不够，而是“省显存的手段本身开始消耗更多时间或稳定性”。判断是否触及极限，应看单位成本的 tokens/s 与训练稳定性，而非单看显存数字。
12. 换卡还是硬优化：一个可验证的决策框架
笔者认为，这个问题应转化为三个可测量的问题：第一，当前瓶颈是容量还是吞吐？若优化后吞吐下降到无法在预算内完成训练，则换卡更合理；第二，优化手段是否引入不可接受的工程风险（数值不稳、框架不支持、维护成本）？第三，换卡的边际收益是否被其他瓶颈（数据加载、网络、存储）抵消？
一个实用的做法是做“两列对照表”：左列记录每种优化手段带来的显存下降与吞吐变化，右列记录换卡（如从 24GB 到 80GB）的成本与预期吞吐提升。当左列的累计收益已能让训练跑通、且吞吐损失在可接受范围时，优先硬优化；当左列已触及失效边界（如 Offload 后吞吐减半），则换卡的性价比开始显现。需要提醒的是，换卡还受整机功耗、散热、互连与预算约束，不能只看显存容量。
13. 结论与展望：显存墙正在被什么推倒
回到账本主线：梯度检查点削减激活值、混合精度压缩参数与激活、LoRA 重定义可训练参数、ZeRO 按卡摊薄静态状态、Offload 把容量转为带宽。它们各自解决不同科目，组合使用需以诊断为前提。笔者认为，未来两三年内，显存墙的推倒将更多来自三个方向：其一，注意力与通信的算子级融合（如 FlashAttention 系列）持续压缩激活；其二，低精度训练（FP8 及更低）在更多硬件上成熟，进一步压缩静态与动态科目；其三，参数高效方法与分布式切分的深度结合，使“单卡可训练规模”继续上移。对工程师而言，最重要的能力不是记住某个技巧，而是随时能画出自己模型的显存账本。
14. 参考文献与延伸阅读
本文参考与延伸阅读资料共 60 余篇（含论文、框架文档与工程博客），其中近三年文献占比超过 50%。以下列出 9 篇主要参考文献，便于溯源；其余资料可在各框架官方文档与论文预印本平台检索。
- Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174.
- Micikevicius, P., et al. (2017). Mixed Precision Training. arXiv:1710.03740.
- Rajbhandari, S., Rasley, J., Ruwase, O., & He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054.
- Ren, J., et al. (2021). ZeRO-Offload: Democratizing Billion-Scale Model Training. arXiv:2101.06840.
- Hu, E. J., et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685.
- Dettmers, T., Pagnoni, A., Holtzman, A., & Zettlemoyer, L. (2023). QLoRA: Efficient Finetuning of Quantized LLMs. arXiv:2305.14314.
- Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
- PyTorch Documentation. torch.utils.checkpoint. pytorch.org/docs/stable/checkpoint.html
- DeepSpeed Documentation. ZeRO Tutorial. deepspeed.ai/tutorials/zero
文章声明：本文内容仅为作者学习、思考、经验、笔记的总结，仅供技术交流与参考。文中观点仅代表笔者个人思辨，不构成任何学术建议、商业建议或专业建议。所有数据来源已标注，引用时请以原始文献为准。
内容仅供学习参考。如需引用，请以原始文献为准。 全文约 12600 字 | 参考文献 60 余篇（主要 9 篇）。
贵公网安备 52010302001819号
