引言

在生成式 AI 浪潮的推動(dòng)下，推薦系統(tǒng)領(lǐng)域正經(jīng)歷深刻變革。傳統(tǒng)的深度學(xué)習(xí)推薦模型 (DLRMs) 雖已展現(xiàn)出一定效果，但在捕捉用戶興趣偏好和動(dòng)態(tài)行為序列變化時(shí)，常面臨可擴(kuò)展性挑戰(zhàn)。生成式推薦系統(tǒng) (Generative Recommenders, GRs) 的出現(xiàn)，為這一領(lǐng)域帶來(lái)了全新思路與機(jī)遇。

本文將介紹NVIDIA recsys-examples中針對(duì)生成式推薦場(chǎng)景設(shè)計(jì)的高效實(shí)踐參考。博客內(nèi)容共分為上下兩篇，本篇將整體介紹 recsys-examples 的設(shè)計(jì)和功能，下篇將對(duì)核心模塊進(jìn)行深入的解析。

GitHub repo:https://github.com/NVIDIA/recsys-examples

生成式推薦系統(tǒng)的崛起

Meta Research 的最新研究表明，生成式推薦系統(tǒng)通過(guò)將推薦問(wèn)題重構(gòu)為生成式建?？蚣芟碌男蛄?transduction 任務(wù)，展現(xiàn)出顯著優(yōu)勢(shì)：

更個(gè)性化的推薦：能夠深入挖掘用戶獨(dú)特的行為模式與偏好。

更強(qiáng)的上下文感知能力：能夠更好地捕捉上下文序列信號(hào)，滿足序列建模的模型需求

以 Meta 提出的 HSTU (Hierarchical Sequential Transduction Units) 為例，其在推薦場(chǎng)景中性能超越傳統(tǒng) Transformer 模型，且推理速度更快。

大規(guī)模訓(xùn)練的挑戰(zhàn)與應(yīng)對(duì)方案

然而，基于類 Transformer 架構(gòu)的生成式推薦系統(tǒng)，因上下文長(zhǎng)度增加和計(jì)算需求提升，在大規(guī)模訓(xùn)練與部署時(shí)面臨嚴(yán)峻的計(jì)算和架構(gòu)挑戰(zhàn)。為解決這些難題，NVIDIA 開(kāi)發(fā)了 recsys-examples 參考實(shí)現(xiàn)，旨在展示大規(guī)模生成式推薦系統(tǒng)中訓(xùn)練和推理的最優(yōu)實(shí)踐。

NVIDIA recsys-examples 中的深度優(yōu)化

NVIDIA recsys-examples 目前主要包含以下特性：

混合并行分布式訓(xùn)練：基于 TorchRec（處理 sparse 部分的模型并行）和NVIDIA Megatron Core（適用于 dense 部分的數(shù)據(jù)并行與模型并行），優(yōu)化多 GPU 分布式訓(xùn)練流程，實(shí)現(xiàn) sparse 和 dense 部分多種并行的高效協(xié)同。

高效 HSTU 注意力算子：通過(guò)NVIDIA CUTLASS實(shí)現(xiàn)高性能的 HSTU 注意力算子，提升計(jì)算效率。

動(dòng)態(tài) embedding 功能：結(jié)合NVIDIA Merlin HKV和 TorchRec，支持無(wú)沖突哈希、embedding eviction 及 CPU offloading 等動(dòng)態(tài) embedding 能力，適配大規(guī)模訓(xùn)練場(chǎng)景。

當(dāng)前，recsys-examples 中提供了基于 HSTU 排序和召回模型的大規(guī)模訓(xùn)練示例，方便用戶快速使用和參考。

圖 1：NVIDIA recsys-examples 的軟件架構(gòu)，其中綠色部分是 NVIDIA 開(kāi)源組件，紅色部分是社區(qū)開(kāi)源組件

一、高效的 HSTU 內(nèi)核

HSTU (Hierarchical Sequential Transduction Unit) 注意力結(jié)構(gòu)是論文中提出的一種針對(duì)推薦系統(tǒng)優(yōu)化的高效注意力機(jī)制。與標(biāo)準(zhǔn)的多頭注意力 (Multi-head Attention) 相比，HSTU 注意力做了以下關(guān)鍵改進(jìn)：

