Kubernetes-native 弹性分布式深度学习系统

9月11日，蚂蚁金服在 Google Developer Day Shanghai 2019 上宣布开源了基于 TensorFlow 2.0 eager execution 的分布式深度学习系统 ElasticDL。基于 TensorFlow 的支持弹性调度的深度学习系统，据咱们所知，ElasticDL 是第一 个。项目负责人王益和咱们分享了 ElasticDL 项目的设计意图和现状，尤为是 ElasticDL 与 TensorFlow 2.0 以及 Kubernetes 的技术关联。

分布式深度学习的技术思路

基于 TensorFlow 的分布式训练系统大体能够分为如下四类：

其中，ElasticDL 位于田字格的右上角。之因此选择这条技术思路，是为了利用 Kubernetes 实现容错和弹性调度。

高性能计算和云计算

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

随着深度学习技术的大规模使用，不少工程师和研究员共用一个集群，经过商量 来协调调度显然不可行了，你们开始使用集群管理系统调度分布式做业。这其中， Kubernetes 近年来一枝独秀，已经在各大公有云中普遍使用。

云计算和弹性调度

在 Kubernetes 上启动分布式 TensorFlow 做业的经常使用方式是使用 Google Cloud 开源的 Kubeflow。Kubeflow 是 Kubernetes 的一个”插件“，它询问 Kubernetes 计划分配哪几台机器来运行一个分布式做业中的各个进程，随后告 知每一个进程，全部其余进程的 IP 地址和 port。从而保证一个做业里各个进程 之间互相知道对方。

为何须要让全部进程互相知道对方呢？这是 TensorFlow ps-based distribution 方式（上述表格中的左上）要求的。TensorFlow 1.x 原生的分布 式训练功能让一个做业中全部进程都执行 TensorFlow 1.x runtime 程序。这些 进程互相通讯，互相协调成为一个“分布式 runtime“，来解释执行表示深度学习 计算过程的计算图（graph）。在开始分布式训练之初，graph 被 TensorFlow runtime 拆解成若干子图；每一个进程负责执行一个子图 —— 任何一个进程失败 （多是被更高优先级做业抢占），则整个大图的执行就失败了。因此 TensorFlow 原生的分布式训练能力不是容错的（fault-tolerant）。不过， 它是能够从错误恢复（fault-recoverable）—— TensorFlow API 提供 checkpoint 的能力；若是一个做业失败了，能够重启做业，从最近的 checkpoint 开始继续执行。

Kubeflow 能够在 Kubernetes 上启动基于 TensorFlow 原生的分布式计算能力的做业。可是 由于后者并不能容错，因此 Kubeflow 并不能无中生有。不能容错，也意味着不 能弹性调度。

对弹性调度的诉求

在不少人共用计算集群的状况下，支持弹性调度意味着极大提高团队效率和集群 的整体利用率。前者支持快速迭代以保持技术领先；后者决定企业成本和云计算 业务的盈利能力。

一个展现弹性调度效果的例子以下。假设一个集群里有 N 个 GPU，一个做业包 括一个进程，占用了 N/2 个 GPU。第二个做业须要 N/2+1 个 GPU；可是此时机 群里空闲 GPU 只有 N/2 个。若是没有弹性调度能力，那么第二个做业被迫等待， 直到第一个做业结束释放资源。这个等待时间极可能和第二个做业的运行时间同 量级。此时，集群的利用率很低，是 50%。若是有弹性调度，那么第二个做业可 以立刻启动，用 N/2 个 GPU 作计算。往后若是有更多空闲资源了，调度系统可 以增长其进程数量，充分利用资源。

另外一个例子是，假设有一个做业已经在执行了，此时一个新的更高优先级的做业 须要资源，因此调度系统杀掉了（preempt）了第一个做业的几个进程来腾出资 源启动第二个做业。若是没有弹性调度和容错，那么第一个做业会失败，全部进 程都结束。直到有足够资源重启它，而且沿着最近的 checkpoint 继续。若是有 弹性调度，则第一个做业的剩下的进程能够继续执行，只是由于可用的进程 （GPU）少了，因此速度慢一些而已。

以上两个例子都展现了弹性调度对集群利用率的提高，以及对团队工做效率的保 障。须要注意的是：容错和弹性调度互为因果。容错的意思是，做业不受其 中进程数量变化影响。弹性调度时，做业里的进程数量会随集群 workload 状况 增减，因此做业必须是容错的，才能和调度系统配合，实现弹性调度。也由于如 此，弹性调度依赖 分布式编程框架和调度系统配合。

今天，不少分布式编程框架均可以和 Kubernetes 配合实现容错和弹性调度。比 如 用于离线数据处理的 Spark、用于在线数据处理的 Storm、在线 流数据引擎 Flink、分布式存储系统 Redis 和 HBase。其中适合深度学习的框 架有 Paddle EDL。基于 TensorFlow 的支持弹性调度的深度学习系统，据咱们 所知，ElasticDL 是第一个。

