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