Normalization改進(jìn)：用 SiLU 替代 softmax，提升模型表達(dá)能力。

引入相對(duì)位置偏置：通過(guò) Relative Attention Bias (RAB) 捕獲序列中的相對(duì)位置/時(shí)間信息。

在 recsys-examples 中，我們基于 NVIDIA CUTLASS 庫(kù)實(shí)現(xiàn)了高性能的 HSTU 注意力算子，并針對(duì)訓(xùn)練和推理場(chǎng)景分別進(jìn)行了優(yōu)化。目前實(shí)現(xiàn)也已經(jīng)合并到 FBGEMM 中，用戶可以直接通過(guò) FBGEMM 使用。

1、訓(xùn)練優(yōu)化技術(shù)

Kernel Fusion 計(jì)算融合：借鑒 Flash Attention 的思想，將多個(gè)連續(xù)操作融合為單個(gè) GPU 內(nèi)核

靈活掩碼和 RAB 機(jī)制：支持可定制的 mask tensor 以及 RAB tensor，適應(yīng)不同推薦場(chǎng)景下的序列建模需求

2、推理優(yōu)化技術(shù)

簡(jiǎn)化計(jì)算邏輯：使用 RAB (Relative Attention Bias) 作為負(fù)無(wú)窮替代傳統(tǒng)的 mask 操作，減少計(jì)算復(fù)雜度和內(nèi)存訪問(wèn)，提升推理速度

稀疏目標(biāo)優(yōu)化：針對(duì)推薦系統(tǒng)中常見(jiàn)的稀疏 target 計(jì)算模式進(jìn)行優(yōu)化，減少內(nèi)存占用，并支持大規(guī)模目標(biāo)的批量推理

在 NVIDIA Hopper 架構(gòu)上，我們的 HSTU 注意力算子相比與 Triton 實(shí)現(xiàn)的版本，在各個(gè)問(wèn)題尺寸上都有超過(guò) 3.5x 的加速比，并且在序列增長(zhǎng)的情況下，加速比進(jìn)一步提升。

圖 2：CUTLASS Kernel 在 NVIDIA Hopper 架構(gòu)上與 Triton 的前向性能對(duì)比

圖 3：CUTLASS Kernel 在 NVIDIA Hopper 架構(gòu)上與 Triton 的后向性能對(duì)比

二、動(dòng)態(tài) embedding 與TorchRec 的無(wú)縫集成

TorchRec 目前對(duì)動(dòng)態(tài) embedding 的支持有兩種，分別是 contrib / dynamic_embedding 通過(guò)外掛 CPU redis 集群和在 ManagedCollision 模塊中通過(guò)額外的排序步驟來(lái)支持，兩者都會(huì)在原有 TorchRec 訓(xùn)練流程的基礎(chǔ)上，增加額外的訓(xùn)練時(shí)間開(kāi)銷。

在 recsys-examples 中，我們引入 NVIDIA Merlin HierarchicalKV 作為底層存儲(chǔ)，并與 TorchRec 團(tuán)隊(duì)合作基于 TorchRec 官方插件接口，直接替換 TorchRec 中原本的 FBGEMM 靜態(tài)存儲(chǔ)，支持了動(dòng)態(tài) embedding 支持能力。這一方案可在大規(guī)模推薦系統(tǒng)訓(xùn)練場(chǎng)景中：

支持無(wú)沖突哈希映射

支持基于頻率或時(shí)間或自定義的 embedding 淘汰策略

提供 CPU offloading 機(jī)制來(lái)實(shí)現(xiàn)超大規(guī)模 embedding 存儲(chǔ)

支持 incremental dump 功能，根據(jù)用戶的需求只 dump 在過(guò)去一段時(shí)間內(nèi)訓(xùn)練過(guò)的 embedding

保持與原生 TorchRec 分布式訓(xùn)練流程的無(wú)縫集成

相比 contrib / dynamic_embedding 中的實(shí)現(xiàn)，NVIDIA recsys-examples 能夠大幅度減少 CPU 上的操作開(kāi)銷，在大規(guī)模訓(xùn)練中能有超過(guò) 20 倍的加速效果。

更多詳細(xì)內(nèi)容您可觀看 "RecSys Examples 中的訓(xùn)練與推理優(yōu)化實(shí)踐——以 HSTU 模型為例"：

