ElasticDL：蚂蚁金服开源基于TensorFlow的弹性分布式深度学习系统

9 月 11 日，蚂蚁金服在2019谷歌开发者大会上海站上开源了 ElasticDL 项目，这是业界首个基于 TensorFlow 实现弹性深度学习的开源系统。

开源地址为：elasticdl.org

开源中国采访了ElasticDL项目负责人王益，对该深度学习系统的技术细节进行了全面介绍。

1、基于 TensorFlow 2.0 和 Kubernetes 实现弹性深度学习

这个基于 Eager Execution 模式的开源项目名为“ElasticDL”，它是一个 Kubernetes 原生深度学习框架，根据介绍，ElasticDL 主要有四大特色：

• 容错性
• 弹性调度
• 易用性
• 高效

其中又以容错与弹性调度特性最具特点。

ElasticDL 实现了容错和弹性调度的分布式深度学习，能够极大提高集群的整体利用率，同时显著减小用户提交做业以后等待做业启动的时间（pending time）。

王益介绍：“ElasticDL 是咱们知道的第一个基于 TensorFlow 实现弹性深度学习的开源系统。具体地说，ElasticDL 是基于 TensorFlow 2.0 和 Kubernetes 实现弹性深度学习的。”

（一）集群效用从 1/N 到 N/N
在深度学习技术研发的早期，公用一个计算集群的人相对少， 计算做业之间的协调能够经过口头交流实现。开发者更关心缩短运行时间，也就是从做业启动到结束的这段时间。高性能计算技术（HPC）是解决这个问题的有效途径，好比 NVIDIA 的 cuBLAS 和 cuDNN 优化高性能数学计算、NCCL 优化 GPU 之间的通讯效率。

随着深度学习技术的大规模应用，在许多工程师和研究员公用一个集群的状况下，经过商量来协调调度显然不可行，因而你们开始使用集群管理系统调度分布式做业。

Kubernetes 近年来已经逐渐成为集群管理的重要工具，目前已经在各大公有云中普遍采用。所以，让 TensorFlow 能更好地运行在 Kubernetes 集群上，同时提高利用集群进行深度学习的效率和资源利用率（效用），显得很是具备实际意义。

关于提高集群资源利用率，王益举了一个比较极端的例子：假设一个集群有 N 个 GPU，而一个任务只使用其中一个，如今有一个任务占用了一个 GPU。当没有弹性调度机制时，一个要求全部 N 个 GPU 的任务须要等待前一个任务结束才能开始，这个等待时间可能高达数天甚至数周，在等待期间，集群的效用是 1/N；而拥有弹性调度能力以后，新的任务能够在 N-1 个 GPU 上当即运行，而且 Kubernetes 能够在第一个任务完成后将占用的 GPU 赋予这个任务，这种状况下，集群总体效用是 100%。

ElasticDL 在容错与弹性调度上都有不错的表现，它的现实意义即是高效解决集群效用问题。

（二）ElasticDL 如何实现？
前边讲到集群资源利用率提升的前提其实就是 ElasticDL 的“弹性调度”特性带来的，而弹性调度依赖于容错能力。

容错是指做业不受其中进程数量变化的影响，在弹性调度过程当中，做业里的进程数量会随集群 workload 状况相应增减，因此做业必须是容错的，才能配合调度系统，实现弹性调度。

在这个过程当中，容错一般由分布式框架实现，好比 Spark 和 ElasticDL 均可以作到当有进程挂掉，或者新的进程加入时，做业不会暂停或者重启，而是平滑地继续。而弹性调度是由分布式框架和分布式操做系统（集群管理系统）一块儿实现的。好比，当有进程挂掉的时候，分布式框架应该通知集群管理系统新启进程来补位 —— 至于集群管理系统能不能启动起来，取决于用户剩余 quota 和集群的忙碌状况。

1. 基于 Kubernetes-native
一般使用 Keras 的 model-fit API 和 Estimator，开发者只须要调用 API 便可进行分布式训练或预测，然而 ElasticDL 不依赖于 TensorFlow runtime 实现分布式计算，它的实如今 runtime 以外。

ElasticDL 经过 Kubernetes-native 机制来完成分布式计算，而这也为其带来了容错性与弹性调度的能力。

所谓 Kubernetes-native 指的是一个程序调用 Kubernetes API 来起止进程，它与 Google MapReduce 的机制相似。MapReduce 是一个 Borg-native 的分布式计算框架，用户经过运行一个 Borg 客户端程度启动一个 MapReduce 做业；Borg 客户端调用 Borg API 提交做业，而且启动一个 master 进程；这个 master 调用 Borg API 启动其它 workers 进程。

