RECom是一個旨在加速GPU上深度推薦模型推理的機器學習編譯器，他解決了已有機器學習編譯器無法有效優化推薦模型中重要的嵌入層計算的問題。他提出了一種新穎的面向子圖間并行的融合方法，以在單個GPU內核中生成處理大量嵌入列的高效代碼，大幅減少了內核啟動的開銷并顯著提升了GPU的利用率。他通過一個形狀計算簡化模塊來解決動態形狀下形狀計算帶來挑戰，并通過一個嵌入列子圖優化模塊來消除冗余計算。我們在阿里巴巴的四個真實的內部生產推薦模型和兩個生成模型上評估了RECom。 實驗表明，RECom在端到端推理延遲和吞吐量方面分別優于TensorFlow基線6.61倍和1.91倍。

背景

基于深度學習的推薦模型在各大企業的業務中變得越來越重要。通過為用戶推薦合適的商品，它們能夠幫助提升客戶體驗，增加銷售額，以及提高客戶保留率。典型的推薦模型包括兩部分，即嵌入層和深度神經網絡層（DNN）。 在生產環境中，嵌入層通常由對應不同特征字段的大量嵌入列（即通過查表將輸入特征轉換為嵌入向量的子圖，如下圖中虛線方框所示）組成，它能將輸入的特征（比如用戶ID，商品ID等）映射到低維的連續向量。

為了獲得更高的模型精度，開發人員通常會生成數千個統計特征，并使用不同的嵌入列來處理它們。 然而，處理如此大量的嵌入列是昂貴的。 例如，我們在阿里巴巴的模型上進行的實驗表明，嵌入列可以占據GPU上99%以上的端到端推理延遲。 因此，業務中迫切需要對推薦模型中的海量嵌入列進行優化。

挑戰

雖然目前已有一些加速嵌入層計算的手寫算子庫[1]，但它們很難在真實業務中被廣泛應用，因為它們不可能在庫中遍歷嵌入層中所有可能的算子組合。此外，出于隱私考慮，實際業務中經常要求根據模型的計算圖IR（比如TensorFlow GraphDef）而非源代碼來進行優化。然而，一個公司的業務中通常會包含大量具有不同嵌入列結構的模型，并且一個模型可以包含上百萬行的IR。因此，基于這些IR再手工地使用這些算子庫來重新搭建所有的業務模型是不現實的。

另一方面，機器學習編譯器，如XLA[2]等被廣泛地應用于自動加速機器學習模型的計算。然而，目前的機器學習編譯器主要關注于DNN的計算，而無法為嵌入層中包含的大量嵌入列提供有效的優化。它們主要存在以下三個問題：

首先，在 GPU 上為數千個嵌入列中的大量算子生成有效的代碼是一項挑戰。現有編譯器的算子融合策略既無法有效消除模型中大量的非計算開銷（包括內核啟動與算子調度），也無法充分利用模型中子圖間的并行性。我們在一個特定模型的實驗上發現，XLA 可以為該模型生成超過10,000個GPU內核，這會引入大量的非計算開銷，進而導致GPU的活躍時間僅占33%。并且，線上推理時輸入的batch size通常不會太大，導致大多數生成的GPU內核的Waves per SM都非常小（小于 0.06），意味著大部分的GPU SM在執行中都沒有被利用。

其次，推薦模型通常具有動態形狀（dynamic shapes）。傳統的機器學習編譯器通常依賴張量的形狀來進行各種優化，而推薦模型的計算圖中的張量形狀一般無法在編譯期獲取。一方面，動態形狀造成的張量間形狀關系信息缺失阻礙一些編譯優化的pass。另一方面，動態形狀場景下，部分算子（如SparseReshape）需要形狀的輸入，繼而引入了大量形狀計算。形狀計算與張量計算的耦合使得計算圖變得復雜，阻礙了后續的圖優化。

下圖(a)展示了一個級聯SparseReshape的子圖，其中藍色的算子表示形狀計算。由于SR1輸出的sparse tensor的消費者只有SR2，因此第一次reshape是冗余的。理想情況下，編譯器檢測到這樣的情況后可以直接把SR1的輸入張量（即張量1）連接到SR2，越過冗余的SR1，如(b)中所示。然而，由于形狀計算與張量計算的耦合，導致實際轉化后的圖會如(c)中所示，SR1并不能被正確的優化掉。

