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

总第481篇
2021年第051篇

文章插图

美团内部深度定制的TensorFlow版本，基于原生TensorFlow 1.x架构与接口，从大规模稀疏参数的支持、训练模式、分布式通信优化、流水线优化、算子优化融合等多维度进行了深度优化。在推荐系统场景中，分布式扩展性提升10倍以上，单位算力性能也有显著提升，并在美团内部业务中大量使用，本文介绍了相关的优化与实践工作。1 背景
2 大规模训练优化挑战
2.1 业务迭代带来的挑战
2.2 系统负载分析
3 优化实践
3.1 大规模稀疏参数介绍
3.2 分布式负载均衡优化
3.3 通信优化
3.4 延迟优化
3.5 单实例PS并发优化
3.6 单位算力吞吐优化
4 大规模稀疏算法建模
5 总结与展望
1 背景
TensorFlow（下文简称TF）是谷歌推出的一个开源深度学习框架，在美团推荐系统场景中得到了广泛的使用。但TensorFlow官方版本对工业级场景的支持，目前做得并不是特别的完善。美团在大规模生产落地的过程中，遇到了以下几方面的挑战：
所有参数都是用Variable表达，对于百亿以上的稀疏参数开辟了大量的内存，造成了资源的浪费；只支持百级别Worker的分布式扩展，对上千Worker的扩展性较差；由于不支持大规模稀疏参数动态添加、删除，增量导出，导致无法支持Online Learning；大规模集群运行时，会遇到慢机和宕机；由于框架层不能处理，导会致任务运行异常。
以上这些问题，并不是TensorFlow设计的问题，更多是底层实现的问题。考虑到美团大量业务的使用习惯以及社区的兼容性，我们基于原生TensorFlow 1.x架构与接口，从大规模稀疏参数的支持、训练模式、分布式通信优化、流水线优化、算子优化融合等多维度进行了深度定制，从而解决了该场景的核心痛点问题。
首先新系统在支持能力层面，目前可以做到千亿参数模型，上千Worker分布式训练的近线性加速，全年样本数据能够1天内完成训练，并支持Online Learning的能力。同时，新系统的各种架构和接口更加友好，美团内部包括美团外卖、美团优选、美团搜索、广告平台、大众点评Feeds等业务部门都在使用。本文将重点介绍大规模分布式训练优化的工作，希望对大家能够有所帮助或启发。
2 大规模训练优化挑战
2.1 业务迭代带来的挑战
随着美团业务的发展，推荐系统模型的规模和复杂度也在快速增长，具体表现如下：
训练数据：训练样本从到百亿增长到千亿，增长了近10倍。稀疏参数：个数从几百到几千，也增长了近10倍；总参数量从几亿增长到百亿，增长了10~20倍。模型复杂度：越来越复杂，模型单步计算时间增长10倍以上。
对于大流量业务，一次训练实验，从几个小时增长到了几天，而此场景一次实验保持在1天之内是基本的需求。
2.2 系统负载分析
2.2.1 问题分析工具链
TensorFlow是一个非常庞大的开源项目，代码有几百万行之多，原生系统的监控指标太粗，且不支持全局的监控，如果要定位一些复杂的性能瓶颈点，就比较困难。我们基于美团已经开源的监控系统CAT[2]，构建了TensorFlow的细粒度监控链路（如下图1所示），可以精准定位到性能的瓶颈问题。

文章插图

图1 TensorFlow PS架构全链路监控

同时，在性能优化的过程中，会涉及到大量的性能测试和结果分析，这也是一个非常耗费人力的工作。我们抽象了一套自动化的实验框架（如下图2所示），可以自动化、多轮次地进行实验，并自动采集各类监控指标，然后生成报告。

文章插图

图2 自动化实验框架

2.2.2 业务视角的负载分析
在推荐系统场景中，我们使用了TensorFlow Parameter Server[3]（简称PS）异步训练模式来支持业务分布式训练需求。对于这套架构，上述的业务变化会带来什么样的负载变化？如下图3所示：

文章插图

图3 TensorFlow PS架构大规模训练负载分析

总结来看，主要包括通信压力、PS并发压力、Worker计算压力。对于分布式系统来说，通常是通过横向扩展来解决负载问题。虽然看来起可以解决问题，但从实验结果来看，当PS扩展到一定数量后，单步训练时间反而会增加，如下图4所示：