How ai-hypercomputer/recml works
RecML is a recommender-modeling library focused on scalable Keras/JAX workloads, especially Cloud TPU and SparseCore. The principal workflow is to prepare and batch recommendation data, apply a recommendation model, train and evaluate it with a framework-specific trainer, then retain checkpoints or use a model for inference. The README also documents reusable layers, embedding support, and standardized metrics. Some connections are documented capabilities rather than verified calls: sampled excerpts confirm the JAX trainer’s dataset, evaluation, and checkpoint responsibilities, but not every example’s wiring.
Main entry points: jax_trainer.py JAX trainer · training Trainer APIs · DLRM_DCNv2 DLRM models · sparsecore.py SparseCore embeddings
Components 19
Data preparation
- Preprocessing
recml/core/data/preprocessing.py - Dataset factory
recml/core/data/tf_dataset_factory.py - Data iterator
recml/core/data/iterator.py
Recommendation models
- Sequential models
recml/layers/keras - HSTU models
recml/examples/DLRM_HSTU - DLRM models
recml/inference/models/jax/DLRM_DCNv2
Training and evaluation
- Trainer APIs
recml/core/training - JAX trainer
recml/core/training/jax_trainer.py - Keras trainer
recml/core/training/keras_trainer.py - SPMD partitioning
recml/core/training/partitioning.py - Checkpoint files
Reusable components
- Metrics
recml/core/metrics - Recommendation layers
recml/layers - SparseCore embeddings
recml/layers/linen/sparsecore.py - Embedding operations
recml/core/ops/embedding_ops.py
Inference
- DLRM inference
recml/inference/models/jax/DLRM_DCNv2/dlrm_main.py
Other components
- Practitioner
- Input records
- Cloud TPU
Connections 22
- Practitioner → Input records: provides data
- Input records → Preprocessing: prepares features
- Preprocessing → Dataset factory: feeds examples
- Dataset factory → Data iterator: batches data
- Practitioner → Trainer APIs: configures training
- Trainer APIs → JAX trainer: selects JAX
- Trainer APIs → Keras trainer: selects Keras
- JAX trainer → Data iterator: consumes datasets
- JAX trainer → Checkpoint files: saves state
- JAX trainer → Metrics: reports evaluation
- Keras trainer → Metrics: evaluates with
- JAX trainer → SPMD partitioning: uses partitioning
- Practitioner → Sequential models: chooses model
- Practitioner → HSTU models: chooses model
- Practitioner → DLRM models: chooses model
- Sequential models → Recommendation layers: uses building blocks
- HSTU models → Recommendation layers: uses building blocks
- DLRM models → SparseCore embeddings: uses embeddings
- SparseCore embeddings → Embedding operations: performs lookup
- Cloud TPU → SparseCore embeddings: accelerates lookup
- Practitioner → DLRM inference: runs inference
- DLRM inference → DLRM models: loads model