TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs

要約

正確なハードウェア パフォーマンス モデルは、コードの最適化において重要な役割を果たします。
これらは、コンパイラーがヒューリスティックな決定を下すのを支援したり、自動チューナーが特定のプログラムの最適な構成を特定したりするのを支援します。
たとえば、機械学習コンパイラである XLA のオートチューナーは、Google で大量の運用トラフィックを処理する最先端のモデルで 10 ～ 20% の高速化を発見しました。
プログラムのパフォーマンスを予測するためのデータセットはいくつか存在しますが、それらは基本ブロックやカーネルなどの小さなサブプログラムを対象としています。
この論文では、Tensor Processing Unit (TPU) 上で実行される、計算グラフとして表現される完全なテンソル プログラムのパフォーマンス予測データセットである TpuGraphs を紹介します。
データセット内の各グラフは、トレーニング エポックや推論ステップなどの機械学習ワークロードの主な計算を表します。
各データ サンプルには、計算グラフ、コンパイル構成、およびその構成でコンパイルされたときのグラフの実行時間が含まれています。
データセット内のグラフは、ResNet、EfficientNet、Mask R-CNN、Transformer などの一般的なモデル アーキテクチャを特徴とするオープンソースの機械学習プログラムから収集されています。
TpuGraphs は、機械学習プログラムの既存のパフォーマンス予測データセットと比較して、最大のグラフ プロパティ予測データセット (同等のグラフ サイズ) の 25 倍のグラフと、平均 770 倍の大きなグラフを提供します。
大きなグラフに対するこのグラフ レベルの予測タスクでは、スケーラビリティ、トレーニング効率、モデル品質に至るまで、学習における新たな課題が生じます。

要約(オリジナル)

Precise hardware performance models play a crucial role in code optimizations. They can assist compilers in making heuristic decisions or aid autotuners in identifying the optimal configuration for a given program. For example, the autotuner for XLA, a machine learning compiler, discovered 10-20% speedup on state-of-the-art models serving substantial production traffic at Google. Although there exist a few datasets for program performance prediction, they target small sub-programs such as basic blocks or kernels. This paper introduces TpuGraphs, a performance prediction dataset on full tensor programs, represented as computational graphs, running on Tensor Processing Units (TPUs). Each graph in the dataset represents the main computation of a machine learning workload, e.g., a training epoch or an inference step. Each data sample contains a computational graph, a compilation configuration, and the execution time of the graph when compiled with the configuration. The graphs in the dataset are collected from open-source machine learning programs, featuring popular model architectures, e.g., ResNet, EfficientNet, Mask R-CNN, and Transformer. TpuGraphs provides 25x more graphs than the largest graph property prediction dataset (with comparable graph sizes), and 770x larger graphs on average compared to existing performance prediction datasets on machine learning programs. This graph-level prediction task on large graphs introduces new challenges in learning, ranging from scalability, training efficiency, to model quality.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google