# 再マテリアライゼーション ## 定義 再マテリアライゼーション(re-materialization)は、逆伝播に必要な中間活性化(intermediate activations)を事前に保存しておく代わりに、逆伝播時に順伝播を部分的に再計算して復元する手法である。「活性化再計算(activation recomputation)」「勾配チェックポイント(gradient checkpointing)」とも呼ばれる。 計算速度とメモリのトレードオフを制御する手法であり、メモリを削減する代わりに追加の計算コストを支払う。 ## 横断的知見 ### GPipe での採用(2019) [[@2019__NeurIPS__GPipe Easy Scaling with Micro-Batch Pipeline Parallelism]] はパイプライン並列化の文脈で再マテリアライゼーションを組み合わせて採用した。 - **保存するもの**: セル(ステージ)境界の出力活性化のみ。 - **破棄するもの**: 各セル内の中間活性化。 - **再計算するタイミング**: 逆伝播時に第 k 加速器が合成順伝播関数 $F_k$ を再実行して中間活性化を復元する。 ピーク活性化メモリは $O(N \times L)$(保存なし・分割なし)から $O(N + K \times L/M)$ に削減される。ここで N はミニバッチサイズ、L は総層数、K は分割数、M はマイクロバッチ数。 **実効効果:** - AmoebaNet 単一 TPUv2(8GB): ピーク活性化 6.26GB → 3.46GB 削減により、対応モデルが 82M → 318M パラメータに拡大(単一加速器で 3.87 倍)。 - Transformer 単一 TPUv3(16GB): モデルサイズ 282M → 785M(2.7 倍)。 ### 先行研究との接続 再マテリアライゼーションのアイデア自体は GPipe に先行する。Griewank & Walther(2000)の REVOLVE アルゴリズム(ACM TOMS)と Chen+(2016)の「Training deep nets with sublinear memory cost」(arXiv:1604.06174)が起源として引用されている。GPipe はこれをパイプライン並列化のコンテキストに適用し、ステージ単位での再計算と組み合わせた点が新規性である。 ### 逆伝播スケジューリングとの関係 GPipe では逆伝播時の再計算が、上位レイヤーからの勾配を待たずに前倒しでスケジュールできる。これにより再計算の計算コストがバブルオーバーヘッドの一部と重なり、実効的なバブルをさらに削減できる。 ## 未解決の問い - 再計算コスト(追加の FLOPS)対メモリ削減量のトレードオフは、モデルの層構成(演算種別・次元)によりどう変化するか。アテンション層 vs FF 層で再計算コストに差が大きい。 - Flash Attention などの効率的な注意機構と再マテリアライゼーションの相互作用: Flash Attention は既に活性化をタイル単位で再計算しているが、これと独立に適用するべきか統合するべきか。 - Full recomputation(全中間活性化を破棄)vs. selective checkpointing(計算コストが高い演算だけ再計算)の最適境界はどこか。 ## 関連 - 概念: [[パイプライン並列化]] / [[並列化戦略]] / [[LLM分散学習]] / [[チェックポイント]] - エンティティ: [[GPipe]] / [[Google Brain]] - ソース: [[@2019__NeurIPS__GPipe Easy Scaling with Micro-Batch Pipeline Parallelism]] ## 出典 - [[@2019__NeurIPS__GPipe Easy Scaling with Micro-Batch Pipeline Parallelism]](§2.3 Performance Optimization: 再マテリアライゼーションの定式化・ピーク活性化 $O(N+K \times L/M)$・実効効果の実測)