NVIDIA Ampere 架構(gòu)的結(jié)構(gòu)化稀疏功能

及其在搜索引擎中的應(yīng)用

深度學(xué)習(xí)徹底改變了我們分析、理解和處理數(shù)據(jù)的方式，而且在各個(gè)領(lǐng)域的應(yīng)用中都取得了巨大的成功，其在計(jì)算機(jī)視覺、自然語言處理、醫(yī)療診斷和醫(yī)療保健、自動(dòng)駕駛汽車、推薦系統(tǒng)以及氣候和天氣建模方面有許多成功案例。

在神經(jīng)網(wǎng)絡(luò)模型不斷變大的時(shí)代，對計(jì)算速度的高需求對硬件和軟件都形成了巨大的挑戰(zhàn)。模型剪枝和低精度推理是非常有效的解決方案。

自 NVIDIA Ampere 架構(gòu)開始， 隨著 A100 Tensor Core GPU 的推出，NVIDIA GPU 提供了可用于加速推理的細(xì)粒度結(jié)構(gòu)化稀疏功能。在本文中，我們將介紹此類稀疏模型的訓(xùn)練方法以保持模型精度，包括基本訓(xùn)練方法、漸進(jìn)式訓(xùn)練方法以及與 int8 量化的結(jié)合。我們還將介紹如何利用 Ampere 架構(gòu)的結(jié)構(gòu)化稀疏功能進(jìn)行推理。

騰訊機(jī)器學(xué)習(xí)平臺部門 (MLPD) 利用了漸進(jìn)式訓(xùn)練方法，簡化了稀疏模型訓(xùn)練并實(shí)現(xiàn)了更高的模型精度。借助稀疏功能和量化技術(shù)，他們在騰訊的離線服務(wù)中實(shí)現(xiàn)了 1.3 倍～1.8 倍的加速。

NVIDIA Ampere 架構(gòu)的結(jié)構(gòu)化稀疏功能

NVIDIA Ampere 和 NVIDIA Hopper 架構(gòu) GPU 增加了新的細(xì)粒度結(jié)構(gòu)化稀疏功能，該功能主要用于加速推理。此功能是由稀疏 Tensor Core 提供，這些稀疏 Tensor Core 需要 2:4 的稀疏模式。也就是說，以 4 個(gè)相鄰權(quán)重為一組，其中至少有 2 個(gè)權(quán)重必須為 0，即 50% 的稀疏率。

這種稀疏模式可實(shí)現(xiàn)高效的內(nèi)存訪問能力，有效的模型推理加速，并可輕松恢復(fù)模型精度。在模型壓縮后，存儲(chǔ)格式只存儲(chǔ)非零值和相應(yīng)的索引元數(shù)據(jù)（圖 1）。稀疏 Tensor Core 在執(zhí)行矩陣乘法時(shí)僅處理非零值，理論上，計(jì)算吞吐量是同等稠密矩陣乘法的 2 倍。

圖 1. 2:4 結(jié)構(gòu)化稀疏模式及其壓縮格式

（結(jié)構(gòu)化稀疏矩陣具有 2:4 的稀疏模式。在 4 個(gè)相鄰權(quán)重當(dāng)中，至少有 2 個(gè)值為零。在模型壓縮后，僅存儲(chǔ)非零值和相應(yīng)的索引元數(shù)據(jù)。）

結(jié)構(gòu)化稀疏功能主要應(yīng)用于能夠提供 2:4 稀疏權(quán)重的全連接層和卷積層。如果提前對這些層的權(quán)重做剪枝，則這些層可以使用結(jié)構(gòu)化稀疏功能來進(jìn)行加速。

訓(xùn)練方法

由于直接對權(quán)重做剪枝會(huì)降低模型精度，因此在使用結(jié)構(gòu)化稀疏功能的時(shí)候，您需要進(jìn)行訓(xùn)練來恢復(fù)模型精度。下面，我們將介紹一些基本訓(xùn)練方法和新的漸進(jìn)式訓(xùn)練方法。

