グラフ最適化のためのエンドツーエンドの転送可能な深層強化学習

[[425806]]

多様なアクセラレータ セットでトレーニングされた大規模で複雑なニューラル ネットワークによって駆動されるアプリケーションが増えています。このプロセスは、高レベルの計算グラフを低レベルのデバイス固有の実行可能ファイルにマッピングする ML コンパイラによって促進されます。これを実現するために、ML コンパイラは、グラフの書き換え、デバイス上の操作分散、操作の融合、テンソルのレイアウトとタイリング、スケジューリングなど、多くの最適化問題を解決する必要があります。たとえば、デバイス配置問題では、コンパイラは、トレーニング ステップ時間などの目的関数を最小限に抑えることができるように、計算グラフ内の操作とターゲットの物理デバイス間のマッピングを決定する必要があります。配置パフォーマンスは、デバイス間のネットワーク帯域幅、ピークデバイスメモリ、コロケーション制約などの複雑な要因の組み合わせによって決まります。そのため、通常は高速だが最適ではないソリューションに落ち着くヒューリスティックまたは検索ベースのアルゴリズムにとっては困難です。さらに、特に新しいモデル アーキテクチャが登場すると、ヒューリスティックの開発と維持が難しくなります。

学習ベースのアプローチを使用した最近の試みは有望な結果を示していますが、実際に展開するのを妨げるいくつかの制限があります。まず、これらの方法は、特に新しいモデル アーキテクチャによって生成された未知のグラフには簡単に一般化できません。また、サンプル効率が低いため、トレーニング中にリソース消費が高くなります。最後に、単一の最適化タスクしか解決できないため、コンパイル スタック内の密結合された最適化問題間の依存関係をキャプチャできません。

最近 NeurIPS 2020 で口頭発表された「ML コンパイラ向けの転送可能なグラフ オプティマイザー」では、上記のすべての制限を克服する計算グラフ最適化 (GO) 向けのエンドツーエンドの転送可能な深層強化学習アプローチを提案しています。 TensorFlow のデフォルトの最適化と比較して、3 つのグラフ最適化タスクで 33% ～ 60% の高速化を実証します。 Inception-v3、Transformer-XL、WaveNet など、最大 80,000 個のノードで構成される代表的なグラフの多様なセットにおいて、GO は、専門家による最適化よりも平均 21%、最先端の最適化よりも 18% の改善を達成し、収束速度は 15 倍高速化しました。

ML コンパイラにおけるグラフ最適化問題

では頻繁に発生する 3 つの結合最適化タスクがあり、それらを学習戦略を使用して解決できる意思決定問題として定式化します。各タスクの決定問題は、計算グラフ内の各ノードに対して決定を行うこととして再定義できます。

最初の最適化タスクはデバイスの配置です。その目的は、グラフのノードをそれが実行される物理デバイスに最適に割り当てる方法を決定し、エンドツーエンドの実行時間を最小限に抑えることです。

2 番目の最適化タスクは操作のスケジュール設定です。計算グラフ内の操作は、その入力テンソルがデバイス メモリ内に存在する場合に実行できる状態になります。一般的なスケジューリング戦略は、各デバイスの準備完了操作のキューを維持し、先入れ先出しの順序で操作をスケジュールすることです。ただし、このスケジューリング戦略では、操作によってブロックされる可能性のある他のデバイスのダウンストリーム操作が考慮されず、十分に活用されていないデバイスがスケジュールされることがよくあります。このようなデバイス間の依存関係を追跡できるスケジュールを見つけるために、私たちのアプローチでは、各操作の優先度に応じて準備キュー内の操作をスケジュールする優先度ベースのスケジューリング アルゴリズムを使用します。機器の配置と同様に、操作のスケジュールは、実行時間の観点から報酬を最大化するためにグラフ内の各ノードに優先順位を割り当てるポリシーを学習する問題として定式化できます。

3 番目の最適化タスクは操作の融合です。簡潔にするために、この問題の詳細な説明はここでは省略しますが、優先度ベースのスケジューリングと同様に、操作の融合でも優先度ベースのアルゴリズムを使用して、融合するノードを決定できることだけ注意してください。この場合も、ポリシー ネットワークの目標は、グラフ内の各ノードに優先順位を割り当てることです。

最後に、3 つの最適化問題のそれぞれで行われた決定が、他の問題に対する最適な決定に影響を与えることを認識することが重要です。たとえば、2 つの異なるデバイスに 2 つのノードを配置すると、実質的に収束が無効になり、通信遅延が発生してスケジュールに影響する可能性があります。

RL ポリシー ネットワーク アーキテクチャ

私たちの研究では、上記の最適化問題をそれぞれ個別に、または共同で解決するために適応できる深層 RL フレームワークである GO を提案します。提案されたアーキテクチャには、次の 3 つの重要な側面があります。

まず、グラフ ニューラル ネットワーク (具体的には GraphSAGE) を使用して、計算グラフにエンコードされたトポロジ情報を取得します。 GraphSAGE の誘導ネットワークは、ノード属性情報を活用して、これまで見たことのないグラフに一般化します。これにより、大きなトレーニング コストをかけずに、見たことのないデータに基づいて意思決定を行うことができます。

第二に、多くのモデルの計算グラフには、10,000 を超えるノードが含まれることがよくあります。このような大規模な最適化問題を効率的に解決するには、ネットワークがノード間の長距離依存関係を捕捉できる必要があります。 GO のアーキテクチャは、セグメント レベルの再帰を使用してこのような長距離ノード依存関係をキャプチャするスケーラブルなアテンション ネットワークで構成されています。

