Bert/Transformer模型加速

前言

Bert/Transformer結構及其變體，已成為自然語言處理 (NLP)、語音識別 (ASR)等領域的主流序列建模結構，并且，相比于卷積操作的區域特征提取能力、以及平移不變性，Self-attention的全域Context資訊編碼能力，能夠使視覺模型具備更強的特征表達能力、與領域適應性，因此在計算機視覺領域，Vision Transformer結構也日益流行、方興未艾，然而，Transformer模型的高計算復雜度與引數量，限制了模型的推理部署（包括服務端與移動端應用），其計算、存盤與運行時記憶體等資源開銷都成為限制因素，例如，從標準Transformer layer的Tensor處理結構來看，計算量與序列長度、特征維度的平方成正比：

• Standard Self-Attention:

• Standard FFN:

本文從以下幾個維度舉例說明，決議Bert/Transformer模型的加速方法：

• 模型結構精簡化與知識蒸餾；
• 模型量化 (Quantization)；
• 網路結構搜索 (NAS: Network Architecture Search)；
• 計算圖優化；
• 推理優化引擎 (Faster Transformer)；

模型結構精簡化與知識蒸餾

Lite-Transformer

參考文獻: https://arxiv.org/abs/2004.11886

Lite Transformer是一種融合卷積與Self-attention操作的、高效精簡的Transformer結構，可應用于NMT、ASR等序列生成任務，其核心是長短距離注意力結構 (LSRA：Long-Short Range Attention)，LSRA將輸入Embedding沿Feature維度split成兩部分，其中一部分通過GLU (Gate Linear Unit)、一維卷積，用以提取區域Context資訊；而另一部分依靠Self-attention，完成全域相關性資訊編碼，Lite Transformer核心結構如下：

SAN-M

參考文獻: https://arxiv.org/abs/2006.01713

SAN-M表示Self-attention與DFSMN記憶單元的融合，是一種Transformer ASR模型，DFSMN適合捕獲區域資訊，Self-attention模塊具備較強的長時依賴捕獲能力，因此二者存在互補性，SAN-M通過將兩個模塊的特性融合在一起，實作了優勢互補，Biasic Sub-layer表示包含了SAN-M的Self-attention layer，DFSMN添加在values后面，其輸出與Multi-head Attention (MHA)相加：

MiniLM (知識蒸餾)

參考文獻: https://arxiv.org/abs/2002.10957

針對NLP任務，深度自注意力知識蒸餾 (Deep Self-Attention Distillation)，通過遷移Teacher model最后一層Self-attention layer的Attention score資訊與Value relation資訊，可有效實作Student model的誘導訓練，只遷移最后一層的知識，顯得簡單有效、訓練速度更快，而且不需要手動設計Teacher-student之間的層對應關系，Attention score資訊與Value relation資訊的知識遷移如下：

• Attention score transfer:

• Value relation transfer:

• 若選擇的Self-attention layer為LSRA，除了在multi-head self-attention分支遷移attention score與value relation；在CNN分支需要遷移feature map的資訊，這里主要計算AT loss：

  

  式中Fs,j表示學生網路里第j個網路層的特征輸出，Ft,j表示teacher network里第j個group的特征輸出，

模型量化 (Quantization)

Transformer ASR模型壓縮

參考文獻: https://arxiv.org/abs/2104.05784

針對Transformer ASR，文章提出了聯合隨機稀疏與PTQ量化 (KL量化、ADMM與混合精度設定)的壓縮策略，整體實作了10倍壓縮，且絕對精度損失約0.5%，總體流程如下：

• 模型稀疏化，更新Weight重要性：

• KL量化，計算Activation量化引數：

• ADMM，優化Weight量化引數：

• 混合精度量化設定，減少量化誤差：

此外，針對Transformer模型的量化，需要討論全網路INT8量化與計算的意義：

• 全網路INT8量化：
  • 同時減少計算密集算子與訪存密集算子的開銷；
  • 實作模型壓縮，INT8模型是FP32模型的1/4；
• Fully INT8 Attention Path：
  • 基本的全INT8通路：
    • Dense layer的輸入與輸出均為INT8 Tensor；
    • Reshape與Transpose的操作物件為INT8 Tensor，節省記憶體開銷：
      • Dense+Reshape+Transpose、與Transpose+Reshape+Dense可以實作Op fusion；
    • BatchMatMul、Softmax的操作物件保留為FP32 Tensor，確保模型預測精度；

