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.

该研究介绍了一种新的嵌入式表（embedding table）缓存内存架构及其优化，该架构主要针对推荐系统中的大规模模型训练，通过低精度训练大部分嵌入式表行数据，实现与标准训练精度下相同准确率的情况下，INT8 位精度嵌入式表与缓存大小仅为嵌入式表大小的 5%，最终实现达到 3 倍的减存储效果，并在 GPU 与主机间的数据传输方面达到了 16% 的训练加速。