[論文レビュー] Memory Optimization for Deep Networks

Aashaka Shah, Chao-Yuan Wu+3 more|arXiv (Cornell University)|Oct 27, 2020

Advanced Neural Network Applications参考文献 34被引用数 10

ひとこと要約

MONeT は、チェックポイントスケジューリングとメモリ効率の良いオペレータ実装を同時に最適化することで、ディープラーニング学習におけるメモリ使用量を削減する自動フレームワークである。メモリと計算のトレードオフを整数計画問題としてモデル化することにより、PyTorch よりも 3 倍のメモリ削減を達成し、計算オーバーヘッドはたったの 9–16% にとどまる。これは、手動最適化および自動化されたベースラインを上回る性能を発揮する。

ABSTRACT

Deep learning is slowly, but steadily, hitting a memory bottleneck. While the tensor computation in top-of-the-line GPUs increased by 32x over the last five years, the total available memory only grew by 2.5x. This prevents researchers from exploring larger architectures, as training large networks requires more memory for storing intermediate outputs. In this paper, we present MONeT, an automatic framework that minimizes both the memory footprint and computational overhead of deep networks. MONeT jointly optimizes the checkpointing schedule and the implementation of various operators. MONeT is able to outperform all prior hand-tuned operations as well as automated checkpointing. MONeT reduces the overall memory requirement by 3x for various PyTorch models, with a 9-16% overhead in computation. For the same computation cost, MONeT requires 1.2-1.8x less memory than current state-of-the-art automated checkpointing frameworks. Our code is available at https://github.com/utsaslab/MONeT.

研究の動機と目的

• 5 年間で GPU メモリは 2.5 倍にしか増加していないが、計算能力は 32 倍に増加したという、ディープラーニング学習における増大するメモリボトルネックに対処すること。
• 局所的なオペレータ実装やグローバルなチェックポイントスケジューリングを別々に最適化する従来のアプローチの限界を克服すること。
• 固定されたメモリ予算の下で、局所的およびグローバルなメモリ削減技術を同時に最適化する自動フレームワークを開発すること。
• フォワードおよびバックワードパス中のピークメモリ消費量を最小限に抑えつつ、計算効率を維持すること。
• メモリ使用量をきめ細かく制御し、メモリと計算の最適なトレードオフを可能にする、統一的かつ理論的に裏付けられた最適化フレームワークを提供すること。

提案手法

• 異なるオペレータ実装およびチェックポイントスケジューリングの下で、フォワード、バックワード、再計算フェーズにおけるメモリおよび計算コストの境界を形式化する。
• オペレータ実装とチェックポイントの共同最適化を、線形化されたメモリ制約を持つ 0-1 整数計画問題としてモデル化する。
• 標準の混合整数線形プログラミングソルバを用いて、与えられた予算下でメモリ使用量を最小化する最適な設定を特定する。
• 複数の最適化戦略をサポートする：畳み込みアルゴリズム選択、インプレイス活性化、出力活性化再利用、中間活性化再計算。
• PyTorch と統合可能で、フレームワークの変更を最小限に抑えつつエンドツーエンドの学習を可能にする。
• 理論的分析を用いてピークメモリ消費量のタイトな境界を導出し、正確な最適化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 局所的なオペレータ実装とグローバルなチェックポイントスケジューリングの共同最適化は、それぞれを別々に最適化するのと比べて、より優れたメモリ-計算トレードオフを達成できるか？
• RQ2 バックプロパゲーションにおいて、異なるメモリ節約技術を組み合わせた場合のピークメモリ消費量の理論的上限は何か？
• RQ3 畳み込みアルゴリズム選択とチェックポイントの共同最適化は、全体のメモリおよび計算効率にどのように影響を与えるか？
• RQ4 自動選択によるメモリ効率の良いオペレータ実装は、手作業で最適化されたものやヒューリスティックベースのアプローチをどれほど上回れるか？
• RQ5 メモリ削減量と計算オーバーヘッドの観点から、MONeT は最先端の自動チェックポイントフレームワークと比べてどの程度優れているか？

主な発見

• MONeT は、ResNet、VGG、UNet、GoogleNet、MobileNet-V2 といった複数のモデルにおいて、PyTorch よりも 3 倍のメモリ削減を達成した。
• 同じ計算コストの下で、MONeT は現在の最先端の自動チェックポイントフレームワークよりも 1.2–1.8 倍少ないメモリを必要とした。
• 標準的な PyTorch 学習と比較して、MONeT は計算時間にたった 9–16% の増加にとどまった。
• 畳み込みアルゴリズムとチェックポイントの共同最適化は顕著なパフォーマンス向上をもたらした。畳み込みのワークスペースメモリ特性は、速度とメモリのトレードオフに強く影響する。
• アブレーションスタディの結果、畳み込み、出力活性化、中間活性化のすべての最適化戦略を組み合わせることで、計算オーバーヘッドが最小限に抑えられた。
• ResNet-50 における詳細なケーススタディでは、MONeT がピークメモリを PyTorch の 14.7 GB から 8.0 GB に削減し、学習の安定性を維持したことが確認された。