阿里妈妈基于 TensorFlow 做了哪些深度优化？TensorFlowRS 架构解析

一． 综述

深度学习比传统的逻辑回归有着更强的模型刻画能力，同时也带来了计算力百倍提升的需求。相比图像、语音、视频等领域，搜索、广告、推荐等场景有着独特的场景特点: 样本规模和特征空间通常非常巨大，千亿样本、百亿特征并不罕见，同时存在大量的稀疏特征作为 Embedding 输入。这就要求我们针对此场景下的计算特点对深度学习框架进行设计和优化。

本文所阐述的工作由阿里妈妈基础平台团队与 PAI 团队合作完成，我们基于 TensorFlow 在搜索、广告、推荐场景下进行了深度的优化与增强，内部项目名称为 TensorFlowRS，主要的成果如下：

(1) 解决了原生 TF 水平扩展能力不足的问题。在我们的测试中，绝大多数搜索广告模型的训练性能提升在十倍以上，某些模型的极限性能最高可提升百倍。

(2) 支持完备的在线学习语义，模型变更实时写出；稀疏特征无需做连续 ID 化，可以直接使用原始特征表征进行训练，大幅简化了特征工程的复杂度。

(3) 异步训练的梯度修正优化器（grad-compensation optimizer），有效减少了异步大规模并发引起的训练效果损失。

(4) 集成了高效的 Graph Embedding、Memory Network、Cross Media 等多种高级训练模式。

(5) 模型可视化系统 DeepInSight 提供深度模型训练的多维度可视化分析。

二． TensorFlowRS 分布式架构

在使用 TensorFlow 的过程中我们发现 TF 作为一个分布式训练系统有两个主要的问题：

1. 水平扩展能力差：在大部分模型的性能测试中, 我们发现随着数据并行度的增加，单个 worker 的样本处理 QPS 急剧下降。当 worker 数量增大到一定规模的时候，系统整体 QPS 不再有增长甚至有所下降。

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

针对上述问题，TensorFlowRS 采取的解决方案包括：

• 通过对接独立参数服务器提升水平扩展能力

  在对 TF 做过细致的 profiling 之后，我们发现 TF 原生的 PS 由于设计和实现方面的多种原因（grpc，lock，graph-engine），很难达良好的水平扩展能力。于是我们决定丢掉 TF-PS 的包袱，重新实现一个高性能的参数服务器：PS-Plus。此外我们提供了完整的 TF on PS-Plus 方案，可以支持用户在 Native-PS 和 PS-Plus 之间自由切换，并且完全兼容 TensorFlow 的 Graph 语义和所有 API。用户可以在深度网络代码一行不改的情况下，将参数分布和运行在 PS-Plus 上，享受高性能的参数交换和良好的水平扩展能力。 

• 重新设计 Failover 机制，支持动态组网和 Exactly-Once 的 Failover

  TensorFlowRS 引入了 worker state，在 checkpoint 中存储了 worker 的状态信息，worker 重启后，会从接着上次的进度继续训练。此外 TensorFlowRS 通过 zk 生成 cluster 配置，支持了动态组网的 Failover。新的 Failover 机制可以保证任意角色挂掉的情况下，系统都能在分钟级完成 Failover，并且不多算和漏算数据

TensorFlowRS 的整体架构如图所示：

三． PS-Plus

PS-Plus 相对于传统的 ParameterServer 有如下特点：

(1) 高性能：PS-Plus 通过智能参数分配，零拷贝，seastar 等多项技术，进一步提升了单台 server 的服务能力和系统整体的水平扩展能力。在实测中，在 64core 的机器上单个 server 能轻松用满 55 + 的核心，在 dense 场景下 io 能打满双 25G 网卡，系统整体在 1~4000 worker 的范围内都具有近似线性的水平扩展能力

(2) 高度灵活：PS-Plus 拥有完善的 UDF 接口，用户可使用 SDK 开发定制化的 UDF 插件，并且可以通过简单的 C++ 以及 Python 接口进行调用。

