大模型推論最佳化

前言

大型transformer模型已成為最近的主流，在處理各種不同的任務上都獲得了很好的結果。大型transformer功能強大，但在計算效能和記憶體方面都有極高的要求，因此大型transformer的訓練和使用成本非常高，這也造成使用transformer大規模解決真實世界任務時的一個大瓶頸。為什麼大型transformer模型的推理執行如此困難？除了最先進模型規模的增加外，還有兩個主要因素導致推理挑戰（Pope等人，2022年）：

1. 佔用大量的記憶體。大型transformer在推理時需要在記憶體中同時存儲模型參數和中間狀態。例如，
2. • 在解碼時需要在記憶體中存儲KV快取記憶體；例如，對於批次大小為512和上下文長度為2048的情況，KV快取記憶體總計3TB，是模型大小的3倍。注意力機制的推理成本隨輸入資料長度呈平方成長。
2. 低平行化程度。推理生成以回歸方式執行，使解碼過程難以同時執行。

在這篇文章中，我們將探討幾種使transformer推理更高效的方法。有些是通用的網路壓縮方法，而其他則是專門針對transformer架構的。

蒸餾法

知識蒸餾（KD；Hinton等人，2015年，Gou等人，2020年）是一種直接的方法，通過將預訓練的昂貴模型（「教師模型」）的內容轉移到較小、較便宜的模型（「學生模型」）中來加速推理。知識蒸餾法對於學生模型的建構方式沒有太多限制，除了需要與教師模型有相符的輸出空間來建構合適的學習目標。

一個早期成功的嘗試是DistilBERT（Sanh等人，2019年），它能夠將BERT的參數減少40%，同時在微調的任務上保持BERT 97%的性能，執行速度提高71%。DistilBERT的預訓練損失是軟蒸餾損失、監督訓練損失（即BERT的遮蔽語言模損失）和一個特殊的餘弦嵌入損失的組合，以對齊教師和學生之間的隱藏狀態向量。蒸餾可以很容易地與量化、剪枝或稀疏化技術結合，其中教師模型是原始的全精度、密集模型，而學生模型為了達到更高的稀疏度是量化的、剪枝的或修剪的模型。

量化法

對深度神經網路應用量化有兩種常見方法：

1. 訓練後量化（PTQ）：首先將模型訓練至收斂，然後我們將其權重轉換為更低精度，無需進一步訓練。與訓練相比，這通常實施起來比較便宜。
2. 量化感知訓練（QAT）：在預訓練或進一步微調期間應用量化。QAT能夠獲得更好的性能，但需要額外的計算資源和輸入具有代表性的訓練數據。

我們應該注意理論上最佳量化策略與硬體支援之間的差距。由於缺乏對某些類型矩陣乘法的GPU核心的支持（例如INT4 x FP16），並非所有的方法都能在實際推理中實現加速。

剪枝法

神經網路剪枝是通過修剪不重要的模型權重或連接來減少模型大小，同時保持模型容量。它可能需要或不需要重新訓練。另外剪枝可以是非結構化的或結構化的。

1. 非結構化剪枝允許刪除任何權重或連接，因此不保留原始網路架構。非結構化剪枝通常不適用於現代體，也不會導致實際的推理加速。
2. 結構化剪枝旨在保持密集矩陣乘法形式，其中某些元素為零。它們可能需要遵循某些模式限制以配合硬體核心支援。這裡我們專注於結構化剪枝，以在transformer模型中實現高稀疏性。

建構剪枝網路的常規工作流程有三個步驟：

1. 訓練密集網路直至收斂；
2. 剪枝網路以移除不需要的結構；
3. 可選地重新訓練網路以恢復新權重的性能。

通過網路剪枝在密集模型中發現稀疏結構，同時稀疏網路仍能保持類似性能的想法，是受彩票假說[1]的啟發：隨機初始化的密集前饋網路包含一池子子網路，其中只有一個子集（稀疏網路）是「中獎彩票」，在單獨訓練時可以達到最佳性能。

Reference

1. https://arxiv.org/abs/1803.03635