作者

metaapp-推薦廣告研發(fā)部：臧若舟，朱越，司靈通

1 背景

推薦場景大模型在國內(nèi)的使用很早，早在 10 年前甚至更早，百度已經(jīng)用上了自研的大規(guī)模分布式的 parameter server 系統(tǒng)結(jié)合上游自研的 worker 來實現(xiàn) TB 級別的萬億參數(shù)的稀疏模型。后來，各家平臺也陸續(xù)基于這種方案，開發(fā)了自己的分布式訓(xùn)練系統(tǒng)，普遍特點是大量使用 id embedding，因此參數(shù)量巨大，模型大小也非?？鋸?。當(dāng)然，隨著開源訓(xùn)練工具 TensorFlow/Pytorch 的流行，使用 TensorFlow/Pytorch 作為 worker，結(jié)合自研 ps 的方案也十分流行。究其原因，以 TensorFlow 為例，雖然內(nèi)置了分布式訓(xùn)練系統(tǒng)，但是對于大規(guī)模 id embedding 的支持卻非常糟糕，無法作為完整的平臺使用。而使用 TensorFlow+ 自研 ps 的方案也存在不少問題，比如自研 ps 一般對于特征輸入都有特定的要求、二次開發(fā)成本比較高等。

一個典型的分布式 worker-ps 架構(gòu)

2 業(yè)務(wù)介紹

metaapp- 推薦廣告研發(fā)部，主要負(fù)責(zé) metaapp 拳頭產(chǎn)品 233 樂園的首頁信息流的推薦和廣告系統(tǒng)，是比較傳統(tǒng)的推廣搜組。我們在 2020 年之前也是采用了 TensorFlow+ 自研分布式 ps 的方案，模型大小在接近 TB 級別（業(yè)務(wù)體量較?。麄€方案的迭代和維護(hù)成本都比較高。

在這種背景下，經(jīng)過多方考量，阿里云機器學(xué)習(xí)平臺 PAI 開源的 DeepRec（脫胎于 PAI-TF），作為支持了淘寶搜索、猜你喜歡、定向、直通車等核心業(yè)務(wù)的訓(xùn)練平臺，直接基于 TensorFlow 做二次開發(fā)，針對稀疏模型在分布式、圖優(yōu)化、算子、Runtime 等方面進(jìn)行了深度的性能優(yōu)化，并且完全開源。

而因為我們公司本身跟阿里云有著深度的合作，阿里云也主動介紹了當(dāng)時還是內(nèi)部項目的 DeepRec 給我們嘗試。在近 2 年的工作后，DeepRec 已經(jīng)全量用于我們的模型訓(xùn)練和線上 inference，并且取得了顯著的性能提升和成本下降。

3 稀疏模型訓(xùn)練

3.1 EmbeddingVariable 多級存儲

由于模型參數(shù)量大，一些特征的 embedding 大小達(dá)到了接近 TB 級別，完全基于內(nèi)存存儲對于成本的要求過高，因此自然而然就會想到多級存儲：將最熱的 embedding 放在顯存或者內(nèi)存里，其余的可以分級放在 PMEM、SSD 等成本較低的存儲介質(zhì)中。而 DeepRec 中 提供了基于 EmbeddingVariable 的 Embedding 多級存儲功能。DeepRec 目前對于 embedding 存放在各種存儲介質(zhì)的支持已經(jīng)相當(dāng)完善。

下面介紹下我們團隊升級 DeepRec 在存儲這一塊的過程和經(jīng)驗：

3.1.1 compaction 的性能問題

我們原本基于自研的分布式 parameter server，而當(dāng)時 PMEM 類的存儲介質(zhì)還不普及，因此我們選擇了比較高性能的 SSD 作為多級存儲介質(zhì)。于是我們自然而然采用了類 leveldb（rocksdb）的方案作為 SSD 存儲方案。但這種方案在模型訓(xùn)練時，由于參數(shù)不斷增加和更新，后臺會進(jìn)行頻繁的 compaction，此時會有嚴(yán)重的寫放大問題導(dǎo)致 ps 的讀取時間大大延長，從而導(dǎo)致模型訓(xùn)練的瓶頸幾乎都在 ps 側(cè)。ps：據(jù)說 rocksdb 在 2022 年底的 7.5.3 版本大幅改進(jìn)了 compaction 的性能，在后臺 compaction 時幾乎不會影響讀取的性能。

3.1.2 DeepRec 的方案

而在早期我們試用 DeepRec 時，DeepRec 的 EmbeddingVariable 對于 SSD 存儲的方案同樣是基于 leveldb，因此同樣遇到了跟我們自研的方案類似的問題。后續(xù)我們將此問題的測試結(jié)果反饋給了 DeepRec 相關(guān)的同學(xué)，他們基于此后續(xù)推出了基于 SSDHASH 的存儲方案，大大提升了 compaction 時的讀取性能，因此模型訓(xùn)練基于不再受困于 ps 的讀取性能問題。后續(xù)又進(jìn)一步了基于 SSDHASH 的同步和異步兩種 compaction 的方式。使用同步 compaction 時，向 SSD 寫入數(shù)據(jù)和 compaction 將會使用同一個線程，異步時則各使用一個線程。這里也推薦大家使用這種方案。

