腾讯太极机器学习平台|Light在广告粗排中的特征与Embedding优化

在《腾讯太极机器学习平台|Light在广告粗排中的数据下载与解析优化》一文里，我们介绍了在广告粗排场景中业务模型的特点，与我们在数据下载和解析方面所做的部分优化。

本文将继续介绍在粗排训练场景中，我们在特征Embedding方面所做的优化工作。

在大多数情况下，由于推荐类的任务对延时非常敏感，广告粗排场景的模型结构通常比较“矮胖”，有大量的特征输入然后做Embedding，NN部分的深度相对较低。

因此，特征样本内容的解析和特征的Embedding在整体耗时中的占比相对较高。图1是一种典型模型训练时的Timeline：

图1. 模型timeline (脱敏)

其中，字符串特征由Hash算子获得索引位置后，从变量中Gather出对应的权重。而整数特征首先通过AsString算子转为字符串表示，然后再做Hash。

如图1中，虽然Hash算子有GPU版本实现，但是其输入是一个string Tensor。每个元素的地址空间是离散，且有较远间隔的。GPU无法高效地处理这种数据。

因此，在baseline的实现中，使用CPU将数据组装成了便于GPU处理的格式，然后再拷贝到显存上进行计算。而CPU做数据组装的这一系列操作，占了整个Hash算子耗时的90%以上。

一方面，由于特征Hash在整体耗时中的占比十分显著，因此我们在优化模型训练性能的时候，对特征Hash做优化是不可避免的。

另一方面，当特征从原始的整数/字符串格式做Hash计算转化成了Embedding权重向量的index后，再从对应的权重变量中去取出对应的向量值，并按一定的规则collapse到对应的特征域当中。这部分功能在Tensorflow中仅支持CPU的版本，缺少GPU的实现。

我们实现了对应的GPU版本，相比CPU有更高的性能，且在GPU训练pipeline中避免了在Host和GPU Device之间来回的Tensor拷贝。

整数特征Hash优化

广告业务中，特征的表达形式有字符串和整型。在原Nabu训练框架中，对于整型特征的处理方式是先将整数转化成字符串（AsString，十进制），然后再做Hash。

由于GPU不支持String类型运算，因此使用CPU来进行整型到字符串类型的转化运算。一个string Tensor中，各个string元素的地址相差很大，散落在非连续空间中。

若直接将各个元素直接拷贝到GPU上，则会造成大量Host-to-Device拷贝，存在很大的拷贝启动overhead。

若先将各个元素先拷贝转移到连续的CPU主存空间，再一次性拷贝到显存上，则存在大量的glibc memcpy调用，占用大量CPU资源。而如果直接将各个碎片化的string从Host Memory拷贝到Device Memory，又会引入巨大的Overhead。

原方案采用了后者（图2），其中的AsString使用了Tensorflow中的实现，完成整型到string类型转换的是glibc中的vsnprintf。

在上一期中，我们提到了使用fmt中的to_string来替换vsnprintf的实现，对于CPU部分的int2string计算有显著性能提升。但是，这种方法仍不能避免string Tensor的离散地址空间问题。该问题是Hash过程中的主要性能瓶颈。

图2. 整数特征Hash Baseline方案。整数特征以Tensor的形式存在主存中，其Backend是eigen::Tensor，存储空间是一片连续的内存。整数特征做AsString运算，转变成具有离散内存空间的string Tensor。

string Tensor的各个元素拷贝到一片临时申请的连续内存中，并计算sizes和offsets来从连续内存中访问每个string元素。将连续的strings拷贝到GPU显存做Hash运算。

为了避免在Host内存中碎片化拷贝导致的性能问题，我们选择直接将整数特征拷贝到GPU显存，然后用GPU替代CPU完成AsString和Hash的计算，从而避免碎片化拷贝问题（如图3）。

图3. 优化思路: 整数特征直接拷贝到显存，做int2string转换，并计算Hash结果。

基于这一思路，我们考虑了几种实现的方案。

方案1 CPU计算与Host2Device拷贝并行

例如输入是int64类型的特征Tensor，我们首先将它拷贝到显存上。并同时计算每个特征的size（整数转化为字符串自后的长度，例如1234的长度为4，-13的长度为3）和offset（整数转化为字符串之后首地址的偏移量）。通过GPU的itoa（整型-字符串）转换，将这个int64的Tensor转换成连续显存上的字符串表达形式，如图4。

