TpuGraphs:一个基于大型张量计算图的性能预测数据集
机器学习
2023-12-07 v3 硬件体系结构
社会与信息网络
摘要
精确的硬件性能模型在代码优化中起着关键作用。它们可以辅助编译器做出启发式决策,或帮助自动调优器为给定程序确定最优配置。例如,机器学习编译器 XLA 的自动调优器在 Google 服务于大量生产流量的前沿模型上发现了 10-20% 的加速。尽管存在一些用于程序性能预测的数据集,但它们针对的是基本块或内核等小型子程序。本文介绍 TpuGraphs,一个在张量处理单元(TPU)上运行、以计算图表示完整张量程序的性能预测数据集。数据集中的每个图表示一个机器学习工作负载的主要计算,例如一个训练轮次或推理步骤。每个数据样本包含一个计算图、一个编译配置,以及该图在此配置下编译后的执行时间。数据集中的图收集自开源机器学习程序,具有流行的模型架构,如 ResNet、EfficientNet、Mask R-CNN 和 Transformer。TpuGraphs 提供的图数量比最大的图属性预测数据集(图大小相当)多 25 倍,且平均图规模比现有的机器学习程序性能预测数据集大 770 倍。这一针对大图的图级预测任务在学习上带来了从可扩展性、训练效率到模型质量的新挑战。
引用
@article{arxiv.2308.13490,
title = {TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs},
author = {Phitchaya Mangpo Phothilimthana and Sami Abu-El-Haija and Kaidi Cao and Bahare Fatemi and Mike Burrows and Charith Mendis and Bryan Perozzi},
journal= {arXiv preprint arXiv:2308.13490},
year = {2023}
}