在 ElasticDL 中，用户调用 ElasticDL 的命令行客户端程序启动做业；这个客户端程序调用 Kubernetes API 启动 master 进程，master 进程继续调用 Kubernetes API 启动其它进程。

“ElasticDL 的整个容错和弹性调度机制都依赖于 Kubernetes-native 架构”，王益介绍：“若是 worker 挂了，按照分布式深度学习训练算法的数学特性，能够不用处理， 便可确保训练过程继续。若是一个 parameter server 进程挂了，master 会选择一个 worker 进程，让它转换角色替补上挂掉的 parameter server 进程。”

在这两种状况下，master 都会调用 Kubernetes API，请它再启动一个额外的 worker 进程。若是启动成功，master 会带其加入到与其它进程的协做中。master 进程的状态（主要是三个 task queues：todo、doing 与 done）能够保留在 Kubernetes 集群的 etcd 存储系统中。

“这样，万一 master 挂了，重启的 master 进程能够从 etcd 继承前世的状态。任何进程挂了，master 都会请 Kubernetes 去启动一个新的进程代替挂掉的进程。而 Kubernetes 是否能完成使命取决于用户剩余 quota 和集群剩余资源状况。”
2. 基于 TensorFlow 2.0 Eager Execution
为何 ElasticDL 又基于 TensorFlow 2.0 呢？王益介绍，这是由于 TensorFlow 2.0 带来了 Eager Execution 特性，正是针对这一特性的尝试，让开发团队实现了 Kubernetes-native 的调度方式，从而让 ElasticDL 支持容错和弹性调度。

分布式学习须要了解每一个进程根据局部训练数据计算获得的 gradients，才能汇总这些 gradients 来更新模型。

TensorFlow 1.x 的执行方式被称为 Graph Mode —— 深度学习计算步骤被表示成一个 graph 数据结构，TensorFlow runtime 会解释执行这个 graph。其中，gradients 的计算过程是 graph 的一部分，因此为了获得 gradients，分布式深度学习系统须要 hack 进入 graph 的执行过程“偷取”gradients。

这个作法须要用户写程序的时候写一些帮助“偷取”的代码，增长了程序的复杂度，也增长了对编程者的要求。

TensorFlow 2.0 提供的 Eager Execution Mode 中，经过一个叫 tape 的数据结构，它能够把获取 gradients 的能力以 API 的方式暴露给开发者，ElasticDL 正是以这样的方式将其实现。

经过这种对比，其实也反映了业界基于 TensroFlow 进行分布式深度学习的不一样设计思路。王益介绍，当前基于 TensorFlow 的分布式训练系统大体能够分为四类：

其中须要修改 TensorFlow runtime 的工做主要由 Google TensorFlow 团队完成。由于 TensorFlow runtime 是用 C++ 写的，把网络通讯和同步功能实如今这个层次里，运行效率很高。并且，理论上 C++ 代码能够经过感知 TCP/IP 连接是否中断，来判断进程是否挂掉，从而实现容错。

“可是 TensorFlow runtime 应该是平台无关的，因此不该该包含访问特定集群管理系统，请它重启挂掉的进程的代码，因此不易实现弹性调度”，王益指出了两者的区别：“与之相对应的，经过调用 TensorFlow API 实现分布式计算的思路，通讯性能每每受到 Python 语言性能以及不能在 runtime 内部实现‘微操做’的限制。但它的好处是能够自由调用集群管理系统 API 来管理进程。”

很明显，ElasticDL 经过 TensorFlow 2.0 带来的新特性实现了 TensorFlow runtime 外直接调用集群管理 API 而完成了弹性调度。

2、ElasticDL 替代 Kubeflow 的使用

Kubernetes 原本是一个用来管理无状态应用的容器平台，可是当前愈来愈多公司用它来运行各类各样的工做负载，特别是使用它来运行机器学习相关任务。

Kubeflow 基于 Kubernetes，它把模型训练、超参数训练与模型部署等机器学习任务类型进行组合并以容器化的方式部署，提供了整个机器学习流程各个系统的高可用与便捷性，使用 Kubeflow 就能够进行各类各样的机器学习任务。

目前 Kubeflow 是在 Kubernetes 上启动分布式 TenosrFlow 做业的主流操做，这可能也是开发者更为熟悉的模式。

“具体来说，Kubeflow 会询问 Kubernetes 计划分配哪几台机器来运行一个分布式做业中的各个进程，随后告知每一个进程全部其它进程的 IP 地址和 port，从而保证一个做业里各个进程之间互相知道对方。”

为何须要让全部进程互相知道对方呢？这是 TensorFlow ps-based distribution 方式要求的。（也就是前边提到的对比基于 TensorFlow 的分布式训练系统表格中左上角的类型）

