[論文レビュー] Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Training

Bojian Zheng, Abhishek Tiwari+2 more|arXiv (Cornell University)|May 22, 2018

Advanced Neural Network Applications参考文献 72被引用数 6

ひとこと要約

Echoは、GPUメモリフットプリントを削減するコンパイラベース最適化であり、特徴マップを保存する代わりに選択的に再計算することで、レイヤー固有のランタイム推定とグローバルメモリ利益分析を活用する。平均1.89倍、最大3.13倍のメモリフットプリント削減を達成し、最小限のランタイムオーバーヘッドで、同じハードウェア上でより高速な学習、より大きなバッチサイズ、より深いモデルを可能にする。

ABSTRACT

The Long-Short-Term-Memory Recurrent Neural Networks (LSTM RNNs) are a popular class of machine learning models for analyzing sequential data. Their training on modern GPUs, however, is limited by the GPU memory capacity. Our profiling results of the LSTM RNN-based Neural Machine Translation (NMT) model reveal that feature maps of the attention and RNN layers form the memory bottleneck and runtime is unevenly distributed across different layers when training on GPUs. Based on these two observations, we propose to recompute the feature maps rather than stashing them persistently in the GPU memory. While the idea of feature map recomputation has been considered before, existing solutions fail to deliver satisfactory footprint reduction, as they do not address two key challenges. For each feature map recomputation to be effective and efficient, its effect on (1) the total memory footprint, and (2) the total execution time has to be carefully estimated. To this end, we propose *Echo*, a new compiler-based optimization scheme that addresses the first challenge with a practical mechanism that estimates the memory benefits of recomputation over the entire computation graph, and the second challenge by non-conservatively estimating the recomputation overhead leveraging layer specifics. *Echo* reduces the GPU memory footprint automatically and transparently without any changes required to the training source code, and is effective for models beyond LSTM RNNs. We evaluate *Echo* on numerous state-of-the-art machine learning workloads on real systems with modern GPUs and observe footprint reduction ratios of 1.89X on average and 3.13X maximum. Such reduction can be converted into faster training with a larger batch size, savings in GPU energy consumption (e.g., training with one GPU as fast as with four), and/or an increase in the maximum number of layers under the same GPU memory budget.

研究の動機と目的

• LSTM RNN学習における高いGPUメモリ消費量がバッチサイズと計算リソースの活用を制限する問題に対処すること。
• 従来の再計算技術がLSTM環境下で正確なメモリ利益とランタイムオーバーヘッドを推定できないという限界を克服すること。
• 学習ソースコードの変更なしに自動的かつ透明にメモリフットプリントを削減すること。
• LSTM RNNにとどまらず、TransformerやResNetを含む多様なモデルを、同じGPUメモリ予算内でサポートすること。
• フットプリント削減をもとに、より高速な学習、より大きなバッチサイズ、またはモデルの深さの増加を実現すること。

提案手法

• Echoは、計算グラフ全体を分析し、特徴マップを再計算するか保存するかの両者におけるネットメモリ利益を推定するコンパイラベースのアプローチを採用する。
• 再計算に伴い解放される特徴マップと再計算用に新たに割り当てられるデータの両方を考慮することで、グローバルメモリ節約を実現する実用的なメカニズムを導入する。
• 非全結合層における低コストな計算特性を活用し、非保守的なアプローチで再計算オーバーヘッドを推定する。
• メモリ節約とランタイムコストのトレードオフに基づき、計算負荷の高いレイヤーを単純に除外しない最適な再計算対象レイヤーを同定する。
• モデルのソースコードを変更する必要なく、既存のディープラーニングフレームワークにスムーズに統合可能であり、手動のカーネルチューニングも不要である。
• Echoは、NMT、DeepSpeech2、Transformer、ResNetといった最先端の複数のモデルを実際のGPU上で評価した。

Figure 1 : Left: A single-layer LSTM RNN that scans through an input sequence. Right: Zoom-in view of one LSTM cell. Both diagrams are simplified for clarity.

実験結果

リサーチクエスチョン

• RQ1 選択的再計算は、禁止的なランタイムオーバーヘッドを伴わずに、LSTM RNN学習におけるGPUメモリフットプリントを効果的に削減できるか？
• RQ2 計算グラフ全体にわたり、保存領域の解放と新たな割り当ての両方を考慮した場合、再計算のネットメモリ利益を正確に推定できるか？
• RQ3 コンパイラベースのシステムは、Transformer や ResNet のような複雑なモデルにおいて、最適な再計算ターゲットを自動で同定できるか？
• RQ4 フットプリント削減をどれほど実用的利点（例：より大きなバッチサイズやより深いモデル）に変換できるか？
• RQ5 多様なDNNワークロードにおいて、Echoは従来の技術と比較して、メモリ削減効果とランタイム効率の両面で優れているか？

主な発見

• Echoは、NMT、DeepSpeech2、Transformer、ResNetを含む多様なモデルにおいて、平均1.89倍、最大3.13倍のGPUメモリフットプリント削減を達成した。
• この手法により、1枚のGPUでベースライン設定の4枚GPUと同等の速度で学習が可能となり、エネルギー消費量が顕著に削減された。
• バッチサイズを固定した状態で、1枚のRTX 2080 Ti（11 GBメモリ）上で、Transformerでは最大レイヤー数が1.83倍、ResNetでは4.0倍にまで増加した。
• 最小限のランタイムオーバーヘッドでこれらの削減を達成しており、実世界の学習ワークロードにおいて実用的であることが示された。
• EchoはLSTM RNNにとどまらず、Transformer やResNetのようなCNNにも強く適用可能である。
• この手法により、同じメモリ予算内でより高速な学習収束が可能となり、より深いモデルのサポートが可能になった。これは、ディープラーニングアーキテクチャ設計のトレンドと整合的である。

Figure 2 : NMT (LSTM RNN Encoder (Middle)-Decoder (Right) with Attention (Left))