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

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

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

客戶簡(jiǎn)介

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

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

需有效提升算法迭代

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

SOK 和 HKV

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

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

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

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

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

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

圖 2. 性能加速效果

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

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

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

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

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

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

持續(xù)合作：

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

雙方團(tuán)隊(duì)通過(guò) SOK 和 HKV 對(duì)原方案進(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)一步對(duì)性能熱點(diǎn) multi-head-attention 部分進(jìn)行優(yōu)化，正在將 Flash-Attention 以 tf-plugin 形式進(jìn)行集成，預(yù)計(jì)此項(xiàng)優(yōu)化集成后，整體加速比可達(dá)到 3 倍，同時(shí)整體的優(yōu)化方案使得顯存使用量下降約 70%，可以顯著地緩解顯存緊張的問(wèn)題。

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

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

• 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 日在美國(guó)加州圣何塞會(huì)議中心舉行，線上大會(huì)也將同期開(kāi)放。點(diǎn)擊“閱讀原文”或掃描下方海報(bào)二維碼，立即注冊(cè) GTC 大會(huì)。