JIT in MegEngine

一、背景

什么是天元

旷视天元(MegEngine)是一个深度学习框架,它主要包含训练和推理两方面内容。训练侧一般使用 Python 搭建网络;而推理侧考虑到产品性能的因素,一般使用 C++ 语言集成天元框架。无论在训练侧还是推理侧,天元都担负着将训练和推理的代码运行到各种计算后端上的任务。目前天元支持的计算后端有 CPU、GPU、ARM 和一些领域专用的加速器,覆盖了云、端、芯等各个场景。

天元主要有三大特征:

1.训推一体,不管是训练任务还是推理任务都可以由天元一个框架来完成。

2.动静结合,天元同时支持动态图和静态图,并且动静之间的转换也非常方便。

3.多平台的高性能支持。

JIT in MegEngine


图 1. 天元架构

如图 1 所示,我们可以看到天元提供了 Python 和 C++ 两种接口。在图表示上分为动态图和静态图。运算层组件包括自动求导器、图优化和图编译等。天元的运行时模块包括内存管理和计算调度,其中内存管理包括静态内存管理和动态内存管理,以及亚线性内存优化技术。计算内核层包含了天元支持的所有计算后端,我们后续会开源出更多的计算后端。除此之外,天元还包含了一个高性能异构通信库,它一般会在多机多卡的场景下被用到。

JIT in MegEngine

图 2. 计算图

动态图和静态图是相对的,在动态图下是没有计算图的概念的。但在静态图下,天元会维护一张计算图。如图 2 所示为天元中的计算图表示,图中圆形表示算子(operator),三角形表示输入。在天元框架中,动态图和静态图之间的转换只需要一条简单的语句即可完成,如下方代码所示:

if __name__ == '__main__’:
 gm = ad.GradManager().attach(model.parameters())
 opt = optim.SGD(model.parameters(), lr=0.0125, momentum=0.9, weight_decay=1e-4)
 # 通过 trace 转换为静态图
@trace(symbolic=True)
 def train():
     with gm:
         logits = model(image)
         loss = F.loss.cross_entropy(logits, label)
         gm.backward(loss)
     opt.step()
     opt.clear_grad()
     return loss
 loss = train()
 loss.numpy()

什么是 AOT 和 JIT

AOT(Ahead Of Time) 和 JIT(Just In Time) 都是编译中的概念。以传统的 C/C++ 语言为例,我们写完代码之后,一般会通过编译器编译生成可执行文件,然后再执行该可执行文件获得执行结果。如果我们将从源代码编译生成可执行文件的过程称为 build 阶段,将执行可执行文件叫做 runtime 阶段的话,JIT 是没有build 阶段的,它只有 runtime 阶段。JIT 一般被用在解释执行的语言如 Python 中,JIT 会在代码执行的过程中检测热点函数,随后对热点函数进行重编译,下次运行时遇到热点函数则直接执行编译结果即可。这样做可以显著加快代码执行的速度。

什么是 MLIR

随着各种编程语言的出现,现代编译器也日趋多样化。特别是近年来随着深度学习的兴起,深度学习软件框架和 AI 领域专用硬件呈爆发式增长。不断增加的软件框架和 AI 硬件之间逐渐形成了一个越来越大的沟壑,如何将框架层对深度学习模型的描述精准高效的翻译成适应各类硬件的语言成为难点。MLIR(Multi-Level Intermediate Representation) 是一种可以在统一的基础架构下满足多样化需求的混合 IR。MLIR 可以满足包括但不限于以下的需求:

1.表达数据流图(如静态图模式下的 MegEngine Graph)

2.表达对该图做的优化和变换操作

3.进行各种算子优化如算子融合(kernel fusion)、循环融合、算子分块和内存格式(memory layout)转换等

自动代码生成、显式缓存管理、自动向量化

