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

我们首先解决了有无的问题，使用HashTable来替代Variable ，将稀疏特征ID作为Key ， Embedding向量作为Value 。相比原生使用Variable进行Embedding的方式，具备以下的优势：
HashTable的大小可以在训练过程中自动伸缩，避免了开辟冗余的存储空间，同时用户无需关注申请大小，从而降低了使用成本。针对HashTable方案实施了一系列定制优化，训练速度相比Variable有了很大的提高，可以进行千亿规模模型的训练，扩展性较好。得益于稀疏参数的动态伸缩，我们在此基础上支持了Online Learning 。API设计上保持与社区版本兼容，在使用上几乎与原生Variable一致，对接成本极低。
简化版的基于PS架构的实现示意如下图6所示：

图6 支撑大规模稀疏参数的HashTable方案

核心流程大致可以分为以下几步：
稀疏特征ID（通常我们会提前完成统一编码的工作）进入Embedding模块，借助TensorFlow搭建的Send-Recv机制，这些稀疏特征ID被拉取到PS端， PS端上的Lookup等算子会实际从底层HashTable中查询并组装Embedding向量。上述Embedding向量被Worker拉回进行后续训练，并通过反向传播计算出这部分参数的梯度，这些梯度进一步被位于PS端的优化器拉回。PS端的优化器首先调用Find算子，从HashTable获取到梯度对应的原始稀疏参数向量和相应的优化器参数，最终通过优化算法，完成对Embedding向量和优化器参数的更新计算，再通过ｉｎｓｅｒｔ算子插入HashTable中。
3.2 分布式负载均衡优化
这部分优化，是分布式计算的经典优化方向。PS架构是一个典型的“水桶模型” ，为了完成一步训练， Worker端需要和所有PS完成交互，因此PS之间的平衡就显得非常重要。但是在实践中，我们发现多个PS的耗时并不均衡，其中的原因，既包括TensorFlow PS架构简单的切图逻辑（Round-Robin）带来的负载不均衡，也有异构机器导致的不均衡。
对于推荐模型来说，我们的主要优化策略是，把所有稀疏参数和大的稠密参数自动、均匀的切分到每个PS上，可以解决大多数这类问题。而在实践过程中，我们也发现一个比较难排查的问题：原生Adam优化器，实现导致PS负载不均衡。下面会详细介绍一下。
在Adam优化器中，它的参数优化过程需要两个β参与计算，在原生TensorFlow的实现中，这两个β是所有需要此优化器进行优化的Variabl（或HashTable）所共享的，并且会与第一个Variable（名字字典序）落在同一个PS上面，这会带来一个问题：每个优化器只拥有一个β
和一个β
，且仅位于某个PS上。因此，在参数优化的过程中，该PS会承受远高于其他PS的请求，从而导致该PS成为性能瓶颈。