美團(tuán)應(yīng)用

NVIDIA recsys-examples 實(shí)踐

在過(guò)去幾個(gè)月中，我們與美團(tuán)緊密合作，助力其加速基于 HSTU 架構(gòu)的推薦模型在離線和在線試驗(yàn)中的應(yīng)用。在美團(tuán)外賣場(chǎng)景下，通過(guò)引入 GR 模型結(jié)構(gòu)，CTR 和 CTCVR 指標(biāo)均實(shí)現(xiàn)了顯著提升（詳情參考 MTGR 博客）。

圖 4：美團(tuán)業(yè)務(wù)引入 GR 后的收益。

該圖片來(lái)源于 MTGR：美團(tuán)外賣生成式推薦 Scaling Law 落地實(shí)踐一文，若您有任何疑問(wèn)或需要使用該圖片，請(qǐng)聯(lián)系美團(tuán)

我們的優(yōu)化版 HSTU 算子，在訓(xùn)練中，端對(duì)端吞吐提升 85%；在推理中，通過(guò) TRT plugin 封裝，在 TRT 中引入了 HSTU fp16 算子，相比 TRT fp32 算子時(shí)延降低 50%，端對(duì)端耗時(shí)減少 30%。

總結(jié)與展望

NVIDIA recsys-examples 將生成式推薦（如 “Actions Speak Louder than Words” 論文中提出的技術(shù)）與分布式訓(xùn)練（借由 TorchRec 增強(qiáng)）及優(yōu)化訓(xùn)練推理相結(jié)合，助力開(kāi)發(fā)和部署能夠提供高度個(gè)性化用戶體驗(yàn)的復(fù)雜推薦模型。我們誠(chéng)摯邀請(qǐng)研究人員和從業(yè)者試用該工具，并期待與您共同推動(dòng)生成式推薦系統(tǒng)的技術(shù)演進(jìn)。

作者

劉仕杰

2020 年加入 NVIDIA DevTech 團(tuán)隊(duì)，專注于 NVIDIA GPU 的性能優(yōu)化及推薦系統(tǒng)加速領(lǐng)域。

張俊杰

來(lái)自 NVIDIA DevTech 團(tuán)隊(duì)，從事企業(yè)用戶 GPU 加速計(jì)算支持工作，目前主要負(fù)責(zé)推薦系統(tǒng)訓(xùn)練端到端優(yōu)化工作。

姚家樹(shù)

來(lái)自 NVIDIA DevTech 團(tuán)隊(duì)，從事企業(yè)用戶 GPU 加速計(jì)算支持工作，目前主要負(fù)責(zé)推薦系統(tǒng) Embedding 存儲(chǔ)的開(kāi)發(fā)和性能優(yōu)化工作。

康暉

2022 年加入 NVIDIA DevTech 工程師團(tuán)隊(duì)，目前從事機(jī)器人仿真加速相關(guān)工作，之前參加過(guò) HugeCTR，SOK，recsys example 等項(xiàng)目開(kāi)發(fā)和優(yōu)化。

柴斌

來(lái)自 NVIDIA DevTech 團(tuán)隊(duì)，從事企業(yè)用戶 GPU 加速計(jì)算支持工作。目前主要負(fù)責(zé)搜廣推鏈路的性能調(diào)優(yōu)和 kernel 開(kāi)發(fā)。

陳喬瑞

來(lái)自 NVIDIA DevTech 團(tuán)隊(duì)，從事企業(yè)用戶 GPU 加速計(jì)算支持工作，目前主要負(fù)責(zé) HPC 程序的開(kāi)發(fā)和 kernel 性能優(yōu)化。

孫佳鈺

來(lái)自 NVIDIA DevTech 團(tuán)隊(duì)，從事企業(yè)用戶 GPU 加速計(jì)算支持工作。

張琪

來(lái)自 NVIDIA DevTech 團(tuán)隊(duì)，從事企業(yè)用戶 GPU 加速計(jì)算支持工作，目前主要負(fù)責(zé) CUTLASS 在推薦系統(tǒng)、LLM 等相關(guān)應(yīng)用場(chǎng)景的性能優(yōu)化與開(kāi)發(fā)工作。