Kubernetes-native 的弹性调度

ElasticDL 经过实现一个 Kubernetes-native 的框架，调用 TensorFlow 2.0， 来实现弹性深度学习。

所谓 Kubernetes-native 指的是一个程序调用 Kubernetes API 来起止进程。 Google MapReduce 是一个 Borg-native 的分布式计算框架。用户经过运行一个 Borg 的客户端程度启动一个 MapReduce 做业。Borg 客户端调用 Borg API 提 交做业，而且启动一个 master 进程。这个 master 调用 Borg API 启动其余 workers 进程。ElasticDL 也相似，用户调用 ElasticDL 的命令行客户端程序 启动做业。这个客户端程序调用 Kubernetes API，启动 master 进程。master 进程继续调用 Kubernetes API 启动其余进程。master 进程也能够调用 Kubernetes API 监控其余进程。

若是 worker 挂了，按照分布式深度学习训练算法的数学特性，能够不用处理， 便可确保训练过程继续。若是一个 parameter server 进程挂了，master 会选 择一个 worker 进程，让它转换角色替补上挂掉的 parameter server 进程。在 以上两种状况下，master 都会调用 Kubernetes API，请它再启动一个额外的 worker 进程。若是启动成功，master 要带它入门，加入到与其余进程的协做中。 master 进程的状态（主要是三个 task queues：todo、doing、done）能够保留 在 Kubernetes 集群的 etcd 存储系统中。这样，万一 master 挂了，重启的 master 进程能够从 etcd 继承前世的状态。

以上是一个简化的描述。 ElasticDL 实现了多种分布式计算模式，每种模式实 现 fault-tolerance 的方式略有不一样。咱们会在后续文章中详细介绍。
Kubernetes-native 架构使得 master 进程有机会与 Kubernetes 协做实现容错 和弹性调度。不过，由于 ElasticDL 调用 Kubernetes API，也就意味着 ElasticDL 只能运行在 Kubernetes 上。

TensorFlow 原生的分布式计算能力不是 Kubernetes-native 的。因此 TensorFlow 不是绑定在 Kubernetes 这个平台上的。这是你们若是要用现有技 术在 Kubernetes 运行 TensorFlow 做业的话，须要依赖 Kubernetes 的扩展 Kubeflow 的缘由。

理论上，不调用 Kubernetes API 也是能够实现必定程度的容错的。即便没有 Kubernetes 的通知，master 能够经过检查其余继承的心跳（heartbeat）或者 检查 TCP 连接状态，判断其余进程的生死存亡。可是，不调用 Kubernetes API （或者其余调度系统的 API），master 没法通知调度系统重启进程，也没法得 知新启动的进程的信息，而且帮助它加入做业。这种“非 Kubernetes-native”的 容错方式颇为被动，只能接受资源紧张时一些进程被抢占而挂掉的事实，而不能 在其余做业释放资源后增长进程充分利用空闲资源。

TensorFlow 2.0

如上文解释，为了保证 TensorFlow 最核心的 runtime 是平台无关的，咱们没 法经过修改 runtime 实现完备的主动的容错和弹性调度。因此如文首的田字格 所示，ElasticDL 和 Uber Horovod 都是在 TensorFlow 的 API 上包一 层。

Horovod 基于 TensorFlow 1.x。 一个 Horovod 做业的每一个进程调用单机版 TensorFlow 作本地计算，而后收集 gradients，而且经过 AllReduce 调用汇聚 gradients 而且更新模型。Horovod 也是平台无关的，因此它提供的 AllReduce 操做不支持容错和弹性调度。这一点和 ElasticDL 不同。

和 ElasticDL 同样的是，Horovod 须要从 TensorFlow 偷偷“截获” gradients， 在 TensorFlow 1.x 中，深度学习计算是表示成一个计算图（graph），而且由 TensorFlow runtime 解释执行，因此 Horovod 为了得到每一个进程算的 gradients 而且 AllReduce 它们，就得 hack 进入图执行的过程。为此， Horovod 要求使用者使用特定的 optimizer 代替 TensorFlow 提供的 optimizer，从而能够在优化模型阶段透露出 gradients。

一个调用 Horovod 的用户程序的结构以下。其中标记为 () 和 (*) 的部 分是 Horovod 要求用户写的，帮助 Horovod 截获 TensorFlow 计算获得的 gradients 的代码。若是用户不慎忘记写了，那么程序执行结果就不对了。

ElasticDL 没有这些问题，由于它依赖的是 TensorFlow 2.0。TensorFlow 2.0 主推的 eager execution mode 采用和解释执行图彻底不一样的深度学习计算方式。 相似 PyTorch 的作法，前向计算过程把对基本计算单元（operator）的调用记 录在一个内存数据结构 tape 里，随后，反向计算过程（计算 gradients 的） 能够回溯这个 tape，以此调用 operator 对应的 gradient operator。这个 tape 提供一个操做让用户能够获取每一个参数的 gradient。

