显存优化工具箱：梯度检查点/FlashAttention/ZeRO

一句话定义

显存优化四件套——梯度检查点（重算换显存）、FlashAttention（IO 优化的精确注意力）、ZeRO 分片（多卡切分状态）、8-bit 优化器（压缩优化器状态）——让同样的卡装下更长的序列与更大的模型。

为什么重要

显存是微调最硬的约束。四件套与 QLoRA（kp-007）组合后，「单卡 24G 训 7B」「48G 训 70B（切片）」从幻想变成日常。同时它们是面试与排错高频考点：知道每件省的是显存账里的哪一项，才能在 OOM 时对症下药而不是乱开开关。

前置知识

kp-004 的显存四大件（权重/梯度/优化器状态/激活值）。

核心概念

• 梯度检查点（Gradient Checkpointing）：前向只存部分中间结果，反向时重算被丢弃的部分，激活显存 O(L) 降为 O(√L)，代价约 20%-30% 速度。
• FlashAttention：分块计算注意力并利用 GPU 显存层级减少 HBM 读写，是精确算法（非近似），省激活显存、显著提速，长序列下收益巨大。
• ZeRO-1/2/3：把优化器状态/梯度/参数切分到多卡，单卡显存占用近似降为 1/N，代价是通信开销递增。
• 8-bit 优化器：AdamW 的 m、v 状态以 8-bit 块量化存储，状态显存降约 75%。
• 序列打包（packing）：把多条短样本拼进同一序列，减少 padding 浪费。

原理与机制

每件工具对应显存账（kp-004）的不同科目：梯度检查点与 FlashAttention 省的是激活值；8-bit 优化器省的是优化器状态（全参场景的大头）；ZeRO 在多卡间切分状态+梯度+参数；QLoRA 省的是冻结权重。组合而非互斥。单卡 24G 训 7B 的经典配方：QLoRA（0.5P 权重）+ FlashAttention-2 + 梯度检查点 + BF16，seq 1024、batch 2 时峰值常在 10-14G。

速度与显存的交换关系要记牢：梯度检查点用 20-30% 时间换激活显存；ZeRO-3 通信量大，多机场景可能吃掉一半吞吐；FlashAttention 是少见的「又省又快」，无理由不开。

图示

显存科目      省法
权重     →   QLoRA(4bit) / ZeRO-3(多卡切)
梯度     →   ZeRO-2
优化器状态 →  8bit优化器 / ZeRO-1
激活值   →   梯度检查点 + FlashAttention + 减batch/seq

直观类比

梯度检查点像登山时只记路标不记全程录像，回程凭路标重走（多花时间，省存储）；FlashAttention 是把仓库整理得随取随用（活儿一样，搬运量暴减）；ZeRO 是把行李分给队友各背一份（人多力量大，但集合点头的沟通成本上升）。

实例或案例

同机（24G）同任务（7B、seq 2048、batch 1）四组对照：baseline LoRA → OOM；+梯度检查点 → 峰值 21G，可用；再 +FlashAttention-2 → 峰值 17G，每步时间反降 35%；再 + packing（平均样本 300 token）→ 等效吞吐再翻近 2 倍。结论：四件套组合让原本不可行的配置成为日常默认。

操作步骤

以 HF Trainer/LLaMA-Factory 为例的开启方式：

args=SFTConfig(
    gradient_checkpointing=True,        # 梯度检查点
    optim="paged_adamw_8bit",           # 8-bit 分页优化器
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,
    bf16=True,
)
# FlashAttention: 加载模型时传 attn_implementation="flash_attention_2"

多卡 ZeRO 由 DeepSpeed 配置文件指定 stage（LLaMA-Factory 直接给 deepspeed: examples/deepspeed/ds_z2_config.json）。

排错清单

• 开梯度检查点后报错（如与 use_cache 冲突）：训练时设 model.config.use_cache=False。
• 开了 FlashAttention 报不支持的卡：需 Ampere 及以上（SM80+），旧卡退回 SDPA。
• ZeRO-3 加载慢或通信打满：参数按需重建的开销，改 ZeRO-2 或减少卡间通信拓扑层级。
• 8-bit 优化器训练不稳：换 paged_adamw_32bit 或全量 32-bit 对照。
• packing 后效果异常：确认实现正确处理了样本边界与注意力 mask，避免跨样本「串味」。

常见误区

• 梯度检查点省一切：它只省激活值，权重/梯度/优化器状态一项都不少。
• FlashAttention 是近似注意力：它是精确计算，结果与标准注意力一致，只是 IO 更优。
• 工具全开一定最好：每件都有时间/通信代价，按 OOM 的科目对症开。

与其他知识点的关系

QLoRA 省权重项见 kp-007；BF16 与数值稳定性的精度面见 kp-018；OOM 的完整决策树见 kp-030。

延伸阅读

Dao et al. 2022《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》；Rajbhandari et al. 2020《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》。

自测题

1. 训练中后期 OOM（非第一步），先开哪个？

答：梯度检查点——中后期 OOM 多由长样本激活峰值导致，检查点直接削减激活显存。

1. 为什么 FlashAttention 又省显存又快？

答：它分块计算并重排读写，大幅减少 HBM 与计算单元之间的搬运；数学上与标准注意力完全等价，非近似。

1. 单卡场景 ZeRO 有用吗？

答：基本无用——ZeRO 是多卡状态分片技术，单卡没有「分」的对象，应转向 QLoRA/8-bit 优化器/检查点。