XEngine：深度学习模型推理优化

摘要：从显存优化，计算优化两个方面来分析一下如何进行深度学习模型推理优化。

本文分享自华为云社区《XEngine-深度学习推理优化》，作者： ross.xw。

前言

深度学习模型的开发周期，包括训练阶段和部署阶段。训练阶段，用户需要收集训练数据，定义自己的模型结构，在CPU或者GPU硬件上进行训练，这个过程反复优化，直到训练出满意精度的模型。有了模型之后，我们需要将模型服务部署运行，我们期望服务延迟越低越好，吞吐越高越好。这里会从显存优化，计算优化两个方面来分析一下如何进行深度学习模型推理优化。

1. 显存优化

1.1 显存分布

模型推理需要占用一定量的显存空间（以GPU推理为例），其中主要包括如下4个部分:

不可控制空间
用户数据
模型参数
运行时空间

op计算的激活值
op计算需要的临时空间

其中“不可控制空间”指系统分配的空间，如每个进程CUDA Context所占用的显存空间，一般在100-300MB左右；“用户数据”指用户自行分配的显存空间，如模型输入输出Tensor占用的空间；“模型参数”指训练的深度学习模型的参数所占用的显存空间，我们需要将模型参数加载到显存中，才能进行计算；“运行时空间”是指模型的算子在计算的时候，需要的显存空间。

以ResNet-50模型为例，显存分配空间占比如下。我们可以看到随着Batch Size的增大，运行时空间会线性增长，运行时空间成为显存占用的瓶颈。

不同的模型显存分布也是不一样的。在NLP场景中，transformer类型的模型近几年涌现了许多超大参数量的模型，模型参数空间将成为显存的瓶颈。

接下来，会从激活优化和参数优化两个方面讲解如何进行显存空间优化，并最后扩展到多模型显存空间共享。

1.2 激活优化

激活值优化的中心思想就是显存复用。推理和训练不一样，推理计算只有forward过程，当一个op计算完后，它所占用的输入空间其实就可以被后面的op进行复用了。

1.3 参数优化

参数优化主要是为了解决超大模型的问题，如果模型太大，一个卡装不下就需要多张卡。参数空间和激活不一样，它是固定的值，提前训练好了，而激活值是临时计算出来的。这就使得我们不能用复用的方式。这些参数总要在一个地方保存。可以借鉴多级缓冲的思路，将训练好的参数可以缓存到磁盘和cpu内存中，在需要的时候提前读取上来，这样我们就不需要所有的参数都存储到显存中，将大模型单张卡加载成为可能。

为了减少数据拷贝对推理性能的影响，需要将数据预读取和计算并行起来。在GPU计算里面，我们可以通过cuda stream + event的方式将计算和拷贝并行起来，如下图所示：

1.4 多模型显存共享

除了单模型内部的激活优化和参数优化，在多并发多模型的服务场景，我们还可以进一步进行多模型显存共享。假设如下：

模型：M1, M2, M3...
激活空间：A1, A2, A3...
参数空间：P1, P2, P3...

1.5 显存优化效果

这里选取了CV和NLP两个场景的模型，对比了一下XEngine推理引擎和其他运行时引擎的显存占用，可以看到有明显的优化。如果不考虑性能开启参数优化（ParamOpt），可以进一步的降低显存占用。

2. 计算优化

2.1 计算分析

下图取了ResNet-50模型的一个片段，并抽象到计算和访存的流水线过程。我们可以看到，每个OP计算对应一个CUDA Kernel计算，每个kernel计算会从显存（Global Memory）中读取数据，在片上CUDA Core进行计算。右图为CUDA的存储架构，也符合金字塔访存原则，即越靠近片上，访存速度越快。其中Global Memory是GPU的显存，访问速度最慢；而片上的Shared Memory次之；最快的是片上Register空间。当模型的算子OP比较多的时候，模型推理计算会反复读写显存，效率比较低，为了解决这个问题，常用的方法是图融合技术；当OP计算的速度越快，整个模型推理计算速度也会越快，因此我们需要高性能的算子实现。接下来会从这两个方面分析。

2.2 图融合

图融合技术指将一些OP进行融合计算，由多个OP的Kernel计算转化为一个融合后OP的Kernel计算。通过这个优化，可以减少显存的反复读写，可以通过一次读取数据，在片上进行尽可能多的计算后，再将数据存储到显存中。同时也可以减少Kernel Launch的开销。融合后的算子实现可以通过手写CUDA Kernel或者调用第三方库或者CodeGen方式进行生成。

下图中左边为resnet-50模型优化前和优化后的Graph结构。CV类型的模型常用的融合手段是将线性计算和激活进行融合，如Conv + BN + Relu进行融合计算。

下图为Bert模型的图融合前后的变化。

2.3 高性能算子

高性能计算库

Cublas/cudnn/cutlass

手写CUDA Kernel

Cuda programming guide
Cuda best practice guide

低精度计算

FP16/INT8/TensorCore (half2, dp4a, wmma)

CodeGen

TVM/XLA/MLIR…

2.4 计算优化效果

XEngine针对Bert类型模型进行优化，测试环境为NVIDIA V100，FP32计算。如下是性能对比，相对于原始框架未优化版本，在sequence length比较小的时候，有明显的性能提升。同时也和NVIDIA的SOTA解决方案FasterTransformer进行了对比，有轻微优势。

XEngine和Pytorch推理性能对比：

XEngine和FasterTransformer推理性能对比：

3. 总结

本文主要从显存优化和计算优化两个角度分析了一下模型推理常用的优化思路和技巧，并展示了一下优化的结果。希望对大家做推理工程优化有帮助。

点击关注，第一时间了解华为云新鲜技术~