[論文レビュー] Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory

Julien Herrmann, Olivier Beaumont+4 more|arXiv (Cornell University)|Nov 27, 2019

Stochastic Gradient Optimization Techniques参考文献 29被引用数 13

ひとこと要約

この論文は、限られたメモリ環境下で深層ニューラルネットワークをトレーニングする際の最適なチェックポイント戦略を提案し、動的にどの順方向活性化を保存するか、どのものを再計算するかを選択する。入力のみと完全な操作履歴の両方をモデル化する独自の動的計画法を用い、既存の手法と比較して平均で17.2%高いスループットを達成している。また、任意のメモリ制限とPyTorchにおける複雑な順序付きネットワークをサポートしている。

ABSTRACT

This paper introduces a new activation checkpointing method which allows to significantly decrease memory usage when training Deep Neural Networks with the back-propagation algorithm. Similarly to checkpoint-ing techniques coming from the literature on Automatic Differentiation, it consists in dynamically selecting the forward activations that are saved during the training phase, and then automatically recomputing missing activations from those previously recorded. We propose an original computation model that combines two types of activation savings: either only storing the layer inputs, or recording the complete history of operations that produced the outputs (this uses more memory, but requires fewer recomputations in the backward phase), and we provide an algorithm to compute the optimal computation sequence for this model. This paper also describes a PyTorch implementation that processes the entire chain, dealing with any sequential DNN whose internal layers may be arbitrarily complex and automatically executing it according to the optimal checkpointing strategy computed given a memory limit. Through extensive experiments, we show that our implementation consistently outperforms existing checkpoint-ing approaches for a large class of networks, image sizes and batch sizes.

研究の動機と目的

• 深層ニューラルネットワークのトレーニングにおけるメモリボトルネック、特にIoTシステムのような低メモリデバイスにおける課題に対処すること。
• 従来の自動微分チェックポイント戦略を、任意に複雑な内部層を持つ異種DNNに拡張すること。
• 任意のメモリ制限をサポートし、トレーニングスループットを最適化する完全自動のPyTorch互換実装を開発すること。
• さまざまなネットワークアーキテクチャ、画像サイズ、バッチサイズにおいて、従来の定期的および最適チェックポイント戦略を上回ること。

提案手法

• DNNの順方向および逆方向演算をモジュールのシーケンスとしてモデル化し、どの活性化をチェックポイントに保存するかを動的に選択可能にする。
• 2種類のチェックポイント戦略を導入：メモリ効率の高い層の入力のみを保存する戦略、または完全な操作履歴を記録する戦略（再計算回数が少ない）。
• メモリ制約下で最適な永続的スケジューリングを計算するための動的計画法を採用し、最適解のメモリ永続性特性を活用する。
• 複雑な内部計算を含む任意の順序付きPyTorchモジュールをサポートする。
• autogradにシームレスに統合され、ユーザーが定義したメモリ制限に基づき、自動的に最適戦略を適用する。
• 離散的なセグメントベース戦略とは異なり、連続的なメモリ適応を可能とし、前方および後方フェーズの両方で最小限のオーバーヘッドで動作する。

実験結果

リサーチクエスチョン

• RQ1 任意に複雑な内部層を有する異種DNNチェーンに対して、チェックポイント戦略をどのように最適化できるか？
• RQ2 入力のみを保存するか、完全な操作履歴を保存するかの選択において、メモリ使用量と再計算コストの最適なトレードオフは何か？
• RQ3 動的計画法を用いるアプローチが、メモリ制約下で従来の定期的および最適チェックポイント戦略を上回るトレーニングスループットを達成できるか？
• RQ4 本手法は、ネットワークの深さ、画像サイズ、バッチサイズの変化に対して、PyTorchのデフォルト戦略と比較してどのようにスケーリングするか？
• RQ5 本手法により、より大きなモデル、画像、またはバッチサイズのトレーニングを低メモリデバイスで可能にする程度はどの程度か？

主な発見

• 提案された最適戦略は、すべてのテスト設定において、最良の競合手法と比較して平均で17.2%高いトレーニングスループットを達成した。
• 224×224の画像とバッチサイズ8のResNet-1001では、最適戦略が1秒あたり9.18画像のスループットを達成したが、PyTorchのデフォルト戦略はメモリ制限のため失敗した。
• バッチサイズ8の場合、PyTorchの戦略は225 GiBのメモリを必要とし、最適戦略のピーク使用量の16倍以上にのぼった。
• 最適戦略は、PyTorchのデフォルト戦略が失敗する状況でも、大規模モデルのトレーニングを可能にした。特に、ResNet-1001のバッチサイズ8でも同様に有効であった。
• 本手法は、順次（セグメント化）戦略および以前のrevolveアルゴリズムを上回った。特に、メモリ制限が高い場合に、後者では利用可能なメモリを十分に活用できなかったが、本手法はその欠陥を克服した。
• 本手法は、さまざまなネットワークサイズおよび画像解像度において安定したパフォーマンスを維持し、頑健性と柔軟性を示した。