引言

在生成式 AI 浪潮的推動下，推薦系統領域正經歷深刻變革。傳統的深度學習推薦模型（DLRMs）雖已展現出一定效果，但在捕捉用戶興趣偏好和動態行為序列變化時，常面臨可擴展性挑戰。生成式推薦系統（Generative Recommenders, GRs）的出現，為這一領域帶來了全新思路與機遇。本文將介紹 NVIDIA recsys-examples 中針對生成式推薦場景設計的高效實踐參考。博客內容共分為上下兩篇，本篇將整體介紹 recsys-examples 的設計和功能，下篇將對核心模塊進行深入的解析。

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

生成式推薦系統的崛起

Meta Research 的最新研究表明，生成式推薦系統通過將推薦問題重構為生成式建模框架下的序列 transduction 任務，展現出顯著優勢：

• 更個性化的推薦：能夠深入挖掘用戶獨特的行為模式與偏好。
• 更強的上下文感知能力：能夠更好地捕捉上下文序列信號，滿足序列建模的模型需求

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

大規模訓練的挑戰與應對方案

然而，基于類 Transformer 架構的生成式推薦系統，因上下文長度增加和計算需求提升，在大規模訓練與部署時面臨嚴峻的計算和架構挑戰。為解決這些難題，NVIDIA 開發了 recsys-examples 參考實現，旨在展示大規模生成式推薦系統中訓練和推理的最優實踐。

NVIDIA recsys-examples 中的深度優化

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

• 混合并行分布式訓練：基于 TorchRec （處理 sparse 部分的模型并行）和 NVIDIA Megatron Core（適用于 dense 部分的數據并行與模型并行），優化多 GPU 分布式訓練流程，實現 sparse 和 dense 部分多種并行的高效協同。
• 高效 HSTU 注意力算子：通過 NVIDIA CUTLASS 實現高性能的 HSTU 注意力算子，提升計算效率。
• 動態 embedding 功能：結合 NVIDIA Merlin HKV 和 TorchRec，支持無沖突哈希、embedding eviction 及 CPU offloading 等動態 embedding 能力，適配大規模訓練場景。

當前，recsys-examples 中提供了基于 HSTU 排序和召回模型的大規模訓練示例，方便用戶快速使用和參考。

圖1: NVIDIA recsys-examples 的軟件架構，其中綠色部分是 NVIDIA 開源組件，紅色部分是社區開源組件

一、高效的 HSTU 內核

HSTU（Hierarchical Sequential Transduction Unit）注意力結構是論文中提出的一種針對推薦系統優化的高效注意力機制。與標準的多頭注意力（Multi-head Attention）相比，HSTU 注意力做了以下關鍵改進：

• Normalization 改進：用 SiLU 替代 softmax，提升模型表達能力。
• 引入相對位置偏置：通過 Relative Attention Bias (RAB) 捕獲序列中的相對位置/時間信息。

在 recsys-examples 中，我們基于 NVIDIA CUTLASS 庫實現了高性能的 HSTU 注意力算子，并針對訓練和推理場景分別進行了優化。目前實現也已經合并到 FBGEMM 中，用戶可以直接通過 FBGEMM 使用。

1、訓練優化技術

• Kernel Fusion 計算融合：借鑒 Flash Attention 的思想，將多個連續操作融合為單個 GPU 內核
• 靈活掩碼和 RAB 機制：支持可定制的 mask tensor 以及 RAB tensor，適應不同推薦場景下的序列建模需求

2、推理優化技術

• 簡化計算邏輯：使用 RAB（Relative Attention Bias）作為負無窮替代傳統的 mask 操作，減少計算復雜度和內存訪問，提升推理速度
• 稀疏目標優化：針對推薦系統中常見的稀疏 target 計算模式進行優化，減少內存占用，并支持大規模目標的批量推理

在 NVIDIA Hopper 架構上，我們的 HSTU 注意力算子相比與 Triton 實現的版本，在各個問題尺寸上都有超過 3.5x 的加速比，并且在序列增長的情況下，加速比進一步提升。

圖 2: CUTLASS Kernel 在 NVIDIA Hopper 架構上與 Triton 的前向性能對比

圖 3: CUTLASS Kernel 在 NVIDIA Hopper 架構上與 Triton 的后向性能對比

二、動態 embedding 與 TorchRec 的無縫集成

TorchRec 目前對動態 embedding 的支持有兩種，分別是 contrib/dynamic_embedding 通過外掛 CPU redis 集群和在 ManagedCollision 模塊中通過額外的排序步驟來支持，兩者都會在原有 TorchRec 訓練流程的基礎上，增加額外的訓練時間開銷。

在 recsys-examples 中，我們引入 NVIDIA Merlin HierarchicalKV 作為底層存儲，并與 TorchRec 團隊合作基于 TorchRec 官方插件接口，直接替換 TorchRec 中原本的 FBGEMM 靜態存儲，支持了動態 embedding 支持能力。這一方案可在大規模推薦系統訓練場景中：

• 支持無沖突哈希映射
• 支持基于頻率或時間或自定義的 embedding 淘汰策略
• 提供 CPU offloading 機制來實現超大規模 embedding 存儲
• 支持 incremental dump 功能，根據用戶的需求只 dump 在過去一段時間內訓練過的 embedding
• 保持與原生 TorchRec 分布式訓練流程的無縫集成

相比 contrib/dynamic_embedding 中的實現，NVIDIA recsys-examples 能夠大幅度減少 CPU 上的操作開銷，在大規模訓練中能有超過20倍的加速效果。

美團應用 NVIDIA recsys-examples 實踐

在過去幾個月中，我們與美團緊密合作，助力其加速基于 HSTU 架構的推薦模型在離線和在線試驗中的應用。在美團外賣場景下，通過引入GR模型結構，CTR和CTCVR指標均實現了顯著提升（詳情參考MTGR博客）。

圖 4: 美團業務引入 GR 后的收益。

（該圖片來源于 MTGR：美團外賣生成式推薦 Scaling Law 落地實踐一文，若您有任何疑問或需要使用該圖片，請聯系美團）

我們的優化版 HSTU 算子，在訓練中，端對端吞吐提升 85%；在推理中，通過 TRT plugin 封裝，在 TRT 中引入了 HSTU fp16 算子，相比 TRT fp32 算子時延降低 50%，端對端耗時減少 30%。

總結與展望

NVIDIA recsys-examples 將生成式推薦（如 “Actions Speak Louder than Words” 論文中提出的技術）與分布式訓練（借由 TorchRec 增強）及優化訓練推理相結合，助力開發和部署能夠提供高度個性化用戶體驗的復雜推薦模型。我們誠摯邀請研究人員和從業者試用該工具，并期待與您共同推動生成式推薦系統的技術演進。

相關資源

1. TorchRec 官方文檔:?https://docs.pytorch.org/torchrec/?
2. Meta HSTU 論文:?https://arxiv.org/abs/2402.17152?
3. MTGR博客：MTGR：美團外賣生成式推薦 Scaling Law 落地實踐:?https://mp.weixin.qq.com/s/JiDOqD-ThU0Upp6xnNg3Nw?
4. Add HSTU in fbgemm_gpu/experimental/:?https://github.com/pytorch/FBGEMM/pull/4090?
5. enable customized emb lookup kernel for TorchRec:?https://github.com/pytorch/torchrec/pull/2887?