NvidiaはAIを使用してGPU演算回路を設計します。これは最先端のEDAよりも25％小さく、より高速で効率的です。

膨大な数の演算回路を搭載した NVIDIA GPU により、AI、高性能コンピューティング、コンピューターグラフィックスのかつてない高速化が実現します。したがって、これらの演算回路の設計を改善することは、GPU のパフォーマンスと効率を向上させるために非常に重要です。 AI がこれらの回路の設計を学習したらどうなるでしょうか?最近の NVIDIA の論文「PrefixRL: 深層強化学習を使用した並列プレフィックス回路の最適化」で、研究者らは、AI がこれらの回路をゼロから設計できるだけでなく、AI 設計の回路は最先端の電子設計自動化 (EDA) ツールで設計された回路よりも小型で高速であることを実証しました。

論文アドレス: https://arxiv.org/pdf/2205.07000.pdf

最新の Nvidia Hopper GPU アーキテクチャには、AI 設計の回路インスタンスが約 13,000 個あります。下の図 1 の左側にある PrefixRL AI によって設計された 64 ビット加算器回路は、図 1 の右側にある最先端の EDA ツールによって設計された回路よりも 25% 小さくなっています。

回路設計の概要

コンピュータチップ内の演算回路は、論理ゲート (NAND、NOR、XOR など) とワイヤのネットワークで構成されています。理想的な回路には次の特性が必要です。

小型: 面積が小さいため、チップ上に多くの回路をパッケージ化できます。
高速: レイテンシが低減し、チップのパフォーマンスが向上します。
消費電力が低い。

この Nvidia の研究では、研究者は回路面積とレイテンシに焦点を当てました。彼らは、消費電力が対象回路の面積と密接に関係していることを発見しました。回路面積と遅延は競合する特性であることが多いため、これらの特性を効果的にトレードオフする設計のパレート最適点を見つけることが望ましいです。つまり、研究者は遅延ごとに回路面積が最小化されることを期待しているのです。

そのため、PrefixRL では、研究者は演算回路の一般的なクラスである並列プレフィックス回路に焦点を当てています。アクセラレータ、インクリメンタ、エンコーダなど、GPU のさまざまな重要な回路はプレフィックス回路であり、より高いレベルでプレフィックスグラフとして定義できます。

そこで疑問になるのが、AI エージェントは適切なプレフィックスグラフを設計できるかどうかです。すべてのプレフィックスグラフの状態空間は O(2^n^n) と大きく、ブルートフォースアプローチを使用して探索することはできません。下の図 2 は、4b 回路例を使用した PrefixRL の 1 回の反復を示しています。

研究者たちは回路ジェネレーターを使用して、プレフィックスグラフをワイヤと論理ゲートを含む回路に変換しました。次に、生成された回路は、ゲートサイズ、複製、バッファ挿入などの物理合成最適化を使用する物理合成ツールによって最適化されます。

これらの物理合成の最適化により、最終的な回路プロパティ (遅延、面積、電力) は、元のプレフィックスグラフプロパティ (レベルやノード数など) から直接変換されません。 AI エージェントがプレフィックスグラフの設計を学習するだけでなく、プレフィックスグラフから生成される最終回路の特性を最適化するのはこのためです。

我々は、演算回路の設計を強化学習 (RL) タスクとして定式化し、エージェントが演算回路の面積とレイテンシ特性を最適化するようにトレーニングします。プレフィックス回路の場合、RL エージェントがプレフィックスグラフ内のノードを追加または削除し、次の手順を実行できる環境を設計しました。

プレフィックスマップは、常に正しいプレフィックスの合計計算を維持するために正規化されます。
正規化されたプレフィックスグラフから回路を生成します。
物理合成ツールを使用して、回路の物理合成最適化を実行します。
回路の面積と遅延特性を測定します。

次のアニメーションでは、RL エージェントがノードを追加または削除することで、プレフィックスグラフを段階的に構築します。各ステップで、エージェントは回路面積とレイテンシの改善という報酬を得ます。

オリジナル画像はインタラクティブバージョンです。

完全畳み込みQ学習エージェント

研究者たちは、Q学習アルゴリズムを使用してインテリジェントエージェントの回路設計をトレーニングしました。下の図 3 に示すように、プレフィックスグラフをグリッド表現に分解し、グリッド内の各要素をプレフィックスノードに一意にマッピングします。このグリッドは、Q ネットワークの入力と出力を表します。入力グリッド内の各要素は、ノードが存在するかどうかを示します。出力グリッド内の各要素は、追加または削除されたノードの Q 値を表します。

Q 学習エージェントの入力と出力は両方ともグリッド表現であるため、研究者は完全な畳み込みニューラルネットワークアーキテクチャを使用しました。エリアと遅延の報酬はトレーニング中に別々に観測できるため、エージェントはエリアと遅延属性の Q 値を別々に予測します。

