文章插图
导致这种结果的核心原因是：Worker单步训练需要和所有的PS通信同步完成，每增加1个PS要增加N条通信链路，这大大增加了链路延迟（如下图5所示） 。而一次训练要执行上百万、上千万步训练 。最终导致链路延迟超过了加PS算力并发的收益 。


文章插图
而对于这个系统，优化的核心难点在于：如何在有限的PS实例下，进行分布式计算的优化 。
3 优化实践
3.1 大规模稀疏参数介绍
对于推荐系统模型，绝大多数参数都是稀疏参数，而对稀疏参数来说有一个非常重要的操作是Embedding，这个操作通常也是负载最重的，也是后续优化的重点 。由于我们对稀疏参数进行了重新定义，后续的优化也基于此之上，所以我们先介绍一下这部分的工作 。
在原生的TensorFlow中构建Embedding模块，用户需要首先创建一个足够装得下所有稀疏参数的Variable，然后在这个Variable上进行Embedding的学习 。然而，使用Variable来进行Embedding训练存在很多弊端：
Variable的大小必须提前设定好，对于百亿千亿的场景，该设定会带来巨大的空间浪费；训练速度慢，无法针对稀疏模型进行定制优化 。
我们首先解决了有无的问题，使用HashTable来替代Variable，将稀疏特征ID作为Key，Embedding向量作为Value 。相比原生使用Variable进行Embedding的方式，具备以下的优势：
HashTable的大小可以在训练过程中自动伸缩，避免了开辟冗余的存储空间，同时用户无需关注申请大小，从而降低了使用成本 。针对HashTable方案实施了一系列定制优化，训练速度相比Variable有了很大的提高，可以进行千亿规模模型的训练，扩展性较好 。得益于稀疏参数的动态伸缩，我们在此基础上支持了Online Learning 。API设计上保持与社区版本兼容，在使用上几乎与原生Variable一致，对接成本极低 。
简化版的基于PS架构的实现示意如下图6所示：


文章插图
核心流程大致可以分为以下几步：
稀疏特征ID（通常我们会提前完成统一编码的工作）进入Embedding模块，借助TensorFlow搭建的Send-Recv机制，这些稀疏特征ID被拉取到PS端，PS端上的Lookup等算子会实际从底层HashTable中查询并组装Embedding向量 。上述Embedding向量被Worker拉回进行后续训练，并通过反向传播计算出这部分参数的梯度，这些梯度进一步被位于PS端的优化器拉回 。PS端的优化器首先调用Find算子，从HashTable获取到梯度对应的原始稀疏参数向量和相应的优化器参数，最终通过优化算法，完成对Embedding向量和优化器参数的更新计算，再通过Insert算子插入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成为性能瓶颈 。


文章插图