为君作磐石——人人都能搭建大规模推荐系统

前言

什么是个性化推荐？简单说，就是给用户推荐他喜欢的物品。近 10 年，移动互联网高速发展，个性化推荐扮演了很重要的角色。以运营一款内容类产品为例：用户增长团队通过广告投放等手段为产品拉新，提升 DAU；产品技术团队为用户分发感兴趣的内容，提升留存及停留时长；商业化团队分发用户可能感兴趣的广告，提升单位流量变现效率；商业化收入又用于用户增长，形成正向循环。个性化推荐技术贯穿每个环节，成为了很多公司的高速增长引擎。

怎么做个性化推荐？通常，对一项业务来说，首先会定义出多个优化目标（例如视频的播放时长、点赞、分享，电商的点击、加购、购买等），之后构建一个或多个模型来预估这些目标，最后融合多个目标的预估分来完成排序。对推荐系统来说，最核心的工作，便是构建精准的预估模型。这些年，业界的推荐模型一直朝着大规模、实时化、精细化的趋势不断演进。大规模是指数据量和模型非常大，训练样本达到百亿甚至数万亿，单个模型达到 TB 甚至 10TB 以上；实时化是指特征、模型、候选实时更新；精细化则在特征工程、模型结构、优化方法等多方面有所体现，各种创新思路层出不穷。

大规模推荐系统的落地，工程挑战很大。本文选择大家最关心的 Training 和 Serving 系统，介绍搭建过程中会遇到哪些挑战，我们做了哪些工作。对任何一家公司来说，从 0 搭建这样一套系统都绝非易事，投入非常大。在字节跳动内部，我们也经过了多年的探索与沉淀，有上千名工程师，不断迭代和优化推荐系统。那么，搭建推荐系统一般会遇到哪些问题？我们先来看一个故事：

A公司的故事

A是一家电商公司，他们的产品有300万DAU，有一个10人的算法团队，他们在搭建推荐系统的过程中，遇到了不少麻烦，我们具体来看看。

A公司想训练一个点击率模型，每天有1亿次曝光，100万次点击，他们想用3个月的数据训练模型，样本量级达到90亿。他们设计了200个特征，包含用户ID、商品ID、用户的点击序列等，想为每个特征分配16维的向量来表征，粗略计算下来模型大小为500G。分析之后，他们发现要做分布式训练和模型存储，于是调研了一些开源方案：

• Tensorflow：Google开源的机器学习系统，可以使用Partitioned Variable来分布式地存储Embedding，从而实现大规模训练。但由于table size固定，有hash冲突风险。
• PyTorch：Facebook开源的机器学习系统，使用Ring All Reduce同步参数，要求单机能容纳所有参数，难以训练超大模型。
• XDL：国内开源的机器学习系统，自研PS系统，用TF作为训练引擎，并且内置了一些开箱即用的推荐模型。功能上可以实现大规模训练，但是这套系统开源支持较弱，使用在生产中有风险。
• Angel：国内开源的机器学习系统，其特点是与大数据系统Spark紧密结合，使用Spark完成数据预处理与特征工程。自研Parameter Server，内嵌Pytorch为训练引擎，可以训练超大模型。但是Angel的在线离线特征难以保证一致性，只适合做离线训练平台。

经过对比，A公司选择了Tensorflow来做分布式训练。但是，训练模型的时候发现速度非常慢，即使投入大量资源依然需要5天才能训完3个月的数据。他们花了很多时间研究Tensorflow，profiling训练过程，发现了一些问题：

• TensorFlow 的分布式runtime 性能不好, 对于每个特征都单独产生了一对send/recv op来连接worker 和 PS，这样单个worker 就跟 PS 产生了200个send/recv，造成了TensorFlow Runtime的调度困难，降低了分布式训练的速度。
• 训练过程中CPU的使用率非常不稳定，看起来CPU并没有被充分利用起来。
• 有些算子运算的特别慢，推测可能和内存带宽有关。
• 虽然网络带宽并没有满载，但是增加更多的机器不能够再提升训练速度了。
• 浏览TF官方网站的时候发现TF最近推出了各种不同的分布式策略，它们又分别对应着训练集群不同的拓扑结构。他们非常疑惑，不知道应该选择哪一种。

