基于 MPI Operator 的弹性训练与实践

背景
弹性训练
弹性 Horovod
Kubernetes 上的弹性 Horovod
演示
结论

随着数据集和深度学习模型规模的增加，分布式训练成为业界训练神经网络模型的主流方法。虽然现在使用 Kubeflow 在 Kubernetes 上启动大规模分布式训练作业是可行的，但当我们谈论 Kubernetes 上的深度学习作业时，弹性工作负载等高级功能以及其他成本缓解方法仍受到限制。

为了解决成本和资源利用率方面的问题，TKE (腾讯云容器服务) AI 团队在 Kubeflow 社区设计并实现了 弹性训练。

本文我们将介绍弹性训练是如何在 Kubernetes 上实现的。通过特定环境下的实验验证，弹性训练降低了云上分布式训练的成本。

背景

让我们先回顾一下深度学习模型的训练。当我们谈论“训练”时，通常是指通过梯度下降法迭代优化神经网络模型中的参数。通过 GPU 加速，训练速度可以提高 10 到 100 倍。

当厂商试图将更多计算资源（如 GPU）集成到单台机器中，以使用更多数据或模型参数进行训练实验时，成本呈指数级增长。因此，在由 Mu Li 在 OSDI'14 上首次提出后，当研究人员处理海量数据集或大型模型时，分布式训练取代了单机训练。

对于数据并行分布式训练，Horovod 因其对 TensorFlow 和 PyTorch 等深度学习框架的优秀支持、通信优化和更简单的编程模式而被广泛采用。在 Horovod 中，所有训练进程都是平等的参与者，每个进程都负责梯度计算和通信。

alt_text

由于训练速度的显著提升以及更易于理解的编程模式，以 Horovod 为代表的数据并行分布式训练正受到越来越多的关注。然而，仍然存在一些问题

云上训练的成本仍然是一个障碍。虽然得益于 Kubernetes 和 Kubeflow，研究人员在云上训练时不再面临复杂性，但云上训练的成本让一些用户望而却步。
与单机训练相比，多节点分布式训练增加了训练失败的概率。当任何一个训练进程出错时，整个训练实验就会失败。当某些实验需要几天甚至几周的时间时，这个问题变得更加严重。
当训练任务与（具有更高优先级的）其他工作负载共置时，随着这些其他工作负载的请求可能周期性地变化，资源需求也会波动。这种资源可用性的不平衡为使用混合部署最大化资源利用率的想法泼了冷水。

弹性训练

研究人员和工程师提出了 弹性训练 作为解决这个难题的关键。

传统上，分布式训练作业的资源配置是固定的。弹性训练打破了这一规则，使用户能够更改参与分布式训练作业的实例数量，为运行分布式训练作业的集群带来以下好处

容错性（Fault Tolerance）：只要至少有一个 worker 实例存活，任何 worker 实例都可以失败。
资源利用率（Resources Utilization）：当资源压力堆积时，集群能够减少低优先级工作负载（分布式训练工作负载）的副本，将资源释放给其他工作负载（如预测服务），以确保业务的 SLA；工作负载释放资源后，弹性训练作业能够通过扩缩工作负载副本吸收这些资源。
云上训练（Training on Cloud）：云上有一种资源被称为“Spot”或“抢占式”实例；它的价格非常低，但可能在保证时长到期后被回收。

弹性训练似乎与公有云完美契合。结合 Spot 实例，我们将 GPU 的成本从每小时 ¥16.21 降低到每小时 ¥1.62，将训练作业的总成本降低了近 70%。在相同预算下，弹性训练可以利用更多 GPU，并将训练速度提高 5 到 10 倍。

弹性 Horovod

作为分布式训练框架的主要参与者，Horovod v0.20.0 提供了其弹性训练解决方案，即弹性 Horovod。这里我们引用了 RFC Elastic Horovod 中关于弹性 Horovod 和现有 Horovod 的架构差异描述

alt_text

所有集合操作都在 hvd.elastic.run 函数中进行协调。
在用户训练函数执行之前，worker 之间会同步状态。
worker 失败或 worker 添加事件将导致其他 worker 上发生重置事件。
重置事件充当屏障，用于
- 根据 worker 退出码确定作业是否应继续。
- 将失败的主机列入黑名单。
- 在新主机上启动 worker。
- 更新现有 worker 上的 rank 信息。
重置事件后同步状态。