王益解释：“TenosrFlow 1.x 原生的分布式训练功能让一个做业中全部进程都执行 TensorFlow 1.x runtime 程序。这些进程互相通讯，互相协调成为一个‘分布式 runtime’来解释执行表示深度学习计算过程的 graph。在开始分布式训练之初，graph 被 TensorFlow runtime 拆解成若干子 graph；每一个进程负责执行一个子 graph —— 任何一个进程失败 （多是被更高优先级做业抢占），则整个大 graph 的执行就失败了。因此 TensorFlow 原生的分布式训练能力不是容错的（fault-tolerant）。”

不过，TensorFlow Python API 提供了 checkpoint 的能力：若是一个做业失败了，能够重启做业，从最近的 checkpoint 开始继续执行。因此它能够从错误中恢复（fault-recoverable）。

Kubeflow 能够在 Kubernetes 上发挥 TensorFlow 原生的分布式计算能力，可是由于后者并不能容错，因此 Kubeflow 并不能无中生有。不能容错，也意味着不能弹性调度，而这正是 ElasticDL 的特长。

3、与 SQLFlow 联动

前边介绍了 ElasticDL 的实现机制与现实意义，总结起来主要是由于 ElasticDL 经过 TensorFlow 2.0 提供的新特性 Eager Execution 实现了 TensroFlow runtime 外直接调用集群管理 API，从而实现了 Kubernetes-native 机制来完成分布式计算，继而实现容错与弹性调度，最终达到极大提高集群整体利用率的目标。

除此以外，ElasticDL 还有一个重要的特性——易用性。ElasticDL 的易用性与另外一个工具密不可分。

几个月前，蚂蚁金服开源了一个机器学习工具 SQLFlow，这个工具旨在让开发者调用 AI 像写 SQL 同样简单。据介绍，SQLFlow 可以抽象出端到端从数据到模型的研发过程，配合底层的引擎及自动优化，具有基础 SQL 知识的开发者便可完成大部分机器学习模型训练及预测任务。

经过与 SQLFlow 联动，开发者能够用扩展后的 SQL 语法，很是精炼地描述整个数据流和 AI 流程。SQLFlow 把一个 SQL 程序翻译成一个实现整个 end-to-end machine learning 的程序，这个程序能够调用 xgboost、PyTorch，或者 ElasticDL、TensorFlow 实现训练任务。

王益举例，在有 SQLFlow 以前，若是要为一个电子商务网站构造一个推荐系统，须要开发日志收集、在线数据清洗、特征工程、模型训练、验证与预测等模块，每一个模块可能须要投入一个团队数周甚至数月的时间。

而 SQLFlow 出现以后，这个流程能够用 SQL 语言描述成一个很简短的程序，SQLFlow 能够把它翻译成上述数据和 AI 流。

由于 SQL 是一种只描述意图，不描述过程的语言，因此 SQL 程序一般都很简短。可是也由于这个缘由，SQL 程序里包含的信息量有限。好比，用户不会经过 SQL 指定分布式调度和训练算法。“这些部分须要 ElasticDL 根据模型特色自主决定”，王益补充：“这也是为何说 ElasticDL 也能够反过来为 SQLFlow 提供易用性。”

4、ElasticDL 开源的下一步计划

关于开源后接下来的发展，王益表示，ElasticDL 项目目前处于早期探索阶段，API 还在演化过程当中。“这次开源的版本，尚不包括自动选择分布策略和算法的代码，相比在 TensorFlow runtime 中实现分布式计算，基于 TensorFlow 2.0 Eager Mode 的 Python API 实现的分布式训练性能差距还很大”，他介绍：“ElasticDL 团队在和 Google Brain 团队合做，开发上述 asynchronous SGD + delayed model update 能力、以及 Kubernetes-native AllReduce，但愿在下一个版本中能够提供给你们使用。”

随后王益又具体介绍，上述两种分布式训练策略，一种会用于模型中有较大的参数的状况，好比分布式 embedding table，另外一种用于模型参数较小的状况。而这也是 ElasticDL 自动决断分布式训练算法的一个例子。

另外一方面，在 SQLFlow 中，若是要让用户能提供尽可能少的参数，AI 引擎还须要更加智能，提供包括 AutoML 等功能。

王益感叹：“ElasticDL 项目任重道远。”

5、受访嘉宾

王益，目前在蚂蚁金服负责 AI 基础架构工做。他于 2007 年从清华大学计算机系博士毕业，前后在 Google（中国）、腾讯、LinkedIn（美国总部）与百度硅谷研究院工做，期间在硅谷和北京各有一次创业经历。参加工做以来，王益一直专一于 AI 基础架构工做，参与和领导了多个核心 AI 系统的研发。