[论文解读] Mixed-Precision Embedding Using a Cache

Jie Yang, Jianyu Huang+3 more|arXiv (Cornell University)|Oct 21, 2020

Topic Modeling参考文献 32被引用 15

一句话总结

本文提出一种混合精度嵌入训练方法，通过使用小规模高精度缓存来存储频繁或最近访问的嵌入行，而大部分嵌入行则以低精度（INT8 或 INT4）进行训练。该方法通过最小化 GPU 到主机的数据传输并利用高效的缓存替换策略，在保持完整准确性的前提下，实现了工业级模型高达 7 倍的内存减少和 16% 的训练速度提升。

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.

研究动机与目标

• 解决深度学习推荐系统中大规模嵌入表日益增长的内存占用问题。
• 在大规模训练工作负载中减少 GPU 内存使用量，同时不牺牲模型准确度。
• 探索将低精度训练与高精度缓存相结合在嵌入表中的有效性。
• 评估缓存大小、替换策略和精度级别对模型准确度和训练效率的影响。
• 通过更智能的内存访问模式减少 PCIe 带宽使用，从而提升训练速度。

提出的方法

• 该方法使用小规模高精度缓存（嵌入表大小的 1–5%）来存储最频繁或最近访问的嵌入行。
• 大部分嵌入行以低精度（INT8 或 INT4）进行训练，而缓存则存储活跃行的全精度梯度和参数。
• 采用 LRU 和 LFU 等缓存替换策略，根据访问频率或最近性决定哪些行被提升至高精度。
• 在低精度训练期间应用随机舍入，以减轻系统性舍入偏差并保持准确度。
• 系统与现有深度学习框架集成，对大表使用统一内存（UVM），并通过缓存减少 PCIe 流量。
• 该方法在开源（DLRM 与 Criteo-Kaggle）和工业级模型上均进行了评估，测量了内存节省、准确度和训练速度。

实验结果

研究问题

• RQ1 当大部分嵌入行以低精度（INT8/INT4）训练时，小规模高精度缓存能否有效维持模型准确度？
• RQ2 不同的缓存替换策略（LRU、LFU）在低精度嵌入训练中如何影响模型准确度和缓存命中率？
• RQ3 相对于嵌入表大小，最优缓存大小应如何设定，以在内存减少与准确度恢复之间取得平衡？
• RQ4 在存在或不存在高精度缓存的情况下，随机舍入与向最近舍入相比，在低精度嵌入训练中的表现如何？
• RQ5 该缓存机制在多大程度上能减少 GPU 到主机的数据传输并提升端到端训练速度？

主要发现

• 在使用 INT4 精度和 1% 高精度缓存的情况下，该方法在工业级模型上实现了高达 7 倍的内存减少，同时保持准确度中性。
• 在使用 INT8 精度和 5% 高精度缓存的情况下，该方法在开源 DLRM 模型上实现了 3 倍内存减少，并保持了完整准确度。
• 随机舍入在低精度训练中始终优于向最近舍入，尤其是在结合高精度缓存时，能有效减轻系统性舍入偏差。
• 缓存命中率与训练速度密切相关：高命中率可饱和 GPU 内存带宽，而低命中率仍优于统一内存（UVM）访问性能。
• 通过有效缓存减少 GPU 到主机的数据传输，在工业模型上实现了 16% 的端到端训练速度提升。
• 随着缓存大小增加，准确度提升呈现边际递减趋势，表明小规模缓存（1–5%）已足以实现接近最优的性能。