基本訓(xùn)練方法

基本訓(xùn)練方法可保持模型精度，并且無需任何超參數(shù)調(diào)優(yōu)。了解更多技術(shù)細(xì)節(jié)，請參閱論文 Accelerating Sparse Deep Neural Networks（https://arxiv.org/abs/2104.08378）。

基本訓(xùn)練方法易于使用，步驟如下：

1. 訓(xùn)練一個(gè)常規(guī)稠密模型，不需要稀疏化的特殊處理。

2. 對全連接層和卷積層上的權(quán)重以 2:4 的稀疏模式進(jìn)行剪枝。

3. 按照以下規(guī)則重新訓(xùn)練經(jīng)過剪枝的模型：

   a. 將所有權(quán)重初始化為第 2 步中的值。

   b. 使用與第 1 步相同的優(yōu)化器和超參數(shù)（學(xué)習(xí)率、調(diào)度方法、訓(xùn)練次數(shù)等）進(jìn)行稀疏調(diào)優(yōu)訓(xùn)練。

   c. 保持第 2 步中剪枝后的稀疏模式。

圖 2. 基本訓(xùn)練方法

（基本訓(xùn)練方法就是使用剪枝后的權(quán)重和掩碼后的優(yōu)化器重復(fù)原始稠密模型的訓(xùn)練過程。）

對于復(fù)雜情況，還有一些進(jìn)階的方法。

例如，把稀疏訓(xùn)練應(yīng)用在多階段式的稠密模型訓(xùn)練當(dāng)中。比如對于一些目標(biāo)檢測模型，如果下游任務(wù)的數(shù)據(jù)集足夠大，您只需做稀疏調(diào)優(yōu)訓(xùn)練。對于像 BERT-SQuAD 等模型，調(diào)優(yōu)階段使用的數(shù)據(jù)集相對較小，您則需要在預(yù)訓(xùn)練階段進(jìn)行稀疏訓(xùn)練以獲得更好的模型精度。

此外，通過在稀疏調(diào)優(yōu)之前插入量化節(jié)點(diǎn)，您可以輕松將稀疏調(diào)優(yōu)與 int8 量化調(diào)優(yōu)結(jié)合起來。所有這些訓(xùn)練以及調(diào)優(yōu)方法都是一次性的，即最終獲得的模型只需要經(jīng)過一次稀疏訓(xùn)練處理。

漸進(jìn)式稀疏訓(xùn)練方法

一次性稀疏調(diào)優(yōu)（fine-tuning）可以覆蓋大多數(shù)任務(wù)，并在不損失精度的情況下實(shí)現(xiàn)加速。然而，就一些對權(quán)重?cái)?shù)值變化敏感的困難任務(wù)而言，對所有權(quán)重做一次性稀疏訓(xùn)練會(huì)導(dǎo)致大量信息損失。在小型數(shù)據(jù)集上只做稀疏化調(diào)優(yōu)可能會(huì)很難恢復(fù)精度，對于這些任務(wù)而言，就需要稀疏預(yù)訓(xùn)練（pretraining）。

然而稀疏預(yù)訓(xùn)練需要更多數(shù)據(jù)，而且更加耗時(shí)。因此，受到卷積神經(jīng)網(wǎng)絡(luò)剪枝方法的啟發(fā)，我們引入了漸進(jìn)式稀疏訓(xùn)練方法，在此類任務(wù)上僅應(yīng)用稀疏化調(diào)優(yōu)便可以實(shí)現(xiàn)模型的稀疏化，同時(shí)不會(huì)造成明顯的精度損失。了解更多細(xì)節(jié)，請參閱論文 Learning both Weights and Connections for Efficient Neural Networks（https://arxiv.org/pdf/1506.02626.pdf）。

圖 3. 漸進(jìn)式稀疏訓(xùn)練的概念

