XEngine-深度學習推理優化

前言

深度學習模型的開發周期，包括訓練階段和部署階段。訓練階段，用戶需要收集訓練數據，定義自己的模型結構，在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. 總結

本文主要從顯存優化和計算優化兩個角度分析了一下模型推理常用的優化思路和技巧，并展示了一下優化的結果。希望對大家做推理工程優化有幫助。

深度學習

版權聲明：本文內容由網絡用戶投稿，版權歸原作者所有，本站不擁有其著作權，亦不承擔相應法律責任。如果您發現本站中有涉嫌抄襲或描述失實的內容，請聯系我們jiasou666@gmail.com 處理，核實后本網站將在24小時內刪除侵權內容。