作为一个公用的 IR,MLIR 具有非常优秀的表达能力和可扩展性。MLIR 可以表达图层面的运算,同时可以表达传统编译器中的 IR 信息,也可以表示硬件专用的运算。这种不同属性,不同类型的运算的集合构成了 MLIR 中的方言(Dialect)。MLIR 还提供方便的机制实现不同方言之间的转换(Lowering Down),因此 MLIR 的一个通用优化将会在多个方面产生收益。接入 MLIR 也将有更大可能享受到它的生态好处,包括性能和扩展性等方面。

二、动机

为什么做

众所周知,深度学习模型中有很多 element-wise 操作,例如加减乘除算术运算和神经网络中的激活函数一般都是 element-wise 操作。天元将 element-wise 操作分为一元操作、二元操作和多元操作。一元操作主要有 RELU、ABS、SIN 和 COS 等等;二元操作有加法、减法、乘法和除法以及 MAX 等;多元操作有 FUSE-MUL-ADD3 和 FUSE-MUL-ADD4 等,它们分别计算的是 “a*b+c” 以及 “a*b+c*d”。

JIT in MegEngine

表 1 卷积神经网络中的 element-wise 操作

element-wise 操作在卷积神经网络中所占的地位不可忽视。如表 1 所示,我们选择公开的卷积神经网络训练模型,以纯 device kernel 的执行时间为基准统计卷积神经网络中element-wise 操作的重要性。

首先可以清晰的看到,element-wise 的计算量的占比相比于运行时间占比要低 1-2 个数量级。它的计算量占的非常少,但是它的运行时间占比非常多,这个结论是比较反直觉的。并且随着 batch size 的增加,这个现象也越来越明显。这是因为 element-wise 操作计算量较低但是访存量较高,即计算访存比较低,是一种典型的访存受限 (memory bound) 的操作。以 “a+b” 为例,我们首先要将 a 读到内存中,再将 b 读到内存中,做完一次加法之后,我们将结果 c 再写到内存中。整个过程要经过两次读和一次写才能完成一次计算,所以它的计算反应访存比非常低。针对访存受限的操作,优化计算时间实际上是没有没有太多的意义的,而应该集中精力优化访存,访存优化的常见的优化手段是融合 (fusion)。如果我们能将网络中连在一起的 element-wise 操作融合成一个算子,则将减少 element-wise 操作的访存量,增加计算访存比从而加速网络的整体性能。

为什么用 JIT 做

卷积神经网络有两个鲜明的特征。一个是静态图模式下的模型训练过程中模型的结构一般是不会变的跑;另一个是在模型训练的过程中,一般会经过很多个 iter/min-batch,不同的 iter/min-batch 之间输入张量形状(tensor shape)一般也不会变。基于卷积神经网络的这两个特征,我们决定应用 JIT 技术,原因如下:

1.只需要在首次运行的时候编译一次,随后的不同 iter/mini-batch 可以重用第一次编译出来的结果。

2.JIT 具有较强的可移植性,因为它在运行时获取平台信息,然后生成可以在该平台运行的代码。

3.JIT 可以解决 element-wise 模式组合爆炸的问题。

三、技术方案

我们通过 Element-wise Fusion 可以把多个 element-wise 操作融合成一个,减少了算子数量也就减少了算子之间的读写次数。如图 3 所示计算图算的是 “a*b+c”,它需要 4 次读,2 次写。4 次读分别是乘法在读 a 和 b 两个输入,乘法其实还要写一个隐藏的输出,加法会读乘法的输出作为输入,以及加法读 c 作为输入。两次写分别是乘法和加法对它们结果的两次写操作,总共加起来是 4 次读,2 次写。

我们将其融合成一个算子 FUSE_MUL_ADD3,由于天元现在已经支持 FUSE_MUL_ADD3 这个 element-wise 模式,所以我们可以直接做模型手术将计算图从图 3 左侧形式转到图 3 右侧形式。对于融合之后的计算图,我们只需要 3 次读和 1 次写就可以完成等价计算,相比于融合前减少了 1 次读和 1 次写操作。

