[論文レビュー] TFLMS: Large Model Support in TensorFlow by Graph Rewriting

Tung D. Le, Haruki Imai+2 more|arXiv (Cornell University)|Jul 5, 2018

Topic Modeling参考文献 11被引用数 28

ひとこと要約

本稿では、TensorFlowにおけるグラフリライト手法TFLMSを提案する。この手法により、中間活性化をCPUメモリにオフロードするためのスワップアウトおよびスワップイン操作を挿入することで、GPUメモリの制限を超えて大規模なディーブニューラルネットワークの学習が可能になる。本手法は、プログラム変換規則を用いてこれらの操作を形式的に導出する。ResNet-50および3DUnetでは、それぞれ最大4.7倍および2倍のバッチサイズを達成し、特に3DUnetはタイリングを伴わず$192^3$のフル画像で学習可能となった。

ABSTRACT

While accelerators such as GPUs have limited memory, deep neural networks are becoming larger and will not fit with the memory limitation of accelerators for training. We propose an approach to tackle this problem by rewriting the computational graph of a neural network, in which swap-out and swap-in operations are inserted to temporarily store intermediate results on CPU memory. In particular, we first revise the concept of a computational graph by defining a concrete semantics for variables in a graph. We then formally show how to derive swap-out and swap-in operations from an existing graph and present rules to optimize the graph. To realize our approach, we developed a module in TensorFlow, named TFLMS. TFLMS is published as a pull request in the TensorFlow repository for contributing to the TensorFlow community. With TFLMS, we were able to train ResNet-50 and 3DUnet with 4.7x and 2x larger batch size, respectively. In particular, we were able to train 3DUNet using images of size of $192^3$ for image segmentation, which, without TFLMS, had been done only by dividing the images to smaller images, which affects the accuracy.

研究の動機と目的

• 非常に大規模なディーブニューラルネットワークを、GPUメモリ容量を超えて学習する課題に対処すること。
• スワップアウトおよびスワップイン操作を計算グラフに挿入する形式的かつ証明可能な正しい手法を提供すること。
• グラフリライトを用いてメモリ使用を最適化することで、大規模モデルの学習効率とスケーラビリティを向上させること。
• 3DUnetをタイリングなしで$192^3$のフル画像で学習可能にすること、これにより精度が損なわれるのを回避すること。

提案手法

• 変数を第一級の市民として扱い、形式的意味論を備えた計算グラフを再定式化することで、TensorFlowのグラフモデルを拡張する。
• 形式的プログラム変換規則を用いてスワップアウトおよびスワップイン操作を導出し、学習中のモデルの正しさを保証する。
• 制御依存関係を静的に特定する2つの戦略、chain_ruleおよびdirect_orderを導入し、最適なスワップインタイミングをトリガーする。
• スワップイン操作の統合といった最適化ルールを適用し、オーバーヘッドを低減してパフォーマンスを向上させる。
• コミュニティでの採用を可能にするために、TensorFlowリポジトリにプルリクエストとして実装し、TFLMSと命名する。
• メモリ使用量と学習パフォーマンスのトレードオフを制御するための設定可能なパラメータ（n_tensors, lb, swap_branches）を導入する。

実験結果

リサーチクエスチョン

• RQ1 計算グラフをどのように形式的に再定式化すれば、第一級の変数をサポートし、メモリ管理のための正しいグラフリライトを可能にするか？
• RQ2 学習中のモデル正しさを保つために、スワップアウトおよびスワップイン操作を導出する形式的規則は何か？
• RQ3 制御依存関係を特定するための異なる戦略が、大規模モデル学習のパフォーマンスに与える影響は何か？
• RQ4 パラメータ（n_tensors, lb, swap_branches）の最適な設定は何か？ これはバッチサイズの最大化と学習スループットの向上に寄与する。
• RQ5 TFLMSは、タイリングなしで3DUnetを$192^3$の画像で学習可能にすることができるか？ また、そのパフォーマンスはどの程度か？

主な発見

• TFLMSにより、ResNet-50の最大バッチサイズが832に達し、TFLMSなしのベースライン176と比較して4.7倍の増加を達成した。
• 3DUnetでは、TFLMSによりベースラインと比較して2倍の最大バッチサイズが可能となり、タイリングなしで$192^3$のフル画像での学習が可能になった。
• ResNet-50では、lb=5およびn_tensors=100の設定で最高の学習スループットが達成され、メモリ使用量と計算・通信のオーバーラップのバランスが取れた。
• スワップイン操作の統合によりパフォーマンスが向上したが、中間テンソルのGPUメモリ保持量が増加したため、最大バッチサイズが減少した。
• スワップブランチを有効にし、しきい値を20に設定することで、3DUnetの最大バッチサイズは3から4に増加し、スケーラビリティへの影響が明確に示された。
• 最適な設定により、3DUnetは$192^3$の入力に対して0.17画像/秒の学習スループットを達成した。これは、従来のメモリ不足エラーのため、学習が不可能であった。