本文將分享小紅書推搜場景下，全 GPU 化建設過程中的模型服務、GPU 優(yōu)化等相關工作。

一、前言

近年來，機器學習領域的視頻、圖像、文本和推廣搜等應用不斷發(fā)展，其模型計算量和參數量遠遠超過了 CPU 摩爾定律的增長速度。在此背景下，GPU 的算力發(fā)展和大模型的發(fā)展不謀而合。很多公司都在結合 GPU 的算力發(fā)展，探索出適合自己的機器學習問題解決方案。

與其他公司類似，小紅書各個場景中的模型也在不斷變大，而 CPU 的發(fā)展跟不上其所需的算力。因此，我們在 21 年時開始進行推廣搜模型的 GPU 化改造，以提升推理性能和效率。在遷移過程中，我們也面臨一些困難，比如如何把之前 CPU架構的工作平滑遷到 GPU 架構上；如何結合小紅書的業(yè)務場景和在線架構發(fā)展出自己的解決方案等。同時我們需要做到降本增效，幫助模型持續(xù)迭代。

二、背景

在介紹具體實踐之前，先來介紹一下小紅書的應用場景。小紅書 APP 主頁包括推薦頁、搜索頁，這些頁面還包括了廣告成分。其中精排 CTR model、CVR model、相關性 model 以及部分排序場景、召回場景都會用到一些 GPU。目前精排場景已經全部遷移到 GPU 推理，而精排的主要目的是將 CTR、CVR 或者其他多個目標估計準確。

從計算參數量來說，小紅書的計算規(guī)模從 21 年初到 22 年底擴大了很多。以推薦場景為例，每個請求要花 400 億的 Flops，整個參數量達到了千億量級。

三、模型服務

1、模型特點

在 22 年底 ChatGpt 類模型提出之前，工業(yè)界或者推搜類公司，暫時還沒有特別大的 Dense 模型的應用場景。因此，在 22 年底之前，小紅書主要模型的大參數量主要是進行充分稀疏化。具體而言，以推薦主模型為例，有大量參數需要與 ID 類型進行交叉：比如小紅書筆記與用戶城市交叉，小紅書筆記與用戶 ID 交叉等，構建特征 Embedding 則成為參數稀疏化過程。整個稀疏化因為笛卡爾積問題，參數量可以達到 TB 千億或者萬億級別，很多頭部公司達到了萬億乃至十萬億的參數量。但考慮到成本以及 CTR 模型收益的局限性，我們并沒有把 Dense 部分計算量做得非常大。Dense 部分計算基本控制在 10GB 以內，也就是一張顯卡能容納的狀態(tài)。不過由于場景的不同，也會存在其他的差異之處。例如，在線場景，除了有計算量的限制，還存在高并發(fā)、延遲等限制。

2、訓練&推理框架

接下來介紹小紅書的訓練和推理框架。

我們在使用 GPU 的時候,主要分為兩大類模型，一類是推搜廣類模型，包括前面介紹的稀疏 CTR 類模型，它們有自己的推理框架和訓練框架；而另一類模型包括CV、NLP 類模型，是基于 PyTorch 技術棧搭建的，有另外獨立的訓練推理框架演進。

CTR 類大規(guī)模技術化模型是 TB 級的模型。下面分別介紹推理和訓練部分。

推理框架部分，2020 年之前小紅書使用的是 TensorFlow 技術棧，采用TensorFlow Serving 作為基礎推理框架，并在此開源框架基礎上進行自研。

TensorFlow Serving 的特點是托管了很多模型生命周期的管理，外圍存在大量組件。但是，TensorFlow Serving 自身也有一些額外開銷，尤其是它存在浪費的內存拷貝。原始 TensorFlow 底層的 API，是基于 CTensor 進 TensorFlow 的圖引擎，但 TensorFlow Serving 為了抽象接口，在最外面用 Proto 封裝了一層 TensorProto 結構來響應請求。因此，我們在 20 年時進行了優(yōu)化，把外層的TensorProto 去掉，不直接依賴于 TensorFlow Serving，而是直接基于底層的 CTensor API 也就是 CPI 搭建了自己的 Lambda Service 推理框架。在后續(xù)的迭代中繼續(xù)集成了圖調度器，編譯優(yōu)化等技術。

針對訓練框架部分，同樣在 20 年之前，小紅書沒有大規(guī)模的稀疏訓練框架，我們在 20 年初搭建完畢后，搜推廣場景的訓練框架統(tǒng)一到以 TensorFlow 為底層基礎的自研 Worker & Ps 訓練框架。同時我們也自研了自己的算子等。后來為了解決升級 GPU 后，block 在 CPU、IO 等存在的問題，我們對此自研訓練框架進行了多輪技術升級。

3、機器特性

