讓算力不再成為瓶頸,小紅書機器學(xué)習(xí)異構(gòu)硬件推理優(yōu)化之道
很多公司都在結(jié)合 GPU 的算力發(fā)展,探索出適合自己的機器學(xué)習(xí)問題解決方案。例如,小紅書在 2021 年開始進行推廣搜模型的 GPU 化改造,以提升推理性能和效率。在遷移過程中,我們也面臨一些困難,例如如何平滑遷移到異構(gòu)硬件,如何結(jié)合小紅書的業(yè)務(wù)場景和在線架構(gòu)發(fā)展出自己的解決方案等等。在全球降本增效的趨勢下,異構(gòu)計算成為了一種很有前途的方向,可以通過將不同類型的處理器(如 CPU、GPU、FPGA 等)組合在一起來提高計算性能,從而實現(xiàn)更好的效率和更低的成本。
1.背景
小紅書推薦、廣告、搜索等主要場景的模型服務(wù),統(tǒng)一由中臺推理架構(gòu)承載。隨著小紅書業(yè)務(wù)的不斷發(fā)展, 推廣搜等場景的模型規(guī)模也在不斷增大。以主推薦場景精排的主模型為例, 從 2020 年初開始,算法推出了全興趣建模,用戶歷史行為記錄長度均值擴大了約 100 倍。模型結(jié)構(gòu)也從最初的 muti-task 經(jīng)過多輪迭代,模型結(jié)構(gòu)復(fù)雜度也不斷提升 ,這些變化導(dǎo)致模型推理的浮點運算數(shù)增加了 30 倍,模型訪存增加了約 5 倍。
圖片
2.模型服務(wù)架構(gòu)概覽
模型特點:以小紅書 2022 年底的推薦主模型為例,該模型具有充分的稀疏性, 部分結(jié)構(gòu)由連續(xù)值特征和矩陣運算構(gòu)成, 也存在大規(guī)模的稀疏參數(shù)比如 < 筆記id * 用戶城市 >,單個模型的 sparse 特征多達 1TB,但通過比較有效的模型結(jié)構(gòu)優(yōu)化,dense 部分控制在 10GB 以內(nèi),可放在顯存中。用戶每刷一次小紅書,計算的總 FLOPs 達到了 40B, 超時的控制在 300ms 以內(nèi) ( 除去特征處理,帶 lookup ) 。
推理框架:在 2020 年之前,小紅書采用 TensorFlow Serving 框架作為在線服務(wù)框架,2020 年后,逐漸迭代成基于 TensorFlowCore 自研的 Lambda Service 服務(wù)。TensorFlow Serving 在進圖之前進行一次內(nèi)存拷貝 TensorProto -> CTensor,以確保模型推理的正確性和可靠性。然而,隨著業(yè)務(wù)規(guī)模的擴大,內(nèi)存拷貝操作會對模型性能產(chǎn)生影響。小紅書自研框架通過優(yōu)化免去一次不必要的拷貝,同時保留 Runtime、圖調(diào)度能力、優(yōu)化能力可插拔的特點,并為后期 TRT、BLADE、TVM 等不同優(yōu)化框架的配合使用奠定了基礎(chǔ)?,F(xiàn)在看來,在合適的時候選擇自研是一個明智的選擇, 同時為了最大化減少數(shù)據(jù)傳輸帶來的成本, 推理框架還承擔了一部分特征抽取和轉(zhuǎn)化的實現(xiàn),這里小紅書還在預(yù)估服務(wù)近側(cè)部署自研的邊緣存儲, 解決了遠端拉取數(shù)據(jù)的成本問題。
機型特性:小紅書沒有自建機房, 所有機器采購自云廠商,因此,選擇不同機型的決策很大程度上取決于能夠采購到什么型號的機器。而模型推理的計算并不是純粹的 GPU 計算,合理找到硬件配比,除考慮 GPU\CPU 外,還涉及帶寬、內(nèi)存帶寬、跨 numa 通信延遲等問題。
圖片
GPU 特性
GPU 特性:在這里,小紅書和其它公司遇到的問題是一樣的,GPU kernel 的執(zhí)行可以分為以下幾個階段:數(shù)據(jù)傳輸、kernel 啟動、kernel 計算和結(jié)果傳輸。其中,數(shù)據(jù)傳輸是將數(shù)據(jù)從主機內(nèi)存?zhèn)鬏數(shù)紾PU內(nèi)存;kernel 啟動是將 kernel 代碼從主機端傳輸?shù)?GPU 端,并在 GPU 上啟動 kernel;kernel 計算是實際執(zhí)行 kernel 代碼計算結(jié)果;結(jié)果傳輸是將計算結(jié)果從 GPU 內(nèi)存?zhèn)鬏敾刂鳈C內(nèi)存。如果大量時間都花費在數(shù)據(jù)傳輸和 kernel 啟動上,而交付給 kernel 計算的活不重,實際計算時間很短,則會導(dǎo)致 GPU 的利用率無法提升,甚至出現(xiàn)空跑的情況。
圖片
預(yù)估服務(wù)框架
3.GPU優(yōu)化實踐
3.1 系統(tǒng)優(yōu)化
3.1.1 物理機
在物理機優(yōu)化方面,可以采用一些常規(guī)的優(yōu)化思路, 主要目的是降低除 GPU 以外的其它系統(tǒng)開銷成本, 降低虛擬化的中間商賺差價。一般來說,一套系統(tǒng)優(yōu)化可以提升 1%-2% 的性能,從我們實踐來看,需要結(jié)合云廠商實際能力來進行優(yōu)化。
● 中斷隔離:將 GPU 的中斷單獨分離出來,避免因為其他設(shè)備的中斷而影響 GPU 的計算性能。
● 內(nèi)核版本升級:提高系統(tǒng)的穩(wěn)定性和安全性,提高 GPU 驅(qū)動程序的兼容性和性能。
● 指令透傳:將 GPU 的指令直接透傳到物理設(shè)備上,加速 GPU 的計算速度。
3.1.2 虛擬化和容器
在多卡情況下,將單個 pod 綁定到特定的 NUMA 節(jié)點上,從而提高 CPU 和 GPU 之間的數(shù)據(jù)傳輸速度。
● CPU NUMA Affinity,親和性是指從 CPU 角度來看,哪些內(nèi)存訪問更快,延遲更低。如前所述,與該 CPU 直接相連的本地內(nèi)存是更快的。因此,操作系統(tǒng)可以根據(jù)任務(wù)所在 CPU 來分配本地內(nèi)存,以提高訪問速度和性能,這就是基于 CPU NUMA Affinity 的考慮,盡量讓任務(wù)運行在本地的 NUMA Node 里。在小紅書場景里,CPU 上面的訪存開銷并不小。能夠讓 CPU 直連本地內(nèi)存可以節(jié)約大量的 CPU 上 kernel 執(zhí)行的耗時, 從而給 GPU 留足夠的空間。
● 將 CPU 使用率控制在 70% 下,可以將延遲由 200ms -> 150ms。
3.1.3 鏡像
編譯優(yōu)化。不同 CPU 對指令級的支持能力是有差異的, 不同云廠商采買的機型也有所不同。一個比較簡單的思路是在不同的硬件場景下,編譯鏡像時帶上不同的指令集合。在實現(xiàn)算子時,大量的算子本身已經(jīng)帶有如 AVX512 等指令。以阿里云的 Intel(R) Xeon(R) Platinum 8163 + 2 A10 的機型為例,我們根據(jù)該機型的特點和支持的指令集,編譯優(yōu)化調(diào)整合適的指令集, 整體相比于不進行指令優(yōu)化的情況下,在該機型上的 CPU 吞吐量提高了 10% 。
# Intel(R) Xeon(R) Platinum 8163 for ali intel
build:intel --copt=-march=skylake-avx512 --copt=-mmmx --copt=-mno-3dnow --copt=-msse
build:intel --copt=-msse2 --copt=-msse3 --copt=-mssse3 --copt=-mno-sse4a --copt=-mcx16
build:intel --copt=-msahf --copt=-mmovbe --copt=-maes --copt=-mno-sha --copt=-mpclmul
build:intel --copt=-mpopcnt --copt=-mabm --copt=-mno-lwp --copt=-mfma --copt=-mno-fma4
build:intel --copt=-mno-xop --copt=-mbmi --copt=-mno-sgx --copt=-mbmi2 --copt=-mno-pconfig
build:intel --copt=-mno-wbnoinvd --copt=-mno-tbm --copt=-mavx --copt=-mavx2 --copt=-msse4.2
build:intel --copt=-msse4.1 --copt=-mlzcnt --copt=-mrtm --copt=-mhle --copt=-mrdrnd --copt=-mf16c
build:intel --copt=-mfsgsbase --copt=-mrdseed --copt=-mprfchw --copt=-madx --copt=-mfxsr
build:intel --copt=-mxsave --copt=-mxsaveopt --copt=-mavx512f --copt=-mno-avx512er
build:intel --copt=-mavx512cd --copt=-mno-avx512pf --copt=-mno-prefetchwt1
build:intel --copt=-mno-clflushopt --copt=-mxsavec --copt=-mxsaves
build:intel --copt=-mavx512dq --copt=-mavx512bw --copt=-mavx512vl --copt=-mno-avx512ifma
build:intel --copt=-mno-avx512vbmi --copt=-mno-avx5124fmaps --copt=-mno-avx5124vnniw
build:intel --copt=-mno-clwb --copt=-mno-mwaitx --copt=-mno-clzero --copt=-mno-pku
build:intel --copt=-mno-rdpid --copt=-mno-gfni --copt=-mno-shstk --copt=-mno-avx512vbmi2
build:intel --copt=-mavx512vnni --copt=-mno-vaes --copt=-mno-vpclmulqdq --copt=-mno-avx512bitalg
build:intel --copt=-mno-movdiri --copt=-mno-movdir64b --copt=-mtune=skylake-avx512
3.2 計算優(yōu)化
3.2.1 充分使用算力
● 計算優(yōu)化,首先需要充分了解硬件性能,將其吃透。在小紅書的場景中,如下圖所示,我們遇到了兩個核心問題:
1. CPU 上的訪存較多, 內(nèi)存 page fault 頻率較高,導(dǎo)致 CPU 資源浪費,以及請求 latency 過高
2. 在線推理服務(wù)中,計算通常具有兩個特點:單次請求的 batch size 小,單個服務(wù)的并發(fā)規(guī)模大。小 batch size 會導(dǎo)致 kernel 無法充分利用 GPU 的計算能力。GPU kernel 執(zhí)行時間一般較短,無法充分掩蓋 kernel launch 的開銷,甚至 kernel launch 的時間比 kernel 執(zhí)行時間還長。在 TensorFlow 中,單個 Cuda Stream launch kernel 成為瓶頸,導(dǎo)致推理場景下 GPU 利用率只有 50% 。此外,對于小模型場景(簡單的 dense 網(wǎng)絡(luò)),用 GPU 替換 CPU 完全不劃算,限制了模型的復(fù)雜度。
圖片
● 為解決上述兩個問題,我們采取了以下措施:
1. 針對內(nèi)存 page fault 頻率高的問題,我們使用 jemalloc 庫來優(yōu)化內(nèi)存回收機制,并開啟了操作系統(tǒng)的透明大頁功能。此外,針對 lambda 特殊的內(nèi)存訪問特點,我們設(shè)計專門的數(shù)據(jù)結(jié)構(gòu),并優(yōu)化內(nèi)存分配策略,盡可能地避免內(nèi)存碎片。與此同時,我們直接繞開了 tf_serving 的接口,直接調(diào)用 TensorFlow,減少了一次數(shù)據(jù)的序列化與反序列化。這些優(yōu)化在首頁精排和內(nèi)流精排場景下,提升了 10+% 的吞吐,在廣告大多數(shù)場景下,降低了 50% 的 latency。
圖片
兼容 tensorflow::Tensor 格式,在將特征傳遞給 tensorflow::SessionRun 之前是零拷貝
2. 針對 TensorFlow 單 Cuda Stream 的問題,我們支持了 Multi Streams , Multi Contexts 的功能,避免了互斥鎖導(dǎo)致的性能瓶頸,成功將 GPU 利用率提升到 90+% 。同時,我們利用 Nvidia 提供的 Cuda MPS 功能,實現(xiàn)了 GPU 的空分復(fù)用 (同一時間支持多個 kernel 執(zhí)行),使得 GPU 的利用率進一步提升?;诖?,Search 的排序模型成功在 GPU 上實現(xiàn)。此外,我們也在其他業(yè)務(wù)線上成功落地,包括首頁初排、廣告等等。下表是在搜索排場景下的一個優(yōu)化情況。
圖片
3. Op/Kernel fusion 技術(shù):通過手寫或者圖編譯優(yōu)化工具生成性能更高的 Tensorflow 算子,充分利用 CPU 的 Cache 以及 GPU 的 Shared Memory,提升系統(tǒng)的吞吐。
圖片
在內(nèi)流場景下, 算子進行融合,可以看到單次調(diào)用 12ms -> 5ms
3.2.2 避免算力浪費
1. 系統(tǒng)鏈路上存在優(yōu)化空間
a. 初排前置計算:在處理用戶側(cè)相關(guān)計算時,初排需要計算大量筆記,例如以外流為例,需要計算約 5000 篇筆記,lambda 對其有切片處理。為避免重復(fù)計算,將初排的用戶側(cè)計算前置到和召回階段并行,從而使得用戶向量的計算從多次重復(fù)變成了只需要 1 次,在粗排場景下優(yōu)化了 40% 機器。
2. 圖內(nèi)訓(xùn)練到推理過程中:
a. 計算前置:通過 graph freeze 可以將一部分計算提前處理。在推理時,不需要重復(fù)計算。
b. 產(chǎn)出模型 freeze 優(yōu)化:模型產(chǎn)出時把所有的參數(shù)和圖本身一起生成凍結(jié)圖( frozen graph )并進行預(yù)處理計算,可以將很多預(yù)計算的 Variable 算子轉(zhuǎn)換成 Const 算子( GPU 使用率下降 12% )
c. 推理場景下的合并計算:每個 batch 只包含一個 user , 即用戶側(cè)存在大量重復(fù)計算, 具備合并的可能性
d. CPU/GPU 算子拆分:將 lookup 之后的全部算子移至 GPU ,避免了 CPU 和 GPU 之間的數(shù)據(jù)拷貝
e. GPU 到 CPU 數(shù)據(jù)拷貝:將數(shù)據(jù)打包一次拷貝???????
f. BilinearNet 算子 GPU cuda 實現(xiàn):通過 GPU 加速計算,提升性能???????
g. 部分算子 GPU 化:省去 CPU -> GPU 拷貝???????
h. BatchNorm & MLP 合并:通過實現(xiàn)新的 MLP 層,根據(jù)一個目標減少進 GPU 的次數(shù) ( N -> 1), 增大一次計算的計算量(重復(fù)利用 GPU 小核心的并發(fā)能力)???????
圖片
3.2.3 全天動態(tài)算力
● 動態(tài)計算降級提升全天資源使用效率,秒級別的對 lambda 負載進行自動負反饋調(diào)整,做到對單區(qū)壓測之前不需要人工做降級準備。
● 在外流精排、外流初排、內(nèi)流精排、內(nèi)流初排、搜索等主要業(yè)務(wù)場景均已經(jīng)上線。
● 在多個業(yè)務(wù)線解決了容量問題,有效緩解了業(yè)務(wù)增長導(dǎo)致的資源線性上升,同時大幅提升了系統(tǒng)的魯棒性。在功能上線后的業(yè)務(wù)線中,均沒有出現(xiàn)因為瞬間成功率大幅下降導(dǎo)致的 P3 及以上事故。
● 大幅提升全天資源使用效率,以內(nèi)流精排為例(如下圖所示),五一假期三天的 10:00-24:00 的 CPU 使用核數(shù)均保持 50 核的一條平線(抖動對應(yīng)發(fā)版)
圖片
3.2.4 換更好的硬件
● A10 GPU 的性能是 T4 GPU 性能的 1.5 倍,同時 A10 機型配備的 CPU ( icelake, 10nm ) 比 T4 機型 ( skylake, 14nm ) 更新一代,價格僅為 T4 機型的 1.2 倍。未來我們還會考慮在線使用 A30 等機型。
3.3 圖優(yōu)化
圖片
3.3.1 DL 棧的自動編譯優(yōu)
● BladeDISC 是阿里最新開源的基于 MLIR 的動態(tài) shape 深度學(xué)習(xí)編譯器, 小紅書的自動圖優(yōu)化部分來自于這套框架( Blade 推理加速庫是 Apache 2.0 開源,可以跨任何云使用,無知識產(chǎn)權(quán)風險)。該框架提供了 TF 圖編譯優(yōu)化(包含 Dynamic Shape Compiler ,稀疏子圖優(yōu)化),同時能疊加我們本身做的算子定制化優(yōu)化,可以較好的適配我們的業(yè)務(wù)場景。在壓測單機 inference 中,QPS 能提升 20% 。
● 這套框架關(guān)鍵技術(shù)
(1) MLIR 基礎(chǔ)架構(gòu)
MLIR,即多層次中間表示語言(Multi-Level Intermediate Representation),是由 Google 發(fā)起的開源項目。其目的是提供一個靈活、可擴展的多層 IR 基礎(chǔ)設(shè)施和編譯器實用工具庫,為編譯器和語言工具的開發(fā)者提供一個統(tǒng)一的框架。
MLIR 的設(shè)計受到 LLVM 的影響,但與 LLVM 不同的是,MLIR 主要關(guān)注于中間表示( IR )的設(shè)計和擴展。MLIR 提供了一個多層次的 IR 設(shè)計,可以支持從高層語言到底層硬件的編譯過程,并提供了豐富的基礎(chǔ)設(shè)施支持和模塊化設(shè)計架構(gòu),使得開發(fā)者可以很方便地擴展 MLIR 的功能。此外,MLIR 還具有較強的膠水能力,可以與不同的編程語言和工具進行集成。MLIR 是一個強大的編譯器基礎(chǔ)設(shè)施和工具庫,為編譯器和語言工具的開發(fā)者提供了一種統(tǒng)一的、靈活的中間表示語言,可以方便地進行編譯優(yōu)化和代碼生成。
(2) 動態(tài) shape 編譯
靜態(tài) shape 的限制意味著在編寫深度學(xué)習(xí)模型時需要提前確定每個輸入和輸出的形狀,并且不能在運行時改變它們。這限制了深度學(xué)習(xí)模型的靈活性和可擴展性,因此需要一種支持動態(tài) shape 的深度學(xué)習(xí)編譯器。
3.3.2 精度調(diào)整
● 量化的實現(xiàn)方式之一是使用 FP16
FP16 計算優(yōu)化:在 MLP 層時用 FP16 替換 FP32 計算,能夠較大地減少 GPU 使用率(相對下降 13% )
在調(diào)整 FP16 的過程中,選擇白盒方式進行精度優(yōu)化意味著可以更加精細地控制哪些層使用低精度計算,并能夠根據(jù)經(jīng)驗進行不斷調(diào)整和優(yōu)化。這種方式需要對模型結(jié)構(gòu)有較為深入的了解和分析,可以根據(jù)模型的特性和計算要求進行有針對性的調(diào)整,以達到更高的性價比。
相比之下,黑盒方式則相對簡單,不需要了解模型的內(nèi)部結(jié)構(gòu),只需要設(shè)置一定的容忍閾值即可完成精度優(yōu)化。這種方式的優(yōu)點是操作簡單,對模型同學(xué)的要求也相對較低,但是可能會犧牲一定的性能和精度。
因此,選擇白盒還是黑盒方式進行精度優(yōu)化需要根據(jù)具體情況而定。如果需要追求更高的性能和精度,同時擁有足夠的經(jīng)驗和技術(shù)能力,那么白盒方式可能更加適合。如果操作簡單、快速迭代更加重要,那么黑盒方式可能更加實用。
4.總結(jié)
從 2021 年開始到 2022 年底, 經(jīng)過本項目優(yōu)化,小紅書推理計算算力增加 30 倍,關(guān)鍵用戶指標提升 10%+ ,同時累積節(jié)約集群資源 50%+ 。在我們看來,小紅書在 AI 技術(shù)方面的發(fā)展路徑應(yīng)該是以業(yè)務(wù)需求為導(dǎo)向,平衡技術(shù)和商業(yè)的發(fā)展:實現(xiàn)技術(shù)創(chuàng)新的同時,也要考慮成本、效益和可持續(xù)性。以下是一些優(yōu)化過程中的思考:
優(yōu)化算法和提高系統(tǒng)性能。這是小紅書機器學(xué)習(xí)團隊的核心任務(wù)。優(yōu)化算法和提高系統(tǒng)性可以更好地支持業(yè)務(wù)需求,提高用戶體驗。然而,在資源有限的情況下,團隊需要明確優(yōu)化的重點,避免過度優(yōu)化。
構(gòu)建基礎(chǔ)設(shè)施和提高數(shù)據(jù)處理能力。基礎(chǔ)設(shè)施對于支持 AI 應(yīng)用是非常關(guān)鍵的。小紅書可以考慮進一步投入基礎(chǔ)設(shè)施的建設(shè),包括計算和存儲能力、數(shù)據(jù)中心和網(wǎng)絡(luò)架構(gòu)等。此外,提高數(shù)據(jù)處理能力也是非常重要的,可以更好地支持機器學(xué)習(xí)和數(shù)據(jù)科學(xué)應(yīng)用。
提高團隊人才密度和組織架構(gòu)。一個優(yōu)秀的機器學(xué)習(xí)團隊需要擁有具有不同技能和背景的人才,包括數(shù)據(jù)科學(xué)家、算法工程師、軟件工程師等;優(yōu)化組織架構(gòu)也有助于提高團隊效率和創(chuàng)新能力。
合作共贏和開放創(chuàng)新。小紅書持續(xù)與其他公司、學(xué)術(shù)機構(gòu)和開源社區(qū)合作,共同推進 AI 技術(shù)的發(fā)展,這有助于小紅書獲取更多資源和知識,成為更加開放和創(chuàng)新的組織。
該方案讓小紅書機器學(xué)習(xí)架構(gòu)水平達到了業(yè)界一流水平。未來,我們將不斷推進引擎升級和降本增效, 引入新技術(shù)提高小紅書機器學(xué)習(xí)的生產(chǎn)力, 將更加結(jié)合小紅書的實際業(yè)務(wù)場景, 從單模塊的優(yōu)化升級為全系統(tǒng)優(yōu)化,并進一步引入業(yè)務(wù)側(cè)流量的個性化差異特征, 將降本增效做到極致。期待有志之士,一同加入我們!
5.團隊
張楚嵐(杜澤宇):商業(yè)技術(shù)部
畢業(yè)于華東師范大學(xué), 商業(yè)化引擎團隊負責人, 主要負責商業(yè)化在線服務(wù)搭建。
陸光(彭鵬):智能分發(fā)部
畢業(yè)于上海交通大學(xué) ,機器學(xué)習(xí)引擎工程師,主要負責 Lambda GPU 優(yōu)化。
伊恩(陳建新):智能分發(fā)部
畢業(yè)于北京郵電大學(xué),機器學(xué)習(xí)引擎工程師, 主要負責 Lambda 參數(shù)服務(wù)器和 GPU 優(yōu)化。
赤羽(劉兆宇):智能分發(fā)部
畢業(yè)于清華大學(xué),機器學(xué)習(xí)引擎工程師, 主要負責特征引擎方向的相關(guān)研究和探索。
特別感謝 :智能分發(fā)部 所有同學(xué)