TensorFlow在美团推荐系统中的分布式训练优化实践( 八 ) _云知道

那我们是否可以提前执行稀疏参数的Embedding ，来做到通信和计算的Overlap ，隐藏掉这部分时间呢？从系统实现上肯定是可行的，但从算法上讲，这样做会引入参数Staleness的问题，可能会导致模型精度受到影响。但在实际的生产场景中，大规模异步训练时本身就会带来几十到几百个步的滞后性问题。经过我们测试，提前获取一两步的稀疏参数，模型精度并未受到影响。
在具体实现上，我们把整个计算图拆分为Embedding Graph（EG）和Main Graph（MG）两张子图，两者异步独立执行，做到拆分流程的Overlap（整个拆分过程，可以做到对用户透明）。EG主要覆盖从样本中抽取Embedding Key ，查询组装Embedding向量， Embedding向量更新等环节。MG主要包含稠密部分子网络计算、梯度计算、稠密参数部分更新等环节。

图15 Embedding流水线模块交互关系

两张子图的交互关系为：EG向MG传递Embeding向量（从MG的视角看，是从一个稠密Variable读取数值）；MG向EG传递Embedding参数对应的梯度。上述两个过程的表达都是TensorFlow的计算图，我们利用两个线程，两个Session并发的执行两张计算图，使得两个阶段Overlap起来，以此到达了更大的训练吞吐。

图16 Embedding流水线架构流程图

上图是Embedding流水线的架构流程图。直观来看分为左侧的样本分发模块，顶部的跨Session数据交换模块，以及自动图切分得到的Embedding Graph和Main Graph ，蓝色的圆圈代表新增算子，橙色箭头代表EG重点流程，蓝色箭头代表MG重点流程，红色箭头代表样本数据重点流程。
以对用户透明的形式引入了一层名为Pipeline Dataset的抽象层，这一层的产生是为了满足EG/MG两张计算图以不同节奏运行的需求，支持自定义配置。另外，为了使得整个流水线中的数据做到彼此的配套，这里还会负责进行一个全局Batch ID的生成及注册工作。Pipeline Dataset对外暴露两种Iterator ，一个供EG使用，一个供MG使用。Pipeline Dataset底部共享TensorFlow原生的各层Dataset 。顶部的ExchangeManager是一个静态的，跨Session的数据交换媒介，对外暴露数据注册和数据拉取的能力。抽象这个模块的原因是， EG和MG原本归属于一张计算图，因为流水线的原因拆解为拆为两张图，这样我们需要建立一种跨Session的数据交换机制，并准确进行配套。它内部以全局Batch ID做Key ，后面管理了样本数据、Embeding向量、Embedding梯度、Unique后的Index等数据，并负责这些数据的生命周期管理。中间的Embedding Graph由独立的TF Session运行于一个独立的线程中，通过a算子获得样本数据后，进行特征ID的抽取等动作，并进行基于HashTable方法的稀疏参数查询，查询结果通过c算子放置到ExchangeManager中。EG中还包含用于反向更新的f算子，它会从ExchangeManager中获取Embedding梯度和与其配套的前向参数，然后执行梯度更新参数逻辑。下面的Main Graph负责实际稠密子网络的计算，我们继承并实现一种可训练的EmbeddingVariable ，它的构建过程（d算子）会从ExchangeManager查找与自己配套的Embedding向量封装成EmbeddingVariable ，给稠密子网络。此外，在EmbeddingVariable注册的反向方法中，我们添加了e算子使得Embedding梯度得以添加到ExchangeManager中，供EG中的f算子消费。