激活值檢查點（Activation Checkpointing）是什麼？

在反向傳播時重新計算隱藏層激活值而非存儲，交換計算時間換取記憶體空間，使訓練更大模型成為可能。

核心概念

Activation Checkpointing 基於一個觀察：標準反向傳播需要存儲前向傳播中的所有中間激活值，用於計算梯度。對於深層網路，激活值存儲占用大量記憶體。但這些激活值在反向傳播時可以重新計算（雖然消耗時間）。

Activation Checkpointing 的策略是：在選定的層（檢查點層）存儲輸入，跳過中間層激活值的存儲。反向傳播時，從最近的檢查點層開始，重新計算跳過層的激活值，然後計算梯度。

該技術的key insight是：記憶體節省和計算增加的權衡。如果有 N 層網路，設定 √N 個檢查點，可以將記憶體需求從 O(N) 降至 O(√N)，而計算增加不超過 1 倍。這是一個好的權衡。

運作原理

Activation Checkpointing 的執行流程：

1. 標記某些層為檢查點層（通常是每 √N 層選一個）
2. 前向傳播時：
   • 檢查點層：存儲輸入和層參數
   • 非檢查點層：計算激活值但不存儲（直接通過，或計算後立即丟棄）
3. 反向傳播時：
   • 從最近的檢查點層開始
   • 重新計算該檢查點到下一檢查點間的激活值
   • 計算梯度
   • 丟棄重新計算的激活值
   • 向前一個檢查點層回溯

在 PyTorch 中，這通過 torch.utils.checkpoint.checkpoint 函數實現。例如：

def forward(self, x):
    x = checkpoint(self.layer1, x)  # layer1 是檢查點層
    x = self.layer2(x)
    x = self.layer3(x)
    x = checkpoint(self.layer4, x)  # layer4 是檢查點層
    return x

實際應用

Activation Checkpointing 在大規模 Transformer 訓練中是必需的。Transformer 層數深（100+ 層不罕見），每層有大量注意力頭和前饋網路，激活值累積巨大。應用 checkpointing 後，可以用一半的 GPU 記憶體訓練相同大小的模型。

許多開源 Transformer 實現（Hugging Face、Megatron-LM）都內建 checkpointing 支持。BERT、GPT、T5 的大規模預訓練都使用了 activation checkpointing。

在計算機視覺中，訓練深層 ResNet（如 ResNet-200）時，checkpointing 能增加可用的批量大小。此外，在高分辨率特徵圖的網路（如 semantic segmentation 模型）中，checkpointing 特別有效，因為激活值占用記憶體巨大。

在循環神經網路（RNN）中，checkpointing 幫助訓練長序列。LSTM/GRU 層的激活值隨序列長度增加，checkpointing 能顯著減少記憶體。

在對抗性訓練（如 GAN）中，checkpointing 用於訓練生成器和判別器，特別是對大規模圖像生成模型。

常見誤區

誤區一：認為 Activation Checkpointing 會大幅增加訓練時間。實踐表明，由於現代 GPU 計算和記憶體帶寬的平衡，重新計算的時間通常只增加 20-40%，而記憶體節省可達 50%，是很好的權衡。

誤區二：應用 checkpointing 後忽視其他優化機會。Checkpointing 通常與混合精度、梯度累積等結合使用，各自優化不同的資源瓶頸。

誤區三：在所有層都應用 checkpointing。應該選擇計算密集、激活值大的層作為檢查點層，而不是每層都做。不當使用反而會增加開銷。

與相關技術的比較

• Activation Checkpointing vs 量化：量化降低權重和激活值的精度以節省空間，而 checkpointing 用計算換空間。兩者可結合，量化降低每個激活值大小，checkpointing 降低需要存儲的激活值數量。

• Activation Checkpointing vs 梯度累積：梯度累積延遲參數更新，checkpointing 延遲激活值計算。兩者都是延遲換空間，但作用對象不同。結合使用能進一步減少記憶體。

• Activation Checkpointing vs 模型分片（Pipeline Parallelism）：模型分片將模型分段存放在不同 GPU，checkpointing 在單 GPU 上重新計算激活值。分片針對超大模型（無法放在單 GPU），checkpointing 針對單 GPU 記憶體不足。

• Activation Checkpointing vs Mixed Precision：混合精度用 FP16 加速和節省空間，checkpointing 用計算換空間。兩者互補，通常同時應用。