(3) 完备的在线学习支持：PS-Plus 支持非 ID 化特征训练，特征动态增删，以及模型增量实时导出等支撑在线学习的重要特性。

下面从中选取几点做比较详细的介绍：

1. 智能参数分配

参数分配策略 (variable placement)，决定了如何将一个参数切分并放置到不同的 server 上。placement 策略的好坏在高并发的情况下对 PS 的整体性能有着重大的影响。传统 ParameterServer 的 placement 方案是由系统预先实现几种常见的 placement 算法（比如平均切分 + roundrobin），或者由用户在创建参数的时候手工划分，往往没有综合考虑全局的参数规模、Server 的负载等。

PS-Plus 实现了基于模拟退火算法的启发式参数分配策略，后续也在考虑实现基于运行时负载，动态 rebalance 的 placement 策略。PS-Plus 的 placement 设计有如下优点：

• 综合考虑了全局参数的 shape 信息，在 cpu，内存，网络带宽等限制条件下给出了近似最优的 placement 方案，避免了手工分配造成的不均匀、热点等问题。

• 整个参数分配过程由系统内部自动完成，用户无需配置即可获得接近最优的性能，用户无需了解 PS 底层实现的具体细节。

• Partition 由框架自动完成，在上层算法代码，如 TF 代码中，不需要额外使用 PartitionedVariable 等机制，使用简单方便。

2. 去 ID 化特征支持

目前主流的深度学习框架都是以连续的内存来存储训练参数，通过偏移量（ID 值）来寻址到具体的权重。为了避免内存的浪费，需要对特征做从 0 开始的连续 ID 化编码，这一过程我们称之为特征 ID 化。特征 ID 化是一个非常复杂的过程，尤其是当样本和特征数量非常庞大的时候，特征 ID 化会占用大量的时间和机器资源，给样本构建带来了很大的复杂度。

PS-Plus 内部实现了一个定制化的 hashmap，针对参数交换场景做了专门的优化，在支持特征动态增删的同时提供了超高的性能。通过 hashmap，PS-Plus 直接实现了对非 ID 特征的支持，极大的简化了样本构建的复杂度。

3. 通信层优化

对于 Parameter Server 架构，延迟是影响整体性能的重要原因。尤其是在模型复杂度不高的情况下，模型计算部分往往在 10~100ms 量级，那么总体通信的延迟就成为一个关键因素。

在传统的 pipeline 线程模型，高并发情况下中断和线程上下文切换会导致很大的开销，同时会引起大量的 cache-line miss。此外，高频的锁竞争是带来延迟的最主要原因之一，即便是各类 SpinLock、读写锁等优化也并不能有效消除这个问题。我们认为 polling + run to completion 是一个正确的选择，并且设计了我们的整体通信层架构。在新的通信层中，我们使用了 Seastar 作为底层的框架。对于 Server、Worker 上的 connection，都严格保证 connection 绑定到固定的线程，同时线程与 CPU 核心绑定。Request、response 直接采用 run to completion 的方式在当前线程处理。整体架构如下图所示：

在 Seastar 的基础上，我们做了很多功能、性能的改进和优化，这里做一些简要的介绍。

• 外部线程交互队列。我们借鉴 Seastar 核心之间的交互机制，提供了一个 M:N 无锁生产者消费者队列，用于外部线程与 Seastar 内部线程进行交互。相比传统队列性能有极大的提升。

• 写请求顺序调度。从外部线程 poll 到的写请求，如果直接调用 Seastar 的写接口，会导致写 buffer 无法保证有序。我们通过队列机制的改造，自动保证了写顺序，同时基本不损失多 connection 的并发写的性能。

• 灵活的编解码层。我们提供了一套编解码层的抽象接口，方便用户使用，从而不需要借助 protobuf 等传统的序列化、反序列化的第三方库，同时也避免了 protobuf 的一些性能问题。 

四．性能测试

我们测试了 TensorFlowRS 在 Dense 以及 WDE(Wide-Deep-Embedding) 两种经典模型的性能指标：

1. 模型说明：

•  Dense

• WDE

2. 测试结果：