虽然发现了不少性能问题，但优化起来并不十分容易。经过一段时间的努力，他们优化了部分问题，将训练时间从5天压缩到了3天，勉强可以接受。但是，当训练进行到第40小时的时候，因为一台机器OOM，训练任务挂了。他们多尝试了几次，发现训练成功率比较低, 分析之后发现主要原因是:

• TF 基于静态拓扑配置来构建 cluster，不支持动态组网，这就意味着当某个 ps 或者 worker 挂掉重启之后，如果 ip 或者端口发生变化（例如机器 crash），训练将无法继续。
• TF 的 checkpoint 只包含 PS 存储的参数信息，不包含 worker 端的状态，不是全局一致性的 checkpoint，无法实现 Exactly-Once 语义。

做好容错挑战不小，他们只能先隔离一个独立的集群，让训练尽量稳定一些。不能和其他任务混合调度，资源利用率自然也要低不少。

几经波折，勉强训好了一个500G的模型，他们想把模型推到线上去Serving，于是考虑在线系统的设计。经过一番讨论，他们认为Serving系统必须满足如下要求：

• 分布式：推荐模型的特点是有大量的Embedding，模型很容易达到TB级，考虑未来的模型迭代，必须支持分布式Serving。
• 低时延：单次预估的延时要尽量低，精排模型一般要控制在80ms内。复杂的深度模型，可能需要GPU来Serving，并做一系列的性能优化。
• 高可用：少部分节点挂掉不影响在线稳定性，一般通过多副本解决，需要调度系统的支持。
• 少抖动：模型更新、上线、下线等操作，不会造成延时抖动。
• AB测试：推荐系统迭代很快，算法工程师会开展很多AB实验，实验组的流量会动态调整，在线系统需要能支持模型、服务的动态调度。

目前，没有开源系统能满足上述要求，各大公司都是自研，实际做起来投入也不小。A公司人力有限，经验也不足，只能先通过一些模型压缩的手段，让单机可以Serving，模型也不能做的太复杂。

模型上线之后，A公司又遇到一个新的问题：如何更新模型。定期全量重训成本很高，如果线上有多个同时ABTest的模型，更是会雪上加霜。所以，至少要做到天级的增量更新，实时更新自然更好。但增量/实时更新，实现起来也不太容易。其实，未来还有更多的问题等着A公司，比如：如何保证线上线下特征的一致性；上游数据流不稳定怎么办；如何解决模型越来越大的问题；如何做好多场景数据的混合训练；如何应对大规模候选的问题；如何解决转化事件大幅延迟的问题等等。

我们的工作

通过A公司的故事，大家能看到，开发一套大规模推荐系统，难度确实不小，成本也很高。那么，有没有一款产品可以直接覆盖数据校验、特征工程、模型开发、线上服务、AB测试等全流程，让业务轻松搭建一套一流的推荐系统，不再遭遇A公司的头疼问题呢？有。

字节跳动成立火山引擎之后，我们一直在努力，将字节的推荐技术开放给外部客户。如今，我们已经可以通过火山引擎的智能推荐平台，来帮助大家解决这些难点和痛点。目前这套平台也开放了部分名额供企业免费使用，具体信息可以在文末进行了解。

接下来，再展开介绍一下，智能推荐平台中的大规模Training和Serving方案，我们把它命名为Monolith(磐石)，希望它能成为大家做推荐系统的坚实基础，如下是架构图：

从图中可以看出，Monolith是PS架构，下面看看这套架构是怎样运行的：

批量/增量训练

• Worker/PS启动时会向ZK注册，信息包括(server_type，index)。然后Worker向ZK请求注册信息，生成Cluster信息，实现动态组网，动态组网是容错的基础。
• 训练开始后，Worker会从标准输入或文件中获取数据，同时从PS拉取参数，然后进行forward/backward计算，得到梯度，并将其Push给PS。
• PS获得梯度后，一方面，利用优化器更新内部weight，另一方面，会记录哪些数据更新了。在PS上起一个TF Session，它会定时将更新的参数发送到Online PS，从而实现实时增量更新。此外，特征过滤，特征淘汰等也在PS上进行。
• 在训练过程中或训练结束时，会写checkpoint。为了加速checkpoint，Monolith 没有延用TF 中的saveable，而是利用estimator saving listener，流式多线程地存取，性能大副提升。为了减少checkpoint体积，会将过期特征淘汰。

