Training Memory Accounting

组成

1
2
3
4
5
6
7
model parameters
+ gradients
+ optimizer states
+ master weights / mixed-precision copies
+ activations
+ temporary buffers and communication workspace
+ allocator fragmentation

全参数 Adam 类训练不能只用“参数量 × 2 bytes”估算。不同精度、优化器、checkpointing、序列长度、micro-batch 和并行策略会改变各项。

主要控制杆

  • QLoRA / LoRA：减少 trainable 参数与 optimizer states。
  • gradient accumulation：用更多 step 换更大的有效 batch。
  • gradient checkpointing：用重计算换 activation memory。
  • mixed precision：减少部分张量字节数。
  • ZeRO / FSDP：跨设备切分参数、梯度和优化器状态。

最终以实际 profiler 和峰值显存为准，估算用于排除不可能配置，而非保证可运行。

估算顺序

先按参数状态建立下界，再单独加激活和通信峰值：

1
2
3
4
5
weights + grads + optimizer/master copies
→ 静态训练状态
+ activations(sequence × micro-batch × layers)
+ all-gather/reduce-scatter、checkpoint 与 allocator 峰值
→ 实际 peak memory

这解释了为什么“量化后权重能放下”不等于“能训练”：长序列通常把激活推到主导项，ZeRO/FSDP 可能降低静态项却带来通信 workspace。排查 OOM 时记录发生在 forward、backward、optimizer step 还是 checkpoint 保存，而不是只记录最终显存。

相关节点：LLM Training Lifecycle、Data Parallelism。

先算静态账，再算峰值账

对一个参数量为 N 的全参 Adam 混合精度训练，至少要分别问：参数副本用几字节、梯度用几字节、Adam 的一阶/二阶状态用几字节、是否保留 FP32 master weights。不同框架的状态布局不同，不能用一个固定的“每参数多少字节”覆盖所有配置。这个静态下界算出来后，再加每层 activation、临时 GEMM workspace、通信 buffer 和 allocator 碎片，才接近实际峰值。

activation 大致随 micro_batch × sequence_length × hidden_size × layers 增长，并且 backward 前不能释放；这解释了为什么把序列长度从 4K 加到 8K，显存可能比模型参数翻得更快。gradient checkpointing 通过只保存边界、反向时重跑中间层来交换计算时间；它减少的是 activation 项，不会减少参数和 optimizer state。

用 OOM 的发生阶段定位账目

发生点 更可能的主项 优先检查
第一个 forward 权重、输入、KV/临时 buffer dtype、模型加载、micro-batch
backward 开始 activation、保存的中间结果 sequence、checkpoint、激活函数
optimizer.step gradients、Adam states、master copy optimizer、ZeRO/FSDP、offload
保存 checkpoint 临时 gather、序列化副本 state dict 方式、rank 0 聚合

每张卡的 allocated、reserved 和峰值并不等价：缓存分配器保留的空闲块会造成看起来“还有很多 reserved 却分配失败”。因此同时记录 allocator snapshot、各 rank 峰值和通信 workspace；否则只看 nvidia-smi 很难复现 OOM。

量化或 PEFT 只改变账本中的部分项。QLoRA 可能让冻结权重低到能加载，但激活仍被长序列主导；ZeRO 可能降低每卡静态状态，却增加 all-gather/reduce-scatter 峰值。每次优化都重新做一次账本，而不是沿用“量化后应该能训练”的直觉。