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 并不能无中生有。不能容错，也意味着不 能弹性调度。

对弹性调度的诉求

在很多人共用计算集群的情况下，支持弹性调度意味着极大提升团队效率和集群 的总体利用率。前者支持快速迭代以保持技术领先