随着大模型时代的到来,AI算法的能力上限不断被刷新,算力门槛也在持续飙升,如何在有限的计算资源的条件下优化和训练模型显得尤其重要。面向计算图编译和运行优化场景,昇腾AI异构计算架构CANN(Compute Architecture for Neural Networks)开放GE(Graph Engine)图引擎,通过计算图优化、计算图下沉、内存复用和多流水并行等技术可将计算图执行性能提升20%,同时开放图引擎接口支持开发者自定义图结构,获得最优执行性能。
基于GE图引擎能力,昇腾已支持TensorFlow计算图模式,并原生支持PyTorch框架。开发者只需要使用PyTorch原生的torch.compile接口,昇腾AI处理器后端就会对PyTorch生成的计算图进行接管、转化为AIR,再进行端到端的图编译深度优化,降低内存需求、提升计算性能,同时最大程度减少开发者的修改工作。接下来,我们将分几期内容深度解读CANN的GE图引擎的诸多开放能力和关键技术,从而帮助开发者优化算法编译和运行效率。
1 Eager模式与图模式
当前业界主流的深度学习框架(例如PyTorch、TensorFlow等)都提供了Eager(Eager Execution,即时执行)模式与图模式。Eager模式的特点是每个计算操作下发后立即执行,图模式的特点是所有计算操作构造成一张图,以图的粒度下发执行。Eager模式的痛点在于没有图的全局视角,一些融合、优化手段无法高效开展。成图以后,编译器的视角更广,计算操作可以更好地化简、优化,从而获得更好的执行性能。
GE针对图进行了系列优化,包括通用的图优化技术“公共子表达式消除”、“剪枝”、“死边消除”,以及特有的Shape优化技术、内存优化技术等。下面对GE的图优化技术详细展开介绍。
2 GE支持通用图优化技术
2.1 常量折叠
常量折叠的核心为在编译阶段直接计算并替换常量表达式的值,从而减少运行时的计算负担。例如,表达式“2+3”会在编译时被直接替换为“5”,从而会减少图中节点的数量,这种优化有助于减小模型体积,提高执行效率和性能。
一般来说,常量折叠的处理方式为:若某算子输入都是常量且可进行折叠,折叠后该算子与其Const输入都融合为一个Const,如下图所示。
除此之外,GE还支持如下几类常量折叠相关优化。
2.2 公共子表达式消除
公共子表达式消除识别图中相同的子表达式,将其合并为一份公共表达式,以便在后续计算时直接引用,而不需要重新计算。这样可以减少程序中重复计算的次数,提高程序的执行效率。
2.3 剪枝
GE支持在编译前指定输出节点,只保留对输出数据有贡献的节点,从而缩小模型大小,提升编译性能。如下图所示,若指定输出节点为NetOutput,则仅保留右边紫色底纹的小图。
2.4 死边消除
当图中存在条件控制算子(如If、Case等)时,会根据其输入中的条件(cond)判断走哪个分支。若cond在编译时已确定,即可明确所走的分支,此时可将不会走到的分支删除,从而减少算子编译、图编译耗时,这种删除冗余分支的情况称之为死边消除。
如上图所示,若cond输入为Const,则已经可以确定走右侧的then分支还是else分支。经过死边消除后的图如下所示。其中PartitionedCall代表一个局部调用,用于将子图调度起来。
3 GE特有的增强图优化技术
3.1 Shape优化技术
当整图或局部图Shape为静态时,GE有较大空间可以减少编译耗时,提升调度性能。比如使能更多的算子融合,使能更多的图裁剪优化,以及使能模型下沉调度,其性能相对于Host调度有较大优势,这个后续会有详细介绍。因此GE在增强图优化技术中,会通过下面介绍的各种技术手段将计算图shape优化成静态Shape。
3.1.1 常量折叠与InferShape协同
对于GE来说,除了常规的常量折叠外,还支持对Shape确定的算子进行折叠。在图模式下,Shape推导可以和常量折叠交替进行,直至无可常量折叠的节点。先看下不做协同技术的优化结果,假设先做InferShape再做常量折叠,Gather的indices输入为[1,0,2,3]):
上图的逻辑为对Data取Shape,用于对Data做Reshape操作。
若常量折叠与InferShape协同工作,接上述图示,InferShape再次执行,如下图所示。
3.1.2 对循环算子While的Shape推导
While算子的语义为按某种条件,将body子图中的算子循环计算。但因编译时循环条件往往并不确定,因此While算子的输出shape通常为动态Shape中的Unknown Rank(body中若有维度调整类算子,可能导致输出维度不确定)。
在GE中,对于While算子会多次推导body子图,将上一次推导的子图输出Shape作为子图输入Shape再次推导,直到两次推导的输出Shape一致,则认为Shape已经推导稳定,并将最后一次推导结果作为While算子的输出Shape。这种处理策略的好处是尽可能推导出更多信息。
例如,while算子的输出可能是静态Shape,这样可以促成模型走下沉调度,提升执行性能;也可能是动态Shape中的Unknown Shape(知道维度信息,但不知道轴信息),此时维度信息也是后续Shape推导的重要信息。
3.1.3 动态分档
在模型推理的场景中,有时候模型输入的数据批次大小BatchSize和图像分辨率大小ImageSize是无法固定的。比如智能交通的应用中,需要视频结构化的应用框架,往往需要先做人机非目标检测,再针对每一个行人、机动车、非机动车目标做相对应的属性识别。这类应用场景中,检测模型往往输入的个数不固定,检测到的目标分辨率大小不固定,因此相应的属性识别模型的输入也就不固定。此类场景对于BatchSize和ImageSize有Shape动态变化的需求,但通常其Shape是有变化规律的,如BatchSize一般有倍数关系。
因此,GE提供了动态分档能力,对于shape可变的输入tensor,以若干离散的固定shape进行替代,每次指定其中一个固定shape,并输入该shape大小的数据进行推理。
动态分档模型的优势在于,虽然模型输入shape是动态的,但是用户只需要加载一次模型,通过传入不同动态档位参数就可实现动态需求,在执行时,模型以静态模型下沉的方式执行,从而获得较好的调度性能。
动态分档示意图
1) Data1是网络的原始输入,设置的shape值为最大档位值。
2) Data2作为第二个输入,自动添加到Graph中,表示真实的档位值;
3) 提供新算子MapIndex实现分档映射的作用,对于不同的BatchSize、ImageSize值,输出对应的index值(Scale值,0~N-1),作为case算子的index判断条件。
使用case算子,按照档位配置将原图复制为case的N个子图,根据步骤3的index选择相应Batch分支函数执行。
3.2 内存优化技术
Const节点一般用于保存权重,所需内存往往较大,为了降低Const在Host和Device上的内存占用,GE在图编译的过程中通过下述所示两个阶段对Const做去重。
1. 图优化过程中,通过内存比较的方式对相同value的Const仅保留其中一个。如下图所示,通过多级条件匹配,确定可以被删除的冗余Const。经过比对后,在图上可以删除冗余的Const,从而节省Host内存空间。
2. 经过步骤1后,还存在某些因图结构因素(例如融合后可能导致成环)或者Shape描述不同而无法被删除的冗余Const。因此,在结束图优化后,将Const节点的权重整块拷贝到Device之前,若经过二进制比对后发现某些Const节点的权重相同,则这些Const节点可以共享内存,从而节省Device的内存空间。
4 更多介绍
GE计算图优化的相关技术的介绍就到这里,欢迎大家关注后续技术分享。如需获取更多学习资源请登录昇腾社区(https://www.hiascend.com/)。