機器特性方面，小紅書沒有自建機房，所有機器采購自云廠商。但能夠采購的機器型號有限，無法定制，包括定制帶寬、定制內存、定制卡數等等。因此在技術選型時，需要根據可采購到的機器型號做出最優(yōu)化的架構選擇。

4、GPU 特性

針對 GPU 特性問題，小紅書與其他公司遇到的問題是一樣的，GPU Kernel 執(zhí)行分為幾個階段：數據傳輸、Kernel 啟動、Kernel 計算與結果傳輸。首先數據需要從主機內存?zhèn)鞯?GPU 內存，Kernel 啟動需要將 Kernel 代碼從主機端傳到 GPU 端，并在 GPU 上啟動 Kernel，Kernel 執(zhí)行計算結果后，將結果從 GPU 傳輸到主機端。整個過程有傳輸、計算、傳輸的三段過程，如果大量時間都花在數據傳輸與 Kernel 啟動上，則會導致 GPU 利用率較低，甚至出現(xiàn)跑空的情況。

四、GPU 優(yōu)化實踐

針對上述問題，小紅書提出了以下具體的實踐優(yōu)化方案。

1、系統(tǒng)優(yōu)化

（1）物理機優(yōu)化

首先針對系統(tǒng)優(yōu)化中的物理機優(yōu)化方面。前面介紹到，我們采購的是云廠商機器，云廠商的機器為了屏蔽硬件的復雜性，做了大量的虛擬化。虛擬化本身是好的，在K8S 生態(tài)下面可以讓應用具有彈性，屏蔽掉硬件細節(jié)。但虛擬化會導致性能變差，硬件速度跟不上 GPU 速度。因此，我們與云廠商針對以下幾點進行了合作，降低虛擬化中間過程。

• 中斷隔離：把 GPU 的中斷單獨分離開來，比如 Agent 隔離、其他中斷隔離等，主要目的是減少虛機的抖動。
• 內核版本升級：用于提高 GPU 驅動系統(tǒng)的兼容性和性能。
• 指令透傳：將 GPU 的指令直接透傳到物理設備上，加速 GPU 計算速度。以上幾點整體基于物理機的系統(tǒng)優(yōu)化可以提升 1%-2% 的性能。

（2）多卡優(yōu)化

第二部分，我們進行了多卡優(yōu)化。

由于我們采購的 CPU 是 AMD 的，存在 NUMA 的問題：在線階段跨 NUMA 訪問內存的時間遠高于 CPU 直接訪問本地內存的時間。因此我們在使用過程中對Pod 進行了 NUMA 綁定，綁定后的 Pod 不需要跨 NUMA 內存申請與節(jié)點調用，提高了 CPU 與 GPU 之間的數據傳輸速度，減少了 CPU 階段的內存消耗。CPU 階段速度變快，可以給后面 GPU Inference 留更多的時間。

（3）編譯優(yōu)化

第三部分的系統(tǒng)優(yōu)化也是非常通用的。隨著 CPU 的發(fā)展，不同機型支持的指令集不同，因此 Kernel 實現(xiàn)時，不同機型會存在大量不同的定義。為了解決這個問題，我們針對不同機型的指令集做了交叉編譯，需要找到 Kernel 對應機型的最優(yōu)實現(xiàn)，打開對應的編譯指令。

以我們采購的阿里云的英特爾 8163 加兩張 A10 的機型為例，我們根據該機型的特點定制了它的編譯指令集，因此針對該機型，我們的 Runtime 換掉了它可執(zhí)行 Runtime 的 Binary。

系統(tǒng)階段的編譯指令優(yōu)化，可以帶來約 10% 的性能提升。

2、計算優(yōu)化

計算優(yōu)化方面，仍然是圍繞前面的核心問題：GPU 運行速度快，CPU 內存訪問跟不上的時候如何解決。

CPU 訪存較多，內存 page fault 頻率較高，會導致 CPU 資源浪費、latency 高的問題。Latency 高時，在線推理階段可能會出現(xiàn)超時問題。同時，在線推理服務中，單次請求的 batch size 較小，單個服務的并發(fā)規(guī)模大，上千規(guī)模的 QPS 單次請求較小的 batch，無法充分利用 GPU 的算力，會導致原始 TensorFlow 中，單個 Cuda Stream launch kernel 成為瓶頸，此時推理場景下的 GPU 利用率只有 50%。此外，對于小模型場景，Dense 部分較為簡單，數據傳輸到 GPU 上推理計算再傳輸回 CPU 上的成本比直接 CPU 計算還要高，較為不劃算。

因此，面對不同的場景，我們需要解決這兩個問題：Page fault 、Latency 較高，Batch 約束問題。

（1）針對內存 page fault 頻率高的問題

首先，針對內存 page fault 較高的問題，有兩個解決思路。