3. WDE 模型下 Native-TF 与 TFRS 水平扩展能力比较

五. 在线学习

以 Ftrl 为代表，在线学习近年来在工业界开始被大规模应用，它是工程与算法的深入结合，赋予模型实时捕捉线上流量变化的能力，在一些对时效性要求很高的场景，有十分大的价值。

深度模型和 LR 模型一样，对在线学习有着同样强烈的需求，然而目前主流的深度学习框架都缺乏对在线学习的支持。TensorFlowRS 通过对接 PS-Plus，给出了一套完整的端到端的在线学习解决方案，赋予了 TF 支持千亿规模非 ID 化特征在线训练的能力。

TFRS 针对在线学习的场景做了专门设计和优化，具体包括：

1. 非 ID 化特征支持

在在线学习的场景下做特征实时 ID 化是比较复杂的，需要一个超高性能的全局的 ID 生成器，这给样本生成带来了很大的复杂度。TensorFlowRS 利用 PS-Plus 直接实现了对非 ID 特征的支持，极大的简化了实时样本构建的复杂度。

2. 特征动态增删

在在线训练的场景下，训练任务会以 service 的形式长期运行，在训练过程中，不断会有新特征加入到模型中，为了保证训练可以长期进行而不会因为新特征的不断加入导致 OOM，PS-Plus 在支持特征动态添加的同时，还提供了默认的特征删除策略，可以将低频或者低权重的特征删除掉，用户还可以通过 UDF 定制符合自身业务需求的删除策略

3. 模型增量实时导出

在线学习模型更新的常见方式有全量和增量两种。在模型参数较多的情况下，全量更新的方式会对在线系统的带宽带来巨大的压力，而降低更新频率又会使模型的实效性降低。PS-Plus 支持以任意频率将模型增量部分实时写出到消息队列，在大幅减少网络 IO 的同时实现了真正意义上的模型实时更新。

4. AUC Decay

在在线学习的场景下，我们希望在训练的过程中就可以尽快的发现模型本身的异常，而不是等模型更新到线上之后。因此我们需要有一些方法来评估模型在训练过程中的 AUC 等指标。TF 里默认的 streaming auc 的实现在历史数据累积了一定量的时候，无法及时反应当前模型的状态，反馈有很大的滞后性。因此我们引入了新的 AUC 计算机制：AUC Decay。AUC Decay 本质上是一种特殊的 Moving Average，通过基于时间的减益方式，弱化历史样本和模型在当前 AUC 计算中的比重，以达到更快反应模型变化的目的

六. 大规模训练场景下的收敛效果优化

1. 问题阐述

大数据模型引入了分布式并行训练，同步并行训练受长尾 worker 的制约，并发数容易受限。异步并行是快速训练的主流。异步并行训练打破了普通 SGD 训练的串行性，计算的梯度与更新的模型不是严格一致，引入了梯度 delay 的问题。

 具体来说，在基于 ParameterServer 的训练框架里，系统分为两个角色：worker 和 ps。ps 的职责是模型的切块存储与更新；worker 的职责是加载从 ps 端获得最新的模型，读取数据进行模型训练，最后把学到梯度发送给 ps，由 ps 更新模型。异步并发训练打破了普通 SGD 训练的串行性，引入了梯度 delay 的问题。如图所示，拿到了模型，计算得到了梯度，可是最后回传给 ps 时，却是用在了 ps 上的模型上。因为在计算梯度的同时，有另外 r 个 worker 向 ps 提交了梯度更新，ps 上的模型已经向前走了 r 步。用模型计算得到的梯度，用在了模型上。虽然梯度更新的大方向可能偏差不大，但与模型期望的梯度相比，梯度是存在细微偏差的，因为梯度有些过时了。这是异步训练里梯度 delay 的由来。

2. 梯度补偿

微软在 ICML2017 提出过 DC-ASGD optimizer，使用泰勒展开来近似梯度补偿。我们测试在 50 并发以内收益良好。可是在上百并发训练里，泰勒展开超过了近似收敛半径，导致的误差增大，效果下降。

