Activation Checkpointing
1. 引言：大型模型時代的視訊記憶體困境
當模型引數量邁入千億乃至萬億級別，單張加速器的視訊記憶體容量已不再是寬裕的資源，而成為制約訓練可行性的最直接瓶頸。大語言模型、多模態基礎模型以及大規模推薦系統的訓練，不僅需要儲存海量引數和最佳化器狀態，還必須在前向傳播過程中保留中間啟用張量，以供反向傳播計算梯度。以典型的 Transformer 結構為例，假設模型層數為 L，隱藏維度為 d，序列長度為 s，批次大小為 b，則單個 Transformer 層產生的啟用張量總位元組數約為 b × s × d × 常數。對於 L=96、d=12288、s=4096 的千億引數模型，僅啟用部分即可輕易佔據數百 GB 視訊記憶體，而目前最先進的單 GPU 高頻寬儲存器 （HBM） 容量僅約 80 GB 至 192 GB 不等。即使引入張量並行、流水線並行等模型並行手段將引數切分至多個加速器，單卡的啟用視訊記憶體壓力依然巨大，且在增大批次大小或處理更長序列時呈線性增長。
視訊記憶體與算力增速的剪刀差進一步加劇了這一矛盾。過去十年間，GPU 的浮點運算能力提升了數十倍，而 HBM 容量和頻寬的增速則遠低於此。產業界長期面臨一個現實：算力相對“充裕”，視訊記憶體卻無比“稀缺”。在此背景下，任何一種能夠在不顯著增加總計算量的前提下壓縮啟用記憶體佔用的技術，都將直接轉化為模型擴充套件和工程迭代的自由度。Activation Checkpointing（啟用檢查點，亦常稱作梯度檢查點）正是這樣一種以計算換記憶體的核心技術。它通過在前向傳播中策略性地丟棄部分中間啟用，並在反向傳播時選擇性重新計算，將啟用所需的峰值視訊記憶體降低至亞線性甚至對數級別，代價僅為引入一定比例的額外前向計算。自 1990 年代自動微分領域的理論奠基以來，經深度學習架構的工程化改造，這一技術已成為當前千億/萬億引數大型模型訓練的標配元件，深刻重塑了分散式訓練系統的記憶體管理體系。
在展開論述之前，有必要釐清一個常見誤解：Activation Checkpointing 並非簡單地“少存一些中間結果”。它立足於一個被嚴格證明的圖靈獎級理論——在計算圖中，通過精心選擇保留哪些中間節點的值，可獲得記憶體佔用的最優下界以及對應的重計算複雜度上界。它背後涉及組合最優、層次化分治和現代編譯最佳化等多領域的交叉，也是自動微分、分散式系統和機器學習系統工程三者匯流的典範案例。本節將以此為起點，綜合剖析該技術從數學原理到工業落地的完整圖景。
2. Activation Checkpointing 的核心思想與直觀解釋
Activation Checkpointing 的根本思想可以歸納為一句話：不保留所有中間結果，只保留若干“檢查點”，需要時從最近檢查點快速重演計算，臨時復現被丟棄的啟用。
在深度學習的一次標準訓練迭代中，前向傳播逐層計算輸出，同時將每一層的輸入、輸出以及某些中間快取的張量儲存在視訊記憶體中，形成計算圖的反向依賴。反向傳播時，根據鏈式法則，每一層引數的梯度都依賴於該層的輸入啟用與上游傳遞而來的輸出梯度。因此，若沒有特殊的記憶體管理策略，L 層模型就必須同時在視訊記憶體中駐留 L 組啟用張量，導致啟用記憶體佔用量與網路深度呈嚴格線性關係。
Activation Checkpointing 挑戰了這一樸素邏輯：它允許開發者在計算圖的特定位置設定“檢查點”。前向傳播經過檢查點時，正常儲存該節點的輸出張量；在非檢查點層，啟用張量在完成前向傳遞後即被釋放。反向傳播進行到某個區段時，由於缺少直接可用的中間啟用，梯度計算模組會自動識別當前區段的起始檢查點，並以該檢查點處儲存的輸出為輸入，重新執行該段內若干層的前向計算，從而按需“復現”出該段內每一層所需的中間啟用。梯度計算完成之後，這些臨時啟用隨即被回收，整個過程中視訊記憶體中同時駐留的啟用總量由“全部層”變為“若干檢查點加當前區段內的啟用”。
直觀上，這好比一個登山者在長途跋涉中選擇性搭建營地。傳統方法要求登山者在每一步都留下帳篷（啟用），導致裝備物資（視訊記憶體）快速耗盡；而 Activation Checkpointing 則只在高海拔關鍵處（檢查點）安置大型營地，其餘路段輕裝前進。需要返回某段獲取資訊時，再從最近營地快速沿原路復攀一次，事畢即撤收。顯然，營地越密集則額外攀登越少，但需攜帶的保障物資越多；營地越稀疏則物資需求越低，但重攀代價越大。這一權衡精確地刻畫了啟用檢查點技術的引數配置空間。若將“營地”數目記為 c，則記憶體大致降低至 O(L/c)，而額外計算開銷約為 O(c) 的前向傳播，由此可匯出經典的平方根均衡策略，即在給定記憶體預算下最小化額外計算量。
3. 理論淵源：從自動微分的檢查點儲存到深度學習
Activation Checkpointing 並非深度學習時代的原創發明，其理論基礎可追溯至計算數學領域對自動微分（Automatic Differentiation, AD）記憶體複雜度的長期研究。在反向模式自動微分（即反向傳播的數學本質）中，為了計算標量函式對眾多輸入變數的梯度，系統必須儲存一條記錄所有中間變數的“磁帶”，該磁帶的長度通常與原始計算的操作步數成正比。對於迭代求解、時間序列模擬或深度鏈式神經網路等包含長序列計算的問題，磁帶的記憶體佔用迅速膨脹，成為記憶體瓶頸。
1992 年，Andreas Griewank 在其開創性論文《Achieving Logarithmic Growth of Temporal and Spatial Complexity in Reverse Automatic Differentiation》中，提出了一種革命性的檢查點排程策略：利用“分治遞迴”的思想，將長度為 n 的計算過程劃分為不均勻的段落，並在各階段末端設定檢查點，使得反向傳播過程的總重計算次數僅為原始前向計算的數倍，而記憶體佔用僅需 O(\log n) 個檢查點。這一成果證明：反向模式自動微分的記憶體複雜度可以從線性降至對數級別，而時間開銷僅增加一個常數因子。該策略被後續文獻稱為“ binomial checkpointing”或“ Griewank checkpointing”，奠定了檢查點技術的數學最優性基礎。
在 Griewank 之前，Volin 與 Ostrovskii（1985）以及 Morgenstern（1979）等已研究了類似“蛇形”或“分段”重計算方法，但未達到對數記憶體的理論保證。Griewank 將檢查點問題形式化為有向無環圖上的峰值記憶體最小化問題，並給出了最優檢查點放置的動態規劃演算法。隨著時間的推移，這一理論被吸收進自動微分工具（如 ADIFOR、TAPENADE）中，用於處理大規模數值模擬的梯度計算。進入深度學習時代後，由於神經網路的高度層狀結構和反向傳播的特性，自動微分中的檢查點儲存技術被自然地遷移過來，形成了我們今天所熟知的 Activation Checkpointing。因此，該技術的本質是對計算圖進行記憶體感知的排程，它的數學根基遠比簡單的“丟棄再重算”深厚得多。
4. 形式化定義與計算圖模型
為建立精確的分析架構，我們需要將深度學習模型的前向計算抽象為一個有向無環圖 G = (V, E)，其中節點 v \in V 表示操作（運算元），邊表示張量資料的流動。前向傳播相當於按拓撲序執行每個節點，每個節點在執行結束後會生成一個或多個輸出張量。反向傳播時，梯度沿逆向拓撲序流動，每個節點需要其直接前驅節點的輸出啟用來計算區域性梯度。傳統訓練模式下，所有節點的輸出啟用均須保留在記憶體中，直到反向傳播不再需要為止。因此，記憶體佔用的峰值表現為所有“活躍”張量的大小之和，其下界與圖的關鍵路徑長度相關。
Activation Checkpointing 的核心在於引入一個“檢查點集合” C \subseteq V，並規定：在前向執行過程中，僅當節點屬於 C 時，其輸出才被長期保留；其餘節點執行後立即釋放輸出。反向傳播按層逆序推進，當需要計算某個非檢查點節點 v 的梯度時，必須從一個距離最近的前向檢查點 c 出發，沿前向邊重算 c 到 v 之間所有被丟棄的中間節點，從而臨時再現 v 所需的輸入啟用。這一定義自然地引出了兩個關鍵效能指標：記憶體峰值 M，即任意時刻所存檢查點與當前重計算區段內啟用之和的最大值；重計算開銷因子 R，定義為由於回退重算而引入的總前向計算量與原始前向計算量的比值減一。
設總操作步數為 n，檢查點數目為 k，在不同調度策略下，M 與 R 呈現出不同的權衡函式。均勻檢查點策略中，每隔 n/k 步設定一個檢查點，導致 M = O(n/k)，R = O(k)，在固定記憶體預算下，取 k \propto sqrt(n) 時達到 M = O(sqrt(n)) 和 R = O(sqrt(n)) 的平方根平衡。Griewank 的分治策略則可實現 M = O(\log n) 而 R 僅為 O(\log n) 或 O(1) 部分取決於具體實現，實際上經典 Griewank 演算法的時間開銷約為 R \approx \log_2 n 次重計算。顯然，後者在大規模網路中更具優越性。這一形式化架構使我們能夠將 Activation Checkpointing 視為一個資源受限的圖排程問題，進而藉助動態規劃或機器學習方法自動發現最優檢查點放置。
5. 基本排程策略：均勻檢查點與平方根平衡
最簡單的檢查點排程策略是均勻分佈檢查點。假設模型由 L 個相同的順序塊組成（例如 Transformer 層），我們每隔 s 層設定一個檢查點，總共保留 k = L/s 個檢查點。前向過程中，所有層的輸出初始均被計算，但僅檢查點層的輸出被駐留記憶體，其餘層輸出即時釋放。反向傳播開始時，從最後一個檢查點開始，對每個檢查點區間 [i, i+s)，從該檢查點的儲存輸出重算該區間內所有層的啟用，然後執行該區間所有層的反向傳播，重算的啟用在該區間反向結束後釋放。由此，峰值記憶體為單個區間的啟用加上 k 個檢查點的輸出，即 M = O(s + k)。在檢查點大小與層啟用大小相當的假設下，給定總記憶體預算 M，我們令 s \approx M，k \approx M，則 L = s \times k，即 M \approx sqrt(L)。額外重計算次數為 k \times s = L，恰好等於原始前向計算量，因此額外開銷 R = 1.0（即總前向計算變為原來的 2 倍）。這就是著名的平方根平衡：記憶體從 O(L) 降至 O(sqrt(L))，而計算時間僅翻倍。
均勻檢查點的優點在於實現簡單，且與深度學習架構的分層模型天然契合。PyTorch 的 torch.utils.checkpoint 在早期版本中便採用了固定間隔的檢查點策略（使用者可自定義分段）。對於層數不是極端龐大的網路（如數十到上百層），均勻策略足以在記憶體和計算開銷之間取得實用平衡。然而，當網路進一步加深或單層本身包含複雜子結構時，平方根的縮減可能不再足夠，因為 O(sqrt(L)) 對於數萬步的序列可能仍然過於龐大。這促使人們尋求更優的排程策略，即 Griewank 所提出的對數記憶體方案，它允許記憶體佔用隨 L 的增加僅呈對數增長。
6. 遞迴分治與對數級視訊記憶體：Griewank 最優排程
Griewank（1992）提出的遞迴分治檢查點策略本質上是一種“二分”或“兩進”排程，但常被稱為“binomial checkpointing”。假設計算圖是一個長為 n 的序列，目標是在記憶體中最多同時保留 d 個檢查點的條件下完成反向傳播，並且使重計算步驟總數儘量小。Griewank 證明，通過如下遞迴方式部署檢查點，可實現記憶體 d = \lceil \log_2 n \rceil，且總重計算步數約為 n \log_2 n / 2（在原始論文中略有不同，但本質為對數級別）。
其思想可簡述為：將序列劃分為不等長兩部分，前一部分遞迴處理，後一部分在反向時需要前一部分的最終狀態作為檢查點。具體執行時，前向過程並不一次性執行完整個序列，而是配合反向過程交錯進行。經典 Griewank 排程在實際深度學習架構中實現較為複雜，因為它需要反覆在不同區間內切換前向與反向的執行粒度，難以直接對映到現有的層式執行迴圈上。後來的研究提出了一些近似等效的實現，例如“動態檢查點間隔調整”或“樹形依賴記錄”。
儘管完全的對數策略在工業界尚未成為主流，但其理論價值極其深遠：它證明了 Activation Checkpointing 不存在根本性的擴充套件瓶頸。即使模型層數增長一千倍，記憶體需求僅增加數十倍（對數增長），且計算成本增長非常溫和（對於最深度的網路，額外計算因子通常低於 2 或 3）。這為大型模型開發者提供了信心——視訊記憶體不足絕不是放棄加深網路的理由。近年來，一些高階訓練架構（如 JAX 的 jax.checkpoint 結合 XLA 編譯最佳化）已嘗試在計算圖級別實現近乎對數的檢查點排程，利用編譯器整體分析重計算區間，進一步逼近理論最優。
7. 深度學習架構中的實現機制
現代深度學習架構將 Activation Checkpointing 實施為一種自動微分上下文與自定義前向/反向鉤子的組合。以 PyTorch 為例，torch.utils.checkpoint.checkpoint 函式接受一個任意模組的 forward 函式及輸入，內部不做常規的中間張量儲存，而是定義了一個自定義的 Function，其 forward 方法僅保留輸入和輸出（作為檢查點），中間啟用全部丟棄；backward 方法被呼叫時，使用儲存的輸入重新執行該模組的 forward，在臨時開啟梯度計算的條件下重新生成中間啟用，再依次計算梯度。這個過程對使用者透明，只需將希望進行檢查點的網路片段包裹即可。
TensorFlow 的 tf.recompute_grad 或 tf.GradientTape 結合自定義迴圈提供了類似功能，但需要顯式控制梯度磁帶的範圍。JAX 的 jax.checkpoint（又名 jax.remat）則允許對純函式進行裝飾，並在 XLA 編譯器層面將重計算指令嵌入 HLO 圖，從而獲得更大的最佳化空間，如與運算元融合相結合，減少重計算的記憶體和計算開銷。JAX 的設計使得檢查點策略可以通過變換函式組合輕鬆定製。
關鍵實現細節包括：如何確保在重計算階段原始子圖所依賴的輸入張量仍然存活（即檢查點的保持）；如何處理不可重算的隨機操作（如 Dropout）；如何避免因重計算導致 batch normalization 統計量變化等。實踐中，Dropout 層通常被要求在檢查點區域內關閉或使用固定種子重放，BatchNorm 則常被排除在檢查點區段外，或使用全域性統計量替代。這些工程化細節使得 Activation Checkpointing 從理論走向可靠的生產環境。
8. 與混合精度訓練、梯度累積的協同作用
在大型模型訓練中，Activation Checkpointing 很少單獨使用，而是與混合精度訓練（FP16/BF16）和梯度累積等技術協同，以進一步壓縮記憶體。混合精度訓練通過將啟用與引數儲存為半精度（或 BF16），使得啟用張量位元組數減少一半；梯度累積允許在極小micro-batch下前向與反向傳播，累積梯度後再更新引數，從而降低啟用對批次大小的線性依賴。當三者結合時，總體視訊記憶體縮減呈乘積效應。
例如，在訓練 GPT-3 規模的模型時，採用 BF16 混合精度、序列並行、啟用檢查點（每 1 或 2 層一個檢查點）以及 32 路梯度累積，單卡僅需滿載數百 GB 的 HBM 即可完成原本需要 TB 級視訊記憶體的訓練任務。值得注意的是，混合精度下 Activation Checkpointing 的重計算部分一般以與原始相同精度執行（通常為半精度），因為重計算中的數值誤差可通過損失縮放與動態調整來控制，幾乎不影響收斂性。但某些敏感操作（如 LayerNorm、Softmax）可能保留高精度進行重計算，這需要架構提供靈活的策略配置。此外，梯度累積帶來的多步前向/反向交錯執行，也為檢查點排程提供了更細粒度的排程自由度：可在多個 micro-batch 之間共享檢查點，或利用累積階段的空閒記憶體增加檢查點密度，進一步平衡計算與記憶體。
9. 模型並行體系下的啟用檢查點設計
當模型引數量超出單卡容量，模型並行（包括張量並行和流水線並行）將引數切分至多卡，此時啟用檢查點的設計需要考慮並行劃分的邊界。對於張量並行，一層的前向計算被橫向切分到多個裝置，啟用張量同樣被分片儲存。若在該層設定檢查點，需確保所有裝置同時儲存對應的分片，或者僅在某一裝置保留完整複製（需額外通訊）。更常見的方式是，在張量並行區域內部不設檢查點，而在張量並行單元的入口和出口（即通訊集合點）設定檢查點，重計算時只需從入口的完整啟用重算整個並行區域，無需跨裝置聯合重計算，從而避免了複雜的通訊協議。
流水線並行則天然將模型按層切分為多個階段，檢查點可自然地嵌入各階段內的微批次排程中。1F1B（one-forward-one-backward）或交錯排程中，檢查點的存留與放棄需與微批次的流動相匹配。先進流水線並行系統（如 Megatron-LM、DeepSpeed）提供了整合的啟用檢查點配置，允許使用者指定哪些流水線階段或哪些層啟用檢查點，自動掛載重計算邏輯並管理跨階段的啟用傳輸。在 3D 並行（資料、張量、流水線）的超大規模訓練中，檢查點成為溝通三者的“記憶體控制器”，使各維度並行的區域性記憶體開銷均能被獨立約束和最佳化。
10. 效能模型與開銷量化分析
為了充分理解 Activation Checkpointing 的利弊，需建立效能模型。設單層前向計算時間為 F，反向時間為 B（通常 B ≈ 2F），頻寬傳輸開銷暫忽略。若無檢查點，總迭代時間 T_0 = L(F + B)，峰值記憶體 M_0 = L \cdot A，其中 A 為單層啟用大小。採用均勻檢查點，間隔 s，檢查點數 k = L/s，則前向額外開銷為每個區間重算 s 層，共 k 個區間，故總前向時間 F_total = L \cdot F + k \cdot s \cdot F = 2LF；反向時間 B_total = L \cdot B 不變，因此總時間 T_ckpt = 2LF + LB ≈ 2.66 T_0（假設 B=2F）。記憶體峰值 M_ckpt = s \cdot A + k \cdot A。在記憶體約束 M_budget 下，令 M_ckpt = M_budget 可反解最優 s 與 k。若 M_budget 遠小於 M_0，則 k 會很大，加倍比例接近 2。實際上，由於重計算無需儲存中間啟用的梯度圖，重計算塊的峰值記憶體較原始前向略低，但核心量級一致。
遞迴分治策略下，額外計算因子近似為 \log_2 L，而記憶體為 O(\log L)。當 L 非常大時（如數千層），對數策略的時間優勢明顯。例如 L=1024，均勻策略 R=1（總時間 2x），記憶體 sqrt(L)≈32；對數策略 R≈10（總時間 11x），記憶體 ≈10。但若記憶體極為緊張只能容納 10 層啟用，則對數策略成為唯一選擇。實踐中，千億引數模型的層數通常在 100 層左右，均勻策略 2x 時間是可接受的，因此工業界大多采用每 1~4 層一個檢查點的均勻配置，保持重計算開銷在 10%~30% 之間，滿足效率要求。
11. 檢查點的動態選擇與最優版面配置
給定一個具體的計算圖，檢查點的放置可以形式化為一個最最佳化問題：在記憶體上限約束下，最小化額外計算量。對鏈式圖，有動態規劃演算法可以在 O(n^2) 時間內解出最優檢查點位置。對於更復雜的 DAG（如帶有殘差連線的網路），問題變得 NP-hard，但現實中常採用貪心或啟發式方法選擇瓶頸節點作為檢查點：優先儲存在網路中被多次使用、且重算代價高的節點（如特徵融合後的張量）。這類方法可視為“選擇性檢查點”（selective checkpointing），它與自動微分中的“ checkpoints on a tape”問題密切相關。
另一個維度是自適應排程：在訓練過程中根據即時記憶體利用率動態插入或丟棄檢查點。例如，當最佳化器狀態更新佔用較高記憶體時，系統可以臨時將某些快取的檢查點釋放，並在反向時從更早的檢查點重算，從而在不違反記憶體限制的前提下最大化計算效率。DeepSpeed 的 ZeRO-Offload 和 ZeRO-Infinity 部分結合了這種動態性，將檢查點與 CPU/NVMe 解除安裝協同，進一步拓寬了可能性邊界。然而，過度的動態性可能破壞編譯器最佳化（如圖融合），因此多數靜態圖系統傾向編譯期固定檢查點策略。
12. 針對 Transformer 架構的定製化設計
Transformer 模型的獨特結構為檢查點最佳化提供了豐富機會。每一層由 Multi-Head Self-Attention（MHA）和 Feed-Forward Network（FFN）組成，中間有 LayerNorm 和殘差相加。Attention 模組內部產生的 QKV 矩陣以及 attention score 矩陣是啟用記憶體的巨大來源，其大小與序列長度的平方成正比。為此，很多實現選擇在 Attention 計算內部設定檢查點，即儲存 QKV 投影後的結果，而丟棄 attention score 和 softmax 輸出。反向重算時僅需從 QKV 快速還原 attention score，且避免了儲存 O(s^2) 的大張量。對於 FFN，啟用函式（如 GeLU）的輸入通常被設為檢查點，因為其重算代價極低。
自 Megatron-LM 起，這個策略被規範化：將每個 Transformer 層的前向函式通過 checkpoint 包裝，對注意力塊和 FFN 塊均丟棄中間啟用，僅保留輸入隱藏狀態和部分必要張量。在此基礎上，針對序列長度極長的情況（如 32k 以上），FlashAttention 等 I/O 感知演算法可結合檢查點更精細地管理 SRAM 與 HBM 之間的資料移動，但本質是另一種層次的重計算：FlashAttention 在 kernel 內部進行分塊重算，類似線上排程。未來，這些不同層級的記憶體最佳化將走向統一，編譯器可將高層次檢查點與底層 kernel 重計算融合排程，獲得更極致的視訊記憶體利用率。
13. 檢查點與長序列、多模態等複雜場景
長序列處理是檢查點技術的一大用武之地。當序列長度增至百萬 token 時，注意力啟用達到 TB 級，標準檢查點即使每層只保留一個隱藏狀態也依然過大。因此，需要發展更細粒度的或層次化的檢查點策略：在序列維度上也進行分段檢查點。例如，Blockwise Parallel Transformer 或 RingAttention 將序列切分為多個塊，每個塊分別計算注意力並通訊，此時可在塊間設定檢查點，塊內重算。這使得啟用記憶體與序列長度解耦，更靈活地適應長上下文。
在多模態大型模型中，不同模態的編碼器、解碼器以及融合模組的啟用分佈極不均衡，檢查點策略需根據運算元和張量大小進行異構配置。例如，影像編碼器（ViT）可能注重 patch 間注意力的啟用，而文本骨幹則更關注序列長度的平方項。通過剖析（profiling）靜態或動態的計算圖，可自動識別高記憶體運算元並施加選擇性重計算。此外，引入 CPU 記憶體或 NVMe 作為檢查點的備份儲存，也能在不增加重計算的前提下降低 GPU 視訊記憶體佔用，即檢查點與解除安裝的聯合排程，這在千卡訓練中已有初步驗證。
14. 當前侷限與工程挑戰
儘管 Activation Checkpointing 已極為成功，但它仍面臨多重挑戰。第一，額外的計算開銷始終存在，即便最佳化的均勻策略也使訓練時間延長 30% 左右，對於需要數週乃至數月訓練的超大型模型，這部分開銷轉化為巨大的能源和硬體成本。第二，確定性重計算難題：當計算包含不可逆操作（如隨機 Dropout、隨機深度、某些池化）時，必須小心控制偽隨機種子以保證重算結果完全一致，否則會引入梯度誤差導致收斂問題；在分散式環境下，模型並行的通訊次序也可能因重算而改變，帶來不確定。第三，架構實現的效能瓶頸：基於 Python 的 torch.utils.checkpoint 在呼叫時產生了額外的主機端開銷，且打斷運算元融合，降低 GPU 利用率；雖然通過 TorchScript 或 torch.compile 可部分緩解，但自動化的端到端最佳化仍不完美。第四，異構計算裝置（如 TPU、IPU）的記憶體模型與 GPU 差異大，原有的檢查點策略需重新評估與適配，缺乏統一抽象。克服這些限制將是下一代訓練系統的重要方向。
15. 總結與未來展望
Activation Checkpointing 以極其優雅的“時間換空間”理念，從自動微分的理論被引入深度學習系統，已成為訓練超大規模模型不可或缺的基石。它深植於計算圖的記憶體排程理論，以 Griewank 的對數級記憶體結果為最優邊界，在實踐中通過均勻分段、遞迴分治以及選擇性保留等策略，將啟用記憶體需求降低數十倍，而額外計算代價可控。當前，它與混合精度、模型並行、梯度累積、FlashAttention 等技術深度融合，共同支撐起萬億引數時代的工程落地。
展望未來，檢查點技術將繼續向自動化、異構化和協同化發展。基於機器學習的計算圖最佳化可能自動探索任意模型結構下的最優檢查點版面配置；在異構記憶體池（HBM + DDR + SSD）場景中，檢查點將被統一管理為多級快取，形成以重計算為最內層、解除安裝為外層的層次化記憶體子系統；在分散式訓練排程器中，檢查點策略將作為第一類約束納入全域性最佳化，平衡各裝置的記憶體與計算負載。同時，伴隨著神經架構搜尋和動態網路結構的興起，檢查點機制需支援可變計算圖和動態控制流，這對編譯器與執行時系統提出了更高要求。我們有理由相信，這個起源於 1990 年代的簡潔思想，將在未來引領更深遠的大規模機器學習系統革新，讓人類在有限的物理視訊記憶體中持續突破模型智慧的邊界。