一是在引擎?zhèn)葍?yōu)化數據結構。把 Tensor 數據結構給到圖調度引擎之前，需要減少額外開銷，做到零拷貝，也就是把用戶側的請求數據直接搬到圖引擎的輸入。因此需要針對 TensorFlow 底層 CPI 定制數據結構，該定制化的數據結構接請求側，直接傳到底層圖引擎，降低成本開銷。

二是在操作系統(tǒng)層面開啟透明大頁功能減少 page fault。同時我們也使用了 jemalloc 庫，通過把垃圾回收時間設長、不斷調整到最優(yōu)狀態(tài)來降低額外開銷，解決 page fault 高的問題，。

（2）針對 TensorFlow 單 Cuda Stream 的問題

此外，針對 TensorFlow 單 Cuda Stream 問題，我們支持了 Multi Cuda Streams、Multi Contexts 的功能，可以避免互斥鎖的性能瓶頸，提高并發(fā)率與 GPU 利用率。同時，針對 GPU 多卡并發(fā)使用時存在的問題，我們利用了 Nivida 提供的 Cuda MPS 功能，實現(xiàn)了 GPU 的空分復用。原本在使用 Nvidia 進行 Cuda 運算時，正常情況下一個時刻只能運行一個 Context（一個 CPU 進程下進行 Cuda 程序調用時，GPU 端申請資源與數據的唯一單位）。但如果有些卡支持Hyper-Q 功能，則可以開啟 Nvidia 的 MPS 服務，實現(xiàn) GPU 的空分復用（同一時間支持多個 Cuda 執(zhí)行）來進一步提高 GPU 利用率。但空分復用可能會導致進程故障隔離問題，因此，GPU 的空分復用建議是在相對固定的場景下開啟，比如推搜等具體場景、計算任務固定的情況下使用。

（3）計算合并

計算優(yōu)化的第三部分是在整個計算過程中找到一些計算合并的可能性。

除了 Kernel 級別可以進行 fusion，上游 Op 級別也可以進行 fusion 合并。我們通過手寫或者圖編譯優(yōu)化工具生成了性能更高的 Tensorflow 算子。如 Mat 加 Bn 與 Relu 的計算，可以用一個優(yōu)化后的 FusedMatmul 算子實現(xiàn)相同的功能。集成了 Compile 級別的優(yōu)化之后可以把算子重新編譯，用新的算子執(zhí)行來降低它的計算開銷。

在右圖所示的我們的一個場景中，核心計算成本可以降到 50%。

（4）冗余計算優(yōu)化

計算優(yōu)化的一個較為明智的思路是找到冗余的計算量進行優(yōu)化。

推搜類的 CTR 模型存在了大量的冗余計算：推搜場景較為突出的特點是 1 個 user + N 個 item 的預估，而 Tensor 結構在原始計算 input 階段需要展平，也就是 user 側計算從數據側會拆分成 N 份，而這部分計算是冗余的。第二個冗余是在使用原始 TensorFlow API 搭建整個 graph 時，一次請求數據會多次進GPU 計算處理，會有部分計算在 CPU 結束后傳給 GPU，而這部分是不夠高效的。

因此我們在此進行了算子合并、優(yōu)化推導，利用 CUDA 的 API 統(tǒng)一實現(xiàn)，讓一次請求只進一次 GPU，減少多次進 GPU 的冗余計算功能。

此外，還可以計算前置，減少冗余計算。比如在產出模型時，需要對變量算子進行凍結，把變量算子轉換成常量算子，具體舉例來說，在圖凍結過程中，有+1+2 計算，可以靜態(tài)地直接轉換成 +3。

圖凍結化在產出中讓整個計算使用率下降了 12%，非常有效。合并計算部分我們需要盡量做到用 GPU 的 CUDA 實現(xiàn)，減少外部數據計算完再進 GPU 的拷貝問題，會帶來較高的性能提升。

（5）換更好的硬件

計算優(yōu)化中一個非常直接的優(yōu)化就是換更好的硬件。

在線推理階段沒必要用到 A100 這種特別貴的卡，可以用 A10 替換 T4，提升了 1.5 倍的性能，但只花了 1.2 倍的錢。未來我們還會考慮在線使用 A30 等機型進行推理。同時，隨著硬件廠商的版本迭代，我們也會直接去換更好的硬件，取得線上的成本收益。

但換了硬件后，前置放在 CPU 上的計算部分會更加雞肋，總會存在一些 Vlookup 特征，合并交叉哈希計算等不適合在 GPU 上處理。這時候則需要推進到第三步，也就是我們會浪費一些時延，把前面 CPU 計算部分拆分出去。Inference 架構會變?yōu)檎埱笙冉涍^ CPU 計算，后進行 GPU 在線推理。后面隨著硬件 GPU 不斷升級，CPU 實在跟不上的情況下，則會繼續(xù)拆分。

（6）DL 棧自動編譯優(yōu)化