• ?????Encoder的Self-attention：
  • X=Y；
• Decoder的Self-attention：
  • X=Y，Batch_size=Batch_size * Beam_size，Seq_len=1；
  • k_dense與v_dense的輸出會添加到Cache，按照INT8型別搬運資料，可節省訪存開銷：

PTQ for Vision Transformer

參考文獻: https://arxiv.org/abs/2106.14156

針對Vision Transformer的PTQ量化，分別針對FFN與Self-Attention，提出了Similarity-aware Quantization與Ranking-aware Quantization：

• Similarity-aware Quantization: 基于優化方式確定Weight與Activation量化的最優量化引數，并以Pearson相似度作為Loss；

• Ranking-aware Quantization: 在相似度Loss基礎之上，疊加Attention map的Ranking loss作為約束，以準確感知不同Attention map的重要性，確保Self-attention量化的有效性；

網路結構搜索 (NAS)

AdaBert

參考文獻: https://arxiv.org/abs/2001.04246

針對NLP任務，AdaBert通過可微分搜索、與知識蒸餾，實作卷積型別的NLP模型搜索：

• 搜索空間：類似于DARTS的Cell設計，以構造搜索空間；

  

• 搜索策略：基于Gumbel Softmax實作網路結構的隨機采樣（類似于FBNet）、與可微分搜索；并基于FLOPS與Model size構造了Efficiency-aware Loss，作為搜索訓練的資源約束；
• 知識蒸餾：多層次、任務相關的知識遷移；

  

DynaBert

參考文獻: https://arxiv.org/abs/2004.04037

針對NLP任務，DynaBert模型壓縮策略可實作多個維度的壓縮搜索（主要是Width與Depth維度），通過動態可伸縮性質的One-shot NAS：首先基于預訓練Bert模型（或其變體），進行多維度正則化訓練；正則化訓練之后，按照不同的資源約束需求，能夠進行網路引數的裁剪、以獲得子網路，并進行子網路的微調訓練；從正則訓練、到子網路微調，可協同知識蒸餾訓練，主要步驟描述如下：

• 第一步：基于Neuron與Attention head的重要性，執行引數重排（Weight Re-wiring）：

• 第二步：多階段微調訓練，實作寬度、與深度方向的自適應正則化：

Evolved Transformer

參考文獻: https://arxiv.org/abs/2004.11886

針對序列生成任務，基于NAS搜索獲得的Transformer結構:

• 搜索空間：包括兩個Stackable cell，分別包含在Transformer encoder與Transformer decoder，每個Cell由NAS-style block組成, 可通過左右兩個Block轉換輸入Embedding、再聚合獲得新的Embedding，進一步輸入到Self-attention layer，
• 搜索策略：基于EA (Evolutional Aligorithm)的搜索策略；

網路結構如下，融合了一維卷積與Attention的特點：

HAT: Hardware-aware Transformer

參考文獻: https://arxiv.org/abs/2005.14187

對Transformer模型進行網路結構搜索時，通常會面臨兩個關鍵問題：

• FLOPS不能完全代表Transformer模型的計算速度/實際執行效率，即只能間接反映硬體平臺特性；
• 不同硬體設備上，與計算硬體或計算庫相適配的最優Transformer結構不盡相同;

針對序列生成任務，HAT (Hardware-aware Transformer)從搜索空間、搜索策略與搜索預測三方面加以分析：

• 搜索空間：Encoder-decoder Attention的任意連接方式，以及Transformer layer內部結構（Attention head數目、Width等）；
• 搜索策略：訓練了SuperTransformer作為超網路，并在超網路預訓練之后，結合資源約束，通過進化搜索演算法（EA: Evolutional Algorithm）尋找最優子網路；不同于權重共享型NAS (如FBNet、SPOS等)，HAT是一種動態可伸縮型別的One-shot NAS；

• 搜索預測：訓練了MLP模型作為Cost model或Predicter，用于預測不同子網路結構、在指定硬體平臺上的執行速度 (Latency)；通過搜索預測，可直接、有效獲取硬體平臺特性，作為超網路預訓練的資源約束；

AutoFormer

參考文獻: https://arxiv.org/abs/2107.00651

