Abstract:Sharding a large machine learning model across multiple devices to balance the costs is important in distributed training. This is challenging because partitioning is NP-hard, and estimating the costs accurately and efficiently is difficult. In this work, we explore a "pre-train, and search" paradigm for efficient sharding. The idea is to pre-train a universal and once-for-all neural network to predict the costs of all the possible shards, which serves as an efficient sharding simulator. Built upon this pre-trained cost model, we then perform an online search to identify the best sharding plans given any specific sharding task. We instantiate this idea in deep learning recommendation models (DLRMs) and propose NeuroShard for embedding table sharding. NeuroShard pre-trains neural cost models on augmented tables to cover various sharding scenarios. Then it identifies the best column-wise and table-wise sharding plans with beam search and greedy grid search, respectively. Experiments show that NeuroShard significantly and consistently outperforms the state-of-the-art on the benchmark sharding dataset, achieving up to 23.8% improvement. When deployed in an ultra-large production DLRM with multi-terabyte embedding tables, NeuroShard achieves 11.6% improvement in embedding costs over the state-of-the-art, which translates to 6.6% end-to-end training throughput improvement. To facilitate future research of the "pre-train, and search" paradigm in ML for Systems, we open-source our code at <a class="link-external link-https" href="https://github.com/daochenzha/neuroshard" rel="external noopener nofollow">this https URL</a>

Column storage layout optimization method based on deep reinforcement learning

One Model Packs Thousands of Items with Recurrent Conditional Query Learning

Layout-Conscious Optimization: Beyond Hybrid Row-Column Storage Model

Wide Table Layout Optimization Based on Column Ordering and Duplication

Learning to Optimize Join Queries With Deep Reinforcement Learning

A Scalable Learned Index Scheme in Storage Systems

Mitigating Data Stalls in Deep Learning with Multi-times Data Loading Rule.

Accelerating Column Generation Algorithm Using Machine-Learning-Based Column Elimination

Column ordering for input/output optimization in tabular data

Qd-tree: Learning Data Layouts for Big Data Analytics

Adaptive Cache Management for Complex Storage Systems Using CNN-LSTM-Based Spatiotemporal Prediction

Workload time series prediction in storage systems: a deep learning based approach

Column sequence used for input/output optimization in form data

A Deep Reinforcement Learning Approach To The Optimization Of Data Center Task Scheduling

Query Join Order Optimization Method Based on Dynamic Double Deep Q-Network

COOOL: A Learning-To-Rank Approach for SQL Hint Recommendations

Machine Learning-Enhanced Ant Colony Optimization for Column Generation

Pre-train and Search: Efficient Embedding Table Sharding with Pre-trained Neural Cost Models

Cost-Based Lightweight Storage Automatic Decision for In-Database Machine Learning

Dynamic Data Layout Optimization with Worst-case Guarantees

Deep Learning Compiler Load Balancing Optimization Method for Model Training