MegEngine 动态执行引擎 Imperative Runtime 架构解析

2023-06-16 22:58:31

来源：哔哩哔哩

在之前的文章中我们介绍过 MegEngine的 Imperative Runtime以及它与 MegBrain、MegDNN的关系，这篇文章中我们将介绍 Imperative中包含的常用组件。

在 MegEngine中，从用户在 python层编写代码到在 interpreter层发生计算经过了下面的流程：

(资料图)

用户在 python层编写网络结构代码，执行时向 C++层发射算子执行指令

Imperative的 dispatcher对部分算子做计算前的预处理（transformation）

Imperative的 interpreter执行计算操作（复用 MegBrain的相关组件）

我们将分别介绍这几个阶段系统所做的工作。

MegEngine 的 Python 层

在主流的深度学习框架中，用户往往不需要自己手写算子的具体实现、处理计算图的执行逻辑、或者与复杂的体系结构打交道。一切都被封装为 Python层的接口。

在 MegEngine的 Python层中用户接触较多的模块主要有：data、functional、module、optimizer、quantization、tools，下面简单介绍一下各个模块的功能。

构建数据处理 Pipeline —— data 模块

Data模块，顾名思义就是对数据进行处理的模块。

没有数据就没法训练，在 MegEngine中，通常会借助一个 Dataset 结构来定义数据集。数据集一般分为 Map-stype和 Iterable-style两种。前者叫作 ArrayDataset，这种数据集支持随机访问；后者叫作 StreamDataset，因为是流式的数据集，只支持顺序访问。

有了数据集，我们还需要一个结构来把数据“喂”给模型训练，这样的一个结构叫作 dataloader。

实际上，只给 dataloader一个 dataset有时无法准确地描述加载数据的整个过程，我们可能还需要定义加载数据过程的抽样规则（Sampler），或者定义一些数据变换的规则（Transform），或者是定义抽样后的数据的合并策略（Collator）。

Python 层计算接口 —— functional 模块

深度学习模型通常包含一些基础的计算操作，比如 convolution、pooling等，在 python 层，这些基本计算操作都定义在 functional模块中。

functional中实现了各类计算函数，包含对很多 op的封装，供实现模型时调用。

模型结构的小型封装版本 —— module 模块

使用 functional提供的接口已经足够编写神经网络模型的代码，但随着模型结构的复杂程度加深，多次反复编写相似的结构会使开发和维护成本迅速提高。

考虑到神经网络模型通常是由各种层（layer）组成，我们通常使用 Module来封装模型的部分结构或者层，用户实现算法时往往使用组合 Module的方式搭建模型计算的 pipeline。定义神经网络时有些结构经常在模型中反复使用，将这样的结构封装为一个 Module，既可以减少重复代码也降低了复杂模型编码的难度。

使用 optimizer 模块优化参数

MegEngine的 optimizer模块中实现了大量的优化算法，同时为用户提供了包括 SGD、Adam在内的常见优化器实现。这些优化器能够基于参数的梯度信息，按照算法所定义的策略对参数执行更新。

降低模型内存占用利器 —— quantization 模块

量化是一种对深度学习模型参数进行压缩以降低计算量的技术。它基于这样一种思想：神经网络是一个近似计算模型，不需要其中每个计算过程的绝对的精确。因此在某些情况下可以把需要较多比特存储的模型参数转为使用较少比特存储，而不影响模型的精度。

MegEngine 相关工具汇总 —— tools 模块

用户进行开发时有时需要一些工具进行错误调试或者性能调优，tools下就提供了一些这样的工具。比如对训练程序进行记录并在浏览器上可视化的 profiler、方便用户查看 MegEngine显存占用的 svg_viewer 等。

一般来说，用户会基于上面的模块搭建算法模型，其中定义了非常多的 op 的计算过程，下面我们看一下 c++ 是怎么进行这些 op 的真正的计算的。

Dispatcher 会对 op 做哪些处理？

从 Python层往下的部分用户往往是感知不到的，脱离了“前端”，我们抽丝剥茧，进入到了框架“后端”对 tensor和 op处理的细节。

前面我们提到在 functional模块中封装了很多算子，并以 python接口的形式提供。实际上这些算子需要向下发射指令对 tensor进行操作并返回操作完成后的 tensor，这些发射的 op指令就会到 dispatch层，在进行实际计算之前，dispatcher会对 tensor做一些处理，我们把这些处理叫作 Transformation。

在 imperative中真正执行算子进行计算是在 interpreter层做的，与 tensor处理相关的操作被解耦出来放在 dispatch层，这样更便于维护。

在 MegEngine中，一些重要的 transformation有：

DimExpansionTransformation：某些 op计算时对输入 tensor的 shape有要求，在这里做处理。

DtypePromoteTransformation：某些 op要求计算的 tensor拥有相同的类型，会将所有的输入的类型提升为同一类型之后再进行计算。比如 int类型 tensor和 float类型 tensor进行计算，需要把 int类型的 tensor转换为 float类型 tensor。

InterpreterTransformation：顾名思义，这类 Transformation将指令转发到 Interpreter层（Interpreter可以认为是 Imperative中所有计算操作的入口）进行计算，并获取指令的计算结果。Transformation通常是叠加的，InterpreterTransformation是最后一层，其后不再跟其他的 Transformation处理。

