【云驻共创】MindSpore图算融合,打破算子边界
目录
1.图算融合的背景
1.1计算图和算子的认识
1.2图算融合的优势
2.设计初心
2.1图算融合概述
2.2图算融合的设计初心
3.关键特性简介
3.1统一图算前端表达
3.2统一图算中间表达:算子白盒化
3.3复用图层AD算法
3.3.1自动生成反向算子
3.3.2正反向算子的自动融合
3.4 Bert上的图算融合
4.启方法汲编译流程演示
4.1使用方法
4.1.1样例脚本
4.2自定义组合算子
4.2.1样例脚本
5.实现过程小结
5.1图算融合流程图
5.2融合算子自动生成流程
5.3小结
6.总结
1.图算融合的背景
1.1计算图和算子的认识
包括MindSpore在内的AI框架,不约而同都采用计算图实现网络结构的表达。所以也间接证明了计算图在AI框架中强大的生命力。
可以引用[tutorialspoint]中的一段话来介绍计算图:
A computational graph is defined as a directed graph where the nodes correspond to mathematical operations. Computational graphs are a way of expressing and evaluating a mathematical expression.
For example, here is a simple mathematical equation –
P = x + y
We can draw a computational graph of the above equation as follows.
总结来说,计算图是一个以算子(Operators)为节点的有向图所表达的计算函数。在AI框架里面,这个计算函数对输入tensor依次调用执行有向图中的算子节点,并得到最终的输出tensor。
引用Wikipedia来介绍一下算子:
In mathematics, an operation is a function which takes zero or more input values (called operands) to a well-defined output value. The number of operands is the arity of the operation.
所以,算子本质还是一个计算函数。在AI框架里面,这个计算函数还是对输入tensor依次进行若干内部处理,然后得到最终的输出tensor。
算子的内部处理:通过对大量的算子进行分析,会发现,绝大部分算子都是可以拆分为若干更为基本的算子组成的的子图,比如Sigmoid,其内部计算可以用如下计算图表示:
通过以上分析,不难发现,计算图和算子在计算本质上是一致的。算子是打包后的计算图,计算图是拆包后的算子。所以理论上,可以定义一个小规模的“基本算子”集合,然后通过一个或多个基本算子组合去等价表达任何的现有算子。从而进一步表达任意的现有计算图。
用小规模的“基本算子”集合就可以表达任意现有计算图,但是现有主流AI框架还要定义大量的、本来可以用基本算子等价实现的复合算子,其主要原因是相比多个基本算子组成的计算图,使用复合算子可以具有更好的计算和内存局部性,从而获得更好的执行性能。
因为对于基本算子计算图来说,相邻算子之间只能通过全局内存(或显存)进行数据传递。而对于复合算子来说,相邻的基本计算之间则可以通过局部内存或者寄存器进行数据传递。除了性能之外,在一些场景下,通过算子融合也能有效减少对全局内存的的实际占用。
1.2图算融合的优势
基于以上分析可见,相比复合算子,基本算子具有更好的通用组合表达能力,同时使得AI框架在算子支持方面更为精简。但相比基本算子,复合算子在性能方面存在优势。可以简单图示如下:
对于以上问题,通过某些技术同时兼顾计算性能和组合表达能力,或者更具体地,使得上图的“计算性能”线变为水平。
另一方面,在通过AI框架提供的算子表达计算网络时,通常是混合使用不同复合程度的算子。比如LayerNorm由12个基本算子组成,而Relu则只有1个基本算子。所以,对于这种不同复合程度的算子组成的计算图,从整体执行性能角度,可以将他们组合成更为合理的复合粒度算子,来获得最佳的整体计算图性能。
在TVM、XLA等自动算子编译技术出现之前,AI框架主流采用手工融合的方式解决如上问题。主要思路是:
- 识别常见的热点算子组合子图,比如:Add(Mul(x,y))。然后针对该算子子图手工实现对应融合算子;
- 将融合算子注册到AI框架,并在AI框架中增加对应的优化pass,将匹配的算子子图替换为融合算子节点。
这种方式的缺点是显而易见的。因为它只能针对若干热点场景进行融合,所以是无法做到通用和泛化的。基于以上问题背景,并结合MindSpore自身需求,我们提出图算融合解决方案。其主要思路是:
- 以通用的pattern识别计算图中的融合场景,并生成融合算子子图;
- 将生成的融合算子子图通过自动算子编译技术(AKG)生成对应的融合算子。
通过图算融合之后,用户在进行网络定义时,不管使用组合表达能力更强的基本算子,还是复合算子。最终都将通过图算融合重新组织和融合,并替换为最优复合粒度的融合算子。
2.设计初心
2.1图算融合概述
图算融合是MindSpore特有的网络性能优化技术。它可以通过自动分析和优化现有网络计算图逻辑,并结合目标硬件能力,对计算图进行计算化简和替代、算子拆分和融合、算子特例化编译等优化,以提升设备计算资源利用率,实现对网络性能的整体优化。相比传统优化技术,图算融合具有多算子跨边界联合优化、与算子编译跨层协同、基于Polyhedral的算子即时编译等独特优势。另外,图算融合只需要用户打开对应配置后,整个优化过程即可自动完成,不需要网络开发人员进行其它额外感知,使得用户可以聚焦网络算法实现。
图算融合的适用场景包括:
- 对网络执行时间具有较高性能要求的场景;
- 通过拼接基本算子实现自定义组合算子,并希望对这些基本算子进行自动融合,以提升自定义组合算子性能的场景。
MindSpore的三大特性:
开发态友好:
- 自动微分,网络+算子统一编程,函数式/算法原生表达,反向网络算子自动生成
- 自动并行,模型自动切分实现最优效率的模型并行
- 自动调优,动态图+静态图同套代码
运行态高效:
- On-Device执行,充分发挥异腾大算力
- Pipeline优化,最大化并行线性度
- 深度图优化,自适应硬件算力和精度
部署态灵活:
- 端-边-云按需协同计算,更好的保护隐私
- 端-边-云统一架构,实现按需部署
2.2图算融合的设计初心
图算融合的设计初心就是为了极简的算子表达和高效的自动融合,以下是设计初心的具体表述:
- 统一前端表达:当前算子与图层没有统一的前端表达,用户新增算子需要额外学习TBE DSL,门槛较高。希望提供种通过ME原语在Python像写网络一样定义算子的能力。
- 统一中间表达:当前算子与图层没有统一的中间表达,框架采取的均是提前写好硬件实现的算子并向框架注册算子。对于MSIR而言,每个算子的内部细节是不可知的,这样会丧失部分在图层进行细粒度跨算子编译优化的机会。同时也使得算子无法复用图层的通用优化。
- 泛化融合能力:当前框架采用的均是采用手写融合pattern,手写融合算子人工注册的方式实现算子融合,这样的方式效率低,需要消耗大量人力(~天/pattern)且不具备面尚用户多样化网络需求的泛化能力。因此需要提供可自动形成融合pattern和自动生成高性能融合算子的能力(~分/pattern)。
3.关键特性简介
3.1统一图算前端表达
前端:以定义原语的方式注册算子
后端需要:算子信息、算子DSL
使用图算方式定义算子
只需描述算子的计算意图,所有性能优化的问题(format、schedule、tiling、fusion)都交给编译框架。让定义算子这种让人秃头的工作也变得简单生动起来。
3.2统一图算中间表达:算子白盒化
在IR中原黑盒算子原语被函数子图替
打破算子和图层的信息边界
- 图层知道算子间的依赖关系
- 图层知道算子内的计算逻辑
产生更多优化机会
- 算子可复用图层的通用pass
- 跨算子的细粒度的融合优化
3.3复用图层AD算法
3.3.1自动生成反向算子
因为算子和图层统一使用MSIR表达,因此可以利用框架已有的自动求导的能力,对白盒算子进行自动求导。开发人员不用手动编写反向算子。
3.3.2正反向算子的自动融合
深度学习神经网络中算子的正反向经常有些相同的子计算,算子使用白盒表达后,就可以在图层自动分析消除正反算子中的相同的子计算。
自动发现细粒度的算子融合机会,更高效的利用片上内存,减少外存数据的搬移。3.4 Bert上的图算融合
在bert网络中虽然已经有很多定制的网络的pattern产生了定制融合算子,但是依然还有许多散落的基础黑盒算子,根据我们的融合算法产生了如下的融合pattern:
➢白加黑融合pattern:
LambNextMV + Assign |
398 times |
LambUpdateWithLR + Assign |
398 times |
BiasAddGrad + Cast + Mul |
146 times |
|
|
➢黑加黑融合pattern:
RealDiv + Reshape + Assign |
402 times |
||
Mul + TensorAdd |
150 times |
Cast + TransData |
152 times |
Cast + Mul |
48 tims |
ReduceSum + Cast + Mul |
48 times |
Mul + ReduceSum |
3 times |
ReduceSum + Mul |
1 times |
TensorAdd + TensorAdd |
1 times |
Mul + ReduceSum + Cast + Mul 1 times |
1 times |
ReduceSum + TensorAdd |
1 times |
Cast + ExpandDim+Cast+Sub+Mul |
1 times |
Cast + RealDiv + Minimum + RealDiv + Sub + Pow + Mul + Sub + Mul +TensorAdd + Mul + Mul + TensorAdd |
1 times |
➢融合后基础算子个数变化:
Mul:413->160 |
RealDiv:403->2 |
Assign:1609->13 |
Cast:649->250 |
ReduceSum:56->2 |
|
4.启用方法及编译流程演示
4.1使用方法
- MindSpore中的图算融合优化分布于网络图层编译和执行的多个步骤中,默认关闭状态,我们可以在训练脚本中为context指定参数enable_graph_kernel=True从而启用图算融合
-
from mindspore import context context.set context(enable graph kernel=True)
- 通过定义继承与GraphKernel的子类定义图算融合的算子
4.1.1样例脚本
为了说明图算融合优化场景,构造了一个简单网络MyNet,包含一个乘法和加法计算。
在打开图算融合进行优化之后,这两个计算便会自动合成一个融合算子:
import numpy as np
import mindspore.context as context
from mindspore import Tensor
from mindspore.nn import Cell
import mindspore.ops as ops
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
# save graph ir to view fusion detail.
context.set_context(save_graphs=True)
# enable graph kernel optimization.
context.set_context(enable_graph_kernel=True)
class MyNet(Cell):
def __init__(self):
super(MyNet, self).__init__()
self.add = ops.Add()
self.mul = ops.Mul()
def construct(self, x):
a = self.mul(x, 2.0)
res = self.add(a, 1.0)
return res
x = np.ones((4, 4)).astype(np.float32) * 0.5
net = MyNet()
result = net(Tensor(x))
print("result: {}".format(result))
输出结果:
result: [[2. 2. 2. 2.]
[2. 2. 2. 2.]
[2. 2. 2. 2.]
[2. 2. 2. 2.]]
该计算图的融合结果如图1所示,其中左图为未使能图算融合时的对应计算图,右图为使能图算融合后的对应计算图。可以看到该网络中的加法和乘法被融合成一个算子。该融合过程可以通过查看中间IR,或者通过Profiling等工具跟踪算子执行过程进行验证。
4.2自定义组合算子
基于图算融合技术,用户可以很方便地实现高性能的自定义组合算子。相比其它自定义算子方式,这种方式具有对框架无侵入、简单易用等优点。
其主要流程为:
(1) 在脚本中用基本算子组合的方式实现自定义算子定义和使用;
(2) 打开图算融合配置;
(3) 图算融合对自定义组合算子中的基本算子自动进行算子融合,并生成高性能融合算子。
4.2.1样例脚本
构造一个简单网络MyNet,并在其中使用了自定义算子MyOp。
代码样例如下:
import numpy as np
import mindspore.context as context
from mindspore import Tensor
from mindspore.nn import Cell
import mindspore.ops.operations as P
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
# enable graph kernel optimization.
context.set_context(enable_graph_kernel=True)
class MyOp(Cell):
""" my first custom OP composited by basic OPs """
def __init__(self):
super(MyOp, self).__init__()
self.sub = P.Sub()
self.mul = P.Mul()
def construct(self, x, y):
a = self.sub(x, y)
return self.mul(a, x)
class MyNet(Cell):
def __init__(self):
super(MyNet, self).__init__()
self.mul = P.Mul()
self.pow = P.Pow()
self.my_op = MyOp()
def construct(self, x, y):
a = self.mul(x, 2.0)
b = self.pow(a, 3.0)
res = self.my_op(b, y)
return res
x = np.ones((4, 4)).astype(np.float32) * 0.2
y = np.ones((4, 4)).astype(np.float32) * 0.3
net = MyNet()
result = net(Tensor(x), Tensor(y))
print("result: {}".format(result))
输出结果:
result: [[-0.015104 -0.015104 -0.015104 -0.015104]
[-0.015104 -0.015104 -0.015104 -0.015104]
[-0.015104 -0.015104 -0.015104 -0.015104]
[-0.015104 -0.015104 -0.015104 -0.015104]]
该计算图的融合结果如图2所示,其中左图为未使能图算融合时的对应计算图,右图为使能图算融合后的对应计算图。可以看到不仅自定义算子MyOp中的基本算子进行了融合,并且与主图中的其他算子也进行了更大范围融合。该融合过程可以通过查看中间IR,或者通过Profiling等工具跟踪算子执行过程进行验
5.实现过程小结
5.1图算融合流程图
5.2融合算子自动生成流程
5.3小结
图算融合提供了一种简洁的算子前端表达方式。使用户可以在Python像写网络一样定义算子,仅需描述计算诉求,无需考虑性能优化问题。
- 定义算子效率提升10X,代码量:百行—>十行
图算融合提供了图层和算子的统的中间表达方式。使得算子可以复用图层的通用优化,同时打破了图层与算子之间的信息边界,创造出跨算子的优化机会。
- 自动生成反向算子
- 自动实现跨算子优化
图算融合提供了一种泛化的算子融合能力。对多样化的用户网络可提供种通用的融合能力,更充分地释放AI算力。
- 融合算子效率提升100X,人工定义—>自动发现、自动融合
6.总结
在MindSpore中,通过白盒化算子,初步打开了图算这一扇门。既提供了图算一体的DSL表达能力,又打开了图算混合优化的大门,借此打破了算子边界。可能多年以后,当AI遍布生活生产的每个角落,MindSpore这样的深度学习框架会成为一种必不可少,且每个人都能用的AI操作系统,它也就成为我们生活的一部分。
MindSpore致力于AI开发生态的繁荣,开源开放可扩展的架构,有助力于开发者灵活扩展第三方框架、第三方芯片支持能力,让开发者实现各种定制化需求。MindSpore将在门户网站、开源社区提供更多学习资源、支持与服务。
我相信未来,也相信MindSpore!
本文参与华为云社区【内容共创】活动第19期。
https://bbs.huaweicloud.com/blogs/370132
任务26:MindSpore图算融合,打破算子边界
- 点赞
- 收藏
- 关注作者
评论(0)