DL 棧自動編譯優(yōu)化，也是一個核心的優(yōu)化點。

我們核心用的是阿里開源的 BladeDISC 以及 TensorFlow 官方的 XLA 技術棧。阿里的 Blade 框架在 MRI 基礎上支持了動態(tài) Shark 的深度學習編譯功能。傳統(tǒng)編譯器，以 C++ 來說，具體的 C++ code 會通過編譯優(yōu)化編譯成機器指令，也就是 Machine code 被機器對應的系統(tǒng)運行。在深度學習 ML 編譯器里，需要把圖編譯成 DL 技術棧，找一層中間實現(xiàn)。比如說 MLIR，基于這個中間層表達去編譯出硬件上最終的執(zhí)行態(tài)，這個硬件態(tài)是 GPU 上真正的 Kernel 的執(zhí)行計算。這部分我們集成了 blade、XLA 技術棧，取得了非常大的收益。在這里要特別安利一下 Blade 框架，它給我們帶來了很好的效果，它也是 Apache 2.0 開源的，所有公司都可以使用。

3、訓練優(yōu)化

（1）訓練優(yōu)化 1、2 期

第三部分主要介紹我們訓練早期的一些優(yōu)化。

訓練與推理遇到的問題會有一點差別，訓練是沒有 latency 約束的，不要求在多少毫秒內返回；但它吞吐更大，希望能夠拉取更多的數據。數據 Embedding Layer 過程能夠在 IO 階段完成，同時做到后向梯度反傳成本盡量低。把更多的時間和數據傳給 GPU 計算。在這個過程中，我們遇到了如下問題。

首先是數據 IO 問題，原始 TF 的訓練數據是基于 TFRecord 行存的，也就是數據展平后放在了一行，這會導致 IO 非常大。因此我們首先把 input 數據行轉列，降低 IO 的拉取數據維度。對于一些可枚舉的特征 input 如城市、ID 等，通過行轉列后，IO 拉取數據維度的規(guī)模則會大幅下降。數據行轉列存儲之后，訓練框架需要對應進行升級，適配列存場景。同時列存與行存有較大的數據分布差異，行存的打散會較好，因此，列存的時候盡量要做到列存，然后恢復數據分布，進 worker 進行訓練，然后再進 GPU 進行處理。

拉取完數據后，需要進行數據的前向 lookup。我們?yōu)榱颂岣?GPU 的利用率，較為激進地做了完全 prefetch 化，可以解決訓練過程中，因 CPU 側 lookup 阻塞造成的訓練等待。但此種更新變成了全異步更新，存在更新帶差，影響效果的穩(wěn)定性，需要機制設計進行強保障。在后向我們也進行了升級，如解綁了訓練對反向更新本身的依賴，升級了 Lookup 反向更新異步策略，對數據梯度進行了累計更新：下一個 batch 訓練梯度更新時，本次更新則可以寫入。

（2）訓練推理優(yōu)化

這部分介紹了我們整體的訓練推理優(yōu)化。

我們整個機器學習是基于訓練和推理的應用層實現(xiàn)的，主要包括應用層、軟件層、硬件容器層。圖中列了我們從 21 年到 22 年所做的工作，每個工作有自己獨立的約束與應用場景。如果大家有做相似的推搜大模型，存在沒有嘗試過的方向，可以參考其中的內容。舉例來說，如精度壓縮，可以用 Float 16 或者 Int 8 等更低級別的精度。但在實際的使用場景中，存在著對效果影響把控的問題。因此，我們使用精度壓縮時，更多的是做白盒，也就是在導出模型時，并非對所有層都進行降精度，而是挑選進 GPU 的部分來降精度，同時精度下降程度要通過后面的優(yōu)化邏輯確保整個 AUC 或者實際指標無損，也就是在白盒情況下合理選擇精度壓縮。

4、未來

最后介紹一下小紅書整個機器學習引擎未來的演進方向。

對于稀疏大模型訓練部分，我們會做幾件事情：一是在訓練端會區(qū)分不同的參數規(guī)模，對于特別小的參數規(guī)模，我們會 HPC 化，盡量放在一臺機器內搞定，不產生額外的開銷。對于達到 1TB 級別參數量的，我們會使用 A100+全高速互聯(lián)。對于特別大參數量，達到幾十 TB 或者幾百 TB 的，我們會做基建通信，采用小 GPU（A10）+ 基建通信的能力來實現(xiàn)。

對于推理方面，未來我們會升級哈希、多級緩存以及模型的輕量化技術。同時我們也會跟著 Nvidia 的升級節(jié)奏，更新整個驅動、硬件的版本。

最后為了優(yōu)化整個公司的迭代效率，我們后續(xù)會做拖拽式、畫布式的一站式機器學習平臺。同時也會做 DSL 化的特征管理，把特征管控升級成抽象算子處理。