JIT in MegEngine

图 3 融合优化减少访存次数

我们无法预测用户将搭出来怎样的一张计算图,考虑图 4 所示的计算图,其中 element-wise 的个数和顺序都不固定,显然我们不可能提前将各种 element-wise 模式的组合都写进天元的。在这种情况下,天元会创建一个虚拟的算子来表示整个可被融合的子图。有了虚拟算子的存在,接下来我们还要解决两个问题,一个是用虚拟算子替换原始计算图中可以被融合的子图,这个工作会在图优化阶段做;另一个是我们要动态生成虚拟算子的代码并执行。如果我们解决了这两个问题,我们就解决了整个问题。

JIT in MegEngine

图 4 子图融合优化

图优化

为了将一张计算图中的可被融合的子图融合成一个算子,天元将进行检测(detection)和融合(fusion)两步操作,如下步骤 1-3 属于检测,步骤 4 则属于融合:

1.对原始计算图进行检测后生成 internal graph generator,一个 internal graph generator 对应一个唯一的子图

2.internal graph generator 稍后会生成 internal graph

3.由 internal graph 创建 JITExcutor 算子

4.将 JITExcutor 写回原始的计算图

检测

检测算法的主要功能是找出可以被融合的子图。为了方便描述,设 G 是计算图,opr 是图 G 中的算子,var 是 opr 的输入和输出。检测算法的输入是原始的计算图 G,输出是一个哈希表 M,表中存放的是检测出的可被融合子图的输出 var(记作 endpoint)与其对应的 internal graph generator。算法步骤如下:

按照逆拓扑序列遍历图 G 中的算子 opr

1.如果 opr 不是
Elemwise/PowC/TypeCvt/Reduce/Dimshuffle/JITExecutor,返回步骤1

2.如果 opr 的 input/output 数据类型不是 float32/float16,返回步骤1

3.

4.process_opr(opr)

5.转到步骤 1

JIT in MegEngine

图 5 process_opr 流程图

拓扑序列要求所有的子节点要先于它的父节点被访问到,与之对应的,逆拓扑序列就是所有的父节点要先于它的子节点被访问到。算法第 1 步中我们之所以按照逆拓扑序列遍历计算图,是因为要保证遍历到某个 opr 时,它的父节点都已经被遍历到了。这样算法可以查看该 opr 的所有的父节点是不是都在同一张子图中,如果是,那么当前 opr 就有很大的可能也在该子图中。算法的第 2 步和第 3 步实际上说明了天元中的 JIT 的限制。目前天元 JIT 仅支持
Elemwise/PowC/TypeCvt/Reduce/Dimshuffle 这几种 opr,而且只支持输入输出是 float32/float16 的数据类型。第 4 步详细流程如图 5 所示。需要注意的是算法会经过如下三个判断语句:

1.该 opr 的输出的父节点是不是都已经在当前的这张子图中了?

2.该 opr 的输出的计算节点(compute node)是不是跟子图匹配?天元支持跨计算节点的计算图,例如计算图中一些 opr 可以运行在 CPU 上,一些 opr 可以运行在 GPU上。但目前天元不支持跨计算节点融合。

3.该 opr 的输出的 shape 是不是跟子图匹配?因为最终生成的代码本质上是一个大的循环,循环的维度就是 opr 输出的 shape,所以如果 shape 不匹配是不能被融合的。

JIT in MegEngine

图 6 检测算法检测出的可被融合的子图

图 6 中虚线框出来的即为检测算法检测出的两个可被融合的子图。

融合

融合算法的主要功能是将检测出来的子图融合成一个算子。融合算法的输入是原始的计算图和检测算法输出的那张哈希表 M,它的输出是经过融合的计算图 G‘。算法流程如下:

1.按照拓扑序列遍历图 G 中的算子 opr

2.若 opr 的输入 var 不是 endpoint, 返回步骤 1

3.从 M 中拿到 var 对应的 internal graph generator, 生成 internal graph

