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

通过上面的设计，我们就搭建起了一套可控的EG/MG并发流水线训练模式。总体来看， Embedding流水线训练模式的收益来源有：
经过我们对多个业务模型的Profiling分析发现， EG和MG在时间的比例上在3:7或4:6的左右，通过将这两个阶段并行起来，可以有效的隐藏Embedding阶段，使得MG网络计算部分几乎总是可以立即开始，大大加速了整体模型的训练吞吐。TensorFlow引擎中当使用多个优化器（稀疏与非稀疏）的时候，会出现重复构建反向计算图的问题，一定程度增加了额外计算，通过两张子图的拆分，恰好避免了这个问题。在实施过程中的ExchangeManager不仅负责了Embedding参数和梯度的交换，还承担了元数据复用管理的职责。例如Unique等算子的结果保存，进一步降低了重复计算。
另外，在API设计上，我们做到了对用户透明，仅需一行代码即可开启Embedding流水线功能，对用户隐藏了EG/MG的切割过程。目前，在美团某业务训练中， Embedding流水线功能在CPU PS架构下可以带来20%~60%的性能提升（而且Worker并发规模越大，性能越好）。
3.5 单实例PS并发优化
经过2.2章节的分析可知，我们不能通过持续扩PS来提升分布式任务的吞吐，单实例PS的并发优化，也是非常重要的优化方向。我们主要的优化工作如下。
3.5.1 高性能的HashTable
PS架构下，大规模稀疏模型训练对于HashTable的并发读写要求很高，因为每个PS都要承担成百乃至上千个Worker的Embedding压力，这里我们综合速度和稳定性考虑，选用了tbb::concurrent_hash_map[10]作为底层HashTable表实现，并将其包装成一个新的TBBConcurrentHashTable算子。经过测试，在千亿规模下TBBConcurrentHashTable比原生MutableDenseHashTable训练速度上快了3倍。
3.5.2 HashTable BucketPool
对于大规模稀疏模型训练来说， Embedding HashTable会面对大量的并发操作，通过Profiling我们发现，频繁动态的内存申请会带来了较大性能开销（即使TensorFlow的Tensor有专门的内存分配器）。我们基于内存池化的思路优化了HashTable的内存管理。
我们在HashTable初始化时，会先为Key和Value分别创造两个BucketPool ，每个池子都会先Malloc较大一块内存备用，考虑到可能会有对HashTable进行中的Key和Value进行Remove的场景（如Online Learning训练时），需要对从HashTable中删除的Key和Value所使用的内存进行回收，因此每个BucketPool还有一个ReuseQueue来负责维护回收的内存。每次向内部的哈希表数据结构中ｉｎｓｅｒｔ Key和Value的时候， Key和Value内存和释放分配都进行池化管理。用这种方式降低了大规模稀疏训练中遇到稀疏内存分配开销，整体端到端训练性能提升了5%左右。

图17 HashTable内存优化

3.6 单位算力吞吐优化
经过2.2章节的分析， Worker的计算压力也非常大，如果不优化Worker ，同时要保持吞吐，需要横向扩展更多的Worker ，给PS带来更大的压力。而对于用户来说，如果能在有限的计算资源下带来性能提升，对业务价值更高。我们通过CAT统计出了一些高频算子，并进行了专项优化。这里选取Unique&DynamicPartition算子融合案例进行分享。