通過 Merlin 大幅提升大規(guī)模深度多目標(biāo)精排模型訓(xùn)練性能

本案例中，NVIDIA 團(tuán)隊(duì)與陌陌推薦系統(tǒng)團(tuán)隊(duì)深度合作，共同使用 NVIDIA GPU 和 Merlin 軟件解決方案替代其原有解決方案。

通過使用 Merlin TensorFlow Plugin （即 Sparse Operation Kit，SOK） 和 HierarchicalKV（HKV)，相較于原方案在相同規(guī)模模型和 GPU 下，顯著提高了陌陌大規(guī)模深度多目標(biāo)精排模型的訓(xùn)練性能。在不影響模型效果的前提下，模型整體吞吐提升了 5 倍以上，再結(jié)合通信和 IO 等進(jìn)一步優(yōu)化后，極限情況下可以提升 12 倍吞吐。

客戶簡介

摯文集團(tuán)于 2011 年成立，2014 年 12 月 11 日在美國納斯達(dá)克交易所掛牌上市（NASDAQ: MOMO），擁有陌陌、探探等多款手機(jī)應(yīng)用，以及電影制作發(fā)行、節(jié)目制作等多元業(yè)務(wù)。陌陌是摯文集團(tuán)于 2011 年 8 月推出的一款基于地理位置的移動視頻社交應(yīng)用，是中國領(lǐng)先的開放式社交平臺之一。

訓(xùn)練速度面臨挑戰(zhàn)，

需有效提升算法迭代

陌陌的原始解決方案本質(zhì)是基于 PS-Worker 的 CPU + GPU 混合訓(xùn)練方案，可支持大規(guī)模稀疏參數(shù)的訓(xùn)練。然而，隨著用戶規(guī)模的增加和業(yè)務(wù)的發(fā)展，對于推薦算法的準(zhǔn)確度也有了更高的要求。這導(dǎo)致模型的復(fù)雜性和訓(xùn)練樣本量顯著增加，對單次模型訓(xùn)練速度和新模型算法探索效率都有更大的挑戰(zhàn)。盡管原方案在功能上支持了大規(guī)模稀疏參數(shù)的訓(xùn)練，但在性能上難以滿足業(yè)務(wù)日益增長的需求。因此，陌陌亟需對訓(xùn)練速度進(jìn)行優(yōu)化，加快算法迭代，以提高業(yè)務(wù)效果。

SOK 和 HKV

為推薦系統(tǒng)提升性能與靈活性

NVIDIA Merlin HugeCTR 是 NVIDIA 推出的可以高效利用 GPU 來進(jìn)行推薦系統(tǒng)訓(xùn)練的解決方案，為了使它能直接被其他 DL 用戶，比如 TensorFlow 所直接使用，NVIDIA 開發(fā)了 Merlin TensorFlow Plugin （以下簡稱 SOK），將 HugeCTR 中的高級特性封裝為 TensorFlow 可直接調(diào)用的形式，從而幫助用戶在 TensorFlow 中直接使用 HugeCTR 中的高級特性來加速他們的推薦系統(tǒng)。

圖 1. Merlin TensorFlow Plugin（SOK）模型并行示意圖

Merlin TensorFlow Plugin 以數(shù)據(jù)并行的方式接收輸入數(shù)據(jù)，將稀疏參數(shù)以模型并行的方式分布在多個 GPU 上，將稠密參數(shù)以數(shù)據(jù)并行的方式分布在多個 GPU 上，內(nèi)部實(shí)現(xiàn)“數(shù)據(jù)并行-模型并行-數(shù)據(jù)并行”的轉(zhuǎn)換流程。整個使用方式上盡可能的與原有 TensorFlow 算子對齊，減少對用戶已有的代碼的修改，以更方便、快捷地在多個 GPU 上進(jìn)行擴(kuò)展。此外，SOK 針對 embedding vector 的拷貝和 combiner 進(jìn)行了高度優(yōu)化和內(nèi)核融合，使整個 lookup 的前后向過程擁有更好的性能。

Merlin HierarchicalKV （以下簡稱HKV）是 Merlin 下的針對于推薦系統(tǒng)訓(xùn)練設(shè)計的 KV 加速庫。為兼容大模型訓(xùn)練支持了層次化動態(tài) Embedding 存儲（CPU+GPU），靈活的 eviction（淘汰） 機(jī)制和豐富的 API。目前已經(jīng)集成入 SOK, 協(xié)同加速推薦系統(tǒng) Embedding 的相關(guān)計算。