启动弹性训练作业时，horovodrun 需要一个 discover_hosts.sh 脚本来实时检测可用的主机和 slot。在下一节中，我们将此脚本称为 discover_hosts.sh。尽管如此，该脚本不必命名为 discover_hosts.sh。一个 discover_hosts.sh 脚本的示例可以在此处找到。

Kubernetes 上的弹性 Horovod

MPI-Operator 旨在在 Kubernetes 上部署 Horovod 作业。尽管该 Operator 发布了多个版本，但其核心思想保持不变。它包括

alt_text

MPIJob Controller 根据 MPIJob 中的副本配置创建 launcher pod 和 worker pod
对于每个 MPIJob，controller 创建一个 ConfigMap，其中包含两个文件：hostfile 和 kubexec.sh
所有 worker pod 就绪后，launcher pod 上的 mpirun（被授予 pod/exec 权限）使用 kubexec.sh 在 worker pod 上启动进程

启动弹性 Horovod 作业不可行，因为弹性 Horovod 和 MPIJob Controller 之间存在一些不兼容性。我们以 controller-v1 为例

launcher pod 上没有内置的 discover_hosts.sh 可用
worker 副本数量减少后，controller 不会删除不再需要的 worker pod，导致分布式训练的规模保持不变
worker 副本数量增加后，controller 不会更新绑定到 launcher pod 的 Role 中的规则，阻止 launcher pod 在新创建的 pod 上创建进程

为了解决这些兼容性问题，我们提交了多个关于 Horovod 和 MPI-Operator 的 pull request，包括 mpi-operator-pull-335 和 horovod-pull-2199。由于为弹性 Horovod 的 launcher pod 提供特定于 MPI-Operator 的 discover_hosts.sh 是最关键的，我们考虑了两种将 Running 状态的 worker pod 转换为 discover_hosts.sh 脚本的场景。

由 MPIJob controller 组合并通过 ConfigMap 同步到 launcher pod 的动态 discover_hosts.sh
- MPIJob controller 有一个 podLister，可用于轻松列出 worker pod
- controller 过滤 status.phase == Running 的 worker pod，并将结果编码到 discover_hosts.sh 中
- 当 discover_hosts.sh 被修改时，ConfigMap 会更新，并且更改将由 Kubernetes 传播到 launcher pod
在 launcher pod 中使用 kubectl 从 APIServer 列出所有正在运行的 worker pod 的静态 discover_hosts.sh

场景 2 更改的是交付镜像而不是 controller。然而，由于我们无法限制用户执行 discover_hosts.sh 脚本的频率，这可能对 APIServer 造成潜在威胁，尤其是在 worker pod 数量巨大时。

场景 2 的一个修正方法是将 kubectl 替换为一个 podLister 进程，从而减轻 APIServer 的额外压力。这样一来，我们在 launcher pod 中安装了两个进程，但缺乏一种合适的机制来保持 podLister 存活。一旦 podLister 死亡，训练作业的弹性就不复存在了。

因此，我们选择了第一种场景，并将 disocver_hosts.sh 映射到 /etc/mpi/ 下。我们还修复了 worker 副本配置更改后的其他兼容性问题。对于选择非弹性模式的用户，只需忽略 /etc/mpi/discover_hosts.sh 即可。

场景 1 也存在一些担忧。ConfigMap 与 horovodrun 在 launcher pod 中从 discover_hosts.sh 看到的内容之间存在延迟。这种延迟一方面可以由集群管理员调整，另一方面与训练耗时或弹性 Horovod 处理 worker 变化的时间相比，可以忽略不计。

演示

我们将展示一个演示，说明如何使用 MPI Operator 操作弹性 Horovod 作业。

bash-5.0$ kubectl create -f ./tensorflow-mnist-elastic.yaml
mpijob.kubeflow.org/tensorflow-mnist-elastic 
createdbash-5.0$ kubectl get po
NAME    READY   STATUS    RESTARTS  AGE
tensorflow-mnist-elastic-launcher   1/1     Running   0          14s
tensorflow-mnist-elastic-worker-0   1/1     Running   0          14s
tensorflow-mnist-elastic-worker-1   1/1     Running   0          14s

