[論文レビュー] Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning

Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan+3 more|arXiv (Cornell University)|Jun 20, 2019

Ferroelectric and Negative Capacitance Devices参考文献 23被引用数 30

ひとこと要約

Placeto は、配置を段階的改善としてモデル化し、グラフニューラルネットワークを用いて構造的情報をエンコードすることで、分散ディープラーニングにおける一般化可能なデバイス配置ポリシーを学習する強化学習フレームワークを導入する。これにより、再訓練なしに未観測のグラフに一般化可能であり、最大6.1倍の高速化を達成し、従来のRNNベースの手法を速度および配置品質の両面で上回る。

ABSTRACT

We present Placeto, a reinforcement learning (RL) approach to efficiently find device placements for distributed neural network training. Unlike prior approaches that only find a device placement for a specific computation graph, Placeto can learn generalizable device placement policies that can be applied to any graph. We propose two key ideas in our approach: (1) we represent the policy as performing iterative placement improvements, rather than outputting a placement in one shot; (2) we use graph embeddings to capture relevant information about the structure of the computation graph, without relying on node labels for indexing. These ideas allow Placeto to train efficiently and generalize to unseen graphs. Our experiments show that Placeto requires up to 6.1x fewer training steps to find placements that are on par with or better than the best placements found by prior approaches. Moreover, Placeto is able to learn a generalizable placement policy for any given family of graphs, which can then be used without any retraining to predict optimized placements for unseen graphs from the same family. This eliminates the large overhead incurred by prior RL approaches whose lack of generalizability necessitates re-training from scratch every time a new graph is to be placed.

研究の動機と目的

• 各新しい計算グラフに対して再訓練を必要とする従来のRLベースのデバイス配置手法が引き起こす高い学習オーバーヘッドに対処すること。
• 同じファミリーに属する未観測のグラフに対しても再訓練なしに学習済み配置ポリシーを一般化できること。
• 分散トレーニングにおける高品質なデバイス配置を特定するための配置評価回数を削減すること。
• 任意のノードインデックスやトポロジカル順序に依存しないポリシーを設計することで、耐性および転送性を向上させること。
• RNNベースのベースラインを上回る速度およびパフォーマンスを達成する、スケーラブルで効率的かつ一般化可能なデバイス配置アプローチを構築すること。

提案手法

• Placeto は、デバイス配置を、現在の部分的配置に基づいて1ノードずつデバイスを予測する段階的改善ステップの系列としてモデル化する。
• 計算グラフの構造をエンコードするためにグラフニューラルネットワーク（GNN）を用い、ノードの順序やラベルに依存せず、親子関係および並列関係をメッセージパッシングによって捉える。
• ポリシーネットワークは、ノードの局所的近傍（親、子、並列ノード）からの特徴を、それぞれのサブネットワークを介して集約し、その後デバイス決定を下す。
• GNNベースのアーキテクチャにより、ノードインデックスやトポロジカル順序が異なるグラフに対しても一般化が可能であり、順序処理に依存するRNNベースの手法とは対照的である。
• ポリシーの訓練を高速収束へ誘導するために、中間報酬を用いた強化学習目的関数を採用する。
• フレームワークは、1つのグラフファミリーに対して1つの一般化可能なポリシーを学習し、そのファミリーに属する新しいグラフに対して再訓練なしに適用可能である。

実験結果

リサーチクエスチョン

• RQ1 分散ディープラーニングにおける未観測の計算グラフに一般化可能な強化学習ポリシーを訓練できるか？
• RQ2 配置を段階的改善としてモデル化することで、学習の複雑さが軽減され、ポリシーの一般化性能が向上するか？
• RQ3 ノードラベルやトポロジカル順序に依存せずに、GNNが計算グラフの構造的情報を効果的にエンコードできるか？
• RQ4 Placeto の一般化能力は、配置品質およびトレーニング効率の観点でRNNベースの手法と比べてどの程度優れているか？
• RQ5 ポリシーのアーキテクチャの選択が、異なるグラフファミリー間での一般化性能にどの程度影響を及ぼすか？

主な発見

• Placeto は、RNNベースの手法と比較して、配置評価回数を最大6.1倍削減しながら、最良の先行研究と同等またはそれ以上の配置品質を達成する。
• NMTモデルでは、Placeto がRNNベースの手法よりも16.5%高速な配置を発見した。
• 再訓練なしに一般化ポリシーを適用する Placeto Zero-Shot は、数千回の評価を要する Placeto Optimized とほぼ同等の性能を達成した。
• RNNベースの手法はゼロショット設定で著しく性能が低く、RNN Zero-Shot はランダムベースラインとほぼ同等の性能に留まり、一般化能力の欠如が顕著に現れた。
• Placeto のポリシーは、異なるノード走査順序に対しても強固に一般化され、平均して配置実行時間の変動が5%にとどまる一方、RNNベースの手法は30%の変動を示した。
• アブレーションスタディの結果、親、子、並列ノードからの局所的近傍特徴の集約とプーリングが一般化に不可欠であることが判明。これらのコンponentsを除去すると、性能が13–20%低下した。