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

AI框架动静态图统一的思考 - 康康OSC的个人空间 - OSCHINA - 中文开源技术交流社区

作者：金雪锋

作者知乎：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的方式，目前看还很难解决动态的问题，还需要进一步探索。