4.从 internal graph 创建 JITExecutor

5.写回原始的计算图 G

6.转到步骤 1

步骤 2 中如果一个 opr 的输入 var 不是 endpoint 则表示它是一个子图中的中间节点而不是子图的输出节点。步骤 3 中从 internal graph generator 到 internal graph 需要将子图的输入 var 替换为一个新的 opr JITPlaceholder。JITPlaceholder 中会存诸如子图的输入顺序这些额外信息,因为某些 element-wise 操作是对输入顺序敏感的。例如 a 对 b 取余和 b 对 a 取余显然具有不同的语义。

JIT in MegEngine

图 7 融合后的计算图

图 7 即为经过融合算法之后的计算图,截止到目前为止,我们已经完成了图优化方面的所有工作。

图编译

经过图优化之后,我们成功的将计算图中可被融合的子图融合成为一个新的算子,剩下的工作就是为这个新的算子生成代码了。JITExecutor 算子的运行时代码非常简单,先判断一下当前的可执行对象是不是已经存在,如果不存在则先编译出一个可执行对象,如已存在则直接运行。这段代码在运行时才会被执行到,所以称之为 JIT。当前天元支持三种 JIT 编译器后端,分别是 NVRTC(支持英伟达 GPU),Halide 和 MLIR。其中后两个编译后端支持的平台众多,但是 MLIR 具有更优秀的表达能力和扩展性,所以我们接下来以 MLIR 为例介绍代码生成、编译和执行的过程。

要想使用 MLIR 作为编译后端,首先我们需要定义和实现天元自己的方言(MGE Dialect),随后我们将 MGE Dialect 转换到 MLIR 既有的 Dialect 上,接下来的绝大部分工作都可以复用 MLIR 中的基础组件和工具完成。图 8 描述了 CPU 和 GPU 上大概的执行流程。

JIT in MegEngine

图 8 JIT 编译器工作流

天元首先将 JITExecutor 算子内部的 internal graph 翻译成 MGE Dialect。在 CPU 上,MGE Dialect 会先 Lowering 到 Affine Dialect 上,然后会通过 LLVM 的组件 Lowering 到 LLVM Dialect 上,LLVM Dialect 可以被直接翻译成 LLVM IR。在这一步之后,其他优化工作都可以直接复用 LLVM 的基础组件。最后天元使用 MLIR ExecutionEngine 执行 LLVM IR 生成的代码。在 GPU 上,天元会先将 MGE Dialect Lowering 到 GPU Dialect上,随后 Lowering 到 NVVM Dialect,NVVM 会被翻译成 PTX 汇编。最后通过英伟达提供的 CUmodule 和 CUfunction 两个机制运行。

四、实验和分析

首先参考这篇文档 在天元中开启 JIT 支持。本次实验选了 resnet50, mobilenetV2 和 vgg16 三个业界广泛使用的模型,batch size 分别设置了 1, 8 和 16。测试硬件环境为 NVIDIA T4,软件环境为 MegEngine v1.2.0。

JIT in MegEngine

图 9 打开 JIT 相比于不开 JIT 的加速比

由图 9 可知,和不打开 JIT 支持相比,打开 JIT 支持后 resnet50 最高可以获得 16% 的加速比,mobilenet V2 则能获得 6% 到 7% 的加速比,而 vgg16 其实上没有明显加速效果。这是因为 vgg16 模型很大,可以被优化的 element-wise 操作比较少。JIT 的优化效果跟具体的模型是有紧密关系的。

JIT in MegEngine

图 10 JIT 编译耗时

如果打开了 JIT 支持,那么天元首次运行的时候会有一次 JIT 编译的过程。JIT 编译耗时跟具体的编译的后端以及模型有关,如图 10 所示 resnet50 耗时 2.7 毫秒,mobilenetV2 耗时 3.9 毫秒。

五、总结和展望