（漸進(jìn)式稀疏訓(xùn)練方法將稀疏率分為幾個(gè)步驟，以輕松恢復(fù)精度。漸進(jìn)式稀疏訓(xùn)練方法的核心思想是將目標(biāo)稀疏率進(jìn)行若干次切分。）

?

如上述公式和圖 4 所示，對于目標(biāo)稀疏率 S，我們將其分為 N 份，這將有助于在稀疏調(diào)優(yōu)過程中快速恢復(fù)信息。根據(jù)我們的實(shí)驗(yàn)，在相同的調(diào)優(yōu)迭代次數(shù)內(nèi)，使用漸進(jìn)式稀疏訓(xùn)練相比一次性稀疏訓(xùn)練，可以獲得更高的模型精度。

圖 4. 漸進(jìn)式稀疏訓(xùn)練方法 （以 50% 稀疏率的 2:4 結(jié)構(gòu)化稀疏模式為例）

（漸進(jìn)式稀疏訓(xùn)練方法的示例：計(jì)算權(quán)重掩碼以達(dá)到 25% 稀疏率，再進(jìn)行稀疏調(diào)優(yōu)恢復(fù)性能，最后重新計(jì)算掩碼使之達(dá)到 50% 稀疏率并對網(wǎng)絡(luò)進(jìn)行調(diào)優(yōu)。）

我們以 50% 稀疏率的 2:4 結(jié)構(gòu)化稀疏為例，將稀疏率分為兩份，然后逐步稀疏和調(diào)優(yōu)模型中的權(quán)重參數(shù)。如圖 4 所示，首先計(jì)算權(quán)重掩碼以實(shí)現(xiàn) 25% 的稀疏率，然后執(zhí)行稀疏調(diào)優(yōu)以恢復(fù)模型精度。接下來，重新對剩余權(quán)重計(jì)算權(quán)重掩碼以達(dá)到 50% 的稀疏率，并對網(wǎng)絡(luò)進(jìn)行調(diào)優(yōu)，以獲得一個(gè)精度無損的稀疏模型。

Sparse-QAT：稀疏化與量化、蒸餾相結(jié)合

為了獲得更輕量的模型，我們進(jìn)一步將稀疏與量化、蒸餾相結(jié)合，即 Sparse-QAT。

量化（PTQ 和 QAT）

下方的公式表示一個(gè)通用的量化過程。對于 32 位浮點(diǎn)數(shù)值 x，我們使用 Q [x] 表示其具有 K-bits 表示的量化值。

通常情況下，我們首先將原始參數(shù)量化到特定范圍，并將其近似為整數(shù)。然后，可以使用這個(gè)量化比例 scale (s) 來恢復(fù)原始值。這樣就得到了第一種量化方法，即校準(zhǔn)，也稱為訓(xùn)練后量化（post-training quantization, PTQ）。在校準(zhǔn)中，一個(gè)關(guān)鍵的因素是要設(shè)置一個(gè)適當(dāng)?shù)牧炕壤╯cale）。如果這個(gè)比例值過大，量化范圍內(nèi)的數(shù)字將不太準(zhǔn)確。相反，如果這個(gè)比例值過小，會(huì)導(dǎo)致大量的數(shù)字落在 lmin 到 lmax 的范圍之外。因此，為了平衡這兩個(gè)方面，我們首先獲得張量中數(shù)值的統(tǒng)計(jì)分布，然后設(shè)置量化比例以覆蓋 99.99% 的數(shù)值。許多工作已經(jīng)證明，這種方法對于在校準(zhǔn)過程中找到合適的量化比例非常有幫助。

然而，盡管我們已經(jīng)為校準(zhǔn)設(shè)置了一個(gè)合理的量化比例，但是對于 8 bit 量化來說，模型精度仍然會(huì)顯著下降。因此，我們引入量化感知訓(xùn)練（quantization-aware training, QAT），以進(jìn)一步提高校準(zhǔn)后的精度。QAT 的核心思想是以模擬量化的方法來訓(xùn)練模型。

