Google Releases TpuGraphs Dataset for ML Performance Prediction

Google has released TpuGraphs, a new dataset designed to enable researchers to build better cost models for machine learning programs running on Google’s custom Tensor Processing Units (TPUs). The dataset was presented in connection with NeurIPS 2023.

Scale and Scope

The TpuGraphs dataset represents a significant leap in scale compared to existing benchmarks. It provides 25 times more graphs than the largest prior graph property prediction dataset with comparable graph sizes. Beyond quantity, the average graph size in TpuGraphs is 770 times larger than in existing performance prediction datasets for ML programs, enabling researchers to work with substantially more complex program structures.

Graph Segment Training Innovation

Google has also introduced Graph Segment Training (GST), a method developed by Google DeepMind that speeds up end-to-end GNN training time by 3 times compared to standard full graph training. This represents a practical improvement for researchers working with large-scale graph neural networks.

Community Engagement Through Kaggle Competition

Google’s commitment to advancing the field extended to a Kaggle competition titled ‘Fast or Slow? Predict AI Model Runtime’, conducted over the TpuGraphs dataset. The competition drew significant global participation, with 792 participants across 616 teams from 66 countries submitting a total of 10,507 entries.

The competition also served as an entry point for new researchers, with 153 participants competing in their first Kaggle competition, including 47 in the top 100.

Source: Google Research Blog