本篇文章介绍了天元使用 JIT 实现将任意多个相邻的 element-wise 算子融合成一个算子的优化。我们在 T4 上用 MegEngine v1.2.0 实验,相比于优化前,resnet 50 最高可以获得 16% 的加速比。

以此为基,展望未来我们可能做的事情如下:

1.将 JIT 编译的结果先离线保存,线上直接将线下编译好的可执行对象读进内存。这种做法可以解决线上第一次运行慢的问题,但它可能会损失一部分可移植性,因为在一种设备上编译产生的可执行对象一般不能适配所有线上设备。

2.JIT 支持更多的算子。

3.JIT支持更多的数据类型,天元 JIT 优化暂时只支持 float32/float16 这两种数据类型。

4.动态图 JIT,也就是传统意义上的检测热点代码,重编译后再执行。


MegEngine Meetup 下期预告

时间:2021 年 3 月 25 日 19:00 - 21:00

内容:

Topic 1 :《MegEngine 开源一周年架构回顾》

讲师:许欣然 | 旷视天元技术总监

Topic 2:《利用 MegEngine 的分布式通信算子实现复杂的并行训练》

讲师:周亦庄| 旷视研究院基础模型组研究员

直播地址:
https://live.bilibili.com/22436423

欢迎大家提前加入天元技术交流 QQ 群(群号:1029741705)

活动开始前小助手会在群内提醒观看,并发送直播链接,群内还有第一手分享PPT、视频回放链接、以及讲师在线答疑 :)



理论卷积神经网络神经网络深度学习
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

激活函数技术

在 计算网络中, 一个节点的激活函数定义了该节点在给定的输入或输入的集合下的输出。标准的计算机芯片电路可以看作是根据输入得到"开"(1)或"关"(0)输出的数字网络激活函数。这与神经网络中的线性感知机的行为类似。 一种函数(例如 ReLU 或 S 型函数),用于对上一层的所有输入求加权和,然后生成一个输出值(通常为非线性值),并将其传递给下一层。

调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源,如线程、进程或数据流;也可以指硬件资源,如处理器、网络连接或扩展卡。 进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态,允许多位用户有效地同时共享系统资源,或达到指定的服务质量。 see planning for more details

基准技术

一种简单的模型或启发法,用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

卷积神经网络技术

卷积神经网路(Convolutional Neural Network, CNN)是一种前馈神经网络,它的人工神经元可以响应一部分覆盖范围内的周围单元,对于大型图像处理有出色表现。卷积神经网路由一个或多个卷积层和顶端的全连通层(对应经典的神经网路)组成,同时也包括关联权重和池化层(pooling layer)。这一结构使得卷积神经网路能够利用输入数据的二维结构。与其他深度学习结构相比,卷积神经网路在图像和语音识别方面能够给出更好的结果。这一模型也可以使用反向传播算法进行训练。相比较其他深度、前馈神经网路,卷积神经网路需要考量的参数更少,使之成为一种颇具吸引力的深度学习结构。 卷积网络是一种专门用于处理具有已知的、网格状拓扑的数据的神经网络。例如时间序列数据,它可以被认为是以一定时间间隔采样的一维网格,又如图像数据,其可以被认为是二维像素网格。

算术技术

算术(英语:arithmetic)是数学最古老且最简单的一个分支,几乎被每个人使用着,从日常生活上简单的算数到高深的科学及工商业计算都会用到。一般而言,算术这一词指的是记录数字某些运算基本性质的数学分支。

分块技术

将标注好词性的句子按句法结构把某些词聚合在一起形成比如主语、谓语、宾语等等。

自动代码生成技术

自动代码生成(Code Generation),指生成程序的程序,更详细的是指通过读取工程中设计的各种文档或者模型,按照一定的领域规则,生成规范的计算机能理解的源代码的过程。

量化技术

深度学习中的量化是指,用低位宽数字的神经网络近似使用了浮点数的神经网络的过程。

推荐文章
暂无评论
暂无评论~