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

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

数据并行分布式训练

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

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

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

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

  • 在发现某些节点失联后自动剔除这些节点;
  • 在更高级别的任务需要使用资源时缩减 worker 数量以释放资源;
  • 在资源存在空置时自动扩展。

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

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

才云分布式训练方案 FTLib

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

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

  • 针对故障 Container 的自动重启;
  • 根据多维度指标对 Pod 数量进行自动的伸缩。

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

为了向不同用户提供不同级别的 API,避免对训练框架的侵入式修改,我们将 FTLib 作为一个库引入到 Python 中,因此,它可以针对不同需求提供不同的 API。针对大多数进行数据并行分布式训练的用户,FTLib 提供接口,将梯度的同步统一到 wait_gradients_ready API:

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

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

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

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