Mixed-Precision Embedding Using a Cache

TL;DR

This work tackles the memory challenge of training large embedding tables in recommender systems by introducing a cache-augmented mixed-precision approach. Embedding tables are stored in low precision while a small FP32 cache holds frequently accessed rows, with fetch/update semantics and configurable policies. Empirical results show substantial memory reductions (3× on open datasets and 7× on industrial-scale models) with minimal to neutral accuracy loss, along with notable speedups (up to 16% end-to-end) and throughput gains at high cache hit rates. The method offers a practical path to scalable training of embedding-heavy models and suggests avenues for hardware and policy refinements."

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.

Figures (10)

• Figure 1: Deep Learning Recommendation Model (DLRM) DLRM19
• Figure 2: Diminishing return of accuracy recovered vs. cache sizes
• Figure 3: CDF of sorted row access counts vs. fraction of rows
• Figure 4: Direct-mapped LFU vs. LRU
• Figure 5: 32-Way Associative LFU vs. Direct-mapped LFU vs. LRU