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

一、前言

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

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

二、背景

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

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

三、模型服務(wù)

1、模型特點(diǎn)

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

2、訓(xùn)練&推理框架

接下來介紹小紅書的訓(xùn)練和推理框架。

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

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

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

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

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

3、機(jī)器特性

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

4、GPU 特性

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2、計(jì)算優(yōu)化

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

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

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

（1）針對內(nèi)存 page fault 頻率高的問題

首先，針對內(nèi)存 page fault 較高的問題，有兩個(gè)解決思路。

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

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

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

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

（3）計(jì)算合并

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

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

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

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

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

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

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

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

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

（5）換更好的硬件

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

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

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

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

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

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

3、訓(xùn)練優(yōu)化

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

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

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

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

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

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

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

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

4、未來

最后介紹一下小紅書整個(gè)機(jī)器學(xué)習(xí)引擎未來的演進(jìn)方向。

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

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

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