才云开源云原生分布式训练项目 FTLib：支持弹性伸缩、自动容错

FTLib 是才云科技开源的支持弹性伸缩和自动容错的云原生分布式训练解决方案，目前正计划落地各类生产环境，帮助企业应对因训练数据量激增产生的意外情况，真正部署和运行大规模的分布式训练。

近年来，随着数据的大量累积和深度学习的不断发展，无论是训练数据还是模型，它们在体量上都增长迅速，而单节点的算力提升却变得愈发昂贵。面对这种情况，分布式模型训练应运而生。

所谓分布式深度学习训练，即按照并行的方式把训练过程分为模型并行和数据并行。模型并行针对庞大的神经网络设计，将一个神经网络拆分成多个部分，分置于不同的计算节点；同步更新的数据并行则是将 mini batch 拆分成若干份，利用各个计算节点上完整的同一份模型执行前向和反向传播，在更新梯度前先同步梯度。

数据并行分布式训练

随着人工智能的发展，数据并行的分布式训练在社区上涌现了许多优秀的工具，如 Horovod，各个深度学习框架也在不断完善自身对数据并行模式的支持。然而，随着训练时长增加和训练节点增多，两类问题渐渐浮现。

问题一：同步更新使得计算和通信与每次梯度同步绑定。一旦出现设备、算法或算法实现、网络等意外，整个分布式训练任务就会中断。随着计算节点的增加，这种意外的发生也越来越频繁。

问题二：随着模型变得更加复杂，收敛速度下降，分布式训练任务往往需要长时间占用宝贵的 GPU 资源，而诸如模型推理等任务在使用 GPU 时，会呈现一定的周期性特征，存在使用峰值和谷值。这意味着在当前模式下，很多用户不得不一方面为训练任务增加 GPU 采购，另一方面却面临由 GPU 空置带来的巨大浪费（无法动态地调度资源供其他任务使用）。

可以想象，如果一个数据并行的分布式训练任务可以随意调节其 worker 的数量：

• 在发现某些节点失联后自动剔除这些节点；

• 在更高级别的任务需要使用资源时缩减 worker 数量以释放资源；

• 在资源存在空置时自动扩展。

那么上述问题即可迎刃而解。

随着用户训练的数据量不断增大，一个可预见的事实是用户对上述可容错、可伸缩的分布式训练的需求会愈发强烈。在我们与广大开发者和客户的接触中，我们也深切感受到了他们对更细粒度划分调度 GPU 资源、减少资源浪费的渴求。

才云分布式训练方案 FTLib

才云一直致力于优化分布式训练在 Kubernetes 上的使用体验。继 tf-operator、Katib（了解详情） 之后，我们对可容错、可伸缩的分布式训练进行了探索和尝试。

Kubernetes 作为容器编排平台，对应用容错和弹性伸缩已经有了成熟的实现：

• 针对故障 Container 的自动重启；

• 根据多维度指标对 Pod 数量进行自动的伸缩。

为了进一步拉近分布式训练和 Kubernetes 之间的距离，才云推出了 FTLib 这一项目，为 AI 分布式训练提供可容错、可伸缩的基本功能。

为了向不同用户提供不同级别的 API，避免对训练框架的侵入式修改，我们将 FTLib 作为一个库引入到 Python 中，因此，它可以针对不同需求提供不同的 API。

针对大多数进行数据并行分布式训练的用户，FTLib 提供接口，将梯度的同步统一到 wait_gradients_ready API：

然而，上面提到的 wait_gradients_ready 对大多数算法工程师而言存在学习成本。他们需要先了解数据并行分布式训练和同步更新策略，才能理解为什么在执行梯度更新之前，需要等待 wait_gradients_ready 返回成功。

另一方面，这种后向传播、同步梯度、全局梯度更新依次进行的方式会使计算和网络通信无法同时进行，易造成 GPU 在等待中闲置。

为了解决这两类问题，FTLib 引入了使用更为广泛的 execute 接口：

它使用起来也比较贴近原有代码：

