面向大规模推荐模型推理的HBM-DRAM嵌入向量存储系统

摘要

存储大规模推荐模型的嵌入向量特征需要大量的内存(DRAM),嵌入向量的高频查询和跨计算层的传输已成为推荐模型推理的性能瓶颈.GPU的HBM(High Bandwidth Memory)具备TB/s级高带宽，采用HBM来存储、访问嵌入向量可以显著提高推荐模型推理的性能，但HBM昂贵且容量有限，无法存放所有的嵌入向量.本文结合推荐场景中存在数据倾斜的特性，设计了一种混合存储系统，使用HBM作为一级存储存放热点嵌入向量加速推理，使用DRAM作为二级存储降低推理成本，实现了面向大规模推荐模型推理的嵌入向量存储系统.实验结果表明，相较于常见的纯DRAM方案，嵌入向量部分的吞吐率提升了14倍；相较于其他使用GPU来存储嵌入向量的实现，采用本系统实现的嵌入向量表，嵌入向量部分的吞吐率有3.8倍的提升.

关键词

推荐系统 / 高性能计算 / AI大模型 / GPU加速

基金资助

中国科学院战略性先导科技专项(A类)项目(CCS54WHZ1123116926096472)资助