万物皆可嵌入--embedding在GPU中的实现

发布时间:2022-03-21 07:00

摘要

 

Embedding技术自从谷歌推出word2vec的工作后得到迅速的应用,典型应用之一是在广告推荐场景中,从word2vec演进到item2vec,embedding技术的出现也使深度学习进入广告推荐的场景成为可能。广告推荐模型动辄几十GB甚至TB的模型大小,高效地进行embedding table的操作成为影响性能的关键,本文将介绍embedding的相关知识背景,并以Nvidia推出的面向广告场景的HugeCTR框架为基础,介绍在GPU中关于embedding操作的实现。

 

Embedding介绍

 

一般而言,嵌入表(Embedding table)是指将一系列不同的元素,比如单词、短语或者句子,使用不同的方法提取到的特征向量矩阵。从数学表达上看,embedding是一种空间映射,将高维的向量表达映射到低维的向量表达。以单词的embedding表达为例,将每个单词使用one-hot(向量中只有一个元素非0)的方式表达,那么就是将单词的one-hot空间表达映射到新的特征向量空间。假设有一个句子“The cat sat on the mat”, 一共5个单词,对于cat的表达选择使用向量[1, 0, 0, 0, 0],经过embedding的操作变为图1中右图的维度为4的嵌入表,cat变为向量[1.2,-0.1, 4.3, 3.2],称之为特征向量,转换的过程即为训练的过程。总而言之,embedding就是用一个低维的向量(数学的方式)表示一个物品,可以是商品、人或者单词,所以从表达范畴上讲万物皆可采用embedding表示[1]。

 

    图1 one-hot encoder到嵌入表的转换[2]

 

显而易见,使用one-hot的编码带来的问题之一是数据量巨大且稀疏,如图1,使用embedding的方式除了可以带来数据量的压缩外,还可以引入特征向量之间的关系表达。以对城市和国家的表达为例,看到北京自然会想到是中国的首都,二者的关系可以体现在向量空间距离很近,如图2所示。

 

图2 国家和首都在嵌入表中表达的空间距离示意图[3]

 

正是得益于embedding表达中潜在的关系,使得使用机器学习的方式来探究数据内部的关系成为可能。在具体应用中,通过embedding table的查找(本质是全连接的操作),使得原本高维稀疏的矩阵变为低维稠密矩阵,然后进入到深度学习的计算中,典型的应用就是广告推荐。

 

广告推荐中的embedding

 

虽然嵌入层将高维稀疏矩阵转换为低维稠密矩阵,降低了参数的数量,然而广告推荐场景面对的是千亿级别的商品,十亿级别的用户,因此嵌入表的规模越来越大,达到TB不足为奇。因此,在广告推荐模型中对嵌入层的优化成为影响模型质量的重要因素。规模巨大的嵌入表导致广告推荐模型的计算变为访存密集型任务,如何充分利用GPU上的内存容量减少数据搬移,同时又发挥GPU的通信高带宽的性能优势成为研究的目标。

 

图3 one-hot编码和embedding的对比

 

在广告推荐模型中,采用和word2vec类似的嵌入表示方法,而用户和商品很多,如果使用one-hot编码会导致嵌入表极具膨胀,以图3中亿级别的用户为例,使用embedding编码后,数据量可以从10的9次方压缩到10的3次方。图4表示的是用户和商品ID及对应的嵌入表查询过程。图中的特征向量的维度为4,但在实际应用中,不同特征的特征向量不一定相同。

 

图4 特征向量维度为4的嵌入表[3]

 

在广告推荐场景中,经常使用feature (categorical feature),categorical data,feature field(slot)术语,本文将其翻译为特征,特征数据及特征域。

 

以电影推荐场景为例,电影名称,用户性别,用户年龄,手机等等均可以看作是一类特征。特征数据是具体的特征描述,比如手机分为苹果,华为和小米等。特征域指的是将相关的特征聚合在一起,作为一个域,也就是说在一个feature field会包含几类特征。比如在经典的YouTube视频推荐模型中,关于language的分析,可以看成是一个feature field,里面聚合了user language和video language。对于一个feature filed中的特征向量,会进行combiner的操作,可以是求平均或者是求和等。

 

  图5 YouTube视频推荐模型[4]

 

稀疏Embedding的分配方式

 

关于embedding的分配,一般是根据embedding 的key均匀地分配到CPU或者GPU中,比如根据embedding的key对GPU的模值,将embedding分配到各个GPU。但在HugeCTR中,关于embedding的分配以slot(slot等同于前文的feature field)为单位(slot是一等公民),每个slot的特征可以单独embedding,然后再合并为一个嵌入向量,这样的收益是确保同一个特征域的特征向量在一台GPU上,可以较为高效地进行combiner的操作,如后文的图8的本地化模式所示。

 

HugeCTR是Nvidia推出的针对点击率(Click-through Rate)场景的分布式推荐框架,支持多GPU及多GPU节点的训练模式,底层则采用GPU的优化库实现,比如NCCL,CuDF[5]等。HugeCTR也可以作为TensorFlow的插件使用,大致的自顶向下的层次如图6所示。

 

图6 HugeCTR的层次图

 

稀疏embedding的分配分为分布式(distributed)模式和本地化(localized)模式。在分布式模式下,如图7所示,一个slot可以位于多个GPU上,比如slot0中的特征可以位于GPU0和GPU1上,slot1的特征也可以位于GPU0和GPU1上。分布式的方式适合slot的大小超过了GPU的内存。分配到GPU的计算规则为: feature_key % GPU_NUM,和上文提到的按照embedding的key进行均匀分配相同。

 

图7 分布式模式的稀疏embedding的分布

 