図 3: 4b プレフィックスグラフ表現 (左) と完全畳み込み Q 学習エージェントアーキテクチャ (右)。

分散トレーニングのためのRaptor

PrefixRL は計算負荷が高く、物理シミュレーションには GPU あたり 256 個の CPU が必要で、64 ビットのタスクをトレーニングするには 32,000 時間以上の GPU 時間が必要です。今回、NVIDIA は、NVIDIA ハードウェアの利点を最大限に活用して、この産業レベルの強化学習を可能にする、社内分散型強化学習プラットフォーム Raptor を開発しました (下の図 4 を参照)。

Raptor は、ジョブスケジューリング、カスタムネットワーク、GPU 対応データ構造などのトレーニングモデルのスケーラビリティとトレーニング速度を向上させます。 PrefixRL のコンテキストでは、Raptor は CPU、GPU、スポットインスタンス間でのハイブリッド割り当てを可能にします。この強化学習アプリケーションのネットワークは多様であり、次のような利点があります。

Raptor はピアツーピア転送のために NCCL を切り替え、モデルパラメータを学習 GPU から推論 GPU に直接転送します。
Redis は、報酬や統計などの非同期で小さなメッセージに使用されます。
JIT コンパイルされた RPC の場合、エクスペリエンスデータのアップロードなど、大量かつ低レイテンシのリクエストを処理するために使用されます。

最後に、Raptor は、複数のワーカーからの経験を受け取り、データを並列にバッチ処理して GPU に事前ロードするためのマルチスレッドサービスを備えたリプレイバッファーなどの GPU 対応データ構造を提供します。

下の図 4 は、PrefixRL フレームワークが同時トレーニングとデータ収集をサポートし、NCCL を活用して最新のパラメータを参加者 (下の図のアクター) に効率的に送信していることを示しています。

図 4: 研究者は Raptor を使用して並列トレーニングと報酬計算を切り離し、回路合成の遅延を克服します。

報酬計算

面積とレイテンシの目標を組み合わせるために、トレードオフ重みw（範囲[0,1]）を採用します。さまざまなエージェントを異なる重みでトレーニングして、面積とレイテンシのトレードオフのバランスをとるパレート境界を取得します。

RL 環境での物理合成の最適化により、面積とレイテンシをトレードオフするさまざまなソリューションを生成できます。研究者は、特定のエージェントをトレーニングする場合と同じトレードオフの重みを使用して、物理合成ツールを駆動します。

報酬計算のループ内で物理合成の最適化を実行すると、次の利点があります。

RL エージェントは、ターゲットテクノロジノードとライブラリの最終的な回路プロパティを直接最適化することを学習します。
RL エージェントは、物理合成プロセス中にターゲットアルゴリズム回路の周囲のロジックを組み込み、ターゲットアルゴリズム回路とその周囲のロジックのパフォーマンスを共同で最適化します。

ただし、物理合成の実行は低速なプロセス (64 ビット加算器の場合は約 35 秒) であり、RL のトレーニングと探索が大幅に遅くなる可能性があります。

エージェントは回路合成や以前の報酬を必要とせず、アクションを実行するために現在のプレフィックスグラフ状態のみを必要とするため、報酬の計算を状態の更新から切り離します。 Raptor のおかげで、長時間の報酬計算を CPU ワーカーのプールにオフロードして物理合成を並行して実行できるようになり、アクターエージェントは待機せずに環境内で実行できるようになります。

CPU ワーカーが報酬を返すと、変換をリプレイバッファーに埋め込むことができます。結合された報酬は、状態が再び発生したときに冗長な計算を回避するためにキャッシュされます。

結果と展望

下の図 5 は、PrefixRL と最先端の EDA ツールの Pareto 優位加算回路を使用して設計された 64 ビット加算回路の面積と遅延を示しています。

最良の PrefixRL 加算器は、同じレイテンシで EDA ツール加算器よりも 25% 少ない面積を実現します。物理合成の最適化後にパレート最適加算回路にマッピングされるこれらのプレフィックスグラフは、不規則な構造を持ちます。

図5: PrefixRL を使用して設計された演算回路は、最先端の EDA ツールを使用して設計された回路よりも小型で高速です。

(左) 回路アーキテクチャ; (右) 対応する 64b 加算器回路特性図

私たちの知る限り、これは深層強化学習エージェントを使用して演算回路を設計する最初の方法です。 Nvidia は青写真を描いています。このアプローチにより、AI を現実世界の回路設計問題に適用し、アクション空間、状態表現、RL エージェントモデルを構築し、複数の競合する目的を最適化し、遅い報酬計算プロセスを克服できると期待しています。

<<: 近年、軍事用人工知能スタートアップが人気を集めている理由

>>: 電力業界における人工知能開発の現状