在线推理

• 加载saved_model。Entry本质上是TF Serving，它会从HDFS上加载非Embedding部分，同时向ZK注册，以便上层做负载均衡。Online PS也会先向ZK注册，然后从HDFS中加载参数，并在加载过程中去除优化器辅助参数，将fp32转换成fp16，量化压缩等。
• 对于一次请求，Entry会随机选择一组Online PS，从中获取Embedding，完成预测。Entry/Online PS 是多副本的，只要有一个副本存在，服务就可用。Online PS是多分片的，可以Serving超大模型。可以在一台机器上部署多个分片，也可以Entry/OnlinePS混部。
• 对于一些对模型实时性较高的系统，Training PS会直接通过RPC的方式与Online PS进行通讯，从而将样本反馈到线上模型的时间间隔缩短到分钟级。
• Training PS可以与Online PS通讯，接受Training PS的参数更新；Entry可以自动从HDFS上读取更新参数，从而实现分钟级参数增量更新。

综上所述，Monolith 包括了 Training/Serving/Parameter Sync等，是一套完整的系统。

与业界其它系统相比，Monolith成功应对了多方面的挑战，有如下特色：

解决了TensorFlow PS 通信瓶颈

在工业级的推荐模型中，我们常会使用几百甚至数千类特征，每类特征都需要创建哈希表去存储特征embeddings。直接为每类特征生成一张哈希表，同时对几百张表进行查找会导致两个问题：

1. PS和Worker连接会产生过多的 send/recv op，大大影响分布式 runtime 的运行效率。
2. 这些 ops 导致模型图节点过多，模型图过大，训练初始化时间过长。

针对如上问题，我们在框架层面做了优化：对于配置同构的哈希表（dim 相同、优化器参数相同），在python API 层面合并哈希表来减少表的数量，同时monolith会对通信op进行进一步的合并，从而极大的减少了send/recv ops，解决了原生TensorFlow PS 的通信问题。

针对异步训练，monolith还开发了变量与embedding预取以及梯度异步更新的功能，对于多数模型，能够更加有效的利用带宽与CPU，从而提高训练速度，优化资源利用率。

全方位容错

在服务发现的基础上，无论是Worker还是PS发生错误，都能得到快速恢复。对于Worker，Monolith不同worker节点之间并不直接进行通信，所以一个worker的失败并不会对别的worker产生影响；同时，worker会存储输入的进度，当worker因为意外原因失败时，输入的进度并不会丢失；当PS shard 节点失败，根据离线/在线任务的性质不同，支持部分恢复和全量恢复不同的模式，在正确性以及恢复速度上做一定的取舍。

分布式Serving

Monolith补齐了开源软件在分布式Serving方面的空白，提供了TB级模型的推理服务。支持多副本、高可用，Training PS在训练过程中，分钟级别将刚刚更新过的Embedding同步给Serving PS，从而实现近实时参数更新，提升了产品的推荐效果。

性能优化

除了上面提到的解决 TensorFlow PS 通信瓶颈之外，Monolith 在 Parameter Server 架构、底层 Hash Table 设计、网络传输、多线程加速、OP Fusion、指令集加速等方向也进行了非常细致的优化并取得了可观的性能收益。以异步训练为例，训练时整个过程示意如下：

• 网络通讯优化：通过 embedding prefetch, gradients postpush 将网络 IO 与图的前向/后向计算异步起来，同时支持控制流与数据流分离、压缩传输等优化；
• 内存优化：通过支持特征过滤、特征压缩、特征淘汰等手段，可以极大地节省 training/serving 阶段内存使用；
• 计算优化：Hot spot code 采用 AVX 指令集优化、耗时 Op 精细调优、手工 Op Fusion 等手段加速前向/后向计算过程；
• 其它方面：多线程优化、细粒度锁设计、IO与计算异步起来等。

目前，Monolith已通过推荐平台，成功应用在电商、社区、视频等多个行业的场景上，效果、稳定性、性能均得到了充足的验证。未来，我们也将继续保持高速迭代，不断优化用户体验和平台功能。