アクティベーションの再計算のトレードオフ
アクティベーションの再計算 (勾配またはアクティベーション チェックポイント) は、前方パスで中間アクティベーションを破棄し、後方パスで再計算することで、トレーニング中の GPU メモリを節約します。
概要
It trades extra compute for the ability to train larger models or longer sequences on the same hardware.
ディープダイブ
バックプロパゲーションでは、勾配を計算するためにフォワードパスのアクティベーションが必要であるため、デフォルトですべての層の出力が保存されます。これは、モデル サイズ、バッチ サイズ、シーケンスの長さに応じて増大する膨大なメモリ コストです。アクティベーションの再計算では、少数の「チェックポイント」テンソル (多くの場合、層の境界だけ) のみが保持され、残りは破棄されます。バックワード パス中に、チェックポイント間でフォワード計算を再実行し、破棄されたアクティベーションをオンデマンドで再生成します。典型的な結果として、sqrt(N) レイヤーごとにチェックポイントを配置すると、メモリはおよそ O(sqrt(N)) に低下し、さらに約 1 つの前方パスが追加されます (計算量が約 33% 増加)。選択的バリアントは、安価だがメモリを大量に消費する演算 (アテンションやドロップアウトなど) のみを再計算し、高価な演算をキャッシュして、大幅に少ない再計算オーバーヘッドでメモリのほとんどを節約します。
技術的な洞察
基本的なトレードオフは、メモリと FLOP です。完全な再計算では、おおよそステップごとに 1 つの余分な前方パスが追加されます (約 30 ～ 40% 遅くなります) が、アクティベーション メモリを桁違いに削減できます。賢い方法は、選択的なチェックポイント設定です。メモリは大きくても計算コストが低い操作 (softmax、layernorm、GELU、アテンション スコア) を特定し、それらのみを再計算し、高価な GEMM の結果をキャッシュに保持し、無駄な計算を最小限に抑えます。
戦略的影響
費用と予算
アーキテクチャの決定により、パフォーマンスと運用コストが何年にもわたって推進されます。
より明確な判決
技術教育は、チームが最新のスタックだけでなく、適切なスタックを選択するのに役立ちます。
品質管理
より良いエンジニアリングの選択により、本番環境での信頼性に関するインシデントが減少します。
アクティベーションの将来の再計算のトレードオフ
再計算はますます自動化され、選択的になってきています。フレームワークは、各演算のメモリと FLOP コストをプロファイリングして最適なチェックポイントを選択し、再計算と CPU/NVMe へのアクティベーション オフロードおよび並列処理戦略を組み合わせます。コンテキストの長さとモデルのサイズが増加し続けるため、オペごとの再計算の決定を自動的に選択するコンパイラ駆動のポリシー (PyTorch、JAX/XLA) に加えて、再計算と通信の重複が強化されるため、余分な FLOP が部分的に隠蔽されます。
現実世界の実装
各レイヤーブロックにチェックポイントを設定することで、他の方法では収まらない大きなトランスフォーマーをトレーニングする
PyTorch の torch.utils.checkpoint を使用してトランスフォーマー ブロックをラップし、アクティベーション メモリをカットする
Megatron-LM のアテンション/ソフトマックスの選択的再計算により、最小限の速度低下でメモリを節約
アクティベーションを保存する代わりに再計算することで、固定 GPU バジェットでより長いシーケンス長を有効にする
リスクとガードレール
1 つのベンチマークを最適化すると、より広範なシステムの弱点が隠れる可能性があります。
インフラストラクチャとメンテナンスのコストは過小評価されがちです。
システムが複雑になるにつれて、セキュリティと可観測性のギャップが拡大する可能性があります。
実装ロードマップ
実装前にレイテンシ、品質、コストの目標を定義します。
現実的な負荷とデータ条件でのベンチマーク。
エラー、ドリフト、ユーザーへの影響を計測器で監視します。
スケーリングの前に、ロールバックとインシデント対応のパスを準備します。
探検を続けましょう
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
次のガイド
SmoothQuant とアクティベーション クオンタイズ
よくある質問
What is Activation Recomputation Tradeoffs?
アクティベーションの再計算 (勾配またはアクティベーション チェックポイント) は、前方パスで中間アクティベーションを破棄し、後方パスで再計算することで、トレーニング中の GPU メモリを節約します。追加のコンピューティングと引き換えに、同じハードウェア上でより大きなモデルやより長いシーケンスをトレーニングできるようになります。
メモリを節約するためにアクティベーションの再計算は何を犠牲にするのでしょうか?
再計算では、保存されたアクティベーションが破棄され、バックワード パスで再生成され、メモリ使用量を削減するために余分なコンピューティングが費やされます。
通常、フォワードパスアクティベーションが保存されるのはなぜですか?
バックワード パスは、フォワード アクティベーションを使用して勾配を計算するため、デフォルトでは、バックワード パスが実行されるまで勾配はメモリ内に保持されます。
完全なアクティベーションの再計算では、通常どのくらいの追加の計算量が追加されますか?
完全な再計算では、バックワード パス中にフォワード計算が再実行され、およそ 1 つのフォワード パスが追加され、計算量が約 30 ～ 40% 増加します。
選択的 (完全ではない) 再計算の背後にある考え方は何ですか?
選択的再計算は、無駄な FLOP を最小限に抑えるために高価な GEMM 結果をキャッシュしながら、大量のメモリを使用するが計算量はほとんどない操作 (softmax やlayernorm など) を対象とします。
さらにメモリを節約するために再計算と組み合わせられることが多い補完的な手法はどれですか?
アクティベーション オフロードは一部のアクティベーションを CPU/NVMe ストレージに移動し、メモリをさらに節約するために再計算や並列処理と組み合わせられることがよくあります。