啟動重新計算權衡
激活重新計算（梯度或激活檢查點）透過丟棄前向傳遞中的中間激活並在後向傳遞中重新計算它們，在訓練期間節省 GPU 記憶體。
概述
It trades extra compute for the ability to train larger models or longer sequences on the same hardware.
深入探討
反向傳播需要前向傳遞激活來計算梯度，因此預設儲存每一層的輸出——巨大的記憶體成本隨著模型大小、批量大小和序列長度的增加而增長。啟動重新計算僅保留一些「檢查點」張量（通常只是層邊界）並丟棄其餘部分。在向後傳遞期間，它重新運行檢查點之間的前向計算，以根據需要重新產生丟棄的激活。經典的結果是，在每 sqrt(N) 層放置檢查點時，記憶體會下降到大約 O(sqrt(N))，同時增加大約一次額外的前向傳遞（計算量增加約 33%）。選擇性變體僅重新計算廉價但佔用大量記憶體的操作（例如注意力或丟失），同時快取昂貴的操作，以少得多的重新計算開銷獲得大部分記憶體節省。
技術洞察
基本的權衡是記憶體與 FLOPs。完全重新計算大約每步增加一次額外的前向傳遞（慢約 30-40%），但可以將啟動記憶體減少一個數量級。明智之舉是選擇性檢查點：識別內存大但計算成本低的操作（softmax、layernorm、GELU、注意力分數）並僅重新計算這些操作，同時緩存昂貴的 GEMM 的結果 - 最大限度地減少計算浪費。
戰略影響
成本與預算
多年來，架構決策決定著效能和營運成本。
更明確的決策
技術教育幫助團隊選擇正確的堆疊，而不僅僅是最新的堆疊。
品質管控
更好的工程選擇可以減少生產中的可靠性事故。
啟動重新計算權衡的未來
重新計算變得越來越自動化和選擇性。現在，框架會分析每個操作的記憶體和 FLOP 成本，以選擇最佳檢查點，並將重新計算與啟動卸載到 CPU/NVMe 以及並行策略相結合。隨著上下文長度和模型大小不斷增長，預計編譯器驅動的策略（在 PyTorch、JAX/XLA 中）會自動選擇每個操作的重新計算決策，再加上重新計算與通訊的更緊密重疊，因此額外的 FLOP 被部分隱藏。
現實世界的實施
透過檢查每個層塊來訓練一個大型變壓器，否則該變壓器無法適應
使用 PyTorch 的 torch.utils.checkpoint 包裝變壓器塊並切割啟動內存
Megatron-LM 中註意力/softmax 的選擇性重新計算，以節省記憶體且速度減慢最小
透過重新計算激活而不是儲存它們，在固定的 GPU 預算上實現更長的序列長度
風險與防護欄
優化一項基準測試可以隱藏更廣泛的系統弱點。
基礎設施和維護成本常常被低估。
隨著系統變得更加複雜，安全性和可觀察性差距可能會擴大。
實施路線圖
在實施之前定義延遲、品質和成本目標。
在實際負載和資料條件下進行基準測試。
儀器監控錯誤、漂移和使用者影響。
在擴展之前準備回滾和事件回應路徑。
不斷探索
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Activation Recomputation Tradeoffs quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
常見問題
What is Activation Recomputation Tradeoffs?
激活重新計算（梯度或激活檢查點）透過丟棄前向傳遞中的中間激活並在後向傳遞中重新計算它們，在訓練期間節省 GPU 記憶體。它以額外的運算能力換取在相同硬體上訓練更大模型或更長序列的能力。
啟動重新計算會犧牲什麼來節省記憶體？
重新計算會丟棄儲存的啟動並在向後傳遞中重新產生它們，從而花費額外的計算來減少記憶體使用。
為什麼前向傳遞激活通常會被儲存？
向後傳遞使用前向激活來計算梯度，因此預設情況下它們會保留在記憶體中，直到向後傳遞運行。
完全啟動重新計算通常會增加大約多少額外計算？
完全重新計算會在後向傳遞過程中重新運行前向計算，大約增加一次額外的前向傳遞，計算量增加了 30-40% 左右。
選擇性（非完整）重新計算背後的想法是什麼？
選擇性重新計算的目標是使用大量記憶體但計算量很少的操作（例如 Softmax 或 Layernorm），同時快取昂貴的 GEMM 結果以最大程度地減少浪費的 FLOP。
哪種補充技術經常與重新計算相結合以節省更多記憶體？
啟動卸載將一些激活移至 CPU/NVMe 存儲，並經常與重新計算和並行性相結合，以進一步節省內存。