Mixed-Precision Embedding Using a Cache

Abstract

In recommendation systems, practitioners observed that increase in the number of embedding tables and their sizes often leads to significant improvement in model performances. Given this and the business importance of these models to major internet companies, embedding tables for personalization tasks have grown to terabyte scale and continue to grow at a significant rate. Meanwhile, these large-scale models are often trained with GPUs where high-performance memory is a scarce resource, thus motivating numerous work on embedding table compression during training. We propose a novel change to embedding tables using a cache memory architecture, where the majority of rows in an embedding is trained in low precision, and the most frequently or recently accessed rows cached and trained in full precision. The proposed architectural change works in conjunction with standard precision reduction and computer arithmetic techniques such as quantization and stochastic rounding. For an open source deep learning recommendation model (DLRM) running with Criteo-Kaggle dataset, we achieve 3x memory reduction with INT8 precision embedding tables and full-precision cache whose size are 5% of the embedding tables, while maintaining accuracy. For an industrial scale model and dataset, we achieve even higher >7x memory reduction with INT4 precision and cache size 1% of embedding tables, while maintaining accuracy, and 16% end-to-end training speedup by reducing GPU-to-host data transfers.

References in corpus (16)

• Distributed Representations of Words and Phrases and their Compositionality
• The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks
• Deep Learning Recommendation Model for Personalization and Recommendation Systems
• Training and Inference with Integers in Deep Neural Networks
• Compression-aware Training of Deep Networks
• Training with Quantization Noise for Extreme Model Compression
• Hierarchy-based Image Embeddings for Semantic Image Retrieval
• Distributed Hierarchical GPU Parameter Server for Massive Scale Deep Learning Ads Systems
• Fonduer: Knowledge Base Construction from Richly Formatted Data
• Deep Learning Inference in Facebook Data Centers: Characterization, Performance Optimizations and Hardware Implications
• High-Accuracy Low-Precision Training
• A Study of BFLOAT16 for Deep Learning Training
• Tensorized Embedding Layers for Efficient Model Compression
• Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems
• Post-Training 4-bit Quantization on Embedding Tables
• On Periodic Functions as Regularizers for Quantization of Neural Networks

Cited by in corpus (3)

• Embedding Compression in Recommender Systems: A Survey
• Adaptive Low-Precision Training for Embeddings in Click-Through Rate Prediction
• ECRM: Efficient Fault Tolerance for Recommendation Model Training via Erasure Coding