概述
激活重新計算(梯度或激活檢查點)透過丟棄前向傳遞中的中間激活並在後向傳遞中重新計算它們,在訓練期間節省 GPU 記憶體。它以額外的運算能力換取在相同硬體上訓練更大模型或更長序列的能力。
啟動重新計算權衡是一個技術構建塊,會大規模影響模型品質、基礎設施成本、延遲和可靠性。
深入探討
反向傳播需要前向傳遞激活來計算梯度,因此預設儲存每一層的輸出——巨大的記憶體成本隨著模型大小、批量大小和序列長度的增加而增長。啟動重新計算僅保留一些「檢查點」張量(通常只是層邊界)並丟棄其餘部分。在向後傳遞期間,它重新運行檢查點之間的前向計算,以根據需要重新產生丟棄的激活。經典的結果是,在每 sqrt(N) 層放置檢查點時,記憶體會下降到大約 O(sqrt(N)),同時增加大約一次額外的前向傳遞(計算量增加約 33%)。選擇性變體僅重新計算廉價但佔用大量記憶體的操作(例如注意力或丟失),同時快取昂貴的操作,以少得多的重新計算開銷獲得大部分記憶體節省。
技術洞察
基本的權衡是記憶體與 FLOPs。完全重新計算大約每步增加一次額外的前向傳遞(慢約 30-40%),但可以將啟動記憶體減少一個數量級。明智之舉是選擇性檢查點:識別內存大但計算成本低的操作(softmax、layernorm、GELU、注意力分數)並僅重新計算這些操作,同時緩存昂貴的 GEMM 的結果 - 最大限度地減少計算浪費。
掌握啟動重新計算權衡
激活重新計算(梯度或激活檢查點)透過丟棄前向傳遞中的中間激活並在後向傳遞中重新計算它們,在訓練期間節省 GPU 記憶體。它以額外的運算能力換取在相同硬體上訓練更大模型或更長序列的能力。啟動重新計算權衡是一個技術構建塊,會大規模影響模型品質、基礎設施成本、延遲和可靠性。為了建立深入的理解,請將啟動重新計算權衡視為一種操作模型,而不是單一功能:定義所需的結果,澄清假設,並將系統可以可靠地執行的操作與仍需要專家判斷的操作分開。
在實踐中,使用啟動重新運算權衡的強大團隊可以根據可靠性和成本來優化架構、資料和基礎架構選擇。他們記錄明確的成功標準,根據實際數據和工作流程進行測試,並根據觀察到的失敗模式而不是一次性基準測試勝利進行迭代。這就是理論理解轉變為跨產品、政策和營運的持久能力的地方。
多年來,架構決策決定著效能和營運成本。同時,優化一個基準測試可以隱藏更廣泛的系統弱點。最具彈性的方法是將實驗速度與治理規則結合:運行試點、捕獲證據、發布決策日誌,並隨著模型行為、使用者期望和監管要求的發展不斷更新保障措施。
戰略影響
多年來,架構決策決定著效能和營運成本。
多年來,架構決策決定著效能和營運成本。在高品質部署中,這會轉化為可衡量的操作規則、所有權邊界和定期審查儀式,以便團隊可以增強信心,而不是擴大模糊性。
技術教育幫助團隊選擇正確的堆疊,而不僅僅是最新的堆疊。
技術教育幫助團隊選擇正確的堆疊,而不僅僅是最新的堆疊。在高品質部署中,這會轉化為可衡量的操作規則、所有權邊界和定期審查儀式,以便團隊可以增強信心,而不是擴大模糊性。
更好的工程選擇可以減少生產中的可靠性事故。
更好的工程選擇可以減少生產中的可靠性事故。在高品質部署中,這會轉化為可衡量的操作規則、所有權邊界和定期審查儀式,以便團隊可以增強信心,而不是擴大模糊性。
現實世界的實施
透過檢查每個層塊來訓練一個大型變壓器,否則該變壓器無法適應
使用 PyTorch 的 torch.utils.checkpoint 包裝變壓器塊並切割啟動內存
Megatron-LM 中註意力/softmax 的選擇性重新計算,以節省記憶體且速度減慢最小
透過重新計算激活而不是儲存它們,在固定的 GPU 預算上實現更長的序列長度
實施模式
實踐中啟動重新計算的權衡
透過檢查每個層塊來訓練一個無法適應的大型變壓器。
透過對每個層塊設定檢查點來訓練原本不適合的大型變壓器 團隊在預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移跟踪生產力增益和錯誤成本時通常會獲得更好的結果。
實踐中啟動重新計算的權衡
使用 PyTorch 的 torch.utils.checkpoint 包裝變壓器區塊並剪切啟動記憶體。
使用 PyTorch 的 torch.utils.checkpoint 來包裝變壓器區塊並減少啟動記憶體 當團隊預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移追蹤生產力增益和錯誤成本時,通常會獲得更好的結果。
實踐中啟動重新計算的權衡
Megatron-LM 中選擇性重新計算注意力/softmax,以節省記憶體並最小化速度減慢。
在 Megatron-LM 中選擇性地重新計算注意力/softmax,以最小化速度節省內存 當團隊預先定義質量閾值、為邊緣情況保留人工升級路徑並隨著時間的推移跟踪生產力增益和錯誤成本時,通常會獲得更好的結果。
實踐中啟動重新計算的權衡
透過重新計算激活而不是儲存它們,在固定的 GPU 預算上實現更長的序列長度。
透過重新計算激活而不是儲存它們,在固定的 GPU 預算上實現更長的序列長度 當團隊預先定義品質閾值、為邊緣情況保留人工升級路徑並隨著時間的推移跟踪生產力增益和錯誤成本時,通常會獲得更好的結果。
風險與防護欄
優化一項基準測試可以隱藏更廣泛的系統弱點。
基礎設施和維護成本常常被低估。
隨著系統變得更加複雜,安全性和可觀察性差距可能會擴大。
實施路線圖
在實施之前定義延遲、品質和成本目標。
在實施之前定義延遲、品質和成本目標。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。
在實際負載和資料條件下進行基準測試。
在實際負載和資料條件下進行基準測試。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。
儀器監控錯誤、漂移和使用者影響。
儀器監控錯誤、漂移和使用者影響。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。
在擴展之前準備回滾和事件回應路徑。
在擴展之前準備回滾和事件回應路徑。將每個步驟視為證據門:如果不符合標準,則暫停推出,縮小差距,然後再擴大使用。