メンガー: 大規模分散強化学習アーキテクチャ

過去 10 年間で、強化学習 (RL) は機械学習で最も人気のある研究分野の 1 つになりました。RL を応用すると、チップの配置やリソース管理などの複雑な問題や、囲碁、Dota 2、かくれんぼなどの難しいゲームを効果的に解決できます。簡単に言えば、RL インフラストラクチャは、データ収集とトレーニングのサイクルです。アクターは環境に基づいてサンプル データを収集し、それを学習者に転送してモデルをトレーニングおよび更新します。現在のほとんどの RL 実装では、ターゲット タスクを学習するために、環境内で数万のサンプルを複数回反復する必要があります。たとえば、Dota 2 では、2 秒ごとに数万のフレームのサンプルを学習します。このように、RL アーキテクチャは、大量のサンプルを収集するためにアクターを追加するなどの強力なデータ処理機能を持つだけでなく、トレーニング中にこれらのサンプルをすばやく反復処理できる必要があります。

アクターと学習者間の相互作用のための RL アーキテクチャ。学習者はサンプリングされたデータを使用してモデルをトレーニングし、更新されたモデルをアクター (TF-Agents、IMPALA など) に渡します。

本日、Menger を紹介します。Menger は、複数の処理クラスター (Borg セルなど) にわたる数千のアクターに拡張可能な、ローカライズされた推論を備えた大規模に分散された RL アーキテクチャであり、チップ配置タスクのトレーニング時間を短縮します。次のセクションでは、Menger を Google TPU で構成してトレーニング速度を向上させる方法について説明します。さらに、チップ配置タスクを通じてフレームワークのパフォーマンスとスケーラビリティを検証します。 Menger はベースライン モデルと比較してトレーニング時間を 8.6 倍短縮することがわかります。

メンガーデザインアイデア

現在、Acme や SEED RL など、さまざまな分散 RL システムがあります。ただし、これらのシステムは、分散強化学習システムを 1 つの特定の観点からのみ最適化する傾向があります。たとえば、Acme は頻繁な学習者からモデルを取得し、各アクターがローカル推論を実行できるようにしますが、SEED RL は、バッチ呼び出しを実行するために TPU コアの一部を割り当てることで集中型推論を実行します。通信コストと推論コストの測定は、異なる最適化システム間の差であり、(1)集中型推論サーバーとの間で観測とアクションを送受信する通信コスト、またはLarnerからモデルを取得する通信コスト、(2)アクセラレータ(TPU/GPU)コストと比較したアクターの推論コストの大きさなどが含まれます。 Menger は、観測、アクション、モデル サイズなどのターゲット プログラムの要件を考慮して、Acme と同様のローカル推論を使用しますが、同時にアクターのスケーラビリティを最大限に高めます。優れたスケーラビリティとトレーニング速度を実現するための主な課題は次の 2 点です。

アクターはモデル取得のために学習者に大量の読み取り要求を行うため、学習者に負担がかかります。アクターの数が増えると、モデルのパフォーマンスが大幅に制限されます (収束時間が大幅に増加するなど)。

TPU のパフォーマンスは、トレーニング データを TPU コンピューティング コアに供給する際の入力パイプラインの効率によって制限されることがよくあります。 TPU コンピューティング コアの数 (TPU ポッドなど) が増えると、入力パイプラインのパフォーマンスがトレーニング時間にさらに大きな影響を与えます。

効率的なモデル検索

最初の課題に対処するために、TensorFlow コードでは、Learner と Actor の間に、Reverb によって最適化された透過的な分散キャッシュ コンポーネントを導入しました (Dota で使用されているアプローチと同様)。キャッシュ コンポーネントの主な役割は、アクターからの大量のリクエストと学習者の処理能力のバランスを取ることです。これらのキャッシュ コンポーネントを追加することで、過剰なリクエストによる学習者への負荷が大幅に軽減されるだけでなく、アクターも小さな通信コストで複数の Borg ユニットに分散されます。私たちの調査によると、512 個のアクターと 16 MB のサイズを持つモデルの場合、キャッシュ コンポーネントを導入すると平均読み取りレイテンシが約 4.0 倍短縮され、トレーニングの反復が高速化され、PPO などの戦略アルゴリズムではその効果がより顕著になります。

