这样使用昇思MindSpore大模型编译器，性能提升N倍、速度提升9倍以上

我们观察到深度学习的神经网络结构是由多层Layer组成，在大模型语言模型下，这些Layer都是Tranformer block的堆叠，尤其是开启流水线并行，各个micro batch的Layer层更是完全一样的。所以，我们想能否保留这些Layer结构，不Inline或者不提前Inline，这样就可以倍数级提升编译性能，比如，我们按照micro batch为边界，保留micro batch的子图结构，那么理论上编译时间可以变为原来的Y分之一（Y为micro batch的数量）。

具体到模型编写的代码，我们可以看到重复使用相同Layer的方式一般采用循环或者迭代调用的方式，Layer一般对应迭代过程中顺序结构的一项，往往是一个子图；也就是说，使用循环或迭代来多次调用相同的计算单元，如下面代码所示，block对应一个Layer或者micro batch子图。

class Block(nn.Cell):
    def __init__(self, config):
        .......

    def construct(self, x, attention_mask, layer_past):
        ......

class GPT_Model(nn.Cell): 
  def __init__(self, config):
      ......
      for i in range(config.num_layers):
          self.blocks.append(Block)
          ......
      self.num_layers = config.num_layers
   def construct(self, input_ids, input_mask, layer_past):
       ......
       present_layer = ()
       for i in range(self.num_layers):
           hidden_states, present = self.blocks[i](...)
           present_layer = present_layer + (present,)
       ......

因此，如果我们把循环体看作被频繁调用的子图，通过把它标记为Lazy Inline的方式，告知编译器推迟Inline处理，那么在编译的大部分阶段都可以获得性能收益。比如，神经网络循环调用同一个子图结构的时候，我们编译阶段，不展开子图；然后在编译最后再触发Inline优化，进行必要的优化和转换Pass处理。这样对于编译器来讲，大部分时间都是更小规模的代码，而不是已经Inline膨胀过的代码，从而大幅提升了编译性能。

具体实现的时候，可以在相关的Layer类上打类似@lazy-inline的标记，给编译器提示，打上标记的Layer不论是在循环体内被调用，还是其他方式被调用，在编译期间都不内联，直到执行前，才进行内联展开。

看上去，Lazy Inline的原理和思路并不复杂，但是现有AI图编译机制普遍不是那种支持完备编译特性的编译器，所以要实现这个功能还是很有挑战的。

幸运的是，MindSpore的图编译器在设计IR的时候，已经考虑了通用性，包括子函数的调用、闭包等特性。

① Cell 实例编译为可复用的计算图

Cell是MindSpore神经网络的基本构建块，是所有神经网络的基类，Cell可以是单个神经网络单元，如conv2d, relu, batch_norm等，也可以是构建网络的单元组合。在GRAPH_MODE(静态图模式)中，Cell将被编译成一个计算图。

当你需要自定义网络时，你需要继承Cell类并覆盖__init__和construct方法。Cell类覆盖了__call__方法。当调用Cell类实例时，将执行构造方法。在构造方法中定义网络结构。

在下面的示例中，构建一个简单的网络来实现卷积计算功能。网络中的操作符在__init__中定义，并在construct方法中使用。案例的网络结构为:Conv2d -> BiasAdd。

在构造方法中，x为输入数据，输出为网络结构计算后得到的结果。

import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import Parameter
from mindspore.common.initializer import initializer
from mindspore._extends import lazy_inline

class MyNet(nn.Cell):
     @lazy_inline
     def __init__(self, in_channels=10, out_channels=20, kernel_size=3):
         super(Net, self).__init__()
         self.conv2d = ops.Conv2D(out_channels, kernel_size)
         self.bias_add = ops.BiasAdd()
         self.weight = Parameter(initializer('normal', [out_channels, in_channels, kernel_size, kernel_size]))

      def construct(self, x):
         output = self.conv2d(x, self.weight)
         output = self.bias_add(output, self.bias)
         return output

@Lazy_Inline是Cell::__init__的装饰器，其作用是将__init__的所有参数生成Cell的cell_init_args属性值，self.cell_init_args = type(self).__name__ + str(arguments)。cell_init_args属性在MindSpore 编译中作为Cell 实例的唯一标识。cell_init_args值相同表明Cell 类名和初始化参数值是一样的。