作业最初创建时包含两个 worker。训练开始后，我们将 MPIJob.Spec.MPIReplicaSpecs["Worker"].Replicas 更改为 3，增加一个 worker。让我们看看 discover_hosts.sh 如何变化

bash-5.0$ kubectl exec tensorflow-mnist-elastic-launcher -- /etc/mpi/discover_hosts.sh
tensorflow-mnist-elastic-worker-0:1
tensorflow-mnist-elastic-worker-1:1
bash-5.0$ cat ./patch_r3.yaml
spec:
  mpiReplicaSpecs:
    "Worker":
      replicas: 3
bash-5.0$ kubectl patch mpijob tensorflow-mnist-elastic --patch "$(cat patch_r3.yaml)" --type=merge
mpijob.kubeflow.org/tensorflow-mnist-elastic patched
bash-5.0$ kubectl exec tensorflow-mnist-elastic-launcher -- /etc/mpi/discover_hosts.sh
tensorflow-mnist-elastic-worker-0:1
tensorflow-mnist-elastic-worker-1:1
tensorflow-mnist-elastic-worker-2:1

我们将副本数量减少到 1，回收 2 个 worker 实例。

bash-5.0$ cat ./patch_r1.yaml
spec:
  mpiReplicaSpecs:
    "Worker":
      replicas: 1
bash-5.0$ kubectl patch mpijob tensorflow-mnist-elastic --patch "$(cat patch_r1.yaml)" --type=merge
mpijob.kubeflow.org/tensorflow-mnist-elastic patched
bash-5.0$ kubectl get po
NAME               READY   STATUS        RESTARTS   AGE
tensorflow-mnist-elastic-launcher   1/1     Running       0          4m48s
tensorflow-mnist-elastic-worker-0   1/1     Running       0          4m48s
tensorflow-mnist-elastic-worker-1   1/1     Terminating   0          4m48s
tensorflow-mnist-elastic-worker-2   1/1     Terminating   0          2m21s

弹性训练持续进行。

Thu Mar 11 01:53:18 2021[1]<stdout>:Step #40    Loss: 0.284265
Thu Mar 11 01:53:18 2021[0]<stdout>:Step #40    Loss: 0.259497
Thu Mar 11 01:53:18 2021[2]<stdout>:Step #40    Loss: 0.229993
Thu Mar 11 01:54:27 2021[2]<stderr>:command terminated with exit code 137
Process 2 exit with status code 137.
Thu Mar 11 01:54:27 2021[0]<stderr>:command terminated with exit code 137
Process 0 exit with status code 137.
Thu Mar 11 01:54:57 2021[1]<stderr>:[2021-03-11 01:54:57.532928: E /tmp/pip-install-2jy0u7mn/horovod/horovod/common/operations.cc:525] Horovod background loop uncaught exception: [/tmp/pip-install-2jy0u7mn/horovod/third_party/compatible_gloo/gloo/transport/tcp/pair.cc:575] Connection closed by peer [10.244.2.27]:54432
WARNING:root:blacklist failing host: tensorflow-mnist-elastic-worker-2
WARNING:root:blacklist failing host: tensorflow-mnist-elastic-worker-1
Thu Mar 11 01:54:58 2021[1]<stdout>:Step #50    Loss: 0.207741
Thu Mar 11 01:55:00 2021[1]<stdout>:Step #60    Loss: 0.119361
Thu Mar 11 01:55:02 2021[1]<stdout>:Step #70    Loss: 0.131966

如我们所见，MPI-Operator 上的弹性 Horovod 现在支持动态调整 worker 副本。作为未来的工作，我们计划支持 Horizontal Pod Autoscaler 用于 MPIJob，以及指定 worker 删除等其他功能。

结论

当云原生和分布式训练的概念融合到 Kubernetes 上的弹性训练时，它降低了成本并提供了鲁棒性和灵活性。作为一个团队，我们正与 PyTorch、Horovod 和其他社区合作，推动弹性训练。我们希望进一步介绍我们在 PS/Worker 训练模式的弹性、资源和作业优先级优化以及云原生 AI 其他主题方面的工作。