在前向傳播過程中，我們將權(quán)重量化為 int8，然后將其反量化為浮點(diǎn)數(shù)來模擬真實(shí)量化。在反向傳播過程中，引入 straight through estimation (STE) 的方法來更新模型權(quán)重。STE 的核心思想可以用如下公式表示：

由上述公式可知，閾值范圍內(nèi)的值對應(yīng)的梯度直接反向傳播，超出閾值范圍的值對應(yīng)的梯度被裁剪為 0。

知識蒸餾

除了上述方法外，我們還引入了知識蒸餾（knowledge distilation, KD），以進(jìn)一步確保 Sparse-QAT 模型的精度。我們以原始稠密模型作為教師模型，以量化稀疏模型作為學(xué)生模型。在調(diào)優(yōu)過程中，我們采用了 Mini-distillation，這是一種層級別的蒸餾方法。使用 MiniLM，我們只需要使用 Transformer 模型最后一層的輸出。引入蒸餾作為輔助工具甚至可以獲得比教師模型精度更高的稀疏量化學(xué)生模型。了解更多信息，請參閱 MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers（https://arxiv.org/abs/2002.10957）。

Sparse-QAT 訓(xùn)練流水線

圖 5 顯示了 Sparse-QAT 的訓(xùn)練流水線。稀疏化、量化、蒸餾以并行的方式執(zhí)行，最終獲得一個(gè)稀疏的 int8 量化模型。整個(gè)流水線包括如下三條路徑：

• 在稀疏路徑中，應(yīng)用漸進(jìn)式稀疏化來獲取一個(gè)稀疏權(quán)重張量。

• 在量化路徑中，使用 PTQ 和 QAT 來獲取 int8 類型的權(quán)重張量。

• 在知識蒸餾路徑中，使用 MiniLM 來進(jìn)一步保障最終稀疏 int8 模型的精度。

圖 5. Sparse-QAT 流水線

（將稀疏、量化和知識蒸餾相結(jié)合，以獲得最終的稀疏 int8 模型。）

使用 NVIDIA Ampere 架構(gòu)的

結(jié)構(gòu)化稀疏功能進(jìn)行推理

在訓(xùn)練好稀疏模型后，您可以使用 NVIDIA TensorRT 和 cuSPARSELt 庫來加速基于 NVIDIA Ampere 架構(gòu)結(jié)構(gòu)化稀疏功能的推理。

使用 NVIDIA TensorRT 進(jìn)行推理

自 8.0 版本開始，TensorRT 可以支持稀疏卷積，矩陣乘法 (GEMM) 需要用 1x1 的卷積替代來進(jìn)行稀疏化推理。在 TensorRT 中啟用稀疏化推理非常簡單。在導(dǎo)入 TensorRT 之前，模型的權(quán)重應(yīng)具有 2:4 的稀疏模式。如果使用 trtexec 構(gòu)建引擎，只需設(shè)置 -sparity=enable 標(biāo)志即可。如果您正在編寫代碼或腳本來構(gòu)建引擎，只需按如下所示設(shè)置構(gòu)建配置：

對于 C++：

config->setFlag(BuilderFlag::kSPARSE_WEIGHTS)

對于 Python：

config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)

使用 NVIDIA cuSPARSELt 庫增強(qiáng) TensorRT

在某些用例中，TensorRT 可能因?yàn)檩斎氤叽绮煌鵁o法提供最佳性能。您可以使用 cuSPARSELt 進(jìn)一步加速這些用例。

解決方案是使用 cuSPARSELt 編寫 TensoRT 插件，我們可以為不同的輸入尺寸初始化多個(gè)描述符以及多個(gè) cuSPARSELt 稀疏矩陣乘法 plan，并根據(jù)輸入尺寸選擇合適的 plan。