construct(self, x) 定义网络结构，相同Cell类，网络结构取决于输入参数self和x。Self 包含weights 等Parameters，这些weights 是随机初始化或是训练后的结果，因而这些weights每个Cell的实例是不一样的。其他self属性由__init__参数确定，而__init__参数通过@lazy_inline计算得到Cell实例标识cell_init_args。因此Cell 实例编译计算图construct(self, x)转换为construct(x, self. cell_init_args,self.trainable_parameters() )。

如果同一个Cell类，且cell_init_args参数相同，我们把这些神经元实列叫做可复用的神经元实例，并这种神经元实例对应的计算图命名为可复用计算图reuse_construct(X，self. trainable_parameters())。由此推导出每个Cell 实例的计算图可以转换为：

def construct(self, x)
    Reuse_construct(x, self.trainable_parameters())

引入可复用计算图后，具有相同cell_init_args的神经元Cell(可复用计算图)只需构图和编译一次。如果网络中这样Cell 数量越多，有可能提升性能效果就越好。但事物都有两面性，这些Cell的计算图如果太小，太多，又会导致某些特性编译优化效果不好，例如算子融合，内存复用，整图下沉和多图调用等。

因此MindSpore 版本目前只支持手工标识哪些Cell 编译阶段生成复用计算图。后续版本将规划自动策略生成可复用计算图，如Cell 包含算子多少，Cell 被使用多少次等因素来权衡是否生成复用计算图，同时给出优化建议。

以下使用GPT结构进行抽象简化说明:

class Block(nn.Cell):
    @lazy_inline
    def __init__(self, config):
        .......

    def construct(self, x, attention_mask, layer_past):
        ......

class GPT_Model(nn.Cell):
    def __init__(self, config):
        ......
        for i in range(config.num_layers):
            self.blocks.append(Block(config, None))
            ......
        self.num_layers = config.num_layers

    def construct(self, input_ids, input_mask, layer_past):
        ......
        present_layer = ()
        for i in range(self.num_layers):
            hidden_states, present = self.blocks[i](...)
            present_layer = present_layer + (present,)
        ......

GPT 由多层Block构成，这些Block 的初始化参数都是同一Config，因而这些Block的结构都是一样，内部将被编译器转化为如下结构：

def Reuse_Block(x, attention_mask, layer_past,block_parameters) :
 ......
具体的Block 实例的计算图如下：
def construct(self, x, attention_mask, layer_past):
   return Reuse_Block(x, attention_mask, layer_past,
         self. trainable_parameters())

有了这个结构后，在编译过程的前半段作为独立的计算图不inline到整体的计算图中，只是最后少量Pass优化才Inline 到大的计算图中。

② Lazy Inline与自动微分/并行/重计算等特性的联合

采用Lazy Inline的方案后，对于原有流程会产生一些影响，需要进行相关的适配，主要是自动微分、并行和重计算。

对于自动微分，出现了类似call function的正向节点，需要提供微分处理；

对于并行流程，其中主要是Pipeline并行的pass处理需要对非整图场景进行适配，因为之前的Pipeline切图是依据整图切的，现在需要依据共享的子图去切。具体方案为先根据Stage着色，将共享Cell内的node按Stage切分，只保留当前进程对应Stage的node，并插入Send/ReCV算子，再将共享Cell外的node切分，保留当前进程对应Stage的node，同时将共享Cell内的Send/Recv算子拿到共享Cell外；

对于重计算流程，旧的重计算流程，是在Inline之后的整图上处理算子的，通过搜索重计算的连续算子块，按照用户的重计算配置，确定需要重计算的算子以及重计算的算子执行所要依赖的正向或反向算子。而Lazy Inline后，连续的重计算算子，可能处于不同的子图中，且正向节点与反向节点之间也找不到连接关系了，所以原先基于整图算子的搜索策略失效了。

我们的适配方案是，对于重计算的Cell或者算子，统一在自动微分之后处理。自动微分流程，对于Cell所生产的子图或者单个算子，都会产生一个闭包，该闭包返回正向输出及反向传播函数，并且我们还得到每个闭包跟原有正向部分的一个映射。通过这些信息，再根据用户的重计算配置，以每个闭包为基本单元，对Cell和算子统一处理，将原有正向部分拷贝回原图上，且依赖关系可以通过闭包中的反向传播函数获取，最终能实现一个不依赖于整图Inline的重计算方案。