針對Vision Transformer的結構搜索，AutoFormer提出的Weight entanglement，在不額外增加Block choice的前提下，通過通道寬度、網路深度、Attention head數目等多個維度的調整，實作了Vision Transformer模型的動態可伸縮預訓練與結構搜索，Weight entanglement的做法，類似于BigNAS、FBNet-v2的通道搜索，都不會額外增加通道維度的權重引數量，相比于手工設計的CNN模型（ResNet、ResNext、DenseNet）與Vision Transformer模型（ViT、Deit），AutoFormer模型在相同資源開銷條件下、能夠獲得最好的識別精度，

AutoFormer總體思路如上圖所示，是一種基于Weight entanglement的動態可伸縮搜索方法，其搜索維度包括Attention heads、通道寬度與網路深度，下面從搜索空間、搜索策略與搜索效率這些維度加以分析：

• 搜索空間：Embedding dimension、O-K-V dimension、Attention heads、MLP ratio與Network depth；根據不同的資源約束，分別設定Supernet-tiny、Supernet-small與Supernet-base三個基本的Template model；

• 搜索策略：基于Weight entanglement的One-shot NAS，完成超網路預訓練之后，通過進化演算法執行子網路搜索，包含l個網路層的子網路結構可按下式表示，表示第i層的Block結構，表示相應的權重引數：

而每個Block結構，都是按照Weight entanglement原則從超網路采樣獲得，下式表示n個動態選擇范圍：

• 搜索效率：按照權重共享形式、完成超網路預訓練，訓練效率高、且收斂速度快；另外，由于沒有引入額外的Block choice，因此訓練時的Memory cost較低；

計算圖優化

LINM

LINM (Loop-invariant Node Motion)是一種計算圖等價變換技術，通過將Transformer模型涉及的自回歸解碼 (Auto-regressive Decoding)的重復計算邏輯 (Encoder-decoder Attention的k/v計算)移至while-loop之外，在確保計算功能不變的情況下，實作計算效率的提升：

推理優化引擎

Faster Transformer

GitHub鏈接：https://github.com/NVIDIA/FasterTransformer

Faster Transformer是NVIDIA針對標準BERT/Transformer做的推理優化庫，其發布時間線：

• 1.0版本：
  • 2019年7月，開源了FasterTransformer 1.0，針對 BERT中的Transformer Encoder進行優化和加速；
  • 面向BERT分類、自然語言理解場景；
  • 底層由CUDA和cuBLAS實作，支持FP16和FP32計算，FP16可充分利用Volta和Turing架構的Tensor Core計算單元；
  • 提供C++ API、TF Op與TensorRT Plugin三種介面；
  • 參考資料；
• 2.0版本：
  • 2020年2月，新增對Transformer decoder的優化和加速，包括decoder與decoding兩種加速模式；
  • 面向生成式場景，如NMT、文本內容生成與ASR等；
  • 底層由CUDA和cuBLAS實作，支持FP16和FP32計算模，FP16可充分利用Volta和Turing架構的 Tensor Core計算單元；
  • 提供C++ API、TF Op與TensorRT Plugin三種介面；
  • 參考資料；
• 2.1版本：
  • 2020年6月，引入Effective Transformer優化；通過remove_padding的支持，提高計算與訪存效率；
  • 并新增PyTorch Op介面；
  • 參考資料；
• 3.0版本：
  • 2020年9月，新增BERT encoder的INT8量化加速支持；
  • 僅支持Turing架構GPU；
  • 同時支持PTQ與QAT方法，提供了TF量化工具；
  • 相比于FP16計算，約20~30%加速，但存在精度損失風險；
• 3.1版本：
  • 2020年12月，新增對PyTorch使用INT8推理的支援；
  • 在Turing以后的GPU上，FP16的性能比3.0提升了 10% ~ 20%；
  • INT8的性能比3.0最多提升了70%；
• 4.0版本：
  • 2021年4月，新增對GPT-3等百億/千億級引數規模模型的多機多卡推理加速能力；
  • 新增FP16 fused MHA算子，同時支持Volta與Turing架構的GPU；
  • 以及對解碼端Kernel的優化，可以省略已完成陳述句的計算，節省計算資源；
  • 參考資料；
• 支持矩陣：

Transformer計算圖表示，主要包含三個部分的表示：

• Encoder layer；
• Decoder layer；
• Decoding logic；