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

摘要

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的方式实现嵌入特征的查询以及采用稀疏压缩的方式构建嵌入表。由于篇幅有限，文中略去了实际的代码实现，对部分细节进行了抽象。当然，对于embedding table的操作，采用hash的实现仅仅是一种处理方式，也是技术栈中的一个小环节，在具体的实现中，可以自行设计策略，关键的问题是如何保证embedding table的操作，这里面包括从存储服务器到主机内存（CPU内存）再到GPU内存之间的数据拷贝，查找，更新等操作，里面涉及到很多工程实现问题。

参考文献

[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021

[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] PaulCovington, Jay Adams, and Emre Sargin. 2016. Deep neural networks for YouTuberecommendations. In Proceedings of the Recsys. 191—198

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

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

 往期推荐

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们