[論文レビュー] Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

Paras Jain, Ajay N. Jain+6 more|arXiv (Cornell University)|Oct 7, 2019

Parallel Computing and Optimization Techniques被引用数 51

ひとこと要約

本論文はテンソルリマテリアライゼーションを制約付き最適化問題として定式化し、Checkmateを提案する。Checkmateは MILP または準最適な LP丸め処理を介して最適なリマテリアライゼーションスケジュールを算出し、訓練メモリを削減しつつオーバーヘッドを最小化、より大きな入力サイズを可能にする。

ABSTRACT

We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing strategies. We introduce Checkmate, a system that solves for optimal rematerialization schedules in reasonable times (under an hour) using off-the-shelf MILP solvers or near-optimal schedules with an approximation algorithm, then uses these schedules to accelerate millions of training iterations. Our method scales to complex, realistic architectures and is hardware-aware through the use of accelerator-specific, profile-based cost models. In addition to reducing training cost, Checkmate enables real-world networks to be trained with up to 5.1x larger input sizes. Checkmate is an open-source project, available at https://github.com/parasj/checkmate.

研究の動機と目的

• 任意の DAG に対してテンソルリマテリアライゼーションをメモリ制約付き最適化問題として定式化する。
• メモリ予算の下で訓練時間を最小化する MILP 定式化を開発する。
• 現実的なネットワークへ拡張するための LP 丸めによる高速近似を提供する。
• ターゲットアクセラレータ上でプロファイリングして Checkmate をハードウェアおよびコスト意識に対応させる。
• 実用的な DNN 訓練効率を改善するオープンソースシステムを提供する。

提案手法

• 訓練 DAG を各ノードのメモリ M_v およびコスト C_v を用いたメモリコスト付き計算グラフとしてモデル化する。
• 実行を T 段階に展開し、再計算を表す二値変数 R_t,i とチェックポイント保持を表す二値変数 S_t,i を定義する。
• メモリ制約付き MILP（および線形化された派生版）を定式化して M_budget の下で総計算コストを最小化する。
• 動的メモリ使用を捉えるためにメモリアカウンティング U_t,k および解放指標 Free_t,i,k を導入し、ブール再編成によって多項式制約を線形化する。
• 解法空間を縮小し MILP ソルバーの性能を向上させるため、フロンティア進行型の段階分割を使用する。
• 厳密な MILP が困難な場合に近似的に整数解を得るため、2 段階の LP ベース丸めアルゴリズムを提供する。

実験結果

リサーチクエスチョン

• RQ1 任意のニューラルネットワークグラフに対して、テンソルリマテリアライゼーションをメモリ制約下の最適化問題としてどう定式化できるか？
• RQ2 最適なリマテリアライゼーションを適用する際のメモリ節約と計算オーバーヘッドのトレードオフは何か？
• RQ3 複雑なアーキテクチャに対して実用的な時間内に近似的に最適スケジュールを得られる LP ベースの近似法は可能か？
• RQ4 Checkmate は現実的なネットワークとハードウェアプロファイルにどれだけスケールするか？
• RQ5 最適なリマテリアライゼーションは従来法と比べてより大きな入力やバッチサイズでの訓練を可能にするか？

主な発見

• 最適なリマテリアライゼーションはメモリ使用量を削減し、画像分類と分割アーキテクチャ全体で控えめなオーバーヘッドでより大きな入力サイズを可能にする。
• 本手法は任意の DAG と不均一なレイヤーコストをサポートし、メモリ効率と計算時間の点でベースラインを上回る。
• Checkmate は同じアクセラレータメモリで最大 5.1 倍の大きな入力サイズを可能にする。
• 標準ソルバーを用いた完全な MILP は中規模ネットワークを1時間未満で解くことができ、フロンティア進行型の分割は解く時間を劇的に短縮する。
• 2 段階の丸めは多くのベンチマークで多項式時間で近似的に最適な整数解を提供する。
• システムは TensorFlow 2.0 で実装されており、オープンソースです（GitHub: parasj/checkmate）。