3、深度学习训练内存瓶颈：模型状态内存占用、激活内存、梯度与优化器状态、内存墙问题
在单卡训练一个中等规模模型时，很多人第一次遇到“显存不够”并不是因为模型参数太大，而是因为训练过程需要保存的东西远比推理多。推理时只需保存权重和少量中间张量，而训练时还要保存梯度、优化器状态以及反向传播所需的激活值。理解这些内存构成，是理解 ZeRO-Offload 为何要把部分状态卸载到 CPU 内存的前提。
一、训练内存的四大构成
训练显存占用可以拆成四块：模型参数、梯度、优化器状态、激活值。前三者常被统称为“模型状态”，它们与参数量成正比；激活值则与批量大小、序列长度、网络深度相关，往往在长序列或大 batch 场景下成为主要压力。
- 模型参数（Parameters）：FP16 下每个参数占 2 字节，FP32 主副本再占 4 字节。
- 梯度（Gradients）：通常与参数同精度，FP16 梯度每参数 2 字节。
- 优化器状态（Optimizer States）：以 Adam 为例，需保存 FP32 的一阶动量与二阶动量，共 8 字节/参数。
- 激活值（Activations）：前向产生的中间结果，反向时复用，大小随 batch 与序列长度增长。
二、激活内存：被低估的“隐形大户”
激活值的内存占用与网络结构强相关。以 Transformer 为例，每层需要保存注意力矩阵、前馈中间结果等。其近似关系可表示为：
激活内存 ≈ batch_size × seq_len × hidden_dim × num_layers × 系数
当序列长度从 512 增加到 2048 时，注意力矩阵按平方增长，激活内存可能膨胀十几倍。此时即使模型状态只占 16 GB，激活值也可能再吃掉 10 GB 以上，直接触发 OOM。常见缓解手段包括梯度检查点（重计算）、序列并行、以及将激活值换出到 CPU。
三、内存墙问题的本质
“内存墙”指的是计算能力与内存容量/带宽之间的失衡。GPU 的算力每两年翻倍，但显存容量增长缓慢，单卡显存长期停留在 16 GB、24 GB、40 GB、80 GB 几个档位。模型参数量却从亿级迈向千亿级，导致单卡无法容纳完整训练状态。
| 模型规模 | 参数量 | 模型状态（16B/参数） | 典型单卡显存 | 结果 |
|---|---|---|---|---|
| BERT-Base | 1.1 亿 | 约 1.8 GB | 16 GB | 可训练 |
| GPT-2 | 15 亿 | 约 24 GB | 24 GB | 临界，激活即 OOM |
| GPT-3 级 | 1750 亿 | 约 2800 GB | 80 GB | 单卡完全不可行 |
四、案例：一次典型的 OOM 复盘
某团队在 24 GB 显卡上训练 1.5 亿参数模型，配置为 batch size 32、序列长度 1024、Adam 优化器。训练启动后前向正常，反向传播到第 3 层时抛出 CUDA out of memory。排查发现：模型状态约 2.4 GB，但激活值在序列长度 1024 下达到约 18 GB，两者叠加超过 24 GB。
该团队的解决路径分三步：第一，启用梯度检查点，将激活内存降低约 60%；第二，将 batch size 从 32 降到 8，配合梯度累积保持等效批量；第三，引入 ZeRO-Offload，把优化器状态和梯度卸载到 CPU 内存。三步之后，单卡显存占用从 24 GB 以上降到约 9 GB，训练得以继续。
torch.cuda.memory_allocated() 观察前向、反向、优化器更新三个阶段的峰值。若反向阶段峰值远高于前向，说明激活值是主因；若优化器更新阶段峰值突出，则优化器状态是主因。
五、小结与启示
训练内存瓶颈不是单一因素造成的，而是模型状态、激活值、梯度与优化器状态共同作用的结果。模型状态与参数量线性相关，激活值与 batch、序列长度强相关，二者叠加形成“内存墙”。要突破这堵墙，要么减少单卡保存的内容，要么把部分内容移到容量更大的 CPU 内存——这正是 ZeRO-Offload 的设计出发点。