对开源框架 Tensorflow 的优化总结

一. 阿里云 TensorFlow On PAI

简单介绍：Tensorflow on PAI 是阿里巴巴机器学习产品团队结合阿里自身大数据业务，以及大规模分布式计算的优势，对开源的 TensorFlow 进行优化，并经过产品化封装后为互联网用户提供服务。

优化点：

• 异构计算资源全局调度

• 安全的分布式 TensorFlow：借助于虚拟化技术 ，重构了 TensorFlow 的底层网络拓扑。在用户分布多作业启动前，首先建立安全的虚拟网络通道，保证在多租户场景中，每位用户的作业都在独立的网络内读取自己的数据进行数据交换。

• 增加了高性能的稀疏数据解析 OP，可以直接用来解析稀疏数据，并返回相应的 SparseTensor。

• grpc 多 worker 下的网络通信优化。原生 TensorFlow 中 grpc 底层接收和发送的 handler 都是单线程，在并发请求多的时候需要排队处理，性能会很差。在分布式训练场景中，作业规模越大，对 grpc 性能依赖越高。我们对 grpc 处理线程进行了优化，采用多线程并发处理，并对 rendezvous 锁进行优化，将读写锁换为互斥锁。

二. 阿里妈妈 TensorFlowRS

原生 TF 的缺点：

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

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

• TensorFlowRS 优化点：

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

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

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

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

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

• TensorFlowRS 架构：

• 

TensorflowRS 重新设计了一个高性能的参数服务器：PS-Plus。用户可以在原生-PS 和 PS-Plus 之间自由切换，并且完全兼容 TensorFlow 的 Graph 语义和所有 API。用户可以在深度网络代码一行不改的情况下，将参数分布和运行在 PS-Plus 上，享受高性能的参数交换和良好的水平扩展能力。另外，TensorFlowRS 引入了 worker state，在 checkpoint 中存储了 worker 的状态信息，worker 重启后，会从接着上次的进度继续训练。此外 TensorFlowRS 通过 zk 生成 cluster 配置，支持了动态组网的 Failover。新的 Failover 机制可以保证任意角色挂掉的情况下，系统都能在分钟级完成 Failover，并且不多算和漏算数据。

PS-Plus 优势：

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

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

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

• 与原生 Tensorflow 性能对比：

参考文献：https://zhuanlan.zhihu.com/p/36139277

三. 阿里大规模分布式训练框架 PAISoar

传统的 TF 的 parameter-server 架构缺陷：

• Variable placement 策略，常用的 replica_device_setter 的参数分配策略是 round-robin ，这种策略并没有考虑 Variable 大小，会导致参数分配不均衡，某些 ps 上分配的 Variable size 大就会成为通信瓶颈；

• 多个 Worker 访问同一个 PS 节点时，受 PS 节点带宽限制和 TCP 的拥塞窗口控制，会导致通信效率大幅降低，并且规模越大，效率越差；

• 分布式扩展后，模型需要精细调参才能收敛，需要用户有丰富的调参经验。

• PAISoar 是基于 PAI Tensorflow 实现的分布式训练框架。通过 PAISoar，我们提供了一个从硬件到软件、易用的分布式训练的性能优异框架。

优化点：

• 网络层：搭建了集团内部第一套基于 RoCE 的大规模 RDMA 集群。

• 软件层：接入 RDMA 驱动，使用 verbs 库进行 RDMA 通信，并和思科交换机进行适配调参；对 RDMA 通信的关键路径进行梳理，加速内存拷贝，异步化数据发送，优化通信状态机，提高 RDMA 通信的效率和稳定性；自研了深度优化的 Ring AllReduce 同步算法，通过针对 RDMA 网络的深度适配、多路通信融合等关键点优化。

• API 层：提供了简化用户构建分布式 TF 模型训练的 ReplicatedVarsOptimizer，极大的方便了将单机模型扩展成分布式模型，降低了用户构建分布式 TF 训练代码的难度和学习成本；提供 lr(learning rate) 的动态调节算法 smooth_exponential_decay，算法中，lr 经过一定轮数的 warm up，从一个较小的 lr 平滑增加到目标 lr，再采用指数衰退，每隔固定 epoch 衰退一次，动态调节的 lr 可以加速用户调参过程。

• 性能对比：

参考文献：https://yq.aliyun.com/articles/

本文转载自 Alex-zhai 知乎账号。

原文链接：

https://zhuanlan.zhihu.com/p/75989435