[论文解读] Runtime Prediction of AI Model Operations Using a GRU-Based Neural Network

Mallapragada, Suma|arXiv (Cornell University)|Aug 25, 2023

Parallel Computing and Optimization Techniques被引用 5

一句话总结

该论文介绍了 TpuGraphs，一个用于预测 TPU 上完整机器学习计算图运行时性能的大规模数据集，采用基于 GRU 的神经网络来建模不同布局和分块编译配置下的性能表现。该数据集支持大规模模型的图级性能预测，相较于真实硬件评估，推理速度最快可提升 2400 倍，基线模型在包括 ResNet、BERT 和 Transformers 在内的复杂且多样的工作负载中表现出中等精度。

ABSTRACT

Modern AI models can be represented as computational graphs, where eachnode corresponds to a tensor operation (e.g., matrix multiplication, convolution),and edges represent tensor data flows. Optimizing the executionof these graphs on hardware accelerators such as Tensor Processing Units(TPUs) requires careful selection of compiler configurations that control layoutand tiling strategies.The compilation configuration involves two key types of optimizations:• Layout Configuration: Controls how tensors are arranged in physicalmemory by specifying the dimension order for inputs and outputs ofeach operation node.• Tile Configuration: Controls the tile size of each fused subgraph,impacting data locality and parallelism. Accurately predicting the runtime of AI model graphs under various configurationscan automate and improve the selection of optimal compiler settings,reducing execution time and resource consumption. The Kaggle competition dataset “Google - Fast or Slow? Predict AIModel Runtime” provides runtime data for XLA High Level Optimizer (HLO)graphs running on TPU v3 hardware. This dataset, called TPUGraphs,comprises multiple collections with diverse layouts and tiling configurations,posing a challenging performance prediction task. This work proposes a GRU-based runtime prediction pipeline leveragingopcode runtime features, graph structure dependencies, and configurablenode embeddings. The method consolidates node features, integrates configurationconvolutions, and trains a neural network to predict runtime withmean squared error loss. The approach captures both static graph propertiesand dynamic configuration effects, enabling enhanced runtime estimation toguide compiler heuristics.

研究动机与目标

• 为运行在 TPU 等加速器上的完整机器学习程序，解决缺乏大规模、图级性能预测数据集的问题。
• 实现在不同编译配置（布局和分块）下对计算图执行时间的高效且准确的预测。
• 支持机器学习模型的开发，以替代昂贵的真实硬件测量，从而加速编译器自动调优。
• 提供一个基准数据集，包含 3100 万个布局配置和 1300 万个分块配置，覆盖包括 ResNet、BERT 和 Transformers 在内的多样化、生产规模模型。
• 研究在具有高冗余性和内存限制的大规模、多样化图上训练可扩展、泛化能力强的模型所面临的挑战。

提出的方法

• 该数据集 TpuGraphs 包含来自开源机器学习程序的计算图，每个图均配有编译配置（布局和分块）及其在 TPU v3 上的对应执行时间。
• 提出一种基于 GRU 的神经网络，以建模图结构、配置与执行时间之间的关系，捕捉图操作中的序列依赖性。
• 模型将图节点和边作为序列处理，利用 GRU 层对结构和操作特征进行编码，以实现性能预测。
• 特征工程包括将张量布局顺序和分块大小编码为输入嵌入，使模型能够学习与配置相关的性能模式。
• 训练流水线通过批量处理和内存优化来处理大型图，特征提取在相同图的多个配置间进行分摊。
• 基线模型采用图神经网络（GNN）和监督学习方法，损失函数包括均方误差（MSE）和 ListMLE，用于基于排序的优化。

Figure 1: Scale of TpuGraphs compared to other graph property prediction datasets.

实验结果

研究问题

• RQ1 仅使用图结构和配置特征，学习到的性能模型能否准确预测 TPU 上大规模计算图的执行时间？
• RQ2 在包括 ResNet、BERT 和 Transformers 在内的多样化真实世界机器学习工作负载中，模型性能如何变化？
• RQ3 在处理节点数达 7700 个且配置高度冗余的图时，训练效率、模型可扩展性与预测精度之间的权衡如何？
• RQ4 基于模型的预测相比真实硬件评估快多少？该加速是否可被用于自动调优流水线？
• RQ5 在生产级机器学习环境中，单一模型在未见图结构和配置上的泛化能力如何？

主要发现

• TpuGraphs 数据集包含 3100 万个布局配置和 1300 万个分块配置，图平均包含 7700 个节点——是最大可比图属性数据集的 25 倍，且比现有机器学习性能数据集大 770 倍。
• 模型预测速度最高可达真实硬件评估的 2400 倍，单配置推理时间在 0.004 秒至 0.1 秒之间，而编译和执行时间则为 45 至 475 秒。
• 使用 GRU 和 GNN 架构的基线模型在布局集合上的中位数 top-1 错误率分别为 10.5%（ListMLE）和 19.0%（MSE），在 BERT 和 ResNet 等特定模型上表现显著提升。
• 表现最佳的模型（ListMLE 与 EarlySAGE 结合）在布局集合上的 top-1 错误率为 10.5%，在分块集合上也为 10.5%，展现出在多样化工作负载上的强大泛化能力。
• 模型的特征提取时间在相同图的多个配置间实现分摊，使其在自动调优流水线中（每图需评估多个配置）具有极高效率。
• 本研究证实，学习到的成本模型可显著加速自动调优，将候选配置评估时间从数小时缩短至毫秒级。

Figure 2: Important optimizations in ML compilers include graph-level and kernel-level optimizations. A graph-level optimization requires the context of the entire graph to make optimal decisions and transforms the entire graph accordingly. A kernel-level optimization transforms each kernel (a fused