ElasticDL 经过调用 TensorFlow 2.0 API 能够很直接地获取 gradients：

并且上面这段代码不是须要用户写的，而是 ElasticDL 的一部分。ElasticDL 用户须要写的代码对应上述 Horovod 代码范例中的一行 —— 定义模型。

极简的 API 和使用方式
训练一个模型不仅须要上述模型定义，还须要指定数据、优化目标（cost）、和 优化算法（optimizer）。用户老是但愿能以尽可能精简的方式指定这些信息，以 尽可能少的代码描述训练做业。
ElasticDL 和 TensorFlow 其余的 high-level API，例如 Keras 和 TensorFlow Estimator 同样， 几乎调用一个 API 函数就能够执行一个分布式训练做业。下 面这个程序使用 Keras。Keras 使用 TensorFlow 原生分布式训练能力，不支持容 错和弹性调度。

ElasticDL 的 API 相对更加精简一些。上述范例程序对应的 ElasticDL 版本以下：

主要的区别在于：在 Keras 程序里用户要选择分布式执行策略；而在 ElasticDL 程序里则不须要。这是由于 ElasticDL 自动选择分布式训练算法和 策略。

简单的说，对于有很大参数（须要 model parallelism）的模型，ElasticDL 使 用 asynchrnous SGD。这个方法配合 delayed model update 能把网络通讯量减 少一个数量级。不少 NLP、搜索、推荐、广告的模型都符合这一类。 Asynchronous SGD 对于这类模型的表现比较稳定。对于图像识别和语音识别这 一类参数不太大的模型，ElasticDL 团队在开发一个 Kubernetes-native 的 AllReduce。和 Horovod 使用的 AllReduce 同样，ElasticDL AllReduce 把进 程间通讯的拓扑组织成一个环，从而实现高性能的模型更新。与之不一样的是， ElasticDL AllReduce 是容错的 —— 在有进程失败致使 AllReduce 调用失败的 状况下，master 组织剩下的活着的进程构造一个新的环。

ElasticDL 项目但愿经过这样的分而治之的策略，提供高性能而且易用的深度学习系统。

ElasticDL 和 SQLFlow 的关系

今年早些时候，王益团队 开源了 SQLFlow。用户能够 用扩展后的 SQL 语法，很是精炼地描述整个数据流和 AI 流程。

好比，若是咱们要为一个电子商务网站构造一个推荐系统，须要开发日志收集、 在线数据清洗、特征工程、模型训练，验证和预测等模块。每一个模块可能须要投 入一个团队数轴甚至数月的时间。

最近几年里，不少互联网服务开始把数据直接上传到通用数据库中，好比蚂蚁金 服的不少数据是在 ODPS（也就是阿里云上的 MaxCompute 服务）以及新一代的 智能数据系统 。这促使咱们考虑把数据清洗和预处理放在数据库中作，而特征工程、自动机器 学习、和训练过程在 ElasticDL 这样的 AI 引擎里作。SQLFlow 把扩展语法的 SQL 程序翻译成一个 Python 程序，把两部分连接起来。

在这样的场景中，若是 AI 须要不少参数，则用户也就须要在 SQL 程序中提供 这些参数。好比下面 SQL 语句从数据库中提取用户的年龄、工做部门、和工做 地点，来预测其收入。

其中，TRAIN 从句指定要训练的模型；COLUMN 从句指定如何把数据映射成 特征；LABEL 指定要预测的值；WITH 指定训练过程当中的各类参数，其中 dist_strategy 是调用 Keras/TensorFlow 作训练是须要指定的分布式策略， gpus 指定须要的资源。而这些，在 SQLFlow 调用 ElasticDL 的时候都是不 须要的，由于 ElasticDL 自动选择分布式策略和算法。

从这个例子能够看出，若是要让用户能提供尽可能少的参数，人工智能引擎还须要 更加智能，提供包括 AutoML 和 自动特征工程 的功能。 ElasticDL 项目任重道远。咱们期待把上述 SQL 程序简化为以下形式：

ElasticDL 项目的现状

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

目前 ElasticDL 实现的基于 parameter server 的分布式SGD 训练方法验证了 容错和弹性调度。而且在 Google Cloud 上的 Kubernetes 1.12 集群和阿里 Sigma 3.1（一个 Kubernetes 的高性能实现）上均可以运行。而且，ElasticDL 团队开发了 SQLFlow 生成 ElasticDL 程序的 code generator。

咱们但愿尽早开源 ElasticDL 和尽早分享其设计意图，能汇聚来自不一样公司和 社区的力量，一块儿探索 Google TensorFlow 2.0 和 Kubernetes 的分布式训练 生态，早日实现便捷的端到端的人工智能开发套件。

---

本文做者：缪克卢汉

阅读原文

本文为云栖社区原创内容，未经容许不得转载。