# 7.5 激活重计算：用时间换空间的艺术

在讨论混合精度等技术时，我们反复提到“显存瓶颈”。在大模型训练中，显存不仅被模型参数、优化器状态和梯度占用，还有一个不可忽视的消耗大户——**激活值（Activations）**。

## 7.5.1 为什么激活值占用大量显存

在前向传播过程中，每一层的输出（激活值）都必须保存在显存中。这是因为反向传播计算梯度时需要重用这些前向激活值。

对于标准的 Transformer 层，非 attention 部分的激活值通常会随批次大小（Batch Size）和序列长度（Sequence Length）近似线性增长；但标准 self-attention 在前向中会产生 attention scores/probs 等 $$O(S^2)$$ 中间量（$$S$$ 为序列长度），因此长序列下的 attention 激活显存具有二次增长项。FlashAttention 通过分块和重计算避免保存完整的 $$S \times S$$ 注意力矩阵，可显著降低 attention 的峰值显存。在长序列训练或使用了大批次的情况下，激活值占用的显存甚至会远远超过模型参数本身的显存占用。如果不加干预，这就成为训练更大模型或更长上下文的致命瓶颈。

## 7.5.2 梯度检查点机制

**激活重计算**（Activation Checkpointing，也常被称为 Gradient Checkpointing）是解决激活显存瓶颈的标准方案。其核心思想非常直观：**放弃保存部分前向传播的激活值，在反向传播需要用到时，再重新计算一遍。**

具体做法是：

1. **前向传播**：将模型划分为多个“段”（Segments）。在层与层的边界处保存少量的“检查点”激活值，而丢弃段内部所有的中间激活值。
2. **反向传播**：当反向传播到达某个段时，从该段保存的“检查点”出发，额外执行一次局部的前向传播，重新计算出所需的中间激活值，然后计算梯度，最后再次丢弃这些临时激活值。

## 7.5.3 内存与计算的权衡

通过激活重计算，显存占用可以大幅降低。具体的降低幅度取决于检查点策略：

* **均匀检查点策略**（Uniform Checkpointing）：在理想化的均匀链式模型中，可将 $$L$$ 层分成约 $$\sqrt{L}$$ 个段，在每段边界保存激活值，使显存从 $$O(L)$$ 降低至 $$O(\sqrt{L})$$。这是理解时间-空间权衡的经典结果，而不是所有 Transformer 实现的固定最优解。
* **粗粒度检查点**（Coarse-grained）：只在少数关键位置保存，显存可至 $$O(1)$$，但重计算开销增大。
* **细粒度检查点**（Fine-grained）：保存频繁，显存接近 $$O(L)$$，重计算开销小。

实际应用中，训练框架通常会按层类型、序列长度、micro-batch、并行策略和显存预算选择检查点区域；均匀切分只是常见基线之一。

这种显存节省的代价是额外重计算，开销取决于保存哪些算子和段的长度；粗略口径常会看到 20-50% 的额外前向计算，而不是固定 33%。但在实际的大规模训练中，这通常是值得的：节省下来的显存可以用于更长序列、更大 micro-batch 或更少的模型切分，从而提高整体吞吐。

## 7.5.4 现代重计算策略

如今的模型训练并不满足于简单的“全量重计算”。Megatron-LM 等框架引入了**选择性激活重计算**（Selective Activation Checkpointing）：只对那些“显存占用极大但重新计算很快”的操作（如注意力机制中的 Softmax 和 Dropout）进行重计算，而对那些“显存占用小但计算耗时”的操作（如稠密矩阵乘法）进行正常保存。这种精细化的策略在几乎不增加计算时间的情况下，能实现显著的显存节省。


---

# Agent Instructions: Querying This Documentation

If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question.

Perform an HTTP GET request on the current page URL with the `ask` query parameter:

```
GET https://yeasy.gitbook.io/llm_internals/di-er-bu-fen-xun-lian-pian/07_distributed_training/7.5_activation_checkpointing.md?ask=<question>
```

The question should be specific, self-contained, and written in natural language.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.

Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.