图4. 黑色方框表示在CPU主存上的数据，红色方框表示在GPU显存上的数据，方框相邻表示在同一片内存上。

这里的的HostToDevice拷贝采用异步拷贝，和CPU上的size和offset计算并行。

当图4中的操作完成后，即启动GPU kernel，将整型的features转化为string，并保存在保留段reserved上，其长度等于size之和。之后进行Hash操作时，即可通过offset和reserved段首地址来访问任意一个字符串。该方法采用CPU做sizes和offsets信息的统计，期间并行地做Host2Device拷贝。

该方案的优点是利用了CPU和GPU的异步关系。但通常情况下，在CPU上统计sizes和offsets信息的速度相比于cudaMemcpyAsync的要慢，因此可能造成GPU等待CPU计算结果的情况，限制了GPU的有效利用率。

方案2 全部使用GPU计算

方案一中，通过offset和reserved首地址访问任意字符串，需要在CPU上先算好每个字符串的offset，使用了部分CPU资源，在广告粗排CPU bound的训练中仍可能受到CPU资源不足的约束。对于该问题问题，方案二将全部的size（十进制位数）计算和itoa都放在GPU上进行（图5）。

图5. 整数特征直接拷贝到显存，并在GPU上计算每个整数转成string后的size。

2.1余数法

由于int64表示的最大和最小数字部分位数不超过19，如果是负数，包括负号的长度不超过20位。因此可以在一个定长的buffer来在GPU的local hierachy内存储字符串（不支持动态分配显存）。以 -123 为例（图6）：

图6. 采用余数法做atoi的过程

因为64位整数的长度有限，在thread local显存足够应对所有的64位整数长度。原数求llabs(64位绝对值)之后，依次求对Base=10的模，结果数得到字符后，从尾到头的放到thread local buffer中。

在计算的过程中，可以获取字符串的长度。如-123的长度是4。在buffer中完成itoa转换以后，可以通过buffer首地址 + buffersize - size(字符串长度)，来访问这个string的结果。

2.2先计算size,后做转换

如果采用对数不等式上下界扩展，可将非负整数的位数搜索范围为降至O(1)。即可以O(1)获得itoa转化后的size。从而可以从thread local buffer的头部直接访问字符串，避免循环地计算++size。

可证明：对于二进制数字长度为n的非负整数，其十进制数位数属于元素数量不大于2的集合。因此我们可以从有限的2种情况中直接得到某个整数的位数（含符号位）。

证明: 对于二进制数字长度为n的非负整数，其十进制数位数属于元素数量不大于2的集合。

对于任意非负整数k，假设其二进制位数为n，则其数值范围可以写为：

其中，k的十进制位数满足：

lg表示以10为底的对数。化简后可以将n提出：

由于

恒成立，因此可将不等式上下限扩展：

左右两式相减，得到：

二者差值比1小，说明k的十进制上下限，在数轴上处于同一个整数的两侧，得证。

进一步地，我们可以推出:

因此，只需知道S1和S2中任意一种情况是否满足条件，即得整数的位数。获取size后，依次将abs数值求模去对'0'符号的偏移值填到thread local buffer数值位首部即可。如果是负数，则数值位从第二位开始数。

2.1和2.2两种方法在不同的硬件条件下性能各有优劣。考虑到代码的可读性，最终我们选择了方案2.1来上线。

优化效果

V100卡千次运行耗时Baseline与优化后对比:

字符串特征Hash优化

广告粗排场景中大量使用整形数和字符串来记录和表示用户与物品特征。特征以序列化的形式存在gzip压缩的TFRecord文件中。

我们在训练时需要将文件解压读入，并依次地将这些特征反序列化出来，然后做Hash，得到对应特征权重的索引，最后用索引查询变量中的对应位置权重数值，进行训练和推理计算。

一方面，Tensorflow在反序列化TFRecord格式的example时，采用Tensor来记录各个字段的内容。例如int64特征用int64 Tensor来记录，string特征用string Tensor来记录。其中的string Tensor由于采用非连续内存来记录各个元素值，所以需要将TFRecord的对应string字段内容，一一拷贝到string Tensor的各个元素中，存在大量细碎的短拷贝。（如图7）

