TLP: A Deep Learning-based Cost Model for Tensor Program Tuning

发表时间: 2022-11 · arXiv:2211.03578

原文: https://arxiv.org/abs/2211.03578

TLP：一种用于张量程序调优的基于深度学习的成本模型
作者/机构：
Yi Zhai, Yu Zhang, Shuo Liu, Xiaomeng Chu, Jie Peng, Jianmin Ji, Yanyong Zhang
(University of Science and Technology of China)

速读

一句话结论 本文提出了基于深度学习的张量程序成本模型 TLP 及其多任务学习版本 MTL-TLP，通过将调度原语序列视为自然语言进行特征提取，并利用多任务学习机制，解决了传统模型依赖硬件专家知识和跨硬件平台性能下降的问题，在搜索速度和最终性能上均大幅超越现有基线。

要解决什么问题 现有的张量程序调优成本模型主要卡在特征提取机制和跨硬件泛化能力两个瓶颈上。首先，传统的特征提取严重依赖硬件架构的专家级领域知识，例如需要针对最内层赋值语句手动提取上百个特征。即便耗费大量人工，提取出的特征依然存在极度稀疏的问题，且必须为 CPU 和 GPU 独立设计，通用性极差。此外，部分模型尝试将抽象语法树作为计算流来前向传播特征，但这导致每个张量程序的结构不同，模型只能逐个前向传播而无法并行化，严重拖慢了编译速度。其次，由于不同硬件架构之间存在巨大的领域差异，离线收集的训练数据在不同硬件间往往是无效的。这导致在一个平台上训练好的成本模型，一旦迁移到全新硬件上，其预测性能就会急剧下降，即跨硬件不可用性。为了在目标硬件上重新获得高性能，通常需要重新收集海量实测数据，带来了不可接受的时间和算力成本。

怎么做的 为了绕开上述卡点，本文提出了 TLP 模型，并扩展出 MTL-TLP 模型。核心思路是放弃对复杂的张量程序源代码进行特征工程，转而直接从生成张量程序的调度原语序列中提取特征。因为调度原语序列结构规整、种类少且在不同硬件上高度通用，作者将其视为一种张量语言，从而将预测张量程序延迟的任务转化为了自然语言处理的回归任务。关键设计由预处理、特征提取器和模型主干三个部件构成。预处理器首先剥离原语中的无关字符，仅保留原语类型、数值参数和字符参数这三个包含完整语义的基本元素。随后，特征提取器对这三者分别处理：原语类型被转换为独热向量，数值参数保持原值，字符参数则像自然语言中的单词一样被转换为标记，最后按原始位置拼接。在模型架构上，TLP 采用自注意力机制或长短期记忆网络作为主干来捕捉原语间的上下文特征，后接残差块和线性层输出预测分数。其目标标签定义为归一化延迟：$label = min\_latency/latency$。针对跨硬件不可用性问题，MTL-TLP 引入了多任务学习机制。它将不同的硬件平台设置为不同的学习任务，在模型末端设置多个预测头。其核心机制在于利用共享的主干网络参数来拟合硬件无关的通用特征，同时利用非共享的预测头参数来拟合特定于硬件的特征。总损失函数定义为各个有效任务损失的总和：$$Loss_{total} = \sum_{i=1}^{n} Loss\_i(pred_i, label_i)$$。通过这种设计，MTL-TLP 能够利用其他平台上丰富的数据作为辅助任务，在仅有少量目标硬件数据的情况下，依然能训练出高精度的成本模型。

效果如何 实验在包含 5157 万个张量程序的大规模数据集 TenSet 及额外收集的 TenSet-TLP 数据集上进行，覆盖多种 CPU 和 GPU 平台。测试任务包含 ResNet-50、MobileNet-V2、BERT 等典型模型的子图调优。对比基线有两个：自动搜索框架 Ansor，代表传统的基于手动特征提取和搜索的路线；预训练的 TenSet MLP，代表先进的基于多层感知机的离线深度学习成本模型路线。量化结果表明，由于 TLP 直接从调度原语提取特征而无需实际生成张量程序，其调优执行速度在 CPU 和 GPU 上比 TenSet MLP 平均快 1.7 倍和 1.8 倍。在端到端搜索效率上，为达到 TenSet MLP 调优 2000 次的性能，TLP 在 CPU 和 GPU 上平均可将搜索时间加速 9.1 倍和 3.0 倍；而 MTL-TLP 在仅使用 50 万条目标硬件数据的情况下，依然能加速 4.7 倍和 2.9 倍。与 Ansor 相比，两者更是实现了 10 倍到 16 倍以上的搜索时间加速。该方法也存在局限性：在 GPU 上运行时，TLP 的显存占用从基线的 882MB 增加到了 1634MB；同时，尽管 MTL-TLP 降低了对目标平台数据的依赖，但收集这 50 万条数据依然需要耗费数十小时，跨硬件部署的数据收集成本仍有待进一步降低。