本地模式指的是将slot中的特征完整地放在一个GPU内存中,不会跨GPU存放,比如slot-0放在GPU0,slot-1放在GPU1,slot-2放在GPU0等,如图8所示。使用本地化模式的前提是GPU内存能够完整存放下一个slot的大小,相比于distributed模式,优先推荐使用local模式,具有更好的性能,同一个slot内的特征处理可以在一个GPU上完成。分配到GPU的计算规则为: slot_id % GPU_NUM。

 

图8 本地化模式的稀疏embedding的分布

 

图9展示了slot中特征向量的combiner的过程。一个sample中有7个key,分布在两个slot中,第一个slot有4个key,第二个slot有3个key,第三个slot没有key。在查找的过程中,第一个slot中的4个key分别找到对应的特征向量进行sum操作得到新的特征向量v1,第二个slot找到3个key对应的特征向量进行sum得到新的特征向量v2。组后把v1和v2 concat。注意在hugeCTR中不允许不同的slot中出现相同的key,在hugeCTR中支持的combiner操作为求和和取平均。

 

图9 slot中特征向量的combiner操作

 

 

Embedding表的查询

 

 

在HugeCTR的实现中,为了提高embedding table的查询及插入性能,采用了二级hash的方式进行构建,如图10所示。

1.第一级hash提供的是逻辑hash,背后的实现根据hash计算得到,而不是真的查找表,目的在于节省存储空间。Key相当于稀疏输入(比如one-hot格式输入),经过hash计算后得到value,value表示的是第二级hash table中key的行偏移。

2.第二级是真正存储在GPU内存中hash table,key为嵌入向量的行偏移,value为最终的嵌入向量。图中hash_value_index_tensors_存储的是GPU数量个hash_value_index,hash_value_index中的每个元素表示在当前GPU中嵌入向量的行偏移。整个hash_value_index_tensors_的大小是每个slot的特征数累加和与batch_size大小的乘积,如公式1所示(nnz_per_slot表示每个slot中的特征数量),记录了所有的低维稠密矩阵offset,如此以来,hash table的大小就被固定了(这种设计也限制了支持更大的参数场景)。同理hash_table_value_tensors_表示每个GPU中存放的嵌入向量的值,hash_table_value_tensors_的大小是每个GPU上存放的最大特征数量和特征向量大小的乘积,如公式2所示(max_vocabulary_size_per_gpu表示每个GPU的最大特征数量)。

图10 的上图表示为高维稀疏矩阵经过embedding优化后得到的低维稠密矩阵,储存在GPU内存中。以苹果的稀疏输入为89举例,苹果特征的slot分布在GPU0上,通过一级hash计算得到row_index为2,从GPU0上的hash table中找到row_index=2对应的行偏移为0,如图10的下图所示,然后从hash_table_value[0]的row_offset=0的位置获得苹果的嵌入表向量9898。

 

图10 hash表查询过程示意图

 

第二级hash的实现是完全放在GPU内存中,包装在cuDF(GPU DataFrame)的库中,依赖于concurrent_unordered_map类,是一个GPU加速的哈希实现,支持并发的insert,但不支持并发的insert和get,这是和hugeCTR仅支持同步训练有关系,不会同时进行pull和push操作。

 

图10是简化的示意图,在实际代码实现中引入了bucket的概念,如图11所示,在经过对hash_table_size取模后,得到一个hash bucket,bucket中存有key和value,其中value对应图10中的hash_value_index,根据hash_value_index所对应的行偏移,在hash_table_value(embedding table)中找到嵌入特征。

 

图11 带有bucket的hash查找过程

 

构建embedding table

 

为了压缩稀疏矩阵的存储格式,hugeCTR中采用了CSR(Compressed Sparse Row)的稀疏压缩方法,如图12中的左图所示,针对一个稀疏矩阵,仅存储矩阵的非0数值,以及数值对应的行偏移和列偏移。

 

图12 CSR的数据表示和在slot中的表示

 

图12的右图中,一共有3个slot,每个slot中有不同数量的特征数,比如slot-0中有3个特征,0,1,2表示对应的embedding key。现有3个样本,分别以row-0,row-1和row-2表示。以row-0为例,它的column index是[1, 3, 4, 8],即slot中的key,将row的数据以slot为行进行排布,然后根据CSR的表示方法,可以得到对应的row_index为[0, 1, 3, 4],可以看成是稀疏输入。

 

图13表示构建好的一个embedding table,同样以row-0为例,row-index是前文提到的hash_table_value_index,sparse tensor表示的是稀疏输入,对应前文提到的hash table key。以图12的row-0的row_index[0, 1, 3, 4]作为稀疏输入,分别在不同的slot中找到row-offset(此处的row-offset指的是嵌入特征在内存中的行偏移),然后根据行偏移从GPU内存中找到对应的嵌入特征向量。

 

图13 embedding table示意图

 

结论与思考

 

本文介绍了hugeCTR中关于embedding table在GPU中实现方式,分析了如何利用二级hash的方式实现嵌入特征的查询以及采用稀疏压缩的方式构建嵌入表。由于篇幅有限,文中略去了实际的代码实现,对部分细节进行了抽象。

 

参考文献

 

[1] https://zhuanlan.zhihu.com/p/53194407

[2] https://www.tensorflow.org/text/guide/word_embeddings

[3] https://developer.nvidia.com/blog/using-neural-networks-for-your-recommender-system

[4] Paul Covington, Jay Adams, and Emre Sargin. 2016. Deep neural networks for YouTube recommendations. In Proceedings of the Recsys. 191—198

[5] https://github.com/NVIDIA-Merlin/HugeCTR

[6] https://github.com/rapidsai/cudf

上一个: 用于推荐系统的近存处理器设计

下一个: 比你更懂你- 神经网络与推荐系统的深度结合

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多