Generative Recommenders

Repository hosting code for Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations (ICML'24 paper) and related code, where we demonstrate that the ubiquitously used classical deep learning recommendation paradigm (DLRMs) can be reformulated as a generative modeling problem (Generative Recommenders or GRs) to overcome known compute scaling bottlenecks, propose efficient algorithms such as HSTU and M-FALCON to accelerate training and inference for large-scale sequential models by 10x-1000x, and demonstrate scaling law for the first-time in deployed, billion-user scale recommendation systems.

Getting started

We recommend using requirements.txt. This has been tested with Ubuntu 22.04, CUDA 12.4, and Python 3.10.

pip3 install -r requirements.txt

Alternatively, you can manually install PyTorch based on official instructions. Then,

pip3 install gin-config pandas fbgemm_gpu torchrec tensorboard

Experiments

Public Experiments

To reproduce the public experiments in our paper (traditional sequential recommender setting, Section 4.1.1) on MovieLens and Amazon Reviews in the paper, please follow these steps:

Download and preprocess data.

mkdir -p tmp/ && python3 preprocess_public_data.py

A GPU with 24GB or more HBM should work for most datasets.

CUDA_VISIBLE_DEVICES=0 python3 main.py --gin_config_file=configs/ml-1m/hstu-sampled-softmax-n128-large-final.gin --master_port=12345

Other configurations are included in configs/ml-1m, configs/ml-20m, and configs/amzn-books to make reproducing these experiments easier.

Verify results.

By default we write experimental logs to exps/. We can launch tensorboard with something like the following:

tensorboard --logdir ~/generative-recommenders/exps/ml-1m-l200/ --port 24001 --bind_all
tensorboard --logdir ~/generative-recommenders/exps/ml-20m-l200/ --port 24001 --bind_all
tensorboard --logdir ~/generative-recommenders/exps/amzn-books-l50/ --port 24001 --bind_all

With the provided configuration (.gin) files, you should be able to reproduce the following results (verified as of 04/15/2024):

MovieLens-1M (ML-1M):

MovieLens-20M (ML-20M):

Amazon Reviews (Books):

for all three tables above, the SASRec rows are based on Self-Attentive Sequential Recommendation but with the original binary cross entropy loss replaced with sampled softmax losses proposed in Revisiting Neural Retrieval on Accelerators. These rows are reproducible with configs/*/sasrec-*-final.gin. The BERT4Rec and GRU4Rec rows are based on results reported by Turning Dross Into Gold Loss: is BERT4Rec really better than SASRec? - note that the comparison slightly favors these two, due to them using full negatives whereas the other rows used 128/512 sampled negatives. The HSTU and HSTU-large rows are based on Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations; in particular, HSTU rows utilize identical configurations as SASRec. HSTU and HSTU-large results can be reproduced with configs/*/hstu-*-final.gin.

Synthetic Dataset / MovieLens-3B

We support generating synthetic dataset with fractal expansion introduced in https://arxiv.org/abs/1901.08910. This allows us to expand the current 20 million real-world ratings in ML-20M to 3 billion.

To download the pre-generated synthetic dataset:

pip3 install gdown
mkdir -p tmp/ && cd tmp/
gdown https://drive.google.com/uc?id=1-jZ6k0el7e7PyFnwqMLfqUTRh_Qdumt-
unzip ml-3b.zip && rm ml-3b.zip

To generate the synthetic dataset on your own:

python3 run_fractal_expansion.py --input-csv-file tmp/ml-20m/ratings.csv --write-dataset True --output-prefix tmp/ml-3b/

Efficiency experiments

ops/triton contains triton kernels needed for efficiency experiments. ops/cpp contains efficient CUDA kernels. In particular, ops/cpp/hstu_attention contains the attention implementation based on FlashAttention V3 with state-of-the-art efficiency on H100 GPUs.

HSTU Ultra attention benchmark

ops/hstu_ultra.py provides named, model-derived HSTU Ultra self-attention configurations backed by ops/pytorch/pt_hstu_ultra.py. The accompanying benchmark can compare the PyTorch reference with the Triton forward and backward kernels, including the semi-local attention configuration.

For the model architecture, see the HSTU Ultra paper. modules/hstu_ultra.py builds complete semi-local or target-aware layer stacks from those configurations. Each layer includes input normalization, UVQK projection, HSTU Ultra attention, output normalization and projection, and the residual connection. Pass fp8_addmm_fwd=True to hstu_ultra_stack_configs to use FP8 for both projection GEMMs. For repeated inference, call stack.prepare_fp8_weights() after moving or loading the model to cache the row-wise quantized projection weights.

The end-to-end integration uses the target-aware HSTU Ultra stack inside the existing DLRMv3 data, multitask-loss, optimizer, checkpoint, and eager inference paths. Start a BF16 training run with:

python3 -m generative_recommenders.dlrm_v3.train.train_ranker \
  --dataset kuairand-1k --hstu-ultra

Add --hstu-ultra-fp8 to use FP8 projection GEMMs (it also enables the Ultra stack). Eager inference accepts the same flags and prepares the FP8 weight caches after loading the checkpoint:

python3 -m generative_recommenders.dlrm_v3.inference.main \
  --dataset debug --hstu-ultra-fp8

Cached incremental decoding and scripted inference export are not included yet.

The complete-stack benchmark measures all configured layers, including input normalization, UVQK projection, attention, output projection, and residuals:

python3 -m generative_recommenders.ops.benchmarks.hstu_ultra_stack_bench \
  --config-name hstu_ultra_semi_local --provider all --sequence-lengths 512

Representative forward-only results for the complete configured stacks on one NVIDIA H100 96 GB GPU are shown below. Each result is the median of three trials with 10 warmup iterations and 100 timed iterations per trial.

The FP8 provider uses cached weights, fused LayerNorm-to-FP8 quantization for the UVQK projection, fused output normalization/gating/dropout-to-FP8 quantization, and a persistent Triton row-wise FP8 GEMM. The semi-local and target-aware configurations are measured independently.

python3 -m generative_recommenders.ops.benchmarks.hstu_ultra_bench \
  --config-name hstu_ultra_semi_local --kernel triton --batch-size 1 --mode fwd

Use --config-name hstu_ultra for target-aware full causal self-attention. Pass --sequence-lengths to override the conservative reference defaults; long production lengths may require substantial GPU memory because the PyTorch implementation materializes dense attention matrices.

Representative forward-pass results on one NVIDIA H100 96 GB GPU are shown below. Measurements used BF16, batch size 1, a 1980 MHz clock, a 700 W power cap, 10 warmup iterations, 100 timed iterations, and the median of three runs.

This benchmark measures attention only with precomputed BF16 Q/K/V tensors. It does not include the UVQK or output projection GEMMs, so the FP8 addmm path is not exercised by these numbers.

These figures are illustrative rather than release guarantees; rerun the benchmark on the deployment hardware and software stack when making tuning decisions.

FP8 linear reference and benchmark

ops/fp8.py provides row-wise E4M3 quantization and FP8 addmm. It dispatches to a persistent Triton implementation when requested and retains the portable PyTorch torch._scaled_mm implementation and CPU fallback. Both paths support on-the-fly or prequantized inputs and a full-precision backward for training. HSTU's Triton path also fuses activation quantization into its input LayerNorm and output postprocessing kernels. Set fp8_in_addmm_fwd=True on hstu_compute_uqvk and hstu_compute_output to use the FP8 path for the HSTU input and output projections. STULayerConfig(fp8_addmm_fwd=True) enables both projection GEMMs for an HSTU layer.

python3 -m generative_recommenders.ops.benchmarks.fp8_addmm_bench \
  --provider all --m 4096 --k 512 --n 512

DLRM-v3

We have created a DLRM model using HSTU and have developed benchmarks for both training and inference to faciliate production RecSys use cases.

Run model training with 4 GPUs

LOCAL_WORLD_SIZE=4 WORLD_SIZE=4 python3 generative_recommenders/dlrm_v3/train/train_ranker.py --dataset debug --mode train

Run model inference with 4 GPUs

git clone --recurse-submodules https://github.com/mlcommons/inference.git mlperf_inference
cd mlperf_inference/loadgen
CFLAGS="-std=c++14 -O3" python -m pip install .

LOCAL_WORLD_SIZE=4 WORLD_SIZE=4 python3 generative_recommenders/dlrm_v3/inference/main.py --dataset debug

License

This codebase is Apache 2.0 licensed, as found in the LICENSE file.

Contributors

The overall project is made possible thanks to the joint work from many technical contributors (listed in alphabetical order):

Adnan Akhundov, Bugra Akyildiz, Shabab Ayub, Alex Bao, Renqin Cai, Jennifer Cao, Xuan Cao, Guoqiang Jerry Chen, Lei Chen, Li Chen, Sean Chen, Xianjie Chen, Huihui Cheng, Weiwei Chu, Ted Cui, Shiyan Deng, Nimit Desai, Fei Ding, Shilin Ding, Francois Fagan, Lu Fang, Leon Gao, Zhaojie Gong, Fangda Gu, Liang Guo, Liz Guo, Jeevan Gyawali, Yuchen Hao, Daisy Shi He, Michael Jiayuan He, Yu He, Samuel Hsia, Jie Hua, Yanzun Huang, Hongyi Jia, Rui Jian, Jian Jin, Rafay Khurram, Rahul Kindi, Changkyu Kim, Yejin Lee, Fu Li, Han Li, Hong Li, Shen Li, Rui Li, Wei Li, Zhijing Li, Lucy Liao, Xueting Liao, Emma Lin, Hao Lin, Chloe Liu, Jingzhou Liu, Xing Liu, Xingyu Liu, Kai Londenberg, Yinghai Lu, Liang Luo, Linjian Ma, Matt Ma, Yun Mao, Bert Maher, Ajit Mathews, Matthew Murphy, Satish Nadathur, Min Ni, Jongsoo Park, Colin Peppler, Jing Qian, Lijing Qin, Jing Shan, Alex Singh, Timothy Shi, Yu Shi, Dennis van der Staay, Xiao Sun, Colin Taylor, Shin-Yeh Tsai, Rohan Varma, Omkar Vichare, Alyssa Wang, Pengchao Wang, Shengzhi Wang, Wenting Wang, Xiaolong Wang, Yueming Wang, Zhiyong Wang, Wei Wei, Bin Wen, Carole-Jean Wu, Yanhong Wu, Eric Xu, Bi Xue, Hong Yan, Zheng Yan, Chao Yang, Junjie Yang, Wen-Yun Yang, Ze Yang, Zimeng Yang, Yuanjun Yao, Chunxing Yin, Daniel Yin, Yiling You, Jiaqi Zhai, Keke Zhai, Yanli Zhao, Zhuoran Zhao, Hui Zhang, Jingjing Zhang, Lu Zhang, Lujia Zhang, Na Zhang, Rui Zhang, Xiong Zhang, Ying Zhang, Zhiyun Zhang, Charles Zheng, Erheng Zhong, Zhao Zhu, Xin Zhuang.

For the initial paper describing the Generative Recommender problem formulation and the algorithms used, including HSTU and M-FALCON, please refer to Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations(ICML'24 paper, slides).