A1 主要贡献

本文旨在解决张量程序调优中的两大核心问题。首先，现有的基于深度学习的成本模型在特征提取上严重依赖于硬件架构的专家级领域知识，即便如此，提取的特征仍不理想，且需要为CPU和GPU分别设计，通用性差。其次，在一个硬件平台上训练的成本模型在另一平台上性能会急剧下降，这一问题被称为“跨硬件不可用性”。

为了应对这些挑战，本文提出了两种方法：TLP和MTL-TLP。

本文的主要贡献如下：

  1. 提出了一种简单、高效且通用的张量程序特征提取机制：通过直接从调度原语中提取特征，将问题转化为NLP任务，避免了对张量程序源代码（如AST）进行复杂的分析，减少了对硬件先验知识的依赖。
  2. 提出了一种解决成本模型跨硬件平台不可用性的多任务学习方法：MTL-TLP能够有效利用来自多个硬件平台的数据，显著提升模型在目标硬件上的性能，即使目标硬件数据量很少。
  3. 实现了TLP和MTL-TLP并进行了全面的评估：将所提方法集成到先进的搜索框架Ansor中，并在多种深度学习模型和硬件平台上进行了验证。实验结果表明，与现有最先进技术相比，TLP和MTL-TLP在搜索效率和性能上均取得了显著的提升。据作者所知，TLP是首个直接从调度原语提取特征的张量程序成本模型，而MTL-TLP是首个有效解决跨平台不可用性问题的开源工作。

A3 背景知识

基于搜索的编译器

图1：通用深度学习编译器流程。红色虚线框是 TLP 的工作区。
图1：通用深度学习编译器流程。红色虚线框是 TLP 的工作区。

TenSet

自然语言处理

A3 设计原则

图3展示了集成了TLP的自动搜索框架（auto-tuner）从计算子图生成张量程序的过程。右侧黑框是TLP的流水线，该图描绘了TLP的训练和推理两种状态。

在TLP的流水线中，“具体原语”指自动搜索框架中的调度原语，“抽象原语”指经过预处理后符合TLP特征提取器输入规范的原语。预处理器、提取器和后处理器的细节将在第4节中详细描述。

图3：集成了TLP的自动搜索框架从计算子图生成张量程序的过程。右侧的黑框是TLP的流水线。左侧深绿色的张量程序和延迟（标签）块在训练时为深绿色，右侧黄色的在推理时为黄色。训练时，图中的虚线箭头、虚线框和“仅推理”箭头无效。推理时，图中的“仅训练”箭头无效。
图3：集成了TLP的自动搜索框架从计算子图生成张量程序的过程。右侧的黑框是TLP的流水线。左侧深绿色的张量程序和延迟（标签）块在训练时为深绿色，右侧黄色的在推理时为黄色。训练时，图中的虚线箭头、虚线框和“仅推理”箭头无效。推理时，图中的“仅训练”箭头无效。

A2 方法细节

4 TLP

现有成本模型特征提取的局限性

4.1 TLP的特征提取

从调度原语提取特征

图2：子图，一个融合的 dense + ReLU 激活，应用不同调度原语的组合来生成张量程序。上面计算子图中包含的数学表达式、DAG和朴素张量程序在逻辑上是等价的，但形式不同。调度原语以伪代码编写。红色长虚线框中的调度原语是TLP的特征提取对象，蓝色短虚线框中的张量程序是Ansor [38]、TIRAMISU成本模型 [5]等的特征提取对象。
图2：子图，一个融合的 dense + ReLU 激活，应用不同调度原语的组合来生成张量程序。上面计算子图中包含的数学表达式、DAG和朴素张量程序在逻辑上是等价的，但形式不同。调度原语以伪代码编写。红色长虚线框中的调度原语是TLP的特征提取对象，蓝色短虚线框中的张量程序是Ansor [38]、TIRAMISU成本模型 [5]等的特征提取对象。

