推荐系统训练期间的细粒度嵌入维度优化

现代深度学习推荐模型 (DLRM) 中巨大的嵌入表在训练和推理过程中需要过大的内存。本文提出了 FIITED，一种通过细粒度训练中嵌入维度剪枝自动减少内存占用的系统。利用嵌入向量重要性不均等这一关键洞察，FIITED 在模型训练期间自适应地调整每个独立嵌入向量的维度，为更重要的嵌入分配更大的维度，同时适应数据的动态变化。我们将具有更高频率和梯度的嵌入维度优先视为更重要。为了在模型训练期间高效地剪枝嵌入及其维度，我们提出了一种基于虚拟哈希物理索引哈希表的嵌入存储系统。在两个工业模型和数月的真实数据集上的实验表明，FIITED 能在保持模型质量的同时将 DLRM 嵌入大小减少 65% 以上，优于最先进的训练中嵌入剪枝方法。在公共数据集上，FIITED 能将嵌入表的大小减少 2.1 倍至 800 倍，且精度下降可忽略不计，同时提高了模型吞吐量。