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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: He, Yongjun, Waleffe, Roger, Han, Zhichao, George, Johnu, Yuan, Binhang, Zhang, Zitao, Shan, Yinan, Zhao, Yang, Dutta, Debojyoti, Rekatsinas, Theodoros, Zhang, Ce
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866915223275831296
author He, Yongjun
Waleffe, Roger
Han, Zhichao
George, Johnu
Yuan, Binhang
Zhang, Zitao
Shan, Yinan
Zhao, Yang
Dutta, Debojyoti
Rekatsinas, Theodoros
Zhang, Ce
author_facet He, Yongjun
Waleffe, Roger
Han, Zhichao
George, Johnu
Yuan, Binhang
Zhang, Zitao
Shan, Yinan
Zhao, Yang
Dutta, Debojyoti
Rekatsinas, Theodoros
Zhang, Ce
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2504_01506
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle MLKV: Efficiently Scaling up Large Embedding Model Training with Disk-based Key-Value Storage
He, Yongjun
Waleffe, Roger
Han, Zhichao
George, Johnu
Yuan, Binhang
Zhang, Zitao
Shan, Yinan
Zhao, Yang
Dutta, Debojyoti
Rekatsinas, Theodoros
Zhang, Ce
Machine Learning
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.
title MLKV: Efficiently Scaling up Large Embedding Model Training with Disk-based Key-Value Storage
topic Machine Learning
url https://arxiv.org/abs/2504.01506