RecML: High-Performance Recommender Library

RecML is a high-performance, large-scale deep learning recommender system library. It is specifically optimized for Cloud TPUs (leveraging SparseCore acceleration) and Cloud GPUs, making it suitable for training and serving massive models with large embedding tables on datasets containing millions or billions of items/users.

Key capabilities include:

• Scalable Architecture: Support for massive datasets via efficient data loading (e.g., tf.data) and sharding/SPMD.
• State-of-the-Art Models: Reference implementations for ranking, retrieval, and sequential recommendation tasks.
• Reusable Building Blocks: Access to common recommendation layers (e.g., DCN, BERT4Rec), specialized Embedding APIs (JAX Embedding API for SparseCore), standardized metrics (AUC, NDCG@K, etc.), and common loss functions.
• Unified Trainer: A high-level abstraction for training and evaluation that targets different hardware (TPU/GPU) and frameworks (primarily Keras and JAX).

BERT4Rec uses a transformer-based approach to model user preferences by looking at action sequences in both directions (bidirectional). It is trained using a mask prediction objective where random items are masked and the model learns to predict them based on context.

Architecture Components:

• Embedding Layer: Converts item IDs to dense vectors with optional type embeddings and absolute positional embeddings. It uses a separate embedding for masked features to prevent other tokens from attending to them.
• Multi-Head Self-Attention Layer: Computes attention scores using a separate embedding for masked features.
• Feed-Forward Network: Two linear layers with GeLU activation.
• Residual Connections and Post-Layernorm: Applied around self-attention and feed-forward sub-layers.
• Masked Prediction Head: Projects masked sequence embeddings and decodes them using the item embedding layer, computing categorical cross entropy loss against masked item IDs.

SASRec uses self-attention to predict a user's next action based on past activities. It is designed to balance long-term user interests with recent actions by adapting focus based on history length. It is built with efficient attention blocks, avoiding RNN or CNN structures for faster training.

Architecture Components:

• Embedding Layer: Converts item IDs to dense vectors with learnable absolute positional embeddings and dropout.
• Multi-Head Self-Attention Layer: Computes attention scores within a sequence window using causal masking to prevent information leakage.
• Feed-Forward Network: Two linear layers with GeLU activation applied to each embedding.
• Residual Connections and Pre-Layernorm: Applied around self-attention and feed-forward sub-layers for stability.
• Prediction Head: Decodes sequence embeddings into logits using the input item embedding table and computes causal categorical cross entropy loss.

HSTU is a generative modeling architecture designed for high-cardinality, non-stationary streaming data. It treats recommendation as a sequential transduction task and is highly scalable.

Architecture Components:

• Embedding Layer: Converts action tokens into dense vectors; optionally includes learnable absolute positional embeddings.
• Gated Pointwise Aggregated Attention: Uses a multi-head gated pointwise attention mechanism with Layernorm on outputs to capture interaction intensity.
• Relative Attention Bias: Uses T5-style relative attention bias based on positions and timestamps.
• Residual Connections and Pre-Layernorm: Applied around pointwise attention blocks.
• No Feedforward Network: Unlike other models, the standard feedforward network is removed.
• Prediction Head: Decodes sequence embeddings into logits using separately learned weights and computes causal categorical cross entropy loss.

For development on a TPU VM, it is recommended to use a Python virtual environment to prevent conflicts with system-level packages.

Prerequisites:

• Python 3.12 or higher.

Steps:

1. Create the virtual environment from the repository root.
2. Activate the environment.
3. Install the required dependencies and the specific jax-tpu-embedding wheel.
4. Crucial: Force a specific version of protobuf using the --no-deps flag to ensure compatibility with the TPU stack and prevent accidental downgrades by other libraries.

# 1. Create and Activate
python3 -m venv venv
source venv/bin/activate

# 2. Install Dependencies
pip install ./jax_tpu_embedding-0.1.0.dev20260226-cp312-cp312-manylinux_2_31_x86_64.whl
pip install -r requirements.txt

# 3. Force Protobuf version (Required)
pip install "protobuf>=6.31.1" --no-deps

# 4. Run DLRM training
python dlrm_experiment_test.py

Follow these steps via gcloud alpha compute tpus tpu-vm ssh to prepare the TPU environment for RecML benchmarks. These commands are executed on all workers simultaneously.

1. Clone the RecML repository
2. Install requirements from requirements.txt
3. Install JAX and JAXLib nightly (including libtpu and requests)
4. Install JAX Sparsecore (jax-tpu-embedding)
5. Install additional dependencies (tensorflow, dm-tree, flax, google-metrax)

# Clone the RecML repository
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE} --worker=all  --command="git clone https://github.com/AI-Hypercomputer/RecML.git"

# Install requirements
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE} --worker=all  --command="cd RecML && pip install -r requirements.txt" 

# Install jax and jaxlib nightly
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE} --worker=all  --command="pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html --force"

# Install JAX Sparsecore (jax-tpu-embedding)
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE} --worker=all  --command="pip install -U https://storage.googleapis.com/jax-tpu-embedding-whls/20250604/jax_tpu_embedding-0.1.0.dev20250604-cp310-cp310-manylinux_2_35_x86_64.whl --force"

# Install other dependencies
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT} --zone ${ZONE} --worker=all  --command="pip install -U tensorflow  dm-tree flax google-metrax"