图7. 特征从样本解析到Tensor的过程

另一方面，在用GPU做特征Hash的过程中，由于string特征离散内存的string Tensor，需要将离散的string Tensor转换成连续内存表示，并且在连续内存中记录每条string元素的长度，从而又需要做大量的短拷贝和赋值操作。在整数特征Hash中，我们可以用GPU来实现atoi的过程，从而避免离散拷贝。但是，string特征本来的地址已经是离散的，无法使用和整数特征相同的方式来做优化。

优化方案

在TFRecord文件中，使用了一定的格式来存字符串特征，数据用Protobuf编码（如图8）。可以使用protobuf::io::CodedInputStream来从样本中流式地读取特征。特征在样本的结构中包含一个32bit的头部，它记录了该特征在样本中所占的字节数。

之后是，数个特征的元素，每个特征的元素之间有一个16bit的分隔符。每个特征元素前32bit是该元素字符串部分的长度。工作线程在解析样本中的字符串特征时，首先从样本中将该特征中的元素读到std::string中，然后再放到对应Tensor的元素。

图8. TFRecord字符串特征格式

为了避免string Tensor中不同元素的地址空间离散问题，我们自定义了一种Tensor类型。

当工作线程从TFRecord中读出string后，拷贝到我们自定义的Tensor元素中。自定义的string Tensor具有连续的内存空间，从而在做Hash的时候，可以直接拷贝到GPU显存，不再需要做大量细碎拷贝。

我们在Tensorflow framework中，利用Variant来构造自己的Tensor元素：

// 自定义的string类型
struct ConsecString {
  DynamicBuffer* buf;  // 实际存string的地方
  explicit ConsecString();
  explicit ConsecString(DynamicBuffer* buf);
  void Encode(VariantTensorData* data) const;
  bool Decode(const VariantTensorData& data);
  string TypeName() const {return "ConsecString";}
  string DebugString() const {return "DebugString: ConsecString";}
  bool is_empty() {}
};

// ConsecString仅仅是一个字符串的handle在OP之间作为Tensor的元素传递，
// 这样可以避免大块数据的拷贝。使用Variant元素的get<T>()方法可以获取到
// 自定义的ConsecString。
ConsecString* dummy_ele = const_cast<ConsecString*>(
    input_ptr->get<ConsecString>());

使用Variant来构造自定义Tensor类型的好处是能够以custom OP的方式来实现，可以避免对Tensorflow框架内部的改动。其中ConsecString所持有的DynamicBuffer是实际存字符串的地方。

// 自定义的Buffer
struct DynamicBuffer {
  // Since a batch is processed in many minibatches, so use a vector to
  // store all buffer in minibatches which avoids merging cost.
  std::vector<string> buf_list;
  std::vector<std::vector<int32>> sizes_list;
  void merge() {}
  void reset() {}
};    

DynamicBuffer是一个持久化的结构，它不随OP计算的结束而销毁。每一个字符串特征Tensor都有一个对应的DynamicBuffer。当开始读TFRecord时，如果对应的字符串特征dense/sparse key没有对应的DynamicBuffer，那么就创建一个新的。之后每次解析样本都使用这个buffer。当Tensor在OP之间传播时，只传递ConsecString。

为了能够并行地解析样本，我们将一个batch拆成多个minibatch，使用多个线程来解析数据。每个线程处理一个minibatch。buf_list中的多个string和sizes_list中的数组分别记录对应minibatch的结果。在所有线程处理完成之后，再将多个minibatch合并。

这样就可以在做Hash的时候，直接将连续内存中的字符串拷贝到GPU显存做进一步运算，而不用做大量细碎拷贝。

优化效果

我们使用广告粗排litecvr中的样本对该优化方案在单个特征上做了性能测试，结果如下：

平均字符串长度: 17。仅包含运算functor的耗时。

关于我们

本期编辑：carol/suzie

技术分享：关注微信公众号 【鹅厂架构师】【腾讯架构师】

【往期精彩】

腾讯太极机器学习平台|Light在广告粗排中的数据下载与解析优化

腾讯太极机器学习平台|大规模训练加速框架Light 在广告粗排场景的落地

腾讯太极机器学习平台荣获“云原生应用优秀案例”大奖