应对上百并发的训练，我们引入了与 g 的相关因子来 boost 主流的 sgd-based optimizer。用与 g 的相关性来衡量梯度 delay 的严重程度。在每一维度上，如果与 - g 是正相关，说明大部分 worker 都在往同方向更新，模型 w 在这个方向上走的挺远了，继续前进需要谨慎，所以我们保持 g 的方向不变，但减少 g 的绝对值。与之相反，如果与 - g 是负相关，说明大部分 worker 都在往反方向更新，此时 g 是一个比较强烈的转折信号，揭示了模型 w 的更新方向要发生变化，我们需要重视这个信号，所以我们保持 g 的方向不变，但增大了 g 的绝对值。

引入相关因子是依据如下分析前提的：

(1) 异步训练时，存在隐式的梯度动量加速情况参见《Asynchrony begets Momentum, with an Application to Deep Learning》, 并发越大，隐式动量越大，造成梯度往一个方向过度前进的情况。

(2) 如果不是很老的 w, 相关因子是转折信号，暗示模型在多个 worker 的动量累积推动下前进的有些过度了。

(3)存在着 tradeoff, 太老的 w，信号准确率会下降，这时要控制 (调小) 系数 lambda。

因为与 g 的相关性具备普适性，所以可以和主流的 sgd-based optimizer 结合，适应不同场景的不同优化器并发训练需求。

3. 实验结果

我们用相关性因子 boost 了 SGD、Momentum、AdaGrad 三种算法，并在生产环境和公开数据集上做了实验，实验结果如下

• WDE 模型

• Cifar10 Alexnet 模型

八．高级训练模式

TFRS 中集成了多种高阶训练模式，例如 Graph Embedding，Memory Network，Cross Media Training 等。在本文我们会简要的介绍一下，在以后的文章中做详细的阐述。

Graph Embedding 图是一种表征能力极强的数据结构，但是无法直接作为神经网络的输入。TFRS 支持样本以图的形式进行输入，并支持多种随机游走算法动态生成正负样本。目前 Graph Embedding 已经应用在了搜索直通车的向量化召回等多个项目里，通过在 User-Query-Item 三种节点的异构有向图中随机游走，生成深度神经网络能够处理的稀疏特征。最终学习出 User，Query 和 Item 三者的高维向量化的表示，用于线上广告的向量化召回。值得一提的是，除了 Graph Embedding，我们同样支持对图的结构进行学习，例如在训练过程中反馈调整图中的边的权重等。

Memory Network 记忆网络最早由 Facebook 在 2015 年提出，用于 QA 系统中。在本模型出现之前，机器学习的模型都缺乏可以读取和写入外部知识的组件。对很多任务来说，这是一个很强的限制。比如，给定一系列事实或故事，然后要求回答关于该主题的问题，虽然原则上这可以用 RNN 等模型进行处理，然而它们的记忆（隐藏状态和权重编码）通常太小，并且不能精确地记住过去的事实。在阿里妈妈搜索广告场景下，我们使用记忆网络对用户行为进行建模。

相比一般的在样本组织阶段进行记忆体生成的方式，TFRS 通过在训练过程中引入动态记忆存储模块，支持长短期记忆，大幅提高了序列化类行为数据的训练效率。

九. 可视化模型分析系统 DeepInsight

DeepInsight 是一个深度学习可视化质量评估系统，支持训练阶段模型内部数据的全面透出与可视化分析，用以解决模型评估、分析、调试等一系列问题，提高深度模型的可解释性。

下面我们通过一个过拟合的例子来说明 DeepInsight 在模型质量分析和问题定位方面发挥的作用：

上图是通过 DeepInsight 生成的特征权重分布, 从图中我们可以看到右侧过拟合模型的边权重大小分布很不均匀，出现了大量权重极大的边，且集中在一条带状区域内，其为某一组特征输入所连接的所有边，这表明模型过度拟合了该组特征的信息。在使用正则项和 dropout 之后，过拟合的问题仍然没解决，因此我们最终定位到问题出现在该组特征的输入上。