本篇博客，我们会从第一性原理角度出发，以面向对象和面向方法的双推理视角，拆解出四大显存消费对象对应的六大类优化策略的完整映射

---

一、训练时，显存都被谁吃了？（分析面向的对象）

在谈优化之前，我们先定位优化的对象，也就是显存都是被谁吃了？训练一个神经网络时，GPU 显存中驻留着四大类数据：

显存消费对象	内容	生命周期
模型参数（Parameters）	网络权重 W	全程常驻
梯度（Gradients）	∂L/∂W	反向传播时产生，更新后可释放（重要，非全程常驻）
优化器状态（Optimizer States）	一阶矩 m、二阶矩 v等（这里以Adam为例）	全程常驻
激活值（Activations）	前向传播的中间结果	前向时产生，反向时消费

然后来定量地模拟个具体的数值以便后续分析。假设模型有 Ψ = 100B（1000 亿） 个参数：

若使用混合精度（FP16 参数 + FP32 Adam）的典型方案：

• 模型参数：FP16 存储，每个参数 2 字节 → 2Ψ = 200 GB
• 梯度：与参数同形，FP16 → 2Ψ = 200 GB
• 优化器状态：Adam 需要为每个参数维护 FP32 的一阶动量 m（4 字节）和二阶动量 v（4 字节），再加上一份 FP32 参数主副本（4 字节，用于精确的参数更新）→ 12Ψ = 1200 GB
• 激活值：取决于 batch size、序列长度和模型架构，通常量级为数百 GB 到 TB 级

仅参数 + 梯度 + 优化器状态就需要 16Ψ = 1.6 TB。单张 H100 只有 80 GB 显存。这意味着即使不考虑激活值，你也需要 20 张以上的卡才能装下这些静态数据。

这就是核心矛盾：模型训练过程中的"记忆需求"远超单卡的物理容量。

而这四大消费对象的优化思路又各不相同。接下来，我们逐个分析。

---

二、优化器状态：最大的隐形杀手

2.1 为什么优化器状态是最大的显存消费者？

以最主流的 Adam 优化器为例，它需要为训练过程中的每个参数维护对应的：

• 一阶动量 m（梯度的指数移动平均）
• 二阶动量 v（梯度平方的指数移动平均）
• FP32 参数主副本（当使用混合精度训练时，还需保留一份全精度参数）

如果模型参数本身用 FP16 只需 2Ψ 字节，优化器状态却需要 12Ψ 字节 —— 是参数本身的 6 倍。对于 100B 模型，这就是 1.2 TB。

所以，一个优化的思路就是：我们真的需要用全精度、而且为每个参数都保存这些统计量吗？

2.2 优化思路一：减少优化器的"记忆容量"

• Adafactor（2018, Google） 的具体思路是：

这将二阶动量的存储从 O(mn) 降到了 O(m+n)，相当于做了一次矩阵分解式压缩。

• 8-bit Adam / 8-bit Optimizers（2021, bitsandbytes） 则走了另一条路：直接对优化器状态（m 和 v）做量化：压缩到8-bit。减少量化误差上，则通过了分块量化（block-wise quantization）和动态量化。最终将优化器状态的内存直接砍到原来的 1/4。

• LOMO / CAME 系列（2023-2024） 更加激进：既然优化器状态这么贵，能不能直接不要？LOMO（LOw-Memory Optimization）将梯度计算和参数更新融合在一起，在反向传播的过程中逐层计算梯度、立即更新参数、然后释放梯度，等效于 SGD 但显存开销极低。CAME 则引入了更高效的矩阵分解来进一步压缩二阶信息。

2.3 优化思路二：分散优化器的"存储负担"

即使单卡装不下，但如果你有 N 张卡，为什么每张卡都要存一份完整的优化器状态？

• ZeRO Stage 1（2019, Microsoft DeepSpeed） 它的做法是：

对于 100B 模型使用 64 张卡，每张卡的优化器状态从 1.2 TB 降到约 19 GB。这一步是 ZeRO 系列中收益最高、通信开销最小的一步，因为优化器状态只在参数更新时需要，而参数更新发生在每个 step 的最后，通信可以高效调度。

2.4 总结

优化器状态优化的核心思路可以概括为三个字：不存、少存、分存。

• 不存：换用无状态或低状态优化器（SGD、LOMO）
• 少存：压缩/量化/低秩近似（Adafactor、8-bit Adam）
• 分存：跨设备切分（ZeRO Stage 1）

---

三、模型参数：从精度到分布的全维度优化

训练过程中确实需要较高精度来保证梯度更新的数值稳定性，但前向和反向计算本身对精度的要求远低于此。

