TensorFlow在美团推荐系统中的分布式训练优化实践( 十 )


在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算子存在冗余数据遍历 , 这些操作其实可以在一次数据遍历中全部做完 , 节省掉算子切换、冗余数据遍历的耗时 。

TensorFlow在美团推荐系统中的分布式训练优化实践


图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倍(如果投入更多计算资源 , 可以进一步加速) , 大大提升业务的迭代效率 , 助力外卖广告业务取得了较为明显的提升 。

推荐阅读