深度学习推理引擎-内存共享算法

注意本文的一些算法术语，算法实现方法是作者根据自己的理解提出的，不一定跟业界已有的方法对齐。实际上业界也有不少相关的专利和论文，读者可以进行相应的比较，很多都是相近的思路。

Introduction

深度学习模型通常是用有向无环图(DAG)的方式来表示的。如下图所示的DL模型ONNX表示可视化图，图中每个节点是一个算子，每个边这是前面的算子的输出到后面算子输入的数据传递关系。在推理时，单个深度学习模型的算子通常是以拓扑排序的结果串行执行的。在每个算子执行的时候，该算子的输入数据需要保存在相应的输入张量内存中，而输出的数据需要写入到各个输出张量的内存中。除了Inplace算子（reshape, squeeze, unsqueeze这些算子通常也可以当做Inplace算子处理）直接修改输入张量内容，正常情况下每个算子的输入和输出张量应该是完全不重叠的内存空间，每个算子应该只修改输出张量而不修改输入张量。另外，每个算子执行完成之后，输出张量仍然保持active以作为后续算子的输入，但是该算子的输入张量在其他算子不继续使用时可以被释放掉，或者把这个算子的输入用于作为后续算子的输出内存使用。以上构成了深度学习推理引擎激活张量内存共享的基础：我们可以预先创建一组被共享的基张量，把这些基张量分配给深度学习推理时的各个激活，用作各个算子的输入输出临时使用，而不需要对每个算子的输入输出独立创建一个内存空间的张量。注意DL模型的算子权重通常情况下并不太适合用来进行内存共享，因为一个模型通常需要推理多次。但是如果能够保证一个模型加载后只被推理一次，那么前面算子的权重其实也可以被用于后面算子的激活共享，这种场景非常少见。

内存共享会带来若干好处：在模型推理之前我们就知道了推理过程需要多少内存使用，从而提前分配好，避免内存碎片，避免运行时的内存分配耗时等等。端侧推理引擎如MNN等也重点强调了这一点。云侧的深度学习引擎如Pytorch很多时候并不见得进行了这种提前的内存共享和分配。一方面模型结构和shape可能是动态变化的，这会导致内存共享的改变。另一方面cuda底层提供了cudaMallocAsync和cudaFreeAsync这样的方法在底层提供了内存共享的可能性：

Using the NVIDIA CUDA Stream-Ordered Memory Allocator, Part 1 | NVIDIA Technical Blog

后面介绍了两种可能的内存共享算法实现思路，虽然不一定是全局最优的，但是都简单适用。注意这些算法假定算子是按拓扑排序串行执行的，这个基本上绝大多数推理引擎都能满足。如果对同一个模型的算子采用多stream并发执行，则内存共享算法需要针对性适配才行，从而保证不会发生内存冲突。 

生产-消费producer–consumer模型的内存共享

这里介绍下作者曾实现过的第一种内存共享算法。所谓生产即每个算子产生若干输出张量，而消费则是每个算子会消费多个输入张量。算法实现：

输入：一个（空的）基张量List，一个DAG深度学习图graph表示，graph.node为整个图的拓扑排序后的算子列表。

for node in graph.node:

        for output in node.output:

                根据输出张量的大小从基张量List找到引用次数为0，大小大于等于该大小的最小的张量用于内存共享。如果基张量List没有合适的则创建一个新的张量。根据后续有多少算子使用该张量作为输入的count，那么对该张量引用次数设置为count。

        for input in node.input :

                对input共享使用的基张量List的内存引用次数减1.

返回：基张量List

一些具体的细节还包括整个模型的输入输出张量的内容共享方法（例如可能不想把模型的输入输出内存用于共享，这些也是部分算子的输入输出）和Inplace算子的处理。

基于优先级图着色的内存共享算法

上面提到的producer–consumer模型的内存共享算法实现非常简单，但是缺点是内存使用效率不高。可以预期同一个深度学习模型推理不同内存共享算法最终使用的激活内存量是不同的。这里提供一种通常内存使用更加高效的共享算法：基于优先级图着色的内存共享算法。

首先我们需要获得每个激活张量的生命周期。因为算子按照拓扑排序进行执行，那么N个算子，我们可以假定他们执行时刻分布为0到N-1。而从前面算子输出到后续算子输入传递的激活张量就可以得到相应的生命周期。比如第一个算子的输出只作为第二个算子的输入，那么其生命周期为[0,1]，以此类推。

 

算法输入：激活张量List(按张量内存大小从大到小排序)，每个激活张量具有其相应的生命周期，一个（空的）基张量List，一个DAG深度学习图graph表示

for activation in activation_list：

        从基张量List找到一个足够大的，且生命周期与该激活张量生命周期不重叠的张量用于共享。没有则创建一个。

        把该激活张量的生命周期扩展到用于共享的基张量。