FormatTransformation：由于在不同情况下对不同 format的 Tensor的计算速度不同，因此需要对 NHWC和 NCHW的 Tensor进行转换，为了不让用户感知到这样的转换，这部分的工作由 FormatTransformation完成。

GradTransformation：训练模型时需要通过反向传播更新模型参数，反向传播需要支持 op的自动微分。要实现求导，就需要在前向执行 op的时候记录某些信息，以便之后进行反向求导。Autodiff算法会根据输入的前向图生成一个完整的前向反向图，所谓的前传反传训练过程对 Autodiff来说实际上都是一个计算图的前向过程，grad的数值是在“前向”的过程中就已经拿到的。GradTransformation处理的就是与反向求导相关的操作。

TracingTransformation：

在介绍 Trace之前，我们需要先明确一下计算图的概念。计算图可以认为是对输入的数据（tensor）、op以及 op执行的顺序的表示。计算图分为动态图和静态图。动态图是在前向过程中创建、反向过程销毁的。前向逻辑本身是可变的，所以执行流程也是可变的（因此叫动态图），而静态图的执行流程是固定的。也就是说，动态图在底层是没有严格的图的概念的（或者说这个图本身一直随执行流程变化）。对于动态图来说，graph的 node对应的概念是 function/ 算子，而 edge对应的概念是 tensor，所以在图中需要记录的是 graph中 node和 edge之间的连接关系，以及 tensor是 function的第几个输入参数。

Trace的作用就是将动态图执行转换为静态图执行，这样做的好处就是执行速度更快了，并且占用的显存更少了。因为静态图需要先构建再运行，可以在运行前对图结构进行优化（融合算子、常数折叠等），而且只需要构建一次（除非图结构发生变化）。而动态图是在运行时构建的，既不好优化还会占用较多显存。

Trace中所有的东西都会进行静态优化（加速）。

加了 Trace之后，模型在训练时第一个 iter是动态图执行，Trace会记录下 tensor、op以及 op的执行顺序这些信息（构建静态图）并进行计算，在第二个 iter就跑的是构建好的静态图。

LazyEvalTransformation：类似 TracingTransformation，也会记录 tensor、op等信息构建静态图，不同的是 LazyEvalTransformation在第一个 iter不会跑动态图，但会在第二个 iter开始跑静态图。

ScalarTransformation：用于判断指令的输出是否为 scalar。因为 dispatch的 Tensor要发到 Interpreter层，而 Interpreter层不接受 ndim == 0的 Tensor（在 Interpreter中 ndim为 0表示 Tensor的 shape未知），也就是一个 scalar，因此 ScalarTransformation会将 ndim为 0的 Tensor表示为 ndim不为 0的 Tensor（具体是多少与具体 op有关）发往 Interpreter。

不同的 Transformation之间拥有固定的执行顺序：比如 InterpreterTransformation是执行实际计算并获取计算结果的（需要进入 Interpreter），所以它是在最后一个执行的。TracingTransformation/ LazyEvalTransformation/ CompiledTransformation等属于 Trace相关的操作，因为 Trace需要记录所有指令，所以这些 Transformation是在倒数第二层执行的。如 ScalarTransformation这样只对 Scalar做处理的 Transformation往往在较上层。

因为不同的 Transformation有逻辑上的先后关系，所以开发者往往需要手动规划它们之间的顺序。

不同类型的 Transformation之间是解耦的，这样便于开发与维护。

Interpreter 是如何“解释”算子的？

由于 MegBrain已经是一个非常成熟的静态图框架，因此在开发动态图（Imperative Runtime）深度学习框架 MegEngine的过程中，复用许多静态图中的组件可以大大降低开发成本。

实际上，张量解释器 Tensor Interpreter就是将动态图中的操作——如执行 op、shape推导等操作“解释”为静态图的对应操作，并复用 MegBrain的组件来运行。

这里我们需要先了解一个 MegBrain的静态图“长什么样”。

复用静态图接口的机制 —— proxy_graph

为了复用 MegBrain的静态求导器、静态内存分配器、静态 shape推导器等组件，imperative引入了 proxy_graph。

复用 MegBrain的接口需要实现对应的方法，在 MegEngine/imperative/src/include/megbrain/imperative 目录下可以看到所有需要实现的桥接接口，其中和 proxy_graph相关的接口声明在 proxy_graph_detail.h 中，通常需要实现这几个接口：

infer_output_attrs_fallible

复用 MegBrain的 StaticInferManager进行 shape推导，在执行计算操作前对输入和输出 tensor的 shape进行检查。

apply_on_physical_tensor

根据 infer_output_attrs_fallible推导的 shape结果去分配 op输出的显存，并调用 proxy opr的 execute函数（会转发到 MegDNN的 exec函数）执行计算操作。

make_backward_graph

在求导时，Grad Manager会记录下来一些求导需要的信息（输入 tensor、op以及它们执行的顺序、输出 tensor），make_backward_graph会根据这些信息造一个反向的计算图，供求导使用。

get_input_layout_constraint

一般用来判断一个输入 tensor的 layout是否满足一些限制：比如判断 tensor是否是连续的。

如果不满足限制，则会造一个满足限制的 tensor，供 apply_on_physical_tensor使用。

在实现一个 imperative算子时通常也只需要实现这几个接口，剩下的工作由 MegBrain和 MegDNN完成。

主流框架在 python层的模块封装结构大同小异，关于 MegEngine的 Python层各模块的使用与实现细节以及 transformation和 interpreter实现细节我们会在之后的文章中逐一解析。