本文详细解析了大模型训练中的显存优化策略。内容涵盖 GPU 显存构成分析，包括参数、梯度、优化器状态及中间激活值的占用计算。深入探讨了集合通信操作（Broadcast, AllReduce 等）及其在多卡协同中的作用。重点阐述了数据并行、模型并行、ZeRO（1/2/3 级）及流水线并行的原理、流程与显存收益对比。此外，还介绍了混合精度训练（FP16/FP32）的机制与 Loss Scaling 必要性，以及 Checkpointing 重计算技术在平衡显存与计算时间上的应用。文章最后提供了不同场景下的策略选型建议，帮助开发者有效解决大模型训练中的 OOM 问题。
在大模型微调或预训练过程中，GPU 显存占用是决定能否成功运行的关键瓶颈。显存主要被以下四个部分占用：模型参数、参数梯度、优化器状态和中间激活结果。
对于一个 6B（60 亿）参数量的模型，若使用 FP32（单精度浮点）格式存储，其模型参数占用计算如下：
6 × 10^9 × 4 (Bytes) / 1024^3 ≈ 22 GB
将模型参数视为基准，反向传播时产生的参数梯度占用量通常与模型参数相同，即约 22GB。
优化器方面，主流采用 Adam Optimizer。其核心计算公式包含动量项 m 和二阶矩估计 v：
m_t = β1 * m_{t-1} + (1 - β1) * g_t
v_t = β2 * v_{t-1} + (1 - β2) * g_t^2
由于需要保存 m 和 v 两个状态变量，且它们的规模与参数梯度相同，因此优化器状态通常需要两倍于参数梯度的显存容量，即约 44GB。
此外，在计算中得到的中间结果（激活值）需要保存在显存中，以便反向传播时计算梯度。对于每一个中间结果，其数据形状通常为 [Batch, SeqLen, Dim]。在深层网络中，这部分显存占用可能非常巨大，甚至超过参数本身。
为了节省显存并提升计算效率，可以将模型或者数据分配到不同的显卡上。多卡之间通过集合通信操作（Collective Operations）进行数据交换，常见的操作包括：
广播操作将一张显卡（Root Rank）上的 N 元素缓冲区复制到所有其他显卡（Ranks）。这常用于初始化参数或分发配置信息。
收集所有 Ranks 的数据，合并成一个大小为 k*N 的输出，并分发给所有 Ranks。这在模型并行或 ZeRO 同步参数时常用。
数据并行（Data Parallelism, DP）是最基础的并行策略。它将训练数据分成若干份，装载到不同节点上进行计算，每个节点维护一份完整的模型副本。
在分布式数据并行中，每个设备显存占用情况依然较高，因为每个设备仍需保存完整的模型参数、梯度和优化器参数。当模型参数量超过单卡显存时，此方法不可用。
由于模型越来越大，单个设备保存模型参数、梯度和优化器越来越难。深度学习主要是矩阵计算，而矩阵计算可以分块计算，因此可以将模型参数拆成若干份，每份单独计算，以减少显存占用。
以矩阵乘法为例，假设 $Y = W \times X$，其中 $W$ 为 $A \times B$ 的矩阵。可以将 $W$ 沿行方向切分为 $n$ 个子矩阵 $W^{(1)}, W^{(2)}, ..., W^{(n)}$，分别存储在 $n$ 个设备上。
模型并行后，显存占用显著降低，因为每个设备只需存储部分权重。但由于每个设备处理所有数据，中间激活结果（Activation）仍会保存在所有设备中，这限制了显存的节省比例。
在分布式数据并行中，多个设备中参数相同，梯度相同，优化器状态相同，存在大量冗余。ZeRO 旨在消除这些冗余，进一步降低显存占用。
ZeRO-1 对优化器状态进行分片。在计算梯度后，通过 ReduceScatter 将梯度分片，每个 replica 只负责更新对应分片的参数，从而减少优化器状态的显存占用。
ZeRO-2 在 ZeRO-1 的基础上，进一步对梯度进行分片。在后向传播时，每计算一层梯度，就可以使用 ReduceScatter 进行同步。这样每个 replica 只需要保存部分梯度即可，无需保留完整梯度。
ZeRO-3 在 ZeRO-2 的基础上，将模型参数也进行分片。这是最激进的优化策略。
不同 ZeRO 级别对应的显存占用呈阶梯式下降，ZeRO-3 理论上可将优化器和参数显存占用降至原来的 1/N（N 为卡数），但通信开销也会相应增加。
将模型一层一层分开，不同层放入不同 GPU 进行计算。与模型并行不同的是，模型并行保留从头到尾每一层的部分参数，输入可以计算出结果；流水线并行需要等前一层计算完毕才能进行计算。
流水线并行显存占用较低，因为每块 GPU 只负责部分层。然而，它引入了**气泡（Bubble）**问题，即某些 GPU 在前向传播完成后需等待后续层完成反向传播才能继续工作，导致算力闲置。可以通过微批次（Micro-batching）和流水线调度算法来缓解这一问题。
FP16（半精度）相较于 FP32 计算更快，同时占用更少的显存。但同时 FP16 表示的范围小，可能产生溢出错误。
特别的，在权重更新时 gradient * lr 可能导致下溢出（Underflow），使得梯度变为 0。
混合精度训练的思路是在优化器中保留一份 FP32 格式的参数副本（Master Weights），而模型权重、梯度等数据在训练中都是用 FP16 来存储。优化器中参数更新在 FP32 格式下保证精度，之后转换为 FP16 格式。
为了防止 FP16 梯度下溢出，通常需要使用 Loss Scaling 技术，即在反向传播前放大 Loss 值，反向传播后再缩小梯度。
由于模型反向传播需要中间结果计算梯度，大量中间结果占用大量显存。Checkpointing 思路是保存部分隐藏层的结果（作为检查点），其余的中间结果直接释放。
当反向传播需要计算梯度时，从检查点开始重新前向传播计算中间结果，得到梯度后再次释放。这种方法以计算时间为代价换取显存空间的节省，通常可节省约 50% 的激活显存。
在实际的大模型训练中，通常会组合使用上述策略：
选择合适的并行策略需要综合考虑硬件资源（GPU 数量、显存大小）、网络带宽（影响通信开销）以及模型架构特性。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online