TLP特征提取流程

图4：TLP特征提取规范。
图4：TLP特征提取规范。

特征提取示例

图5：TLP的特征提取示例。
图5：TLP的特征提取示例。

4.2 TLP在TenSet数据集上的特征提取

术语定义

TenSet数据集上的特征统计

图6：TenSet CPU数据集中张量程序序列长度的分布。
图6：TenSet CPU数据集中张量程序序列长度的分布。
表1：TenSet CPU数据集中各调度原语的最大嵌入大小。“RE”、“FU”等是调度原语的缩写，分别指reorder、fuse等原语。
表1：TenSet CPU数据集中各调度原语的最大嵌入大小。“RE”、“FU”等是调度原语的缩写，分别指reorder、fuse等原语。

4.3 TLP特征提取的可行性分析

调度原语序列与张量程序的对应关系

数据验证

调度原语序列的优势

结论：张量语言处理

4.4 模型架构

TLP模型结构

图7：TLP的模型架构。左侧的参数是层的输入和输出形状，其中N表示批量大小，L表示序列长度，El和Eh表示嵌入大小。我们将红色虚线框称为主干（backbone），蓝色虚线框称为头（head）。这将在多任务学习章节中使用。
图7：TLP的模型架构。左侧的参数是层的输入和输出形状，其中N表示批量大小，L表示序列长度，El和Eh表示嵌入大小。我们将红色虚线框称为主干（backbone），蓝色虚线框称为头（head）。这将在多任务学习章节中使用。

5 MTL-TLP

成本模型发展历程

表2：张量程序成本模型的发展。
表2：张量程序成本模型的发展。

5.1 跨硬件不可用性

问题描述

潜在解决方案分析

选择多任务学习的原因

5.2 MTL-TLP

模型结构

图8：MTL-TLP的模型架构。
图8：MTL-TLP的模型架构。

损失函数

$Loss_{total} = \sum_{i=1}^{n} Loss\_i(pred_i, label_i)$, 其中 $label_i$ is not no label.
其中$Loss\_i$是每个任务的损失函数（通常相同），例如MSE损失函数、排名损失函数等。

5.3 MTL-TLP的可行性分析

MTL-TLP之所以能发挥巨大作用，主要归功于以下几点：

A4 实验环境

A5 实验结果

评估指标分为两类：基于数据集的指标（模型在静态数据集上的准确率）和基于搜索的指标（模型集成到搜索算法后的端到端搜索效率或质量）。

6.1 基于数据集的TLP性能评估

使用top-k分数作为评估标准，其表达式如下：
$top-k = \frac{\sum_m \sum_s min\_latency_{m,s} \times weight_{m,s}}{\sum_m \sum_s \min(latency_{m,s,i}) \times weight_{m,s}}, 1 \le i \le k$
$$top-k = \frac{\sum_M \sum_S min\_latency_{m,s} \times weight_{m,s}}{\sum_M \sum_S min (latency_{m,s,i}) \times weight_{m,s}}, 1 \le i \le k$$
其中$min\_latency_{m,s}$是模型$m$的子图$s$所有张量程序中的最小延迟，$weight_{m,s}$是子图$s$在模型$m$中出现的次数，$latency_{m,s,i}$是模型$m$的子图$s$所有张量程序中，成本模型输出分数第$i$大值对应的延迟。

6.2 基于数据集的MTL-TLP性能评估

6.3 基于搜索的端到端评估

A7 补充细节

局限性 (Limitation)

未来工作 (Future Work)

A6 结论

本文提出了TLP和MTL-TLP两种方法，以改进张量程序调优中的成本模型。TLP设计了一种新颖、简单且通用的特征提取机制，直接从调度原语中提取特征，将调优问题转化为张量语言处理任务。MTL-TLP则利用多任务学习技术有效解决了离线成本模型在不同硬件平台间的不可用性问题。本文从理论上分析了这两种方法的可行性和优势，并通过基于数据集和基于搜索的详尽实验验证了它们的有效性。实验结果表明，TLP和MTL-TLP在搜索效率和最终性能上均显著优于当前最先进的实现。