[論文レビュー] Runtime Prediction of AI Model Operations Using a GRU-Based Neural Network

Mallapragada, Suma|arXiv (Cornell University)|Aug 25, 2023

Parallel Computing and Optimization Techniques被引用数 5

ひとこと要約

本稿では、TPU上で実行される完全な機械学習計算グラフの実行時間を予測するための大規模データセットTpuGraphsを紹介する。GRUベースのニューラルネットワークを用いて、レイアウトおよびタイルコンパイル設定のパラメータを考慮したパフォーマンスをモデル化する。このデータセットにより、大規模モデルにおけるグラフ単位のパフォーマンス予測が可能となり、実ハードウェア評価に比べ最大2400倍速い推論が達成される。ベースラインモデルは、ResNet、BERT、Transformersを含む複雑で多様なワークロードにおいて中程度の精度を示す。

ABSTRACT

Modern AI models can be represented as computational graphs, where eachnode corresponds to a tensor operation (e.g., matrix multiplication, convolution),and edges represent tensor data flows. Optimizing the executionof these graphs on hardware accelerators such as Tensor Processing Units(TPUs) requires careful selection of compiler configurations that control layoutand tiling strategies.The compilation configuration involves two key types of optimizations:• Layout Configuration: Controls how tensors are arranged in physicalmemory by specifying the dimension order for inputs and outputs ofeach operation node.• Tile Configuration: Controls the tile size of each fused subgraph,impacting data locality and parallelism. Accurately predicting the runtime of AI model graphs under various configurationscan automate and improve the selection of optimal compiler settings,reducing execution time and resource consumption. The Kaggle competition dataset “Google - Fast or Slow? Predict AIModel Runtime” provides runtime data for XLA High Level Optimizer (HLO)graphs running on TPU v3 hardware. This dataset, called TPUGraphs,comprises multiple collections with diverse layouts and tiling configurations,posing a challenging performance prediction task. This work proposes a GRU-based runtime prediction pipeline leveragingopcode runtime features, graph structure dependencies, and configurablenode embeddings. The method consolidates node features, integrates configurationconvolutions, and trains a neural network to predict runtime withmean squared error loss. The approach captures both static graph propertiesand dynamic configuration effects, enabling enhanced runtime estimation toguide compiler heuristics.

研究の動機と目的

• TPUなどのアクセラレータ上で実行される完全な機械学習プログラム向けに、大規模かつグラフ単位のパフォーマンス予測データセットが不足しているという問題に対処すること。
• レイアウトおよびタイル設定のさまざまなコンパイル設定下での計算グラフの実行時間を効率的かつ正確に予測できることを実現すること。
• コストの高い実ハードウェア測定を置き換えることで、コンパイラのオートチューニングを加速できる機械学習モデルの開発を支援すること。
• ResNet、BERT、Transformersなどの多様で生産規模のモデルをカバーする、3100万件のレイアウトおよび1300万件のタイル設定を含むベンチマークデータセットを提供すること。
• 高い再現性とメモリ制約がある大規模で多様なグラフにおいて、スケーラブルで汎用性のあるモデルを学習する際の課題を調査すること。

提案手法

• TpuGraphsデータセットは、オープンソースのMLプログラムから得られた計算グラフから構成され、それぞれがTPU v3上で実行された際の実行時間と、コンパイル設定（レイアウトおよびタイル）とペアになっている。
• グラフ構造、設定、実行時間の関係をモデル化するため、GRUベースのニューラルネットワークが提案されている。このモデルは、グラフ演算における順序的依存関係を捉えることができる。
• モデルはグラフのノードとエッジをシーケンスとして処理し、GRU層を用いて構造的および操作的特徴を符号化し、パフォーマンス予測を実行する。
• 特徴工学には、テンソルレイアウト順序とタイルサイズを入力埋め込みとしてエンコードする手法が含まれており、これによりモデルは設定固有のパフォーマンスパターンを学習できる。
• 大規模なグラフを処理するためのトレーニングパイプラインは、バッチ処理とメモリ使用量の最適化を実施しており、同じグラフの複数の設定に対して特徴抽出を共有することで、計算コストを低減している。
• ベースラインモデルは、GNN（グラフニューラルネットワーク）と教師あり学習を組み合わせ、MSEやListMLEといった損失関数を用いて順序付け最適化を実現している。

Figure 1: Scale of TpuGraphs compared to other graph property prediction datasets.

実験結果

リサーチクエスチョン

• RQ1 グラフ構造と設定特徴のみを用いて、TPU上での大規模計算グラフの実行時間を正確に予測できる学習済みパフォーマンスモデルは存在するか？
• RQ2 ResNet、BERT、Transformersなどの多様で実世界のMLワークロードにおいて、モデルのパフォーマンスはどのように変動するか？
• RQ3 最大7,700ノードのグラフと高い設定再現性を伴う場合に、トレーニング効率、モデルスケーラビリティ、予測精度の間のトレードオフは何か？
• RQ4 モデルベースの予測は、実ハードウェア評価に比べてどの程度速く、この高速化はオートチューニングパイプラインで活用可能か？
• RQ5 生産環境のMLシステムにおいて、未学習のグラフ構造や設定に対して、1つのモデルがどれほど汎用的に一般化できるか？

主な発見

• TpuGraphsデータセットには3100万件のレイアウト設定および1300万件のタイル設定が含まれており、グラフの平均ノード数は7,700である。これは、最大の類似グラフプロパティデータセットの25倍、既存のMLパフォーマンスデータセットの770倍に相当する。
• モデル予測は、実ハードウェア評価に比べ最大2,400倍速く、1回の設定あたりの推論時間は0.004秒～0.1秒である一方、コンパイルと実行には45～475秒がかかる。
• GRUおよびGNNアーキテクチャを用いたベースラインモデルは、レイアウトコレクションにおいて中央値トップ1誤差が10.5%（ListMLE）および19.0%（MSE）を達成しており、BERT や ResNet などの特定のモデルでは顕著な改善が見られた。
• 最良のモデル（ListMLE + EarlySAGE）は、レイアウトコレクションで10.5%のトップ1誤差、タイルコレクションでも10.5%のトップ1誤差を達成し、多様なワークロードにわたる強力な一般化性能を示している。
• モデルの特徴抽出時間は、同じグラフの複数の設定に対して共有されるため、複数の設定を1回のグラフで評価するオートチューニングパイプラインにおいて極めて効率的である。
• 本研究では、学習済みコストモデルがオートチューニングを顕著に加速できることを確認した。候補設定の評価に要する時間が数時間から数ミリ秒に短縮された。

Figure 2: Important optimizations in ML compilers include graph-level and kernel-level optimizations. A graph-level optimization requires the context of the entire graph to make optimal decisions and transforms the entire graph accordingly. A kernel-level optimization transforms each kernel (a fused