3.1 混合精度训练（Mixed Precision Training）

具体来说：

1. 前向和反向传播用 FP16 进行：计算快（Tensor Core 加速）、显存省（参数和激活值减半）
2. 参数主副本和优化器状态用 FP32 保存：保证梯度更新时的数值精度
3. Loss Scaling：将 loss 乘以一个大数再反向传播，防止小梯度在 FP16 下下溢到零

这个方案将参数和激活值的显存占用直接减半，同时几乎不影响训练精度。FP16 参数只需 2Ψ 字节，对于 100B 模型节省了 200 GB。

• FP8 训练：随着硬件（NVIDIA H100/Blackwell）支持 FP8 Tensor Core（比如 E4M3 和 E5M2 格式），前向传播可以进一步用 FP8 执行，理论上将计算显存再减半（但需要精细的量化策略和缩放因子管理）。

3.2 模型参数的分布式切分

当单卡装不下完整模型参数时，唯一的选择就是把参数"拆开"放到多张卡上。但"拆法"有很大的讲究：

张量并行（Tensor Parallelism, Megatron-LM）

这是最细粒度的切分方式。它的核心思想是：

具体来说，比如对 Transformer 中的 MLP 层：

这样每张卡只需存储 1/N 的参数，但代价是每个 Transformer 层forward或backward时各需要 2 次 All-Reduce 通信（MLP 和 Attention 各一次）。因此张量并行通常只在单机多卡的GPU间（通过 NVLink高速互联）使用，N 一般取 2、4 或 8。

流水线并行（Pipeline Parallelism）

如果说张量并行是"横切"每一层，流水线并行就是"纵切"整个模型：

最朴素的流水线并行有一个致命问题：气泡（Pipeline Bubble）。当设备 1 在计算前向时，设备 2-N 都在空闲等待。

GPipe 的解决方案是将一个 mini-batch 拆分为多个 micro-batch，让不同的 micro-batch 在流水线上交错执行，减少气泡。理论上，如果 micro-batch 数量 M 远大于流水线深度 P，气泡比例约为 (P-1)/M，可以控制在较低水平。

1F1B（One Forward One Backward）调度（PipeDream-Flush）进一步优化：交替执行前向和反向传播，使得流水线更加满载，同时减少了需要保留的激活值数量。

ZeRO Stage 3 / FSDP（Fully Sharded Data Parallelism）

这是另一种思路：不在计算维度上切分模型，而是在存储维度上切分。

这意味着显存占用被削减到 1/N，但通信量增加了——每个前向和反向 step 都需要做一次 All-Gather。ZeRO Stage 3 还将通信与计算重叠（overlap），在 GPU 计算当前层的同时，后台预取下一层的参数，从而隐藏通信开销。

PyTorch 原生的 FSDP（Fully Sharded Data Parallelism） 正是 ZeRO Stage 3 的官方实现。

3.3 总结

模型参数优化的核心思路：用低精度计算，用高精度存储，用切分分担。

• 精度维度：FP32 → FP16/BF16 → FP8，用硬件支持的最低精度完成计算
• 空间维度：张量并行（层内切）、流水线并行（层间切）、ZeRO/FSDP（存储切）

三种并行往往同时使用，形成所谓的 3D 并行。例如 Megatron-LM 中典型的配置是：8 路张量并行（机器内）× 8 路流水线并行（机器间）× 数据并行（机器组间），就可以在数千张 GPU 上协同训练。

---

四、梯度：短命但昂贵的中间产物

4.1 梯度非全程常驻

梯度是反向传播的产物。对于每个参数 W，梯度 ∂L/∂W 在反向传播时被计算出来，用于参数更新 W ← W - lr · g(∂L/∂W)，然后就可以丢弃了。

4.2 梯度累积（Gradient Accumulation）：用时间换空间

其核心想法是：

比如，如果 batch size = 12 对显存压力太大，就把它拆成 3个 mini-batch，每步只跑 batch size = 4，将梯度累加3次后再做一次参数更新。

梯度累积的另一个好处是解耦了计算的 batch size 和统计的 batch size：你可以在小显存上实现大 batch 训练的效果。

4.3 ZeRO Stage 2：梯度分片

与优化器状态类似，梯度也可以在多卡间切分：

ZeRO Stage 2 在 Stage 1（切分优化器状态）的基础上增加了梯度切分。对于 100B 模型 + 64 张卡，梯度从 200 GB/卡降到约 3 GB/卡。

4.4 梯度压缩与通信优化