假設(shè)您在實(shí)現(xiàn)SpmmPluginDynamic插件，該插件繼承自nvinfer1:: IPluginV2DynamicExt，您可以使用一個(gè)私有結(jié)構(gòu)來存儲(chǔ)這些 plan。

 struct cusparseLtContext {
    cusparseLtHandle_t handle;
    std::vector plans;
    std::vector matAs, matBs, matCs;
    std::vector matmuls;
    std::vector alg_sels;
}

TensorRT 插件應(yīng)實(shí)現(xiàn) configurePlugin方法，該方法會(huì)根據(jù)輸入和輸出類型及尺寸設(shè)置插件。您需要在這個(gè)函數(shù)當(dāng)中初始化 cuSPARSELt 的相關(guān)結(jié)構(gòu)。

cuSPARSELt 的輸入尺寸有一些限制，應(yīng)為 4、8 或 16 的倍數(shù)，具體取決于數(shù)據(jù)類型。在本文中，我們將其設(shè)置為 16 的倍數(shù)。了解該限制條件的相關(guān)信息，請查看此文檔（https://docs.nvidia.com/cuda/cusparselt/functions.html#cusparseltdensedescriptorinit）。

for (int i = 0; i < size_num; ++i) {
  m = 16 * (i + 1);
  int alignment = 16;
  CHECK_CUSPARSE(cusparseLtStructuredDescriptorInit(
      &handle, &matBs[i], n, k, k, alignment, type, CUSPARSE_ORDER_ROW,
      CUSPARSELT_SPARSITY_50_PERCENT))
  CHECK_CUSPARSE(cusparseLtDenseDescriptorInit(
      &handle, &matAs[i], m, k, k, alignment, type, CUSPARSE_ORDER_ROW))
  CHECK_CUSPARSE(cusparseLtDenseDescriptorInit(
      &handle, &matCs[i], m, n, n, alignment, type, CUSPARSE_ORDER_ROW))
  CHECK_CUSPARSE(cusparseLtMatmulDescriptorInit(
      &handle, &matmuls[i], CUSPARSE_OPERATION_NON_TRANSPOSE,
      CUSPARSE_OPERATION_TRANSPOSE, &matAs[i], &matBs[i], &matCs[i], &matCs[i], compute_type))
  CHECK_CUSPARSE(cusparseLtMatmulAlgSelectionInit(
      &handle, &alg_sels[i], &matmuls[i], CUSPARSELT_MATMUL_ALG_DEFAULT))
  int split_k = 1;
  CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute(
      &handle, &alg_sels[i], CUSPARSELT_MATMUL_SPLIT_K, &split_k, sizeof(split_k)))
  int alg_id = 0;
  CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute(
      &handle, &alg_sels[i], CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg_id, sizeof(alg_id)))
  size_t ws{0};
  CHECK_CUSPARSE(cusparseLtMatmulPlanInit(&handle, &plans[i], &matmuls[i], &alg_sels[i],
                                          ws))
  CHECK_CUSPARSE(
      cusparseLtMatmulGetWorkspace(&handle, &plans[i], &ws))
  workspace_size = std::max(workspace_size, ws);
}

在enqueue函數(shù)中，您可以檢索適當(dāng)?shù)?plan 來執(zhí)行矩陣乘法。

int m = inputDesc->dims.d[0];
int idx = (m + 15) / 16 - 1;
float alpha = 1.0f;
float beta = 0.0f;
auto input = static_cast<const float*>(inputs[0]);
auto output = static_cast<float*>(outputs[0]);
cusparseStatus_t status = cusparseLtMatmul(
    &handle, &plans[idx], &alpha, input,
weight_compressed,&beta,output,output,workSpace,&stream,1);

搜索引擎中的應(yīng)用

在本部分中，我們將展示在搜索引擎中應(yīng)用了稀疏化加速的四個(gè)應(yīng)用案例：

• 第一是搜索中的相關(guān)性預(yù)測，旨在評估輸入文本和數(shù)據(jù)庫中視頻之間的相關(guān)性。

• 第二是查詢性能預(yù)測，用于文檔召回交付策略。

• 第三是用于召回最相關(guān)文本的召回任務(wù)。

