[論文レビュー] Auto-MAP: A DQN Framework for Exploring Distributed Execution Plans for DNN Workloads

Siyu Wang, Yi Rong+7 more|arXiv (Cornell University)|Jul 8, 2020

Parallel Computing and Optimization Techniques参考文献 31被引用数 5

ひとこと要約

Auto-MAP は、HLO IR 上で動作する DQN に基づく強化学習フレームワークであり、データ並列性、モデル並列性、パイプライン並列性を統合的に探索することで、DNNワークロードの最適な分散実行計画を自動で発見する。NLP および CNN モデルにおいて、2時間以内に最適なスループットを達成し、手動およびヒューリスティック手法を上回る性能を発揮する。

ABSTRACT

The last decade has witnessed growth in the computational requirements for training deep neural networks. Current approaches (e.g., data/model parallelism, pipeline parallelism) parallelize training tasks onto multiple devices. However, these approaches always rely on specific deep learning frameworks and requires elaborate manual design, which make it difficult to maintain and share between different type of models. In this paper, we propose Auto-MAP, a framework for exploring distributed execution plans for DNN workloads, which can automatically discovering fast parallelization strategies through reinforcement learning on IR level of deep learning models. Efficient exploration remains a major challenge for reinforcement learning. We leverage DQN with task-specific pruning strategies to help efficiently explore the search space including optimized strategies. Our evaluation shows that Auto-MAP can find the optimal solution in two hours, while achieving better throughput on several NLP and convolution models.

研究の動機と目的

• 多様な DNN ワークロードにおける効率的な分散学習戦略を手動で設計する課題に対処すること。
• フレームワーク固有またはモデル固有の解決策を回避するために、データ並列性、モデル並列性、パイプライン並列性の探索を統合的に1つのフレームワークで行うこと。
• 低レベルの分散システム設定を排除することでユーザーの負担を軽減し、フレームワークに依存しない、ワークロードに依存しない並列化計画を可能にすること。
• タスク固有のプルーニングと DQN に基づく強化学習を用いて、大規模で複雑な並列化空間における探索効率を向上させること。
• BERT や T5、ResNet といった多様なモデルに対して、自動的で高スループットかつ一般化可能な並列化戦略を提供すること。

提案手法

• Auto-MAP は、XLA によって生成される低レベルでハードウェアに依存しない中間表現である HLO IR を基盤として利用し、オペレータレベルやレイヤーレベルの表現よりも細かく、より表現力のあるモデル解析を可能にする。
• データ並列性、モデル並列性、パイプライン並列性を含む膨大な分散実行計画の探索空間を、タスク固有のプルーニング戦略を用いた Deep Q-Network (DQN) を用いて効率的に探索する。
• 訓練スループットに基づく報酬関数を用いて、DQN エージェントが最適または近似的に最適な並列化戦略へと導かれるようにする。
• XLA コンパイラー スタックを活用することで、TensorFlow や JAX、Trax などの複数のフロントエンドとの互換性を確保し、ユーザーのモデルへのコードインジェクションを回避する。
• DQN エージェントは、実際のモデル分布上で微調整される前の段階として、ランダムな分布を持つ合成環境で事前学習が行われる。
• NLP モデルでは、合成データ上で事前学習した後、エージェントは良好に一般化するが、CNN では、正確なパイプライン分割を達成するために、モデル固有の分布上で微調整が必要となる。

実験結果

リサーチクエスチョン

• RQ1 手動設計を伴わず、多様な DNN モデルにおいて最適な分散学習戦略を自動で発見できるか？
• RQ2 1つの拡張可能な探索空間内で、データ並列性、モデル並列性、パイプライン並列性の戦略を同時に統合的に探索できるか？
• RQ3 HLO IR 上で動作させることで、レイヤーまたはオペレータレベルの抽象化よりも、並列化計画における性能と一般化能力が向上するか？
• RQ4 DQN におけるタスク固有のプルーニングが、探索空間を著しく縮小し、最適解への収束を加速できるか？
• RQ5 事前学習済みの DQN エージェントは、異なるモデルアーキテクチャ間でどの程度一般化可能か？また、微調整が必要となる状況はいつか？

主な発見

• Auto-MAP は、NLP および畳み込みモデルの両方において、最適な分散実行計画を2時間以内に発見し、手動チューニングにかかる時間を顕著に短縮する。
• BERT や T5 モデルでは、合成データ上で事前学習した後、微調整なしでエージェントが最適なパイプライン分割戦略を的確に推論できる。
• T5-11B などの大規模モデルでは、10,000 エピソードにわたるモデル固有の分布上でエージェントを微調整することで、正確な最適パイプライン境界の推論が可能になる。
• 複数の NLP および CNN モデルにおいて、ベースラインのヒューリスティック手法や手動手法と比較して、より高い訓練スループットを達成する。
• HLO IR の使用により、TensorFlow、JAX、Trax の間でフレームワークに依存しない低侵入型の並列化計画が可能になる。
• タスク固有のプルーニングを DQN に統合することで、ハイブリッド並列化戦略の複雑な空間における探索効率と収束速度が向上する。