第三，現有的機器學習編譯器忽視了推薦模型中存在的大量冗余計算。通過對業務模型的分析，我們觀察到現有模型普遍存在大量的冗余計算。例如，使用TensorFlow框架進行嵌入表查找時會引入一些邊界檢查的計算，而這些檢查在整個嵌入列子圖的上下文里可能是冗余的，因為在經過哈希/分桶等操作后，用于嵌入表查找的索引已經滿足了要求。在極端情況下，一個嵌入列中80%的GPU計算時間可能都是冗余的。

破局

我們提出了第一個端到端的推薦模型優化編譯器RECom來解決以上的問題。

針對大量嵌入列的GPU代碼生成，RECom 提出了一種子圖間并行性導向的算子融合方法，將嵌入層中數千個嵌入列包含的大量算子全部融合（fuse）到一個GPU kernel中。一方面，這種方法使得RECom 能夠消除頻繁的內核啟動與框架調度，顯著減少嵌入層執行時的非計算開銷。另一方面，為了生成具有高并行性的代碼，RECom 將獨立的嵌入列映射到不同組的GPU 線程中。每個線程組中的線程協作處理一個嵌入列，并利用GPU 的層次存儲來進行中間數據緩沖和線程間通信。通過這種方式，RECom 可以有效地利用GPU 來挖掘推薦模型嵌入層中算子內部和嵌入列間的并行性。

為了解決動態形狀與形狀計算帶來的問題，RECom提出了一種基于符號表達式的方法。類似于BladeDISC[3]，RECom 通過下圖中的四種pattern構建了嵌入列的全局符號形狀表達式，并利用這些表達式來確定不同張量的形狀間關系。

得到形狀的符號表達式后，RECom進一步將所有形狀計算子圖用統一的算子進行重構。重構后的ShapeConstruct算子僅依賴于生成所需符號的算子，而非依賴于原計算子圖的輸入。

這樣，RECom得以解耦張量計算與形狀計算來簡化計算圖。回到之前冗余SparseReshap的例子，再經過形狀計算重構后，SR2輸入形狀對SR1的依賴被消除，因此在經過原來的編譯優化pass后，我們能得到下圖的結果，成功去除了冗余的SR1。

對于嵌入層中存在的大量冗余計算，我們在RECom中開發了一個嵌入列子圖優化模塊，針對性地為常見的冗余計算設計了優化。具體細節可以參考我們的論文。

RECom的整個編譯優化流程如下圖所示。

與此同時，阿里云PAI團隊與悉尼大學合作的論文《MonoNN: Enabling a New Monolithic Optimization Space for Neural Network Inference Tasks on Modern GPU-Centric Architectures》在OSDI2024中選，論文創新性地提出了MonoNN，這是第一個能夠為單個GPU上的常見靜態神經網絡（NN）推理任務提供新的整體設計和優化空間的機器學習優化編譯器。MonoNN可以將整個神經網絡容納到單個GPU核心中，大大減少了非計算開銷，并從新形成的整體優化空間提供進一步細粒度的優化機會。最重要的是，MonoNN確定了各種NN運算符之間的資源不兼容性問題，這是創建這樣一個整體優化空間的關鍵設計瓶頸。MonoNN通過系統地探索和利用跨不同類型的NN計算的并行補償策略和資源權衡，并提出了一種新的獨立于調度的群體調諧技術來顯著縮小極其龐大的調諧空間。評估表明，MonoNN的平均加速比比最先進的框架和編譯器提高了2.01倍。具體而言，就端到端推理性能而言，MonoNN在TVM、TensorRT、XLA和AStitch方面的表現分別高達7.3倍、5.9倍、1.7倍和2.9倍。

目前，MonoNN已開源，如需獲取MonoNN源代碼請前往：https://github.com/AlibabaResearch/mononn；

論文地址：https://dl.acm.org/doi/10.1145/3623278.3624761。

論文信息

論文標題：《RECom: A Compiler Approach to Accelerating Recommendation Model Inference with Massive Embedding Columns》

論文地址：https://dl.acm.org/doi/10.1145/3623278.3624761

參考文獻

[1] NVIDIA. NVIDIA-Merlin/HugeCTR. https://github.com/NVIDIA-Merlin/

HugeCTR, 2023.

[2] Google. XLA: Optimizing Compiler for Machine Learning. https://www.tensorflow.org/xla, 2023.

[3] Zhen Zheng, Zaifeng Pan, Dalin Wang, Kai Zhu, Wenyi Zhao, Tianyou Guo, Xiafei Qiu, Minmin Sun, Junjie Bai, Feng Zhang, Xiaoyong Du, Jidong Zhai, and Wei Lin. BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach. Proc. ACM Manag. Data, 1(3), nov 2023.