【DL輪読会】MEMORY OPTIMIZATION FOR DEEP NETWORKS (ICLR’21)

DEEP LEARNING JP [DL Papers] MEMORY OPTIMIZATION FOR DEEP NETWORKS (ICLR’21) Presenter: Masanori Misono (Univ. Tokyo) http://deeplearning.jp/ 2020-01-29 1

• http://deeplearning.jp/

書誌情報 タイトル MEMORY OPTIMIZATION FOR DEEP NETWORKS 著者 Aashaka Shah, Chao-Yuan Wu, Jayashree Mohan, Vijay Chidambaram, Philipp Kraehenbuehl (University of Texas) 会議 ICLR’21 (6, 8, 7, 7; spotlight) URL https://openreview.net/forum?id=bnY0jm4l59 https://github.com/utsaslab/MONeT 概要 限られたメモリ制約のなかで大きなNNモデルを効率的に訓練するために，最適な チェックポインティングのスケジュールと各種オペレータを線形計画法を用いて求 める 2020-01-29 2

• https://openreview.net/forum?id=bnY0jm4l59
• https://github.com/utsaslab/MONeT

Motivation • NNモデルの巨大化により，訓練時にメモリ量がボトルネックとなることが多い  過去五年でGPU性能は32xだが，利用可能メモリ量は2.5x程度の増加に留まる • 既存手法  分散学習 (モデル並列，データ並列)  複数のマシンが必要 • 一台(あるいは限られたリソースの中で)必要なメモリがリソース上限を超えるよ うなモデルを訓練したい 2020-01-29 3

そもそも何にメモリを使うか • 各層の重み • 各層の勾配を計算するための途中計算結果 • 訓練データ 2020-01-29 4

NNの使用メモリ量の最適化手法 • 大きく分けて二種類に分類される (それぞれ直交する話) 1. ローカルレベル(レイヤー)単位での最適化  ReLUの結果の量子化 (符号だけ保持する)  in-place activated batchnorm  FP16の利用  etc. 2. グローバル(グラフ全体)レベルでの最適化  チェックポインティング (ネットワークを分割し，分割した内部の中間結果を破棄，必要になったら再計算)  どの部分の再計算するか決定するのをチェックポイントのスケジュールという 2020-01-29 5

この研究のゴール • 自動で最適なチェックポイントスケジュール及び各種オペレーターを見つ ける 2020-01-29 6

Approach • フォワードパス，バックワードパス，及び再計算時におけるメモリの最大 使用量を理論的に求める • 最適なスケジュールを，問題を使用可能なメモリ量の制約下で計算オーバ ヘッドを最小化するような線形計画問題に落とし込むことで求める 2020-01-29 7

Forward / Backward pass 2020-01-29 8

Objective Minimize With 2020-01-29 9

Objective 計算コスト τi,l: オペレータlの計算コスト δi,l: どのオペレータを使うか Minimize With 2020-01-29 10

Objective Minimize With 2020-01-29 メモリ使用量制約 11

Objective Minimize With 2020-01-29 チェックポイント制約 (再計算には保存された値か，再計算された前の値が必要) 12

Evaluation • 10%の計算オーバヘッドを許容した上で，どこまでメモリが削減できるか 2020-01-29 13

Evaluation 2020-01-29 14

Summary • 自動でNNモデル訓練時のメモリを減らすための仕組みMONETを提案 • MONETはローカル(オペレータ)レベルの最適化と，グローバル(グラフ)レ ベルの最適化を線形計画法の形で解き，計算オーバヘッドが低くかつメモ リ効率の良いチェックポイントスケジュールを実現する • メモリ使用量と計算効率のトレードオフでstate-of-the-artを達成 2020-01-29 15

2020-01-29 16