[論文レビュー] Addressing the Memory Bottleneck in AI Model Training

David Ojika, Bhavesh Patel+4 more|arXiv (Cornell University)|Mar 11, 2020

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

ひとこと要約

本論文は、大容量システムメモリを備えた1ノードサーバー上で、Intel最適化TensorFlowを2代目Intel Xeonスケーラブルプロセッサに適用し、1TBのメモリフットプリントを持つディープニューラルネットワークのトレーニングを初めて実現した。このアプローチにより、分散システムを用いずに、医療画像分野を含むメモリ集約的AIモデルのスケールアップトレーニングが可能となり、従来のメモリボトルネックを克服する。

ABSTRACT

Using medical imaging as case-study, we demonstrate how Intel-optimized TensorFlow on an x86-based server equipped with 2nd Generation Intel Xeon Scalable Processors with large system memory allows for the training of memory-intensive AI/deep-learning models in a scale-up server configuration. We believe our work represents the first training of a deep neural network having large memory footprint (~ 1 TB) on a single-node server. We recommend this configuration to scientists and researchers who wish to develop large, state-of-the-art AI models but are currently limited by memory.

研究の動機と目的

• 大規模AIモデルのトレーニングにおけるメモリボトルネックを解消すること。特に医療画像分野を対象とする。
• 大容量システムメモリを備えた1ノードサーバーが、従来は分散システムを必要としていたモデルをトレーニング可能であることを実証すること。
• インfraストラクチャが限られた研究者にとって、実用的でスケーラブルな分散トレーニングの代替手段を提供すること。
• 十分なメモリが利用可能な場合、最新鋭の高メモリ要件を満たすモデルを1台のサーバー上でトレーニング可能であることを検証すること。
• 大規模AIモデルを開発する科学者向けに、ハードウェアとソフトウェアスタックの推奨事項を提示すること。

提案手法

• x86アーキテクチャのサーバーに2代目Intel Xeonスケーラブルプロセッサを搭載し、Intel最適化TensorFlowを活用した。
• 最大1TBの大容量システムメモリを活用し、モデル全体と活性化状態をすべてメモリ上に保持した。
• 通信オーバーヘッドを回避するため、分散トレーニングではなくスケールアップアーキテクチャを採用した。
• 利用可能なRAMを最大限に活用するよう、TensorFlow内のメモリ管理とデータフローを最適化した。
• 3次元畳み込みネットワークの高メモリ要件を鑑み、医療画像を事例として選定した。
• 高いスループットを維持できるよう、十分なメモリ帯域幅とNUMA対応メモリ割り当てをサーバーに構成した。

実験結果

リサーチクエスチョン

• RQ1 1TBのシステムメモリを備えた1ノードサーバーが、同程度のメモリフットプリントを持つディープニューラルネットワークを正常にトレーニングできるか？
• RQ2 分散トレーニングと比較して、1台サーバーでのスケールアップトレーニングは、メモリ効率性とトレーニングスループットの点でどのように異なるか？
• RQ3 モデル並列化を用いずに、安定したメモリ集約的モデルのトレーニングを可能にするハードウェア・ソフトウェアスタックは何か？
• RQ4 Intel最適化TensorFlowは、大規模AIワークロードにおいて、メモリ利用効率をどの程度まで向上できるか？
• RQ5 クラスタインfraストラクチャがなくても、最新鋭の医療画像分野のAIモデルを1台のサーバーでトレーニング可能か？

主な発見

• 著者らは、1TBのメモリフットプリントを持つディープニューラルネットワークを、1ノードサーバー上で正常にトレーニングした。これは、同種の実証において初の事例である。
• トレーニングは、大容量システムメモリを備えた2代目Intel XeonスケーラブルプロセッサにIntel最適化TensorFlowを適用することで達成された。
• スケールアップ構成により、モデル並列化やデータ並列化の必要がなく、トレーニング設定が簡素化された。
• 分散システムの調整を一切不要とすることで、複雑な医療画像分野のモデルを安定してトレーニング可能となった。
• 十分なメモリが確保可能な場合、最新鋭の高メモリ要件を満たすAIモデルを1台のサーバー上でトレーニング可能であることが検証された。
• 本研究は、インfraストラクチャに制限を受ける研究者にとって、より単純で効率的なトレーニングパスを提供する実用的な代替手段を示している。