TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs

要約

正確なハードウェア性能モデルは、コードの最適化において重要な役割を果たす。これらのモデルは、コンパイラが発見的な判断を下すのを支援したり、オートチューナーが与えられたプログラムに対して最適なコンフィギュレーションを特定するのを支援したりすることができる。例えば、機械学習コンパイラであるXLAのオートチューナーは、Googleの大規模な実運用トラフィックに対応する最先端のモデルで、10～20%のスピードアップを発見した。プログラム性能予測のためのデータセットはいくつか存在するが、それらは基本ブロックやカーネルといった小さなサブプログラムを対象としている。本論文では、TpuGraphsを紹介する。TpuGraphsは、Tensor Processing Units (TPU)上で動作する、計算グラフとして表現される完全なテンソルプログラムの性能予測データセットである。データセットの各グラフは、機械学習ワークロードの主要な計算、例えば学習エポックや推論ステップを表す。各データサンプルには、計算グラフ、コンパイル構成、その構成でコンパイルした場合のグラフの実行時間が含まれる。データセットのグラフは、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, DeepL