与此同时，我们也保留了一些底层 API 的暴露，为那些更偏向工程的用户提供接口。

FTLib 架构详解

FTLib 旨在为分布式训练带来可容错和弹性伸缩等功能。当前，我们已经对接了 PyTorch 训练框架，并计划支持更多训练框架。

为了实现这个目标，在设计 FTLib 时我们采用了一个相对松散的架构以便组件更新，如下图所示：

FTLib 由两个主要组件构成：

• 一致性协议（Consensus Protocol）

• 通信库（Communication Library）

它们通过一个“秩分配方案”（Rank Assign Scheme）将一致性协议确认的“成员列表”（member list）转换成秩和总成员数。

当前 FTLib 采用 Gossip 作为其一致性协议来传播 worker 的加入和退出。社区上也有一些支持弹性伸缩的方案，如 PyTorch/Elastic。和这些方案相比，FTLib 的优势之处在于不依赖外部组件（如 Zookeeper 或 etcd）即可确保所有 worker 对于 member list 的一致性确认。

这样的设计一方面有利于用户开箱即用：无需事先安装外部组件，用户就可以在 Kubernetes 集群或其他类型的集群中立刻使用。另一方面，它将一致性协议嵌入到每一个训练的 worker 本身，可以避免因这些外部组件失效而导致的训练中断，更大程度上确保了整体的可容错性。

当前，我们基于 Gossip 对于 member list 的实现，采用了 hashicorp 的 memberlist 模块（github.com/hashicorp/memberlist）。考虑到 worker 退出会直接导致梯度同步失败，因此在后续优化中，我们会降低日常用以确认 liveness 的频率，减少对资源的消耗。

在分布式训练中，通信库的作用主要是在 worker 之间传递梯度或权重。

考虑到不同训练框架往往都有自定义的数据结构用来实现高效的内存、显存管理，且主流训练框架基本都已具备一套高度定制化的通信库，如 pytorch.distributed、tf.distribute 等，我们认为没有必要再去实现或包装一个通信库。FTLib 会在用户选择了对应的训练框架后，对对应框架的通信库做二次封装，以实现可容错功能。

这样轻量级、兼容性好的设计特征使 FTLib 非常适合当下的分布式训练。

未来展望

事实上，真正实现可容错、具备弹性伸缩的分布式训练是一个相当复杂的工程，目前才云的 FTLib 只是提供了一条轻量化的、可高度定制的道路，我们还需要对许多组件进行调整，例如：

• 当 member list 发生变动时，如何对当前 epoch 下尚未处理的数据进行再分配；

• 当 worker 数变动时，如何调节 batch size 或是 learning rate，使得整个训练的收敛不受到过多影响。

除了训练本身，如何在集群中利用可容错、可弹性伸缩的分布式训练，在调度层面也有许多复杂的问题亟待解决。

幸运的是，才云并不是单打独斗。

在支持弹性伸缩的分布式训练上，我们已经和蚂蚁金服的 ElasticDL 团队展开深度合作。此前他们曾推出基于 PS/Worker 模式的弹性伸缩方案 ElasticDL（elasticdl.org），加上才云团队在 AI 上积累的创新实践，双方强强联手将为 FTLib 的未来提供无限可能。

未来，我们会和蚂蚁金服 ElasticDL 团队一起为基于 AllReduce 的弹性伸缩方案努力，下面是双方团队在短期内的一些规划：

• 支持除 PyTorch 外的更多深度学习框架；

• 持续优化基于 Gossip 的 member list 方案；

• 探索动态 batch size 下收敛稳定性的各种方案。

FTLib 致力于拉近分布式训练和 Kubernetes 的距离，为企业和开发者个人在云上高效开发智能模型、实现 AI 普惠化打造基础。

我们真诚地希望能有更多社区开发者可以加入到这个项目中来，一起打造 FTLib，以分布式训练改变人工智能和互联网的未来！项目地址。

本文转载自 才云 Caicloud 公众号。

原文链接：https://mp.weixin.qq.com/s/T3qIAGNLEUiWbeJZb5IBfQ