# MLKV: Efficiently Scaling up Large Embedding Model Training with Disk-based Key-Value Storage

arXiv: https://arxiv.org/abs/2504.01506 · HTML: https://mlsanity.com/paper/2504.01506/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 62 references that are in the map.
- First posted on arXiv 2 Apr 2025.

## Builds on (most cited first, 28 of 28)

- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://mlsanity.com/paper/1412.6575/) (2014)
- [Complex Embeddings for Simple Link Prediction](https://mlsanity.com/paper/1606.06357/) (2016)
- [Deep & Cross Network for Ad Click Predictions](https://mlsanity.com/paper/1708.05123/) (2017)
- [A Comprehensive Survey of Graph Embedding: Problems, Techniques and Applications](https://mlsanity.com/paper/1709.07604/) (2017)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Carbontracker: Tracking and Predicting the Carbon Footprint of Training Deep Learning Models](https://mlsanity.com/paper/2007.03051/) (2020)
- [Deep Learning based Recommender System: A Survey and New Perspectives](https://mlsanity.com/paper/1707.07435/) (2017)
- [PyTorch-BigGraph: A Large-scale Graph Embedding System](https://mlsanity.com/paper/1903.12287/) (2019)
- [Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Training](https://mlsanity.com/paper/2208.06102/) (2022)
- [Software-Hardware Co-design for Fast and Scalable Training of Deep Learning Recommendation Models](https://mlsanity.com/paper/2104.05158/) (2021)
- [DGL-KE: Training Knowledge Graph Embeddings at Scale](https://mlsanity.com/paper/2004.08532/) (2020)
- [Understanding Scaling Laws for Recommendation Models](https://mlsanity.com/paper/2208.08489/) (2022)
- [Petuum: A New Platform for Distributed Machine Learning on Big Data](https://mlsanity.com/paper/1312.7651/) (2013)
- [Distributed Hybrid CPU and GPU training for Graph Neural Networks on Billion-Scale Graphs](https://mlsanity.com/paper/2112.15345/) (2021)
- [Marius: Learning Massive Graph Embeddings on a Single Machine](https://mlsanity.com/paper/2101.08358/) (2021)
- [HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework](https://mlsanity.com/paper/2112.07221/) (2021)
- [Merlin HugeCTR: GPU-accelerated Recommender System Training and Inference](https://mlsanity.com/paper/2210.08803/) (2022)
- [MariusGNN: Resource-Efficient Out-of-Core Training of Graph Neural Networks](https://mlsanity.com/paper/2202.02365/) (2022)
- [Persia: An Open, Hybrid System Scaling Deep Learning-based Recommenders up to 100 Trillion Parameters](https://mlsanity.com/paper/2111.05897/) (2021)
- [BagPipe: Accelerating Deep Recommendation Model Training](https://mlsanity.com/paper/2202.12429/) (2022)
- [BenchTemp: A General Benchmark for Evaluating Temporal Graph Neural Networks](https://mlsanity.com/paper/2308.16385/) (2023)
- [NuPS: A Parameter Server for Machine Learning with Non-Uniform Parameter Access](https://mlsanity.com/paper/2104.00501/) (2021)
- [CoroBase: Coroutine-Oriented Main-Memory Database Engine](https://mlsanity.com/paper/2010.15981/) (2020)

## Abstract (arXiv)

Many modern machine learning (ML) methods rely on embedding models to learn vector representations (embeddings) for a set of entities (embedding tables). As increasingly diverse ML applications utilize embedding models and embedding tables continue to grow in size and number, there has been a surge in the ad-hoc development of specialized frameworks targeted to train large embedding models for specific tasks. Although the scalability issues that arise in different embedding model training tasks are similar, each of these frameworks independently reinvents and customizes storage components for specific tasks, leading to substantial duplicated engineering efforts in both development and deployment. This paper presents MLKV, an efficient, extensible, and reusable data storage framework designed to address the scalability challenges in embedding model training, specifically data stall and staleness. MLKV augments disk-based key-value storage by democratizing optimizations that were previously exclusive to individual specialized frameworks and provides easy-to-use interfaces for embedding model training tasks. Extensive experiments on open-source workloads, as well as applications in eBay's payment transaction risk detection and seller payment risk detection, show that MLKV outperforms offloading strategies built on top of industrial-strength key-value stores by 1.6-12.6x. MLKV is open-source at https://github.com/llm-db/MLKV.