③ 后端处理和影响

前端开启Lazy Inline后生成的IR下发给后端，后端需要对该IR进行切图，才能通过子图下沉的方式在设备上执行。但Lazy Inline后，子图下沉的方式执行还会出现一些问题，例如无法采用最优的方式进行内存复用和流分配，无法在编译时使用图内部的缓存进行编译加速和无法做跨图的优化（内存优化、通信融合、算子融合等）等问题。

为了实现最优性能，后端需要将Lazy Inline的IR处理成适合后端下沉执行的形式，主要做的是将自动微分产生的Partial算子转换为普通的子图调用，将捕获的变量变成普通的参数传入，这样就可以做到整图下沉执行整个网络。

在整图下沉流程中，这些call有两种处理方法：图上Inline和执行序Inline。图上Inline会导致图膨胀，后续编译速度变慢；但执行序Inline会导致内存复用的时候执行序Inline的部分内存生命周期特别长，最后内存不够用。

最终我们采用的处理方式为，在优化pass、算子选择、算子编译等流程中，复用执行序Inline的流程，让图规模尽可能的小，避免图节点过多影响图的后端编译时间。而在执行序优化、流分配、内存复用等流程之前，将这些call做实际的节点Inline，以获取最优的内存复用效果。此外，通过一些图Inline后的内存、通信优化、冗余计算消除等方式，可以做到在内存和性能方面都不产生劣化。

目前还无法做到所有的跨图级别的优化，只能单点识别放到Inline之后的阶段，而且无法节省执行序优化、流分配、内存复用的时间。

④ 实现效果

大模型编译性能优化通过Lazy Inline方案，提升编译性能3~8倍，同样以百亿大模型13B网络为例，应用Lazy Inline方案后，计算图编译规模从13万+节点下降到2万+个节点，编译时间从3个小时下降到20分钟，再配合编译结果缓存，整体效率大大提升。

⑤ 使用限制和下一步计划

1、Cell 由于以Cell的类名和__init__参数值生成Cell实例的标识。这是基于init的参数确定Cell 所有属性，以及construct 构图开始时的Cell 属性和init执行完的属性一致为假设前提，因此Cell与构图有关的属性，在init执行完后不能进行更改。

2、construct函数参数不能有默认值。现有MindSpore版本的对construct 函数参数如果有默认值，每次使用都特化为一个新的计算图；后续版本将对原有特化机制进行优化。

3、Cell有多个共享Cell_X 实例构成，同时每个Cell_X 又有多个共享Cell_Y实例构成。如果Cell_X 和 Cell_Y的init都被装饰为@lazy_inline，则只有最外层的Cell_X能被编译为重用的计算图，里层的Cell_Y的计算图还是Inline了；后续版本计划支持这种多层级的lazy inline机制。

如何辅助客户写出高内聚低耦合构图代码，也是昇思MindSpore框架追求目标之一，例如在使用中存在这种Block:: __init__ 参数包含Layer index,其他参数的值都相同，由于Layer Index 每一层都不一样，导致该Block由于细微的差别不能复用。如在某个GTP 版本代码中存在以下这样的代码：

class Block (nn.Cell):
    """
    Self-Attention module for each layer
    Args:
       config(GPTConfig): the config of network
       scale: scale factor for initialization
       layer_idx: current layer index
    """
    def __init__(self, config, scale=1.0, layer_idx=None):
        ......
        if layer_idx is not None:
            self.coeff = math.sqrt(layer_idx * math.sqrt(self.size_per_head))
            self.coeff = Tensor(self.coeff)
        ......
    def construct(self, x, attention_mask, layer_past=None):
        ......

为了使该Block能够复用，我们可以对其进行优化，把Layer Index 相关的计算提取出来计算，然后作为Construct的参数，输入到原来的构图中，从而使得Block的init参数相同了。

上述代码段修改为以下代码段,删除Init 与Layer Index 相关部分，在construct增加coeff 参数。

class Block (nn.Cell):
    def __init__(self, config, scale=1.0):
        ......
    def construct(self, x, attention_mask, layer_past, coeff):
        ......

在后续昇思MindSpore版本中我们规划特性识别这种有细微差别的Block，对这些Block给予优化建议以便优化改进。