3 番目に、ML コンパイラーは、さまざまなアプリケーション ドメインのさまざまなグラフの最適化問題を解決する必要があります。異種グラフを使用して共有ポリシー ネットワークをトレーニングするという単純な戦略では、クラス固有のグラフの特異性を捉えることはできないでしょう。この問題を克服するために、GO は、パラメータの数を増やすことなく、ネットワークが特定のグラフ タイプに特化できるようにする機能変調メカニズムを使用します。

複数の関連する最適化タスクを共同で解決するために、GO は各タスクに追加の再帰的注意レイヤーを追加し、異なるタスク間でパラメータを共有できます。残余アクション接続を備えた再帰的注意層は、タスク間の依存関係を追跡できます。

結果

次に、実際のハードウェア測定に基づくデバイス配置タスクのシングルタスクの高速化評価結果、さまざまな GO バリアントを持つ未知のグラフへの一般化、および操作の融合、デバイス配置、およびスケジューリングを共同で最適化することによるマルチタスク パフォーマンスを示します。

加速する:

提案されたアーキテクチャのパフォーマンスを評価するために、実際のハードウェア評価に基づいてデバイス配置問題に GO を適用します。まず、各ワークロードでモデルを個別にトレーニングします。 GO-one と呼ばれるこのアプローチは、専門家による手動配置 (HP)、TensorFlow METIS 配置、および最新の強化学習ベースのデバイス配置である階層的デバイス配置 (HDP) よりも一貫して優れています。重要なのは、効率的なエンドツーエンドの単一配置により、GO-one の配置されたネットワークが HDP 上で 15 倍高速に収束することです。

私たちの実験結果によると、GO-one はエキスパート配置、TensorFlow METIS 配置、階層型デバイス配置 (HDP) よりも一貫して優れていることがわかりました。 GO は、8 層の Google Neural Machine Translation (GNMT) モデルなど、80,000 を超えるノードで構成される非常に大規模なグラフに拡張できるように設計されているため、HDP、REGAL、Placeto などの従来の方法よりも優れたパフォーマンスを発揮します。 GO は、GNMT などの大規模グラフに対して最適化されたグラフ実行時間を実現し、HP および HDP よりもそれぞれ 21.7% と 36.5% 高速です。全体として、GO-one は、14 の異なるグラフ セットに対して、HP および HDP と比較して、実行時間をそれぞれ平均 20.5% および 18.2% 短縮します。重要なのは、効率的なエンドツーエンドの単一配置により、GO-one の配置されたネットワークが HDP 上で 15 倍高速に収束することです。

一般化:

GO はオフラインの事前トレーニングを使用して、未知のグラフに一般化し、その後、未知のグラフを微調整します。事前トレーニング中に、トレーニング セット内のグラフの異種サブセットで GO をトレーニングします。次のグラフに切り替える前に、各グラフ バッチで GO を 1000 ステップトレーニングします。この事前トレーニング済みモデルは、その後、ホールドアウト グラフ上で 50 ステップ未満 (通常は 1 分未満) で微調整 (GO 一般化 + 微調整) されます。ホールドアウト グラフの GO 一般化 + 微調整は、すべてのデータセットでエキスパート配置および HDP を一貫して上回り、平均して GO-one に匹敵します。

また、ターゲット保存グラフを微調整せずに、事前トレーニング済みモデルで直接推論を実行し、これを GO-generalization-zeroshot と名付けます。この調整されていないモデルのパフォーマンスは、GO 一般化 + 微調整よりもわずかに劣りますが、エキスパート レイアウトと HDP よりもわずかに優れています。これは、グラフ埋め込みと学習されたポリシーの両方が効果的に転送され、モデルが未知のデータに一般化できることを示しています。

レイアウト、スケジュール、統合の共同最適化 (pl+sch+fu):

配置、スケジューリング、融合を同時に最適化すると、最適化されていない単一 GPU の場合と比較して 30% ～ 73% の高速化が実現され、TensorFlow のデフォルトの配置、スケジューリング、融合と比較して 33% ～ 60% の高速化が実現されます。各タスクを個別に最適化する場合と比較すると、マルチタスク GO (pl+sch+fu) は、すべてのタスクを一度に最適化するシングルタスク GO (p | sch | fu) よりも平均で 7.8% 優れています。さらに、すべてのワークロードにおいて、3 つのタスクすべてを共同で最適化すると、任意の 2 つのタスクを最適化し、3 番目のタスクにデフォルトの戦略を使用する場合と比較して、実行時間が短縮されます。

pl+sch: マルチタスク GO はデフォルトの融合を使用して配置とスケジュールを共同最適化します。 sch+fu: マルチタスク GO の共同最適スケジューリングと人間の配置との融合。 pl | sch | fu: GO はそれぞれ配置、スケジューリング、融合を最適化します。 pl+sch+fu: 配置、スケジュール、融合のマルチタスク GO 共同最適化。

結論は

ハードウェア アクセラレータの複雑さと多様性が増すにつれて、堅牢で適応性の高い ML フレームワークの開発は煩雑で時間がかかり、多くの場合、数百人のエンジニアによる数年にわたる取り組みが必要になります。本稿では、このようなフレームワークにおける多くの最適化問題が、慎重に設計された学習方法を使用して効率的かつ最適に解決できることを示します。