meta-recsys/generative-recommenders

Repository hosting code for "Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations" (https://arxiv.org/abs/2402.17152).

项目内容

Generative Recommenders 是一个研究级代码库,实现了 ICML 2024 论文 Actions Speak Louder than Words: Trillion‑Parameter Sequential Transducers for Generative Recommendations 中的想法。作者证明经典的深度学习推荐模型(DLRM)可以重新定义为生成式建模问题。通过这样做,他们引入了新的训练和推理算法——特别是 HSTU(高速序列转换器)和 M‑FALCON——大幅加速了大规模序列推荐(10×‑1000× 更快),同时提高了准确度。

此仓库包含:

  • HSTU 系列及构建于其上的 DLRM‑v3 的模型定义和训练脚本。
  • 用于在公开数据集(MovieLens‑1M、MovieLens‑20M、Amazon Books)上重现实验的配置文件(Gin)。
  • 通过分形扩展生成合成 30 亿评分数据集的工具。
  • 高度优化的 CUDA/Triton 内核(基于 FlashAttention‑V3),用于超快速注意力与 FP8 矩阵乘法,以及报告 NVIDIA H100 GPU 加速比的基准测试。
  • 用于多 GPU 设置的端到端训练和推理脚本。

简而言之,这是一个完整的研究实现,让您能够大规模训练、评估和基准测试最先进的生成式推荐模型。


如何开始

  1. 安装依赖 – 最简单的方式是 pip3 install -r requirements.txt。README 指出已在 Ubuntu 22.04、CUDA 12.4 和 Python 3.10 上测试过。如果您偏好手动安装,需要 PyTorch 以及 gin-config、pandas、fbgemm_gpu、torchrec 和 tensorboard。
  2. 准备数据 – 运行 python3 preprocess_public_data.py,它会创建 tmp/ 文件夹并下载 MovieLens / Amazon 数据。
  3. 运行实验 – 选择一个 Gin 配置文件(例如 configs/ml-1m/hstu-sampled-softmax-n128-large-final.gin)并启动:
    CUDA_VISIBLE_DEVICES=0 python3 main.py \
        --gin_config_file=configs/ml-1m/hstu-sampled-softmax-n128-large-final.gin \
        --master_port=12345
    
    日志写入 exps/;您可以使用 TensorBoard 可视化它们。
  4. 基准测试内核 – 对于低层性能测试,运行提供的基准测试模块,例如:
    python3 -m generative_recommenders.ops.benchmarks.hstu_ultra_stack_bench \
        --config-name hstu_ultra_semi_local --provider all --sequence-lengths 512
    
  5. 扩展规模 – 仓库包含生成合成 30 亿评分数据集(run_fractal_expansion.py)以及使用 4 个 GPU 训练/推理 DLRM‑v3 模型的脚本。

谁可能觉得这有用

受众 为什么重要
推荐系统或大规模生成式模型的学术研究人员 完整的可重现流程、最先进的结果,以及新颖的算法贡献(HSTU、M‑FALCON)。
构建生产推荐服务的工程师 优化的 CUDA/Triton 内核、FP8 支持,以及可适应真实工作负载的多 GPU 训练脚本。
性能工程师 / 硬件架构师 在 H100 GPU 上的详细基准测试、基于 FlashAttention‑V3 的内核,以及用于衡量加速潜力的 FP8 量化路径。
学习序列推荐的学生 端到端示例代码(数据预处理、模型配置、训练循环),桥接了理论(论文)与实践。

限制 / 开放挑战

  • 硬件需求高 – 最快的内核需要高端 GPU(例如 H100,至少 24 GB HBM)。在一般硬件上运行全规模实验会很慢或不可行。
  • 依赖合成数据 – 30 亿评分数据集是通过分形扩展生成的,可能无法反映真实世界的分布细节。
  • 研究级代码 – 虽然仓库功能完整,但缺乏生产级库的完善度(例如广泛测试、CI 或易于 pip 安装的包)。
  • 范围仅限于序列推荐 – 生成式公式仅针对下一项预测进行演示;扩展到其他推荐任务(例如多模态、基于图的)需要额外工作。

如何融入更广泛的 AI 生态系统

Generative Recommenders 位于大规模语言风格生成式建模与工业推荐系统的交汇点。通过将推荐视为序列到序列的生成问题,此项目利用了 transformer/LLM 领域的进展(例如 sampled softmax、FP8 量化、FlashAttention)来推动推荐吞吐量和模型大小的极限。这反映了更广泛的趋势:最初为 NLP 建立的技术(大参数生成式模型、高效注意力内核)正被重新用于其他领域,如推荐、搜索,甚至机器人技术。这项工作也促进了关于推荐系统计算扩展定律的持续讨论,显示万亿参数模型可以通过正确的算法和系统技巧高效地训练和服务。


快速参考(您最常使用的命令)

# 安装
pip3 install -r requirements.txt

# 预处理公开数据
mkdir -p tmp/ && python3 preprocess_public_data.py

# 在 MovieLens‑1M 上训练 HSTU(单 GPU)
CUDA_VISIBLE_DEVICES=0 \
python3 main.py \
  --gin_config_file=configs/ml-1m/hstu-sampled-softmax-n128-large-final.gin \
  --master_port=12345

# 可视化结果
tensorboard --logdir ~/generative-recommenders/exps/ml-1m-l200/ --port 24001

# 基准测试超快速注意力(H100)
python3 -m generative_recommenders.ops.benchmarks.hstu_ultra_stack_bench \
  --config-name hstu_ultra_semi_local --provider all --sequence-lengths 512

许可证

Apache 2.0 – 免费供商业和学术使用。

相关

  • 项目
  • 项目
  • 项目
  • 项目
  • 项目