3.1.3 壓縮模型大小

進(jìn)一步的，如果能把模型大小控制在數(shù)十 GB，那 ps 的性能就可以進(jìn)一步提升了。因為采用 DeepRec，自定義各種壓縮方式的算子變得非常輕松。我們調(diào)研并實現(xiàn)了了多篇 embedding 壓縮方向的 paper，最后采用了binary code的方式實現(xiàn)了 embedding 的 multihash 方案，可以自由控制 embedding 的大小。我們嘗試在最大的特征 uid embedding 上使用了 multihash，把模型大小從 800GB 降低到 40GB 以下，auc 的損失僅在千分之三左右，線上點擊率下降了 1.5%；進(jìn)一步的，我們通過優(yōu)化序列推薦模型，更好的通過序列特征建模了用戶的個性化，可以發(fā)現(xiàn)在序列模型的基礎(chǔ)上把 uid embedding 換成 multihash 的方案，對于線上點擊率的影響僅有 0.3% 左右，因此可以放心全量 multihash 方案。我們也把基于 multihash 的 embedding variable 算子以 pr 的形式提交給了 DeepRec。

3.2 基于 GPU 的分布式訓(xùn)練

在解決了 ps 的性能瓶頸后，模型訓(xùn)練的速度就和模型 Tensor 計算的算力近似線性相關(guān)了。而近幾年隨著序列模型的發(fā)展，搜廣推的矩陣計算復(fù)雜度也在顯著提升。此時使用 gpu+ 大 batch size 來代替 cpu 作為 worker 的方案，無論在性能還是成本控制上都有巨大的優(yōu)勢。而阿里云機器學(xué)習(xí)平臺 PAI 開源的HybridBackend平臺就支持了基于 GPU 的分布式訓(xùn)練方案，并且深度支持了 DeepRec。

可以看到使用 hb 的方案在訓(xùn)練速度上對比 TF-PS 原生方案的優(yōu)勢。

3.2.1 模型參數(shù)完全放在顯存里

想要充分釋放 gpu 的算力，減少因為數(shù)據(jù)拷貝帶來的性能損耗，最好的方案自然是把所有參數(shù)都放在 gpu 顯存里。上面 2.1.3 提到的壓縮模型大小，為這種方案提供了可能性。調(diào)大 batch size 則可以進(jìn)一步提高顯卡的利用率。經(jīng)過測試，在這種情況下，單張 V100 顯卡的算力可以超過 20 臺 40core worker 節(jié)點的算力。

3.2.2 解決了多卡訓(xùn)練丟失數(shù)據(jù)的問題

在單機多卡訓(xùn)練時，我們發(fā)現(xiàn)和單卡訓(xùn)練相比有近 1/3 的數(shù)據(jù)被丟棄，這是由于 hybridbackend 默認(rèn)使用所有 worker 按照 row group 平分?jǐn)?shù)據(jù)的策略，以提高讀取效率。當(dāng) group 數(shù)目不夠均分時，多余的數(shù)據(jù)會被丟棄，當(dāng) parquet 文件較多且比較小時，該問題尤為嚴(yán)重。我們通過使用每個 worker 加載所有的 group，再按照 batch 平分?jǐn)?shù)據(jù)的策略，極大地緩解了數(shù)據(jù)丟失的情況，讀取壓力也在可接收范圍內(nèi)，后續(xù)可考慮將兩策略結(jié)合降低 worker 的讀取壓力。

4 模型 inference

4.1 痛點

在我們的實際場景里，線上 inference 的痛點大部分來自于維護(hù)成本。因為推薦廣告業(yè)務(wù)場景，需要大量嘗試各種模型在線上分配流量做 AB test，因此線上存在的模型量級大概是 10 倍的基線模型量級。而每次上線一個模型，都需要給對應(yīng)的模型分配相應(yīng)的資源，并且這個資源跟 AB test 的流量正相關(guān)；而每次調(diào)整 AB test 流量（比如模型效果不錯，放大流量觀察）的時候，又需要調(diào)整該模型分配的資源。這個過程比較難實現(xiàn)自動化，往往需要算法工程師手動擴縮容。

4.2 基于 Processer 庫的 inference 方案解決痛點

上面這個圖是我們線上實際的 inference 方案。

4.2.1 單機器運行所有模型

基于上面的痛點，我們給出的方案是使用大規(guī)格機器（比如 128C，512G 內(nèi)存）來做線上 inference，然后每臺機器都會有線上所有的模型實例。每臺機器運行一個 serving-proxy 會自動的管理所有的模型進(jìn)程，包括模型上下線、模型更新等。這種方案的好處是整個維護(hù)成本基本沒有了，所有事情基本都自動化完成了。因為線上整體的流量相對穩(wěn)定（比如擴大 AB test 模型的流量，自然基線模型流量就減少了，整體是穩(wěn)定的），所以各個模型之間資源競爭也不需要重新分配資源。