複数のアクターが異なるボーグセルに配置された分散 RL システム。 異なる Borg ユニット内の多数のアクターからの頻繁なモデル更新要求により、学習者のパフォーマンスと学習者とアクター間の通信ネットワークが制限され、全体的な収束時間が大幅に増加します。破線は異なるマシン間の gRPC 通信を表します。

透過的な分散キャッシュ サービスを備えた分散 RL システムを紹介します。複数のアクターが異なる Borg ユニットに配置され、学習者は更新されたモデルのみを分散キャッシュ コンポーネントに送信します。各キャッシュ コンポーネントは、隣接するアクターとキャッシュ モデルへの更新を要求する必要があります。キャッシュは、モデル更新リクエストを処理するための学習者の負担を軽減するだけでなく、アクターの平均読み取り待ち時間も短縮します。

高スループット入力パイプライン

入力データ パイプラインのスループットを向上させるために、Menger 氏は機械学習アプリケーション専用に設計された新しいオープン ソース データ ストレージ システムである Reverb を使用しました。 Reverb では、オンライン アルゴリズムまたはオフライン アルゴリズムのいずれかを使用してエクスペリエンスの再生を実行できます。ただし、現在、単一の Reverb では、数万のアクターを備えた分散 RL システムに拡張できず、アクターの書き込みスループットは非効率的です。

単一のリプレイ バッファを備えた分散 RL システム。アクターからの書き込み要求が大量に発生すると、リプレイ バッファが調整され、全体的なスループットが低下する可能性があります。さらに、Learner を複数のコンピューティング エンジン (TPU Pod) に拡張すると、これらのエンジンに単一のリプレイ バッファーを供給する効率が非常に低くなり、全体的な収束時間に重大な影響を及ぼしました。

分散 RL システムにおけるリプレイ バッファの効率性をより深く理解するために、さまざまなペイロード サイズ (16 MB ～ 512 MB) とさまざまなアクター (16 ～ 2048) の平均書き込みレイテンシを評価しました。再生バッファとアクターを同じ Borg セルに配置します。アクターの数が増えると、平均書き込みレイテンシが大幅に増加することがわかります。アクターの数を 16 から 2048 に拡張すると、16 MB および 512 MB のペイロード サイズの平均書き込みレイテンシがそれぞれ約 6.2 倍と約 18.9 倍増加します。このような書き込みレイテンシの増加はデータ収集時間に影響し、トレーニングの非効率につながります。

同じ Borg セルに異なるサイズのペイロード (16 MB - 512 MB) と異なる数のアクター (16 から 2048) を配置した場合の、単一の Reverb リプレイ バッファーの平均書き込み遅延。

これを軽減するために、Reverb のシャーディング機能を使用して、アクター、学習者、リプレイ バッファー間のスループットを向上させます。シャーディングにより、単一のリプレイ バッファ サーバーで動作するのではなく、多数のアクターの書き込み負荷を複数のリプレイ バッファ サーバーに分散できます。同時に、少数のアクターが同じサーバーを共有するため、各リプレイ バッファ サーバーの平均書き込みレイテンシを最小限に抑えることができます。これにより、Menger は複数の Borg セルにわたって数千のアクターに拡張できます。

シャードされたリプレイ バッファを備えた分散 RL システム。各リプレイ バッファは、同じボーグ セルにある特定のアクターのデータを保存するために使用されます。さらに、シャードされたリプレイ バッファは、アクセラレータ コアへのスループットが高い入力パイプラインを提供します。

検証例: チップの配置

大規模なネットリストのチップ配置タスクで Menger をテストしました。ベースラインと比較すると、Menger は 512 個の TPU コアを使用し、トレーニング時間の大幅な改善を実現します (最大約 8.6 倍、つまり最適な構成では、トレーニング時間を約 8.6 時間から 1 時間に短縮できます)。 Menger は TPU 向けに最適化されていますが、パフォーマンス向上の鍵となるのはフレームワークであり、GPU で実験した場合にも同様の改善が見られると期待しています。

チップ配置のベースラインと比較して、さまざまな数の TPU コアで Menger を使用した場合のトレーニング時間の改善。

Menger アーキテクチャとチップ配置タスクにおけるその優れたパフォーマンスは、チップ設計サイクルをさらに短縮するための方向性を示していると考えています。同時に、このアーキテクチャを適用して、他の困難な現実世界の問題を実装することもできます。

この記事はLeiphone.comから転載したものです。転載する場合は、Leiphone.com公式サイトにアクセスして許可を申請してください。