TensorFlow在美团推荐系统中的分布式训练优化实践( 七 ) _知识分享

我们在HashTable初始化时，会先为Key和Value分别创造两个BucketPool，每个池子都会先Malloc较大一块内存备用，考虑到可能会有对HashTable进行中的Key和Value进行Remove的场景（如Online Learning训练时），需要对从HashTable中删除的Key和Value所使用的内存进行回收，因此每个BucketPool还有一个ReuseQueue来负责维护回收的内存。每次向内部的哈希表数据结构中Insert Key和Value的时候，Key和Value内存和释放分配都进行池化管理。用这种方式降低了大规模稀疏训练中遇到稀疏内存分配开销，整体端到端训练性能提升了5%左右。

文章插图

图17 HashTable内存优化

3.6 单位算力吞吐优化
经过2.2章节的分析，Worker的计算压力也非常大，如果不优化Worker，同时要保持吞吐，需要横向扩展更多的Worker，给PS带来更大的压力。而对于用户来说，如果能在有限的计算资源下带来性能提升，对业务价值更高。我们通过CAT统计出了一些高频算子，并进行了专项优化。这里选取Unique&DynamicPartition算子融合案例进行分享。
在TensorFlow PS架构中，包括Embedding向量在内的共享参数都存储在PS上，并通过网络与Worker交互，在进行Embedding查询过程中，往往会涉及如下两个环节：
由于稀疏参数的性质，从样本中抽取得到的待查询Embedding ID，它的重复率往往高达70%~90%，如果不进行去重查询，不论是对HashTable的查询还是网络的传输，都会带来不小的压力。因此，通常会在查询前进行Unique操作。在大规模稀疏场景中，为了存储千亿规模的参数，会有多个PS机器共同承载。而Worker端会负责对查询请求按照设定的路由规则进行切分，这里通常会在查询前进行DynamicPartition动作。
通常这两个过程会利用TensorFlow既有的算子进行搭建，但在实际使用中，我们发现它并不是很高效，主要问题在于：
Unique算子原生实现，它内部使用的内存分配策略较为低效。使用了两倍输入参数（Embedding ID）的大小进行内存分配，但由于输入参数较大，而且重复率高，导致HashTable创建过大且非常稀疏。几乎每次插入都会产生一次minor_page_fault，导致HashTable性能下降。我们使用Intel Vtune验证了这一点（参见图18）。Unique和Dynamic Partition算子存在冗余数据遍历，这些操作其实可以在一次数据遍历中全部做完，节省掉算子切换、冗余数据遍历的耗时。

文章插图

图18 Unique算子内部出现DRAM Bound问题

总结来说，HashTable开辟过大会导致大量的minor_page_fault，导致访存的时间增加，HashTable过小又可能会导致扩容。我们采用了基于启发式算法的内存自适应Unique算子实现，通过对训练历史重复率的统计，我们可以得到一个相对合理的HashTable大小，来提高访存的性能；另外Unique算子内HashTable的具体选择上，经过我们的多种测试，选择了Robin HashTable替换了原生TF中的实现。
进一步，我们对围绕Embedding ID的Unique和Partition环节进行了算子合并，简化了逻辑实现。经过上述的优化，Unique单算子可以取得51%的加速，在真实模型端到端上可以获得10%左右的性能提升，算子总数量降低了4% 。
在整个关键算子优化的过程中，Intel公司的林立凡、张向泽、高明进行大量的技术支持，我们也复用了他们的部分优化工作，在此深表感谢！
4 大规模稀疏算法建模
大规模稀疏能力在业务落地的过程中，算法层面还需要从特征和模型结构上进行对应升级，才能拿到非常好的效果。其中外卖广告从业务特点出发，引入大规模稀疏特征完成外卖场景下特征体系的升级，提供了更高维的特征空间和参数空间，增强了模型的拟合能力。重新设计了面向高维稀疏场景的特征编码方案，解决了特征编码过程中的特征冲突问题，同时编码过程去掉了部分冗余的特征哈希操作，一定程度上简化了特征处理逻辑，并降低了特征计算的耗时。
在系统层面，面对百亿参数、百亿样本以上量级的大规模稀疏模型的训练，会带来训练迭代效率的大大降低，单次实验从一天以内，增长到一周左右。美团机器学习平台训练引擎团队，除了上述TensorFlow框架层面的优化、还针对业务模型进行了专项优化，整体吞吐优化了8到10倍（如果投入更多计算资源，可以进一步加速），大大提升业务的迭代效率，助力外卖广告业务取得了较为明显的提升。