【ASPLOS2024】RECom：通過編譯器技術(shù)加速推薦模型推理，論文中選并獲得榮譽(yù)獎項(xiàng)！

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

背景

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

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

挑戰(zhàn)

雖然目前已有一些加速嵌入層計(jì)算的手寫算子庫[1]，但它們很難在真實(shí)業(yè)務(wù)中被廣泛應(yīng)用，因?yàn)樗鼈儾豢赡茉趲熘斜闅v嵌入層中所有可能的算子組合。此外，出于隱私考慮，實(shí)際業(yè)務(wù)中經(jīng)常要求根據(jù)模型的計(jì)算圖IR（比如TensorFlow GraphDef）而非源代碼來進(jìn)行優(yōu)化。然而，一個公司的業(yè)務(wù)中通常會包含大量具有不同嵌入列結(jié)構(gòu)的模型，并且一個模型可以包含上百萬行的IR。因此，基于這些IR再手工地使用這些算子庫來重新搭建所有的業(yè)務(wù)模型是不現(xiàn)實(shí)的。

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

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

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

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

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

破局

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

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

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

得到形狀的符號表達(dá)式后，RECom進(jìn)一步將所有形狀計(jì)算子圖用統(tǒng)一的算子進(jìn)行重構(gòu)。重構(gòu)后的ShapeConstruct算子僅依賴于生成所需符號的算子，而非依賴于原計(jì)算子圖的輸入。

這樣，RECom得以解耦張量計(jì)算與形狀計(jì)算來簡化計(jì)算圖?；氐街叭哂郤parseReshap的例子，再經(jīng)過形狀計(jì)算重構(gòu)后，SR2輸入形狀對SR1的依賴被消除，因此在經(jīng)過原來的編譯優(yōu)化pass后，我們能得到下圖的結(jié)果，成功去除了冗余的SR1。

對于嵌入層中存在的大量冗余計(jì)算，我們在RECom中開發(fā)了一個嵌入列子圖優(yōu)化模塊，針對性地為常見的冗余計(jì)算設(shè)計(jì)了優(yōu)化。具體細(xì)節(jié)可以參考我們的論文。

RECom的整個編譯優(yōu)化流程如下圖所示。

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

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

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

論文信息

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

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

參考文獻(xiàn)

[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.