TpuGraphs: 大規模テンソル計算グラフの性能予測データセット

精密なハードウェア性能モデルはコード最適化において重要な役割を果たします。これらのモデルは、コンパイラがヒューリスティックな決定を行うのを支援したり、オートチューナーが特定のプログラムにとって最適な設定を特定するのに役立ったりします。例えば、GoogleのXLA（マシンラーニング用コンパイラ）のオートチューナーは、最先端のモデルで10-20%の速度向上を達成し、大量のプロダクショントラフィックに対応しました。プログラム性能予測用にいくつかのデータセットが存在していますが、それらは基本ブロックやカーネルなどの小さなサブプログラムを対象としています。本論文では、Tensor Processing Units (TPUs) 上で実行される完全なテンソルプログラムを計算グラフとして表現した性能予測データセットであるTpuGraphsを紹介します。データセット内の各グラフは、マシンラーニングワークロードの主な計算（例：トレーニングエポックや推論ステップ）を表しています。各データサンプルには計算グラフ、コンパイル設定、およびその設定でコンパイルされた場合のグラフの実行時間が含まれています。このデータセットに収録されているグラフはオープンソースのマシンラーニングプログラムから集められており、ResNet、EfficientNet、Mask R-CNN、Transformerなど人気のあるモデルアーキテクチャが特徴となっています。TpuGraphsは最大規模のグラフ特性予測データセット（同程度のグラフサイズを持つもの）よりも25倍多くのグラフを提供し、既存のマシンラーニングプログラム向け性能予測データセットと比較して平均770倍大きなグラフを持っています。大規模なグラフにおけるこのグラフレベルでの予測タスクは、スケーラビリティ、学習効率性、モデル品質など新たな課題をもたらします。