在應(yīng)用了 SOK 和 HKV 后，相同規(guī)模模型和 GPU 下，陌陌精排模型的訓(xùn)練性能相比于原方案，整體吞吐提升了 5 倍以上。除此之外，陌陌推薦團(tuán)隊(duì)在當(dāng)前 SOK + HKV 的架構(gòu)基礎(chǔ)上，基于業(yè)務(wù)場景特點(diǎn)，進(jìn)一步優(yōu)化整體性能，包括梯度合并，減少梯度計算的通信開銷；并行特征數(shù)據(jù)讀取與轉(zhuǎn)換，以及特征數(shù)據(jù)預(yù)取到 GPU 等操作提速特征 IO；使用 XLA 進(jìn)行編譯優(yōu)化，融合 kernel 以減少 kernel launch 時間；設(shè)置 GPU 親和等操作，使得整體性能提升達(dá)到 12 倍。

圖 2. 性能加速效果

在陌陌的實(shí)際應(yīng)用中，動態(tài) Embedding 的特性大大簡化了推薦系統(tǒng)中連續(xù)訓(xùn)練需人工控制顯存中 embedding tab 大小的問題。而 SOK 與 HKV 為陌陌提供了完整的功能和性能支持。

除此之外，我們跟陌陌的合作過程中，也結(jié)合產(chǎn)品部署中的需求進(jìn)一步對產(chǎn)品進(jìn)行了性能優(yōu)化和功能迭代，比如：

SOK 針對 embedding vector 的拷貝和 combiner 進(jìn)行了高度優(yōu)化和 kernel 融合，使整個 lookup 的前后向過程擁有更好的性能。

在陌陌 GPU 高水位線的實(shí)際業(yè)務(wù)中，基于陌陌的測試和反饋，SOK 通過優(yōu)化了性能和功能的平衡點(diǎn)，使得其在保證性能的基礎(chǔ)上，穩(wěn)定性也大大提升。

另外，在陌陌的實(shí)際應(yīng)用中，面臨著模型實(shí)時訓(xùn)練的挑戰(zhàn)，即需要減少對計算資源的占用，因此我們引入了 HKV，它支持了層次化動態(tài)嵌入存儲（包括 CPU 和 GPU），并提供了靈活的 eviction 機(jī)制以及豐富的 API。這種引入在降低資源占用的同時，也提高了系統(tǒng)的靈活性。

陌陌的實(shí)際業(yè)務(wù)場景和 GPU 使用方式對于 SOK 的開發(fā)和迭代提供了非常寶貴的經(jīng)驗(yàn)，同時陌陌的大量測試反饋也幫助 SOK 提升了應(yīng)對復(fù)雜場景的能力，使得 SOK 的 feature 更加穩(wěn)定和貼近客戶。

持續(xù)合作：

優(yōu)化推薦模型性能，降低訓(xùn)練成本

雙方團(tuán)隊(duì)通過 SOK 和 HKV 對原方案進(jìn)行深度優(yōu)化后，成功幫助陌陌提升了 12 倍的訓(xùn)練效率，極大的降低了模型訓(xùn)練的成本和新模型算法嘗試的成本。目前，整體方案已上線，全面支持陌陌推薦系統(tǒng)模型訓(xùn)練。

近期，NVIDIA 團(tuán)隊(duì)還與陌陌進(jìn)行了基于 Transformer 的推薦模型性能優(yōu)化。NVIDIA JOC 團(tuán)隊(duì)和 Merlin 團(tuán)隊(duì)基于客戶的模型做了一系列性能分析，將 XLA+AMP+半精度 allreduce 應(yīng)用到該模型上后，端到端性能實(shí)現(xiàn)了 50% 的加速。在此基礎(chǔ)上，團(tuán)隊(duì)們進(jìn)一步對性能熱點(diǎn) multi-head-attention 部分進(jìn)行優(yōu)化，正在將 Flash-Attention 以 tf-plugin 形式進(jìn)行集成，預(yù)計此項(xiàng)優(yōu)化集成后，整體加速比可達(dá)到 3 倍，同時整體的優(yōu)化方案使得顯存使用量下降約 70%，可以顯著地緩解顯存緊張的問題。

未來，陌陌與 NVIDIA 將繼續(xù)在推薦系統(tǒng)訓(xùn)練和推理等方面持續(xù)合作，持續(xù)推進(jìn) GPU 和 AI 軟件加速計算在陌陌的全面落地，期待能夠?yàn)槟澳暗臉I(yè)務(wù)及場景應(yīng)用帶來更大的價值。

了解更多本案例中相關(guān)的 NVIDIA 產(chǎn)品信息，敬請查閱：

• NVIDIA Merlin：

  https://developer.nvidia.cn/merlin

• Merlin TensorFlow Plugin (SOK) ：

  https://github.com/NVIDIA-Merlin/HugeCTR/tree/main/sparse_operation_kit

• Merlin HierarchicalKV (HKV)：

  https://github.com/NVIDIA-Merlin/HierarchicalKV

GTC 2024 將于 2024 年 3 月 18 至 21 日在美國加州圣何塞會議中心舉行，線上大會也將同期開放。點(diǎn)擊“閱讀原文”或掃描下方海報二維碼，立即注冊 GTC 大會。