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

• URL: http://arxiv.org/abs/2504.01506v1
• Date: Wed, 02 Apr 2025 08:57:01 GMT
• Title: MLKV: Efficiently Scaling up Large Embedding Model Training with Disk-based Key-Value Storage
• Authors: Yongjun He, Roger Waleffe, Zhichao Han, Johnu George, Binhang Yuan, Zitao Zhang, Yinan Shan, Yang Zhao, Debojyoti Dutta, Theodoros Rekatsinas, Ce Zhang,
• Abstract summary: This paper presents MLKV, an efficient, reusable data storage framework designed to address the scalability challenges in embedding model training.<n>In experiments on open-source workloads, MLKV outperforms offloading strategies built on top of industrial-strength key-value stores by 1.6-12.6x.
• Score: 22.848456481878568
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• PLM: Efficient Peripheral Language Models Hardware-Co-Designed for Ubiquitous Computing [48.30406812516552]
  We introduce the PLM, a Peripheral Language Model, developed through a co-design process that jointly optimize model architecture and edge system constraints.<n>PLM employs a Multi-head Latent Attention mechanism and employs the squared ReLU activation function to encourage sparsity, thereby reducing peak memory footprint.<n> evaluation results demonstrate that PLM outperforms existing small language models trained on publicly available data.
  arXiv  Detail & Related papers  (2025-03-15T15:11:17Z)
• 1bit-Merging: Dynamic Quantized Merging for Large Language Models [20.19975755949984]
  texttt1bit-Merging is a novel framework that integrates task-specific routing with 1-bit quantized task vectors to balance performance and storage efficiency. We demonstrate that texttt1bit-Merging achieves comparable or superior performance to existing methods while significantly reducing storage requirements.
  arXiv  Detail & Related papers  (2025-02-15T09:47:50Z)
• A Bring-Your-Own-Model Approach for ML-Driven Storage Placement in Warehouse-Scale Computers [4.849222239746218]
  Storage systems account for a major portion of the total cost of ownership (TCO) of warehouse-scale computers. Machine learning (ML)-based methods for solving key problems in storage system efficiency, such as data placement, have shown significant promise. We study this problem in the context of real-world hyperscale data centers at Google.
  arXiv  Detail & Related papers  (2025-01-10T01:42:05Z)
• Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design [59.00758127310582]
  We propose a novel framework Read-ME that transforms pre-trained dense LLMs into smaller MoE models. Our approach employs activation sparsity to extract experts. Read-ME outperforms other popular open-source dense models of similar scales.
  arXiv  Detail & Related papers  (2024-10-24T19:48:51Z)
• xGen-MM (BLIP-3): A Family of Open Large Multimodal Models [157.44696790158784]
  This report introduces xGen-MM, a framework for developing Large Multimodal Models (LMMs) The framework comprises meticulously curated datasets, a training recipe, model architectures, and a resulting suite of LMMs. Our models undergo rigorous evaluation across a range of tasks, including both single and multi-image benchmarks.
  arXiv  Detail & Related papers  (2024-08-16T17:57:01Z)
• SWIFT:A Scalable lightWeight Infrastructure for Fine-Tuning [4.173156963843178]
  Large Language Models (LLMs) and Multi-modal Large Language Models (MLLMs) have leverage Attention-based Transformer architectures. We develop SWIFT, a customizable one-stop infrastructure for large models. We show that notable improvements on the ToolBench leader-board can be achieved by training with customized dataset on SWIFT.
  arXiv  Detail & Related papers  (2024-08-10T11:00:13Z)
• LLAVADI: What Matters For Multimodal Large Language Models Distillation [77.73964744238519]
  In this work, we do not propose a new efficient model structure or train small-scale MLLMs from scratch. Our studies involve training strategies, model choices, and distillation algorithms in the knowledge distillation process. By evaluating different benchmarks and proper strategy, even a 2.7B small-scale model can perform on par with larger models with 7B or 13B parameters.
  arXiv  Detail & Related papers  (2024-07-28T06:10:47Z)
• Bullion: A Column Store for Machine Learning [4.096087402737292]
  This paper presents Bullion, a columnar storage system tailored for machine learning workloads. Bundy addresses the complexities of data compliance, optimize the encoding of long sequence sparse features, efficiently manages wide-table projections, introduces feature quantization in storage, and provides a comprehensive cascading encoding framework. Preliminary experimental results and theoretical analysis demonstrate Bullion's improved ability to deliver strong performance in the face of the unique demands of machine learning workloads.
  arXiv  Detail & Related papers  (2024-04-13T05:01:54Z)
• Retrieval-based Knowledge Transfer: An Effective Approach for Extreme Large Language Model Compression [64.07696663255155]
  Large-scale pre-trained language models (LLMs) have demonstrated exceptional performance in various natural language processing (NLP) tasks. However, the massive size of these models poses huge challenges for their deployment in real-world applications. We introduce a novel compression paradigm called Retrieval-based Knowledge Transfer (RetriKT) which effectively transfers the knowledge of LLMs to extremely small-scale models.
  arXiv  Detail & Related papers  (2023-10-24T07:58:20Z)
• In Situ Framework for Coupling Simulation and Machine Learning with Application to CFD [51.04126395480625]
  Recent years have seen many successful applications of machine learning (ML) to facilitate fluid dynamic computations. As simulations grow, generating new training datasets for traditional offline learning creates I/O and storage bottlenecks. This work offers a solution by simplifying this coupling and enabling in situ training and inference on heterogeneous clusters.
  arXiv  Detail & Related papers  (2023-06-22T14:07:54Z)
• COMET: A Comprehensive Cluster Design Methodology for Distributed Deep Learning Training [42.514897110537596]
  Modern Deep Learning (DL) models have grown to sizes requiring massive clusters of specialized, high-end nodes to train. designing such clusters to maximize both performance and utilization--to amortize their steep cost--is a challenging task. We introduce COMET, a holistic cluster design methodology and workflow to jointly study the impact of parallelization strategies and key cluster resource provisioning on the performance of distributed DL training.
  arXiv  Detail & Related papers  (2022-11-30T00:32:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.