返回：基张量List

与前面的方法一样，需要注意Inplace算子的共享和模型输入输出的共享细节。

基于内存地址级别更加细粒度的内存共享策略。

这里采用了以激活张量为粒度的内存共享，可以发现，优先对大张量设置了内存共享，后续小张量会把这些大张量拿来作为基张量进行共享，实际上又用不完但占据了这个大的基张量的生命周期使得其他小张量又又不到，这会导致内存的浪费。

因此一种可能的改进策略是，不再基于激活张量级别的共享，而是按照连续内存地址的共享策略。可能的改进是把基张量List改成一个完整的大的基内存内存地址。每次从基内存中找到一段具有合适生命周期的内存地址用于激活张量的共享，没有则增加内存分配。后续从这个大的基内存不同地址偏移创建张量分配给各个激活张量。

动态shape模型内存共享

静态shape指模型的输入shape是固定的常数，并且模型内部所有的算子shape都可以以此经过shape推导唯一确定。通常有两种情况导致动态shape：1是模型输入的shape是可以动态变化的，例如支持不同大小的图像输入；2是存在一些特殊算子，即使输入shape是静态的，但是输出shape也是动态不确定性的，例如torch.nonzero(有些情况下这些算子能够通过修改算法替换掉)。当然也存在一些特殊算子输入是动态而输出是静态，但这个实际上是归为前一类，因为输出能根据输入唯一确定。

针对场景1，作者提出了基于符号shape推导的方法，参考论文Transformer-Lite: High-efficiency Deployment of Large Language Models on Mobile Phone GPUs。也就是采用符号计算表达式来作为模型的输入和激活张量的shape表达和推导。例如sd的vae encoder的输入shape可以表示为[1, 3, "8*hight", "8*width"]。然后内存共享最重要的是要比较激活张量和基张量之间的大小，这个同样可以基于符号计算来实现。一个常用的符号计算工具是Python的sympy，可以从字符串的计算表达式采用Python的ast库解析得到sympy的表达式进行sympy符号计算。

此外，基于符号shape推导的内存共享需要针对性区分张量计算算子和shape计算算子。区别是，张量计算算子是常规意义上的把激活作为输入然后得到输出激活张量，在GPU等加速硬件上执行。而shape计算算子计算的输入为算子的shape信息，需要在CPU上执行。对于静态shape模型只有张量计算算子而没有shape计算算子。此外，为了从模型输入的符号表达式推导出中间每个算子的输入输出shape符号表达式，还需要注意的是，按照拓扑排序对每个算子进行shape推导，对于张量计算算子只需要根据输入shape推导出输出shape，而对于shape计算算子，则是需要对该算子进行计算，从而根据输入的shape信息得到输出的shape信息。上图是Transformer-Lite: High-efficiency Deployment of Large Language Models on Mobile Phone GPUs论文的算子分类示意图。蓝色为shape计算算子，而黄色为张量计算算子。可以看到对于图中的Shape, Gather这样的shape计算算子在shape推导时是把输入的shape进行计算得到输出的shape，最终传递到张量计算算子。那么具体怎么把哪些算子分类为shape计算算子，哪些分类为张量计算算子呢？这些细节可以参考论文。实际上作者在端侧推理引擎完成这个解决方案后续再云侧GPU使用TensorRT时，发现TensorRT内部应该也是使用的这个思路，算是殊途同归吧。

针对场景2的特殊动态shape算子，首先需要在前面解决方案的基础上，要对输出插入额外的shape符号类型和加以大小约束。并且可能要在实际运行时执行完每个特殊动态shape算子后进行一次同步进行后续的shape更新。

 

Nvidia TensorRT引擎内存共享的优缺点观察

1. TensorRT支持动态shape内存共享，这也使得TensorRT比较好的支持了动态shape推理。

2. TensorRT是把所有的激活张量的基张量创建在了一个更大的完整的张量里面，而不是独立进行创建和分配的。也就是根据所有基张量的大小和内存对齐计算得到一个总的张量大小，然后从这个大的张量里面去设置地址偏移创建各个基张量，用于所有激活内存的共享。这个带来一个独特的好处，那就是不同DL模型的激活可以复用激活张量用量最大模型的激活张量分配，只要它们的推理是串行执行的，这样在一个业务里面有多个深度学习模型时，转换为TensorRT可以显著降低激活内存的使用，因为所有的激活内存取决于最大的那个模型的使用量（注意当前TensorRT本身就支持这个特性），而不是每个引擎的使用量之和，这个对于Pytorch等引擎特别有优势。

3. TensorRT的内存共享算法应该做的也并非非常完美，虽然我没有理论分析过，但是实际感受一些模型的激活内存使用量非常大，远超预期，我认为经过恰当的内存共享后应该不需要那么大的激活内存。