万物皆可嵌入--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内存之间的数据拷贝，查找，更新等操作，里面涉及到很多工程实现问题。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述