云原生的弹性 AI 训练系列之一：基于 AllReduce 的弹性分布式训练实践

引言

随着模型规模和数据量的不断增大，分布式训练已经成为了工业界主流的 AI 模型训练方式。基于 Kubernetes 的 Kubeflow 项目，能够很好地承载分布式训练的工作负载，业已成为了云原生 AI 领域的事实标准，在诸多企业内广泛落地。

尽管 Kubeflow 让基于 Kubernetes 的大规模分布式训练变得可行，但是云原生的极致弹性、降本增效等特性在人工智能场景下没有得到很好地释放。

为了解决目前在云原生 AI 场景下的成本高，资源利用率低等问题，TKE AI 团队在 Kubeflow 社区中推动了弹性训练特性的设计与实现。

在这一文章中，我们主要介绍了数据并行的分布式训练任务的弹性能力在 Kubernetes 上的设计与实现。并且通过实验的方式验证了特定的场景下，在保证训练精度的同时，这一特性能够使成本降低 70%。

背景

首先我们简要回顾一下深度学习的模型训练。这里所说的训练，指的是利用数据通过计算梯度下降的方式迭代地去优化神经网络的参数，最终输出网络模型的过程。在这个过程中，通常在迭代计算的环节，会借助 GPU 进行计算的加速。相比于 CPU 而言，可以达到 10-100 倍的加速效果。而分布式的模型训练，最早是由 Mu Li 在 OSDI'14 上提出的。在传统的模型训练中，迭代计算的过程只能利用当前进程所在主机上的所有硬件资源。但是单机的扩展性始终是有限的，在数据集规模特别大或者模型特别复杂的时候，单机训练的速度就会有些捉襟见肘。分布式的训练可以借助不同主机上的硬件资源进行训练加速，大大提高训练速度。

Horovod 是一款基于 AllReduce 的分布式训练框架。凭借其对 TensorFlow、PyTorch 等主流深度学习框架的支持，以及通信优化等特点，Horovod 被广泛应用于数据并行的训练中。在 Horovod 中，训练进程是平等的参与者，每个进程既负责梯度的分发，也负责具体的梯度计算。如下图所示，三个 Worker 中的梯度被均衡地划分为三份，通过 4 次通信，能够完成集群梯度的计算和同步。

依托 AllReduce 的分布式训练由于其简单易懂的编程逻辑和大幅提升的训练速度，逐渐成为分布式训练的主流方式。然而，当前这种模式依然存在一些问题：

• 首先，AI 训练的成本问题显著。借助于 Kubernetes，大规模分布式训练虽然已经不再复杂，但是高昂的训练成本使得这项技术难以真正做到普惠。
• 其次，相比于单机的训练，分布式训练有更大的可能出现任务失败的情况。在分布式训练中，有多个进程同时参与训练，而其中的某个进程出现了问题，整个训练任务都会因此而失败。尤其是当训练任务需要持续几天甚至几个礼拜时，这个问题就会显得尤为严重。
• 同时，由于一些混部的集群存在业务压力周期性波动的特性，在闲时 GPU 占用率通常不到 40%。但是与之相对的是，在任务密集提交时，集群的资源又会出现紧张的情况。资源利用在时间上的不均衡问题非常突出。

弹性训练

为了解决上述问题，更好地向分布式训练释放云原生的红利，业界提出了弹性训练这一概念。

在传统的深度学习分布式训练任务中，通常任务的实例配置是固定的。这很大程度上限制了任务的灵活性和训练速度，对于整个集群的资源利用率而言也不友好。而弹性训练，就是指让训练任务能够在运行时动态地调整参与计算的实例数量。这使得训练更加灵活，同时可以配合集群的负载进行更好的扩缩容和调度。这一特性为训练场景带来了诸多收益：

• 容错性的提升。在这样的选型下，所有实例的失败都是可以容忍的。任务不再会因为某个进程出错而导致任务整体的失败。
• 资源利用率的提升。在集群资源紧张时，通过减少低优先级训练任务的实例数量，能够保证高优先级训练任务的资源配额，