• 第四是文生圖任務(wù)，該任務(wù)根據(jù)輸入的提示詞自動(dòng)生成相應(yīng)的圖片。

搜索相關(guān)性案例

我們使用 PNR (Positive Negative Rate，正負(fù)率) 或 ACC (accuracy，精度) 標(biāo)準(zhǔn)來評估稀疏化加速在這些應(yīng)用案例的效果。在相關(guān)性案例 1 中，我們運(yùn)行 Sparse-QAT 獲得了一個(gè)稀疏 int8 模型，該模型在兩個(gè)重要的 PNR 評估指標(biāo)均優(yōu)于在線 int8 模型。

在相關(guān)性案例 2 中，稀疏 int8 模型可以獲得與 float32 模型接近的 Acc 分?jǐn)?shù)，相比稠密 int8 模型，其獲得了 1.4 倍的推理加速。

? ?

查詢性能預(yù)測案例

在這部分，我們展示了查詢性能預(yù)測 (query performance prediction, QPP) 的四個(gè)用例，其效果使用 NDCG（normalized discounted cumulative gain, 標(biāo)準(zhǔn)化折扣累積增益）評估。如表 3 所示，這些稀疏 float16 模型甚至可以獲得比原始 float32 模型更高的 NDCG 分?jǐn)?shù)，同時(shí)推理速度相比于 float32 模型提高了 4 倍。

文檔查詢案例

表 4 顯示了搜索引擎中文檔查詢案例的結(jié)果，與稠密 int8 模型相比，使用我們推薦的 Sparse-QAT 訓(xùn)練流水線，稀疏 int8 模型可以實(shí)現(xiàn) 1.4 倍的推理加速，準(zhǔn)確度損失可忽略不計(jì)。

文生圖案例

圖 6 展示了文生圖模型的結(jié)果，上面四張圖片是用稠密 float32 模型輸出，下面四張圖片是用稀疏 float16 模型輸出。

從結(jié)果中您會(huì)發(fā)現(xiàn)，輸入相同的提示，稀疏模型可以輸出與稠密模型相當(dāng)?shù)慕Y(jié)果。而且引入模型稀疏化和額外的漸進(jìn)式稀疏調(diào)優(yōu)使得模型從數(shù)據(jù)中學(xué)習(xí)到了更多內(nèi)容，因此部分稀疏模型的輸出結(jié)果看起來更為合理。

圖 6. 搜索引擎中的文生圖案例

（在文生圖的案例中，稀疏模型可能會(huì)產(chǎn)生比密集模型更合理的結(jié)果。）

總結(jié)

NVIDIA Ampere 架構(gòu)中的結(jié)構(gòu)化稀疏功能可以加速許多深度學(xué)習(xí)工作負(fù)載，并且易于結(jié)合 TensorRT 和 cuSPARSELt 稀疏加速庫一起使用。

如需了解更多信息，請觀看 GTC 演講：NVIDIA Ampere（https://www.nvidia.com/en-us/on-demand/session/gtcspring23-s51299/） 架構(gòu)的結(jié)構(gòu)化稀疏功能及其在騰訊微信搜索中的應(yīng)用（https://www.nvidia.com/en-us/on-demand/session/gtcspring23-s51299/）。

下載最新的 TensorRT 和 cuSPARSELt：

TensorRT 下載：

（https://developer.nvidia.cn/zh-cn/tensorrt）

cuSPARSELt 下載：（https://developer.nvidia.com/cusparselt-downloads）。

點(diǎn)擊“閱讀原文”，或掃描下方海報(bào)二維碼，在 8 月 8日聆聽NVIDIA 創(chuàng)始人兼 CEO 黃仁勛在 SIGGRAPH 現(xiàn)場發(fā)表的 NVIDIA 主題演講，了解 NVIDIA 的新技術(shù)，包括屢獲殊榮的研究，OpenUSD 開發(fā)，以及最新的 AI 內(nèi)容創(chuàng)作解決方案。