在分布式训练中，梯度不仅占显存，还要在设备间传输。压缩梯度可以减少通信量：

• 梯度量化：将梯度从 FP32/FP16 压缩到更低精度再传输
• 梯度稀疏化（Top-K / Random-K）：只传输最大的 K% 梯度，其余累积到下一轮
• PowerSGD（2019）：用低秩近似压缩梯度矩阵，通信量从 O(mn) 降到 O((m+n)r)

这些方法本质上都在利用一个事实：梯度信息是冗余的，损失少量精度对训练收敛的影响远小于预期。

4.5 即时释放：Reduce-Scatter + Backward 融合

还有个不可忽视的工程优化就是将梯度的通信与反向传播的计算重叠：

这样，GPU 显存中在任一时刻只需保留正在计算的那一层和正在通信的那一层的梯度，而非所有层的梯度。

4.6 总结

梯度优化的核心思路：快算快丢，边算边传。

• 时间换空间：梯度累积，小 batch 等效大 batch
• 分存：ZeRO Stage 2 梯度分片
• 压缩：量化、稀疏化、低秩近似
• 流水化：计算-通信重叠，即时释放

---

五、激活值：最不起眼却最致命的显存杀手

5.1 为什么需要保存激活值？

这是很容易被忽略的点。源头还是因为反向传播的链式法则：

对于一个简单的两层网络 y = f(g(x))，反向传播需要计算

∂L/∂W_g = ∂L/∂y · ∂y/∂g · (∂g/∂W_g)

注意最后一项 ∂g/∂W_g 的计算需要用到 g 的输入，也就是 x。类似地，计算 ∂y/∂W_f 需要 g(x) 的输出。

这意味着前向传播的每一层中间结果都必须保存到反向传播使用。这些中间结果就是"激活值"。

对于一个 L 层 Transformer，激活值的显存占用大约为：

激活值 ≈ L × batch_size × seq_len × hidden_dim × 常数因子

激活值的可怕之处在于它与 batch size 和序列长度成正比，而且随着长上下文（100K+ tokens）的需求爆发，它正在成为新的显存瓶颈。

5.2 激活重计算（Activation Checkpointing / Gradient Checkpointing）

激活重计算的优化思路是：

这是一个经典的时间-空间权衡（Time-Space Tradeoff）：

• 不做 Checkpointing：保存所有 L 层激活值，显存 O(L)，计算量 1×
• 全量 Checkpointing：只保存网络输入，反向时从头重算，显存 O(1)，但额外计算开销为 O(L)×（因为每一层的反向都需要从头重算到该层）
• √L 策略（最优均衡）：每隔 √L 层设置检查点，显存 O(√L)，计算量约 1.33×

在实践中，最常见的做法是每个 Transformer Block 保存一个检查点（即保存每个 block 的输入），block 内部的激活值不保存，反向时重算。这通常会增加约 30-35% 的计算开销，但能将激活值显存降低一个数量级。

Selective Activation Checkpointing（选择性重计算）：并非所有激活值的重计算代价都一样。矩阵乘法的输出重算很贵（需要再做一次 GEMM），但 LayerNorm、Dropout、GELU 等的输出重算很便宜。因此可以只丢弃便宜的激活值，保留昂贵的激活值，在最小的时间代价下获得最大的显存节省。

5.3 Activation Offloading：卸载到 CPU 或 NVMe

另一条路径是利用异构存储层次：

CPU Offloading 的做法是：前向计算完一层后，将该层激活值异步传输到 CPU 内存；反向传播需要时再传回。关键在于传输与计算的重叠——在 GPU 计算当前层时，后台通过 PCIe 传输上一层的激活值。

PCIe 4.0 x16 的带宽约 32 GB/s，PCIe 5.0 约 64 GB/s。只要每层的计算时间大于激活值的传输时间，offloading 就能几乎无感。

**NVMe Offloading（ZeRO-Infinity, 2021）**更进一步，将数据卸载到 SSD。虽然带宽更低（约 3-6 GB/s），但容量几乎无限。这使得在有限 GPU 上训练超大模型成为可能。

5.4 序列并行（Sequence Parallelism）

随着上下文窗口从 2K 增长到 128K 甚至 1M+，激活值沿序列维度的增长变得不可忽视。

• Megatron-LM 序列并行 将 Transformer 中非张量并行区域（如 LayerNorm、Dropout）的计算沿序列维度切分到多张卡上。这样每张卡只需保存 1/N 序列长度的激活值。

