EasyRec 推薦算法訓練推理優(yōu)化
一、EasyRec 訓練推理架構(gòu)
在介紹 EasyRec 的訓練推理架構(gòu)之前,先來談?wù)勍扑]模型的發(fā)展趨勢和面臨的挑戰(zhàn)。近年來,推薦模型的發(fā)展呈現(xiàn)出以下一些趨勢:首先,特征數(shù)量越來越多,從幾百個增加到上千個,還有許多交叉特征;同時,Embedding 變得越來越大,序列越來越長,Dense 層也越來越復(fù)雜,從簡單的 MLP 發(fā)展為 MMOE、MaskNet、PLE 等復(fù)雜結(jié)構(gòu)。由此帶來的最大挑戰(zhàn)就是算力不足,另外訓練推理成本很高,推理超時嚴重。
EasyRec 推薦算法訓練整體框架主要包括:數(shù)據(jù)層、Embedding 層、Dense 層和輸出層。這個架構(gòu)可以在多種平臺上運行,包括 MaxComput、開源大數(shù)據(jù)平臺 EMR 和深度學習的容器平臺 DLC。
此架構(gòu)的優(yōu)勢是支持配置化、組件化,包括深度支持 Keras 組件,能自定義組件,并通過配置接入各種模型。它還支持大規(guī)模分布式訓練、ODL,以及基于 NNI 的自動調(diào)參,搜索最優(yōu)超參數(shù),和自動特征選擇。支持推薦模型中的常用功能,如 MultiOptimizer,設(shè)定 Embedding 和 DNN 層不同學習率和優(yōu)化器,以及特征熱啟動,大規(guī)模負采樣等。如果模型訓練中斷,可以使用 Work Queue 從斷點恢復(fù)訓練,顯著提升了大型任務(wù)的訓練成功率。此外,在 TF 框架上擴展了分布式 Evaluator,支持大數(shù)據(jù)量的模型評估。
接下來介紹推理框架 PAI-REC 引擎,這是整個推薦鏈路的一個重要部分。PAI-REC 引擎串聯(lián)推薦業(yè)務(wù)的各個階段,常見的階段包括召回、排序、重排和打散。PAI-Rec 引擎基于 go 語言編寫,具有比較高的效率,同時也是模塊化的,因此具有比較強的擴展性,進一步還提供了用戶友好的界面,方便用戶配置 ab 實驗,做特征一致性診斷,分析特征和實驗效果等關(guān)鍵功能。
與 EasyRec 相關(guān)的是 EasyRecProcessor,負責精排和召回模型的在線推理。主要包括三個部分:item 的 Feature Cache,F(xiàn)eature Generator 和 TF Model。EasyRecProcessor 進行了大量的 CPU 和 GPU 推理優(yōu)化,如通過 item 特征緩存減少 item 靜態(tài)特征帶來的網(wǎng)絡(luò)傳輸壓力,通過增量更新加快模型傳輸和部署的速度,在 Feature Generator 和 TF Model 模型推理上也有很多優(yōu)化,下面進行詳細介紹。
EasyProcessor 支持在 PAI-EAS 平臺上一鍵部署。該框架已經(jīng)在阿里云上得到了廣泛應(yīng)用,已服務(wù)數(shù)百家客戶,覆蓋電商、直播、文章分享、視頻分享、廣告和社區(qū)等多種業(yè)務(wù)。同時,在阿里內(nèi)部也有很多客戶在使用該框架。
我們曾服務(wù)過一個電商導(dǎo)購案例,通過優(yōu)化,不僅提升了效果,還顯著降低了成本,我們針對推薦的各個鏈路都進行了升級和優(yōu)化。
二、EasyRec 訓練優(yōu)化
接下來講一下 EasyRec 在訓練方面的優(yōu)化。隨著 sequence 長度的增加,算力、存儲和網(wǎng)絡(luò)開銷顯著增大。我們發(fā)現(xiàn)一次曝光會下發(fā)很多 item,而這些 item 的 SequenceFeature 大多相同。通過去重操作,例如一個 8192 的 batch_size,去重后可能只剩下原來的 5% 到 10%。因此,對 SequenceFeature 進行去重,只存 request_id,再通過 iGraph 查找 SequenceFeature,經(jīng)過 embedding layer 和 deunique 處理,得到 batch seq_embedding。這個優(yōu)化提升了系統(tǒng)吞吐量 20%??紤]到可遷移性,我們目前的 unique 實現(xiàn)基于 Python,若改用 C++,性能將進一步提升。
另一個優(yōu)化是 EmbeddingParallel,即 embedding 分片優(yōu)化。以往多采用 PS-Worker 模式,盡管擴展性好,但存在問題,如 ps 通信量大,算力不足,以及 embedding 劃分不均勻影響訓練效率。算子 placement 不當,如 unique 算子被錯誤地放在 ps 上,也會造成瓶頸。All-Reduce 模式是另一種選擇,所有 Worker 存儲相同參數(shù),避免了 ps 的通信和計算瓶頸。但這種架構(gòu)的問題是 embedding 容量受單機內(nèi)存限制,難以實現(xiàn)多機擴展。
EmbeddingParallel 優(yōu)化中,每個 Worker 獨立存儲 dense 參數(shù),但 Sparse 參數(shù)分片存儲在每個 Worker 上,避免了 All-Reduce 模式的內(nèi)存瓶頸。dense 參數(shù)通過 All-Reduce 更新,小型和桶化的 embedding 也是如此,大型 embedding 則通過 AllToAll 更新。
在 CPU上,我們采用 DeepRec 的 lock-free hash table,比 google 的 dense hash table 效率更高。在 GPU 上,采用 hugectr 的 sok embedding,通過 GPU 緩存的方式加載熱點 embedding,減少 embedding h2d 的開銷。在訓練效果上,MMOE 和 PPNet 模型的對比顯示,PS 模式下每秒約 3.5 步,而 EmbeddingParallel 架構(gòu)顯著提升了訓練速度。由于參數(shù)保存在不同 Worker 上,需額外工作聚合 embedding,導(dǎo)出單機可 serving 的模型。EasyRec 框架已實現(xiàn)這一功能,直接可用。
我們在 CPU 上的另一個訓練優(yōu)化,針對仍使用 CPU 架構(gòu)進行訓練和推理的客戶。推薦模型的 Dense 層越來越復(fù)雜,導(dǎo)致計算量大增。分析模型時間線發(fā)現(xiàn),MatMul 占據(jù) 60% 以上的計算時間。為提升 MatMul 這類算力密集型算子的性能,我們與英特爾合作,利用 AMX 計算能力,進行矩陣 BF16 加速,其算力比普通 CPU 高約 16 倍。在實際模型訓練中,采用 AMX 功能優(yōu)化,顯著提升了訓練速度。
三、EasyRec 推理優(yōu)化
接下來介紹一下 EasyRec 推理方面的優(yōu)化。首先是 Embedding 部分的優(yōu)化,大部分 Embedding 仍然放在 CPU 上。如果用 TF 的 feature column 構(gòu)造 embedding layer,會發(fā)現(xiàn)存在很多小的算子,如 unique 和 SparseSegmentMean,這些小算子帶來大量啟動開銷,影響整體性能.
針對常用的 Embedding 模式,做了一些融合算子優(yōu)化,并通過 AVX 進行并行加速。比如一個 sequence 算子,可能包含幾百個小算子,優(yōu)化后變成一個算子,計算開銷降低且通過 AVX 加速,性能大幅提升。實際應(yīng)用中,算子數(shù)量減少 50% 以上,響應(yīng)時間(RT)也減少一半以上。
我們發(fā)現(xiàn)半精度計算可以加速推理并減少內(nèi)存占用,尤其對大模型的內(nèi)存開銷影響顯著。實驗表明,大部分模型將模型量化為 BF16 對 AUC 基本沒有影響。在 BF16 到 float 的轉(zhuǎn)換中,原生 TensorFlow 的轉(zhuǎn)換速度較慢,我們嘗試用 AVX 進行加速,結(jié)果顯示 QPS 和 RT 顯著提升?;诖?,我們進一步嘗試了 AMX 的矩陣乘法加速,能夠進一步提升約 10% 以上。
接下來介紹一下我們在 Feature 層的優(yōu)化。很多算子用 string 表示,如 look up feature 會解析 string 并構(gòu)建 map,帶來開銷。我們用 AVX 優(yōu)化了 StringSplit。在構(gòu)建 HashMap 時,默認使用 MurmurHash,雖然沖突概率小,但特征解析時,HashMap 規(guī)模不大且用時短。采用更高效的 CrcHash 和 XorHash,均用 AVX 實現(xiàn),替換 MurmurHash 后,RT 降低 5% 以上。
另外是 SequenceFeature 優(yōu)化,使用 item feature cache,減少了遠程網(wǎng)絡(luò)訪問開銷,提高了 sequence 在推理側(cè)的性能,但是帶來了一個新問題:內(nèi)存占用較大。我們設(shè)計了一種緊湊的存儲格式,內(nèi)存開銷相比普通的存儲方式降低了 80% 以上。進一步我們將 Feature 處理算子封裝為 TensorFlow op,支持并行執(zhí)行,復(fù)用 TensorFlow 線程池,實現(xiàn) feature generation 和 embedding look up 的 overlap 執(zhí)行,并減少減少數(shù)據(jù)序列化和網(wǎng)絡(luò)傳輸?shù)拈_銷。整體優(yōu)化后,RT 減少 20%,QPS 顯著增加。
這是實際采集的 timeline,顯示優(yōu)化前的情況,其中很多時間花在 match feature,字符串解析、拼接和 tensor 填充上,開銷較大。優(yōu)化后,這些額外的解析和拼接操作都消除了,主要只剩下 match feature 本身的開銷。
接下來講常用的 user feature tile 優(yōu)化。許多 user feature 和 sequence feature 在一次請求中只需計算一次,但導(dǎo)出模型時,算法同學未考慮這個情況。因此,我們在 processor 側(cè)進行 tile 優(yōu)化(自動 broadcast)。在輸入層補齊 user feature 并做 tile 的效率有限。
進一步提升是在 embedding look up 后進行自動 broadcast,節(jié)省計算開銷。實際測試中,QPS 顯著提升 30% 到 50%。整個優(yōu)化流程是找到需要 broadcast 的算子,很多算子可自動 broadcast,但 select 和 concat 等特殊算子需要對輸入進行 broadcast 處理以確保正確執(zhí)行。找到這些候選算子后進行 top 排序,再對排序后的算子逐一 Tile。Tile 過程中,部分算子 Tile 后使其他算子無需再 Tile,因此只需選擇未 Tile 的算子繼續(xù) Tile,實現(xiàn)自動 broadcast。我們通過分析全圖來將 tile 盡可能后置,以最大化的降低計算量。
接下來講 GPU 上的優(yōu)化,GPU 優(yōu)化最重要的是 Placement 優(yōu)化。GPU 的算力強吞吐高,但啟動開銷高。通常我們會把 embedding 放在 CPU 上,因為 OP 數(shù)目多且單個 OP 計算量小,放在 GPU 開銷大于執(zhí)行時間。這樣用 GPU 反而不如 CPU 效率,加上 CPU 側(cè)有很多 AVX 優(yōu)化,要 GPU 超過 AVX 的效率就更難了。
GPU 主要負責 Dense 計算。Dense 計算量大,OP 執(zhí)行時間超過 kernel launch 開銷,所以用 GPU 性能提升顯著。除 kernel launch 外,還要考慮數(shù)據(jù)拷貝,embedding 到 Dense 的拷貝次數(shù)和數(shù)據(jù)量對性能影響大。我們用 Min-Cut 方法在圖中找到最優(yōu)分割點,將 Embedding Lookup 部分放在 GPU 上,后續(xù) Dense 計算前面在 CPU,后面在 GPU,減少 H2D Memcpy 開銷。
即使進行了 placement 優(yōu)化,但發(fā)現(xiàn)仍有一些模型的 GPU 利用率很高,達到百分之八九十,但整體吞吐仍然不理想。主要原因在于 GPU 的算子,比如 MatMul 和許多 elementwise 算子(如 batch_norm、sigmoid、softmax),在 CPU 上計算效率較高。這些算子屬于訪存密集型算子,訪存和調(diào)度開銷較大,不能充分發(fā)揮 GPU 的計算能力。因此,我們考慮使用 XLA 進行算子融合,減少 kernel launch 開銷,提升系統(tǒng)吞吐。
XLA 主要是 TF to XLA,包含以下流程:自動圈圖(AutoCluster),將目標算子圈出,生成 function library;然后 TF2XLA Compiler 優(yōu)化,轉(zhuǎn)為 HLO 的 XLA 表示;最后通過 LLVM 編譯優(yōu)化到 Cuda。
我們遇到的問題主要是 Dynamic shape,采用的方法是對 XlaRun 的 OP 進行 Padding,執(zhí)行后再剪切出有效的部分,以減少編譯優(yōu)化導(dǎo)致的動態(tài)重編譯問題。優(yōu)化后效果顯著。優(yōu)化前 RT 高,QPS 不高;優(yōu)化后 RT 顯著下降,QPS 提升。即使在一些 GPU 利用率不高的場景下,XLA 融合后 RT 也明顯下降。
剛剛講了 XLA 存在動態(tài)形狀的問題,隨后我們嘗試了 TRT(dense layer optimization)優(yōu)化。TRT 的流程類似,先拆分部分 OP 進行 cast 圈圖,再轉(zhuǎn)成 TRT 表示,最后用 TRTEngineOp 執(zhí)行。TRT 對 BatchNorm、Add、ReLU 等 elementwise 算子進行了深入融合。一個優(yōu)勢是對 dynamic shape 有支持,可以指定 range,在一定范圍內(nèi)避免重編譯。另一個優(yōu)勢是 TRT 支持量化,如 BF16 轉(zhuǎn)換。
我們在算力密集的 Dense 層進行了實驗,QPS 提升明顯。TRT 的缺點是作為閉源系統(tǒng),問題排查較困難。所以我們結(jié)合 XLA 和 TRT 進行模型優(yōu)化。
關(guān)于 dynamic shape,更加優(yōu)雅的解決方案是 blade-disc,現(xiàn)有的使用方式是離線將模型轉(zhuǎn)成 ONNX 后,用 blade-disc 優(yōu)化并加載。實時優(yōu)化尚未實現(xiàn),未來我們會逐步在 EasyRec Processor 中引入 blade-disc compiler 的 dynamic shape 功能。
在廣告場景中,實際的 batch size 較小,即使進行了 XLA 優(yōu)化,吞吐量仍然不理想。單個 Batching 執(zhí)行時,kernel launch 的開銷仍然較大。
我們進行了 batch 優(yōu)化,將多個小 batch 組裝在一起,由 GPU 執(zhí)行。embedding lookup 之前,每個 batch 仍單獨在 cpu 上執(zhí)行,lookup 之后組裝成一個大 batch 提交到 gpu 上執(zhí)行。Batch 模式的一個差異是在 feature tile 層需要進行更多的 broadcast 操作。像 Add、Sum、Mul 這些 OP 在 batch 處理后無法自動 broadcast,因此在 feature tile 層對這些 OP 進行處理,使其在多個小 batch 上也能自動 broadcast。broadcast 完成后,再 concat 并交由 GPU 執(zhí)行大的 batch。
在廣告場景中,這種優(yōu)化顯著提升了 QPS,尤其在 CVR 和 CTR 方面效果明顯。
接下來介紹我們在網(wǎng)絡(luò)直連和請求壓縮方面的優(yōu)化經(jīng)驗。之前在 PAI-EAS 上部署 EasyRec 推理服務(wù)時,通過 Client 請求 Nginx 網(wǎng)關(guān)負載均衡,會增加一次網(wǎng)絡(luò)轉(zhuǎn)發(fā)。改用直連方式后,客戶端定期刷新 pod 的 IP,減少一次網(wǎng)絡(luò)轉(zhuǎn)發(fā),RT 降低約 5 毫秒。
另一個問題是在客戶和我們機房之間做專線連接時,請求流量較大,高 QPS 場景下流量可能達到幾十 Gbps,給專線帶來壓力。我們考慮請求壓縮,嘗試了 gzip、snappy 和 zstd 等方式,最終選擇 snappy 和 zstd,既對 RT 影響小,又顯著降低流量壓力,10Gbps 流量大約減少了五倍,大大減輕了專線壓力。
四、實時學習 Online Learning
接下來介紹我們在 online learning 上的工作。Online learning 現(xiàn)在應(yīng)用非常廣泛,尤其在新品上架和熱點更新等需要及時響應(yīng)的場景中。此外,大促活動時樣本分布變化快,需要 embedding 參數(shù)和 dense 參數(shù)快速更新。online learning 的核心步驟包括:流失樣本、流失訓練和增量參數(shù)更新。
這是在 EasyRec 中使用 online learning 進行實時更新的主要流程。首先,我們通過 PAI-REC 實時回流日志和特征到 SLS 日志系統(tǒng),并通過特征埋點回流到 Datahub 中間件。我們在 Flink 上構(gòu)建了一套完整的樣本聚合和 label 生成流程,支持配置化的方式構(gòu)建流式訓練:從日志生成訓練樣本,聚合 Datahub 埋點特征,最終生成實時訓練樣本并存儲在 Datahub 中間件和實時消息隊列中,推送到實時訓練系統(tǒng)。
實時訓練系統(tǒng)定期從 Datahub 拉取訓練樣本進行訓練。訓練完成后,會定期保存增量參數(shù)到 OSS,并同步到 EAS 的 Processor。我們在穩(wěn)定性和一致性方面做了優(yōu)化,通過特征埋點提高特征一致性,并采用 flink 的 gemini kv 分離方式提升樣本和特征 join 性能。我們還對特征進行 Lz4 壓縮,提高 join 的穩(wěn)定性和效率。
針對實時場景中的異常數(shù)據(jù),我們進行了過濾和去重,如處理延遲或異常上報的 timestamp 和重復(fù)調(diào)用的 callback。對于延遲到達的正樣本,進行延遲下發(fā)校正訓練。這些優(yōu)化在新品和內(nèi)容場景中效果顯著。
這是一些參考文獻,包括 EasyRec 和 Processor 的一些文檔,以及全鏈路推薦系統(tǒng) PAI-REC 和特征工程的相關(guān)文檔,這些都是構(gòu)建整個阿里云上推薦系統(tǒng)的主要組成部分。