AI框架动静态图统一的思考

2021/12/10 14:32
阅读数 1W

作者:金雪锋

作者知乎https://www.zhihu.com/people/jin-xue-feng

文章来源https://zhuanlan.zhihu.com/p/416643687

 

发现好久没更新AI框架分析的专栏,确实有点对不住当初立的flag,赶紧先上一篇,内容准备的还有些匆忙,后面有时间再优化和细化。

本文是AI框架分析专栏的第七篇,总体目录参见:AI框架的演进趋势和MindSpore的构想

MindSpore在第一个版本发布的时候,就开始构建动静态图统一的特性,动静统一的目标非常美好,但是实现起来其实非常困难;本文试图从动静图统一的挑战、趋势以及MindSpore的实践三个维度出发,给读者一些参考。

一、动静态图的发展趋势:

1、什么是动静态图

动态图:AI框架的动态图基本上是指类似PyTorch的实现(非torchscript),原理上比较简单,使用Python的c extension机制注册算子(算子一般性能比较高,采用native的语言开发),正向利用Python进行解释执行,同时通过Tape机制(自动微分)生成反向图,然后基于反向图进行梯度更新,反向执行其实也是图的一种模式,不过不会做图的编译优化。

静态图:与动态图相比,静态图在执行前有一个构图和编译优化的过程,即执行前先生成正反向图并完成编译优化,然后再执行。那如何从Python的神经网络表达翻译到图呢?目前主要有两种方式:

  • Tracing模式:框架把Python代码假执行一遍,记住算子执行序列(Tensor相关操作),作为正向图,并以此进行自动微分生成反向图,并进行正反向图的编译优化,最后以图模式进行执行
  • AST转换:框架获取Python代码的AST,然后通过编译的技术转换成正向图(这里面有Parser、Resolver、Specialize、Optimize等),并基于此生成反向图,同时再进行编译优化,最后以图模式进行执行。

聊一聊AI框架前端

动静态图两者各有优缺点:

动态图:灵活易用,基本上Python的语法都支持,对于动态模型,比如动态shape、有复杂控制流等,尤其有利;当然它的缺点也比较明显,就是性能和部署,动态图的执行性能基本上取决于单算子的性能,缺乏算子间的融合优化,无法充分发挥AI芯片的算力(这个对性能快速提升的AI芯片来说不是个好事情),同时AI在部署的时候,为了性能和功耗(比如端侧部署),一般要需要生成一个部署模型,部署模型本质是静态图的表达,所以这里就存在一个动静态转换的问题,太灵活的动态图训练完后,想转换变成一个部署模型(静态图)同样存在很大的挑战。

静态图:优缺点和动态图刚好相反,表达上有许多限制,不灵活;但是性能好(适合面向芯片的编译优化、适合分布式并行优化),部署能力强。

另外,tracing based和ast based两种静态图模式在表达上/编译优化上也有少许差异,比如tracing based的方法,可以充分利用Python的执行完成推导功能(实现比较完整的Python复杂的数据结构和动态类型到Tensor的映射),问题是很难处理控制流并丢失了scope的信息;ast based的方法原理上功能可以做的比较强大,但是在python语法的翻译需要比较全量的翻译,工程量比较大/方案比较复杂。

大家一个朴素的想法就是我们是否可以做到动静态图灵活转换,从而能同时拿到两者的优点。比如在网络调试或者网络研究的时候,使用动态图,工作效率比较高;在生产环境上,无缝切换到静态图,性能高,部署快。

但是动静态图的无缝转换实际上是一件非常难的事情。

2、动静态图转换的挑战

如果把静态图的表达看成是一种特殊的DSL,那这种DSL实际上是一种静态语言,而Python是一个解释性、动态类型的语言,从本质上讲它是不可能完全无损的转换到一种静态语言。不过AI框架是一种面向领域的框架,其业务特征是有比较固定的范式的,比如以tensor计算和自动微分为中心,最终的计算都会形成tensor流,这一定程度上可以减少动静图转换的难度,但是挑战依然存在,下面分情况来分析:

tracing based的静态图:这种方式拿到的图实际上是平铺的一个执行流,所以很难处理控制流的情况,比如循环,对于tracing来说就是展开,但是有些情况下循环无法展开,如循环条件根据训练的收敛情况/算子的执行结果等。

ast based的静态图:这种方式拿到的图是从python的ast转换而来,好处是控制流和scope信息都可以保留,但是挑战是那么多python的语法和数据结构都要转换到静态图的表达,更难的是python是动态类型的,所以这个ast到静态图的转换中需要一个复杂的类型/值推导过程。

3、动静态图发展的三个阶段

总的来说,个人认为,动静图的发展分为三个阶段:

动静态图分离——>动静图结合——>动静图统一;当前框架主要处于第二阶段。

动静图结合:朴素的想法是,既然动态图不容易转换到静态图,能不能让开发者自己标识模型中的哪部分Layer或者代码需要加速并且可以实现动静转换。基本的实现方式是在需要静态图的代码块上加上装饰符,以MindSpore为例:

@ms_function

def tensor_add_with_dec(x, y):

     z = x + y

     return z

框架会将ms_function修饰的函数进行静态图处理。虽然动静结合是当前框架的主流模式,但是动静图结合的方式对开发者并不友好,需要开发者自己去判断哪里可以转成静态图,意味着开发者能够清楚两个事情:一是哪些地方可以加速;二是哪些代码块可以转成静态图(代码块里面代码的语法符合静态图的约束)。

动静图统一:理想的方式就是动静统一,开发者可以灵活的进行动静态图的切换,这里面可能有两条路径,一条路径是从动态图出发,框架可以在运行过程中自动的JIT,无需用户用修饰符指定,主要的关键技术就是Lazy Tensor;另外一条路径是从静态图出发,在编译过程中如果发现有不支持的语法,可以保留到运行时进行fallback到python,主要的关键技术就是jit fallback。

 

二、动静统一的关键技术

1、LazyTensor

LazyTensor思想也比较简单,该方法建立在动态图的异步执行基础上,遇到tensor的操作(主要是算子),如果不需要查看tensor的具体内容,框架可以把这算子缓存下来,不发到device执行,从深度学习的范式看,大概率会缓存一大串op的执行序列,缓存到一定程度,可以把这个序列进行JIT编译优化,然后通过图的方式执行。

当然这里tensor相关的api需要分为两类,一种是可以被表达成IR图的API以及另一种无法被表达成IR的API。任意一个返回一个或多个Tensor的API,都是可以被转换成IR图,而返回非Tensor类型的API则无法被转换成IR图。这样Layz Tensor机制就可以根据这个标识来确定哪些op序列可以进行JIT优化。

Lazy Tensor机制本质是tracing based的静态图的一个演进,相当于自动打装饰符。

LazyTensor机制的好处:

  • 用户无感,易用性好
  • 语法限制少,理论上只要找到合适的tensor操作序列都可以进行加速。

LazyTensor机制的局限:

  • JIT编译的开销比较大,所以最好是首次的JIT编译结果能够缓存,方便下次重复使用,用在静态网络下比较合适。
  • 动态情况下,比如动态shape/控制流,每次变化都会触发新的编译,从性能上看得不偿失。

LazyTensor的细节可以参考下面的论文分析,同时国内清华大学的计图框架在这一块已经做了非常有意思的尝试。

论文分析-动静态图结合的一种方案:LazyTenor-combining eager execution with domain-specific Compiler

2、JIT fallback

JIT fallback是用静态图的角度出发来考虑动静图的统一,希望静态图能够尽量多的支持动态图的语法,其思路借鉴了传统JIT编译的fallback的思路,传统的JIT编译经常会通过profiling信息,对函数进行多态选择、value推导、分支调度等优化,同时设置guard条件,一旦guard条件发现情况有变,可以去JIT优化,回到原来未优化的函数进行解释执行。

对于AI框架的JIT fallback没有那么复杂,在静态图编译的时候(一般JIT fallback是基于ast based的静态图),如果发现是编译器不支持的Python语法,可以把相关语句保留下来,生成解释节点,然后在后面的处理中,fallback到python去执行相关的语句,从而实现相关语法的支持,这里有几个难点:

1)不支持的Python语法的识别

2)解释节点的推导和执行,解释节点有两个运行时机:首先是编译期的推导,一般而言,解释节点尽量在这个实际执行;其次是运行期的执行。

 

三、MindSpore的实践

MindSpore目前已经支持动静结合的方式,正在完善动静统一的方式。

动静结合的方式:主要是通过ms_function这个装饰符来实现。

动静统一:MindSpore可以通过set_context来实现动静态图的一键式切换,但是不可否认当前确实存在不少动态图无法转到静态图的语法,MindSpore正在做JIT fallback的工作,预计1.6版本先实现编译推导期的fallback。至于LazyTensor的方式,目前看还很难解决动态的问题,还需要进一步探索。

展开阅读全文
打赏
1
1 收藏
分享
加载中
牛,这么好的文章,必须顶一下,传统软件企业要突破行业增速,单靠产品云化概念不能解决增长问题。只有向SaaS转型,才是实现增长的根本途径。做saas这个做的不错,分享一下开源地址: https://www.wei-it.com/
01/10 21:40
回复
举报
更多评论
打赏
1 评论
1 收藏
1
分享
返回顶部
顶部