• Ring Attention & Context Parallelism 将自注意力机制的 KV 在设备间做分块循环传递，每张卡只处理序列的一个片段的 Query，通过 Ring 拓扑依次获取所有 Key-Value。这使得序列长度可以随 GPU 数量线性扩展，是训练超长上下文模型的关键技术。
• FlashAttention 虽然本质上是一个 IO-aware 的精确注意力实现，但它有一个重要的副产物：不需要存储 S×S 的注意力矩阵（S 为序列长度）。传统注意力需要 O(S²) 的显存来存储注意力权重，而 FlashAttention 通过在线 softmax 和分块计算，将显存降到 O(S)。

5.5 总结

激活值优化的核心思路就是：不存就重算，存不下就外搬，太长就切开。

• 时间换空间：激活重计算（Activation Checkpointing）
• 空间换空间：卸载到 CPU/NVMe（Offloading）
• 分治：序列并行、Ring Attention
• 算法改进：FlashAttention 避免 O(S²) 存储

---

六、全局视角：跨维度的系统级优化

6.1 内存碎片化与内存池

GPU 显存的分配和释放过程中会产生碎片。PyTorch 的 CUDA Caching Allocator 通过内存池化减少碎片，但在极端情况下（如动态形状、频繁分配释放）仍会出问题。

预分配 + 内存复用是工程实践中的重要技巧：预先计算整个训练过程的显存使用模式，为激活值、梯度等分配固定的内存区域并循环复用。Megatron-LM 和 DeepSpeed 都有类似的实现。

6.2 ZeRO 全家桶：一张图理清层次

ZeRO 阶段	切分内容	每卡显存 (N 卡)	通信开销
Baseline	无切分	16Ψ（FP32 Adam）	AllReduce（梯度）
Stage 1	优化器状态	4Ψ + 12Ψ/N	同上
Stage 2	+梯度	2Ψ + (2Ψ+12Ψ)/N	Reduce-Scatter + AllGather
Stage 3	+参数	(2Ψ+2Ψ+12Ψ)/N = 16Ψ/N	额外 AllGather（前向/反向取参数）
ZeRO-Offload	+CPU 卸载	极低（GPU 仅保留计算所需）	PCIe 传输
ZeRO-Infinity	+NVMe 卸载	几乎无限	PCIe + NVMe 传输

（注：Ψ 为参数量，此处假设混合精度训练，参数为 FP16 = 2Ψ 字节）

6.3 3D 并行 + ZeRO：现代训练系统的标准架构

当前训练 100B+ 模型的标准做法是组合多种并行策略：

├── 数据并行 (Data Parallelism, DP)        × D 路
│   └── ZeRO Stage 1/2/3 切分优化器/梯度/参数
│
├── 张量并行 (Tensor Parallelism, TP)      × T 路
│   └── 通常 T=8（单机 NVLink 内）
│
├── 流水线并行 (Pipeline Parallelism, PP)   × P 路
│   └── 跨机器，1F1B 调度
│
├── 序列并行 (Sequence Parallelism, SP)     × 与 TP 组合
│   └── 处理非 TP 区域的激活值
│
└── Context Parallelism (CP)               × C 路
    └── 长序列场景下切分 Attention

总 GPU 数 = D × T × P × C

6.4 编译器与算子融合

算子融合（Operator/Kernel Fusion） 是另一个重要的显存优化手段。未融合时，每个算子的输出都要写入显存作为下一个算子的输入。融合后，中间结果保留在寄存器或共享内存中，不写入全局显存。

• torch.compile（PyTorch 2.x）：自动识别可融合的算子序列
• FlashAttention：本质上是 Q·K^T → Softmax → ·V 三步的深度融合
• Fused Adam：将参数更新的多步运算融合为单个 kernel

算子融合不仅减少了中间激活值的显存占用，还减少了 GPU 的全局内存访问次数，同时优化了显存和速度。

类型	中间依赖	存储	HBM 访问	举例
Elementwise Fusion	逐元素，线程私有	寄存器	最少（2 次）	x*2 + x
Block-Level Fusion	跨线程，但 block 内	Reg + Shared Mem	少（1 次读 + 1 次写）	LayerNorm, block reduction
Unfused	跨 block / 全局依赖	HBM	多（读写多次）	Sort, Top-K

---

七、总结：一张全景图

我们形而上的系统总结下，可以发现，核心方法论大概就三条：

1. 能省则省（降精度、压缩、低秩近似）
2. 能分则分（ZeRO、张量并行、流水线并行、序列并行）
3. 能丢则丢（激活重计算、梯度即时释放、offloading）

而将这三条原则应用到优化器状态、模型参数、梯度、激活值这四个维度上的排列组合，就构成了当代大模型训练系统的完整技术图谱。