作者

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

1 背景

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

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

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

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

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

而因?yàn)槲覀児颈旧砀⒗镌朴兄疃鹊暮献?，阿里云也主?dòng)介紹了當(dāng)時(shí)還是內(nèi)部項(xiàng)目的 DeepRec 給我們嘗試。在近 2 年的工作后，DeepRec 已經(jīng)全量用于我們的模型訓(xùn)練和線上 inference，并且取得了顯著的性能提升和成本下降。

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

3.1 EmbeddingVariable 多級(jí)存儲(chǔ)

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

下面介紹下我們團(tuán)隊(duì)升級(jí) DeepRec 在存儲(chǔ)這一塊的過(guò)程和經(jīng)驗(yàn)：

3.1.1 compaction 的性能問(wèn)題

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

3.1.2 DeepRec 的方案

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

3.1.3 壓縮模型大小

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

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

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

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

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

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

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

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

4 模型 inference

4.1 痛點(diǎn)

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

4.2 基于 Processer 庫(kù)的 inference 方案解決痛點(diǎn)

上面這個(gè)圖是我們線上實(shí)際的 inference 方案。

4.2.1 單機(jī)器運(yùn)行所有模型

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

4.2.2 基于 DeepRec 提供的 Processer 庫(kù)

DeepRec Serving Processor 是用于線上高性能服務(wù)的 Library，可以參考文檔。因?yàn)楸旧硎且粋€(gè)獨(dú)立的 so 包，我們可以很方便的對(duì)接到自己的 Serving RPC 框架中。我們采用 golang 語(yǔ)言來(lái)完成了我們自己的 serving rpc 項(xiàng)目，優(yōu)點(diǎn)自然是開發(fā)成本低并且性能不錯(cuò)。

4.2.3 使用 DeepRec 的 Session Group

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

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

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

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

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

4.2.5 子圖優(yōu)化

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

4.2.6 cost model 的設(shè)計(jì)

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

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

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

2、開源我們的 inference 架構(gòu)方案，因?yàn)樵谖覀兊臉I(yè)務(wù)里，基于 DeepRec processor 設(shè)計(jì)的 inference 架構(gòu)帶來(lái)了巨大的便利，并且性能很好，我們預(yù)計(jì)在上半年會(huì)開源我們的 inference 架構(gòu)方案，歡迎大家到時(shí)關(guān)注。