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