ai算子是什么_AI编译优化--计算密集算子优化

作者：PAI团队

进入正题前，还是先打个招聘小广告，欢迎对我们工作感兴趣的同学联系我们，细节参见这里，可以直接邮件muzhuo.yj@alibaba-inc.com。

本文是AI编译优化系列连载的第三篇，总纲请移步：

https://zhuanlan.zhihu.com/p/163717035

针对计算密集算子，我们的工作包括两大部分：

• 围绕GPU硬件上的低精度算子开展了一系列优化工作，以充分发掘NV新硬件提供的以TensorCore为代表的专用硬件加速单元的计算效率。
• 针对多种硬件设备(GPU/CPU/端侧CPU等)，以更具一般性的方式自动完成计算密集算子的codegen支持，这个工作的细节会在未来的文章中进行着重介绍，在本篇不会涉及，感兴趣的同学可以先通过这篇tech report来了解整体的技术思路。

由于最新的AI硬件加速单元都只支持低精度，例如fp16, int8, int4甚至int1，因此想要充分利用这些AI硬件加速单元的前提首先是需要一个低精度的模型，因此在训练侧我们开发了自动混合精度，在推理侧我们支持低精度量化。

在得到低精度模型之后我们便可以基于硬件厂商提供的高性能计算库（例如针对TensorCore有cuBLAS, cuDNN, TensorRT等）利用AI硬件加速单元的算力大幅提升训练和推理的性能，但是硬件厂商极致手工（汇编）优化需要针对各种情况（例如不同算法、不同尺寸、不同硬件架构）分别进行精细调优，工作量极大，因此很难在所有情况下都确保有高效的实现。所以我们需要更通用更自动化的解决方案来进一步提升AI硬件加速单元的计算效率，我们在TVM框架上实现了TensorCore AutoCodeGen. 本文接下来将详细介绍这三部分的工作。

一、自动混合精度

混合精度训练是指在训练过程中在不同的计算中使用不同的数值精度，从而充分挖掘显卡硬件上每一个晶体管的极致性能，目前在PAI上支持的混合精度训练主要是指FP16和FP32两个数值精度。在标准TensorFlow训练任务中，variable及tensor的表示精度目前均默认为FP32的，那么为什么要引入FP16的数值精度呢？我们知道对于神经网络训练过程而言，影响速度的两个关键因素是，计算和访存：

• 从计算而言，NVIDIA在2017年推出了Volta架构，其中重要的模块即为TensorCore（如下图），在TensorCore的加速下，V100在混合精度下的吞吐量比在FP32下有8X的加速。TensorCore主要对matrix-multiply-and-accumulate类计算进行了加速。基于TensorCore我们可以对 MatMul、 Convolution等计算原语进行加速，而二者恰好是神经网络中计算量较大的op

• 从访存而言，若tensor的数值精度由FP32变为FP16，那么我们也可以得到理论上2X的访存加速。 因此我们引入了混合精度对神经网络的训练进行加速。

自动混合精度训练需要解决的问题分为两个方面：

一是自动将用户定义好的单精度模型转换为高效的混合精度模型（单精度和半精度的混合精度模型），即Auto-mixed-precision graph optimization模块；二是通过模型算法策略解决训练过程中数值精度降低所带来的模型训练精度问题，以帮助用户在不改变训练超参（例如learning rate, optimizer等)的情况下保证模型训练的收敛性，即为automatic loss-scale optimizer模块。

易用性是我们PAI平台在追求高性能、高效率的同时同样重视的另一个关键目标，因为我们的用户，即AI Developer，希望专注于模型和算法本身，不希望被各种复杂的性能优化问题分散精力，这要求我们的性能优化能够作为Turnkey Solution，尽可能对用户透明，避免对用户模型代码的入侵。混合精度训练加速优化同样如此。

因此对于第一个问题，图转换问题，我们在计算图优化环节实现了自动的、精细的混合精度转换，对用户完全透明，用户跟正常训练FP32模型一样，无需修改模型代码；同时在计算图优化环节，我们通过相应的性能优化避免混合精度转换过程中伴生的性能不利因素；对于第二个问题，我们实现了比社区对优化更加友好的optimizer封装，辅助用户无缝加入loss scaling策略，使能用户无需关心自动混合精度训练的细节策略，即可基于开关启动混合精度训练，以更好地利用TensorCore的计算加速特性，加速模型训练。

下面我们将依次介绍两个维度的技术解决方案。

自动混合精度图改写

自动混合精度图改写旨在自动帮用户将单精度训练图改写为混合精度训练图，在精度符合模型约束的前提下保证图的高性能。自动图改写可分为以下几个步骤：

White-List:

首先，为了保证模型的高精度以及满足深度学习框架的约束，我们会维护一个white-list，只有white-list中的算子才可能被转换为FP16。

无论对于计算密集算子还是非计算密集算子，我们都期望使得尽可能多的算子转换为FP16或FP16/FP32混合精度进行计算，因为：对于计算密集算子例如Conv2D/MatMul，可以通过混合精度获得TensorCore的硬件加速，而非计算密集算子（即访存密集算子）可以通过FP16获得访存的节省。但在实际图转换中有若干约束使得有些算子不能加入whitel-list。更细节一些，white-list的维护来自两个维度的约束：一方面模型层不适合用FP16进行计算的算法将维护white-list之外，例如reduction类算子（例如batch_norm)，exp/log/softmax类算子等；另一方面TensorFlow框架层不支持的FP16的算子也将维护在white-list之外，例如BatchMatMul等；

Cost-Model:

实际上，white-list中的算子并非都应该转换为FP16，计算图中每个节点需不需要转换，除了和这个计算节点本身的算子类型有关，还和这个节点周围其他节点的连接拓扑等因素紧密相关，所以我们需要一个Cost Model来预测每个节点转换为FP16之后对整体性能的影响