4.2.2 基于 DeepRec 提供的 Processer 庫

DeepRec Serving Processor 是用于線上高性能服務(wù)的 Library，可以參考文檔。因為本身是一個獨立的 so 包，我們可以很方便的對接到自己的 Serving RPC 框架中。我們采用 golang 語言來完成了我們自己的 serving rpc 項目，優(yōu)點自然是開發(fā)成本低并且性能不錯。

4.2.3 使用 DeepRec 的 Session Group

直接使用 TensorFlow 提供的 C++ 接口調(diào)用 Session::Run，無法實現(xiàn)多 Session 并發(fā)處理 Request，導(dǎo)致單 Session 無法實現(xiàn) CPU 的有效利用。如果通過多 Instance 方式（多進(jìn)程），無法共享底層的 Variable，導(dǎo)致大量使用內(nèi)存，并且每個 Instance 各自加載一遍模型，嚴(yán)重影響資源的使用率和模型加載效率。

DeepRec 中 SessionGroup 可配置一組 Session，并且通過 Round Robin (支持用戶自定義策略)方式將用戶請求分發(fā)到某一個 Session。SessionGroup 對不同 Session 之間的資源進(jìn)行隔離，每個 Session 擁有私有的線程池，并且支持每個線程池綁定底層的 CPU Core(numa-aware)，可以最大程度地避免共享資源導(dǎo)致的鎖沖突開銷。SessionGroup 中唯一共享的資源是 Variable，所有 Session 共享底層的 Variable，并且模型加載只需要加載一次。

我們使用 session group 后，實測調(diào)整到合適的 group 數(shù)量，可以提高 50% 的 inference 性能。

4.2.4 基于 oneDNN 的優(yōu)化

DeepRec 集成了英特爾開源的跨平臺深度學(xué)習(xí)性能加速庫 oneDNN（oneAPI Deep Neural Network Library），并且修改 oneDNN 原有的線程池，統(tǒng)一成 DeepRec 的 Eigen 線程池，減少了線程池切換開銷，避免了不同線程池之間競爭而導(dǎo)致的性能下降問題。oneDNN 已經(jīng)針對大量主流算子實現(xiàn)了性能優(yōu)化，包括 MatMul、BiasAdd、LeakyReLU 等在業(yè)務(wù)場景中使用到的常見算子，為業(yè)務(wù)模型提供了強有力的性能支持。更值得一提的是， oneDNN 的算子支持 BF16 數(shù)據(jù)類型，與搭載 AMX(Advanced Matrix Extensions)指令集的第四代英特爾? 至強? 可擴展處理器同時使用，可顯著提升模型訓(xùn)練和推理性能。在 DeepRec Serving Processor 編譯選項中，只需加入“--cnotallow=mkl_threadpool”，便可輕松開啟 oneDNN 優(yōu)化。

4.2.5 子圖優(yōu)化

子圖融合是推理性能優(yōu)化的常用方法。但是對于本模型中左圖所示的子圖結(jié)構(gòu)含有 Reshape 算子，原生 tensorflow 并沒有對應(yīng)結(jié)構(gòu)的圖優(yōu)化器以及算子實現(xiàn)，我們通過手動融合來實現(xiàn)，融合前后的子圖構(gòu)成如下圖所示。這樣減少了多余算子的運行開銷，減少了內(nèi)存訪問，提升了計算效率。再結(jié)合 oneDNN 加速融合算子，最終業(yè)務(wù)端到端加速了 10%，CPU 利用率下降 10%。

4.2.6 cost model 的設(shè)計

由于大機器的 cpu core 數(shù)較多，而我們是一臺機器有所有模型的進(jìn)程，那么所有模型都共享所有 cpu core 顯然會造成不必要的資源競爭等。因此給不同模型設(shè)計合理的 cost model 就很有必要。我們目前采用比較簡單的方式，因為基線模型和需要做 AB test 的模型資源差別較大（流量差距大），我們會給每個基線模型分配對應(yīng)的 core，然后讓所有非基線模型共享一組 core（總體 AB test 的流量有上限）。雖然這個方案很簡單，但是取得了非常好的效果，大概有 30% 的性能提升。

5 后續(xù)規(guī)劃

1、cost model 的優(yōu)化，顯然有更好的方案來動態(tài)的調(diào)整每個模型需要的 core。我們打算開發(fā)更好的 cost model 并提供給 DeepRec。

2、開源我們的 inference 架構(gòu)方案，因為在我們的業(yè)務(wù)里，基于 DeepRec processor 設(shè)計的 inference 架構(gòu)帶來了巨大的便利，并且性能很好，我們預(yù)計在上半年會開源我們的 inference 架構(gòu)方案，歡迎大家到時關(guān)注。