弹性分布式训练在虎牙的实践

分享嘉宾：刘柏芳 虎牙

编辑整理：吴春梅 格灵深瞳

出品平台：DataFunTalk

导读：本次分享题目为弹性分布式训练在虎牙的实践，主要包括以下几方面内容：

为什么要弹性

弹性分布式训练的设计

落地效果

未来展望

首先介绍下虎牙AI平台的发展历程。

虎牙AI平台在2019年前还处于混沌状态：AI开发各自为政、硬件资源不共享、技术栈不统一。为了解决这些问题，2019年我们开始虎牙AI平台的开发，基于k8s云原生平台实现资源统一调度，统一开发、训练、推理流程。从2021年开始，平台逐步走向降本增效，通过提升任务调度编排效率和AI框架层面的优化，提升了训练效率以及资源利用率，降低任务队列等待时长。从2022年虎牙AI平台开始整体转向更加精细化运营，比如AI CI/CD、训练过程可视量化跟踪等。

本次分享的弹性分布式训练，就是在降本增效阶段提出的一个解决方案。接下来详细介绍下为什么要弹性。

我们面对的主要问题包括：

首先虎牙直播平台有一个明显的高低峰流量趋势：从天的角度来看，晚上是高峰期、凌晨白天是低峰期；从周的角度看，周末是高峰期，工作日是低峰期。因为推理应用高低峰的存在，GPU使用也存在很明显的高低峰，在低峰期很多资源是闲置的。

第二个问题是碎片化GPU资源。因为我们调度分配机器给AI算法同事，训练任务的启动和停止完全由算法同事主观来决定。比如有两台8卡机器，每台机器有一张空闲显卡，即总共有2张空闲显卡，假设有一个2卡待训练任务，但因为我们是机器级别分配显卡资源，所以这个2卡的待训练任务也无法使用这2张分布在不同机器上的空闲显卡。如果待训练任务排队时间较长，平台侧可能就会被迫增加机器，进而增加了成本。

第三个问题是如果机器异常宕机，必须要算法同事人工参与才能继续训练，除了耽误时间，也可能会丢失一些训练状态，给算法同事带来不小的麻烦。

为了方便大家理解，我们举例来进一步说明为什么要弹性。

假设任务A正在节点1的两张显卡上执行训练任务，任务B和C在节点2的两张显卡上执行训练任务。一段时间后，任务B和C相继结束。在过去传统训练集群里，如果没有新任务来，节点2的两张显卡将一直是空闲浪费状态。但是如果平台支持弹性分布式训练，那么可以将任务A动态弹性扩容到节点2上，即总共4张显卡上加速训练。一段时间后，更高优先级的任务D来临，任务A可以先缩容自己到最小的2卡，空出节点2给任务D。另外一种情况就是，如果任务A处于4卡训练时，节点1突然宕机，在过去传统训练平台里，任务A就会报错退出中断训练。但是如果支持弹性分布式训练，平台就可以将任务A自动缩至节点2继续运行，训练任务不会被中断。另外整个训练调度过程的扩缩对算法同事都是无感的，不需要他们人工参与。

理解了为什么需要弹性，接下来我们就讨论怎么来设计弹性分布式训练。弹性分布式训练与传统分布式训练的核心改变是当节点扩缩时，每个节点需要能感知到变化，然后重新建立通信，恢复之前训练状态继续训练。

接下来我们演示下这个训练过程：

我们通过ETCD的注册、选举和watch监听功能来感知集群节点的状态变化。每个节点在启动时首先将IP、端口号、GPU卡等信息注册到ETCD，注册完后，每个节点都可以从ETCD获取训练任务相关的所有其他节点信息。

假设某个训练任务初始阶段有3个节点：

1. 启动时，首先每个节点将自己的信息注册到ETCD上去，包括IP地址、端口号、GPU卡信息等等，注册完成之后，节点再从ETCD获取所有跟它相同任务的其它节点的信息，这样每个节点都互相知道相同任务所有节点的信息；

2. 然后利用ETCD进行选举每个节点的角色，比如节点1的角色是rank0，节点2的角色是rank1，节点3的角色是rank2；

3. 有了这些信息后，接下来就可以利用传统的分布式训练方式对所有节点建立通讯，比如通过RingAllReduce建立环状通讯等，建立通讯之后就可以进行分布式训练了；

4. 训练进行一段时间后，新节点4加入到集群，它首先也是将自己注册到ETCD；

5. 其它节点监听到有新节点加入进来，跑完各自当前训练step后，会暂停训练，重新从ETCD获取最新的节点信息；

6. 重复执行步骤2、3，因为节点4新加入，会有一个状态同步，同步后训练任务将会在4个节点上继续执行。

以上步骤是资源扩容的情况，缩容的步骤与其类似。

接下来介绍平台整体模块设计架构图：

平台建立在K8S集群基础之上，通过自定义的k8s operator进行训练任务启停操作。具体训练进程在k8s pod里执行，其中Rendezvous是训练进程的一个核心组件，它负责与ETCD交互，注册容器IP、端口号、GPU卡等信息，并且从ETCD同步获取训练任务的所有节点信息。

然后通过选举赋予每张显卡的角色：rank_0、rank_1……rank_n。接下来通过launcher为每张显卡启动一个训练进程。此时就可以开始分布式训练了。在训练过程中，如果某张显卡上的训练进程因为内存溢出等原因崩溃，其它训练进程可以实时捕获到此异常(通信崩溃)，暂停训练进程，从ETCD获取训练任务的最新所有节点信息，然后重新建立通讯并继续训练任务。

另外平台还有一个组件叫Remote Cache，它缓存训练过程中的一些中间数据，例如训练参数、超参数等。在某些场景下，比如有一个低优先级训练任务使用的全部都GPU被更高优先级任务抢占，此时低优先级训练任务是挂起状态。当有空闲资源时，之前被挂起的低优先级任务会被重新调度，平台会从Remote Cache读取缓存的数据，从而恢复之前被中断的低优先级训练任务的训练状态，继续训练。

平台设计完后，我们进行了大量测试验证效果是否符合预期。下面是我们关心的两个核心指标：

精度：确保弹性分布式训练对算法精度的影响在可接受范围内

训练耗时：期望弹性分布式训练利用闲置可以缩短训练时长

我们测试了模型Resnet50在ImageNet数据集上的表现。下图是测试结果：

其中EFDL弹性数据为在实际集群中不同时间段分别跑了4次的数据。从表格可以看出弹性分布式训练和单机多卡训练相比，精度和总卡时接近，符合预期。从卡数使用趋势图可以看到弹性分布式训练可以充分利用闲置资源，一开始如果集群资源不足，可以利用比较少的卡数将训练任务先启动起来，然后等有空闲资源时，可以动态使用更多显卡以加快训练速度。

测试符合预期，接下来就是实际落地给算法同学来使用了：

算法同学将传统训练任务改成弹性分布式训练任务时，只需要改大概十几行代码。上图提到引入了EFDL包，它就是我们实现的弹性分布式训练框架。EFDL框架适配了虎牙内部常用的几种训练框架：支持Pytorch DDP分布式训练、Pytorch Horovod分布式训练；支持Tensorflow Horovod分布式训练，包括Keras、Estimator这些高层API。

代码修改完后，算法同学在AI平台上新增一个训练任务时，他们只需要修改训练模式为EFDL弹性训练，并设置最小和最大worker数。

下面分享一下我们弹性分布式训练的落地效果。

对于算法同学来说最大的收益是缩短训练时长。平台通过弹性利用低峰空闲GPU资源以及碎片化资源，可无感的成倍的缩短部分训练时间。同时也可以避免因为机器宕机对训练造成的中断。

除了对业务模型训练带来收益，平台的收益也很显著。传统的训练调度策略导致排队等待时间太长，被迫加机器，造成成本浪费。而弹性训练调度策略可以通过更加灵活的调度策略，更易于基于优先级进行全公司的训练任务的编排，从而最大限度的压榨GPU资源池。最终达到降低成本的目的。

最后介绍一下未来展望。

更易用，目前如果将传统训练改成弹性分布式训练，大概需要修改十多行代码，并且还需要对弹性有一定认知。为了降低门槛，我们正在迭代重构一个新的版本，新版本需要算法同学改动的代码会更少，并且也不需要了解弹性概念。

更稳定，这也是每个平台的基本要求。努力做到运行更稳定，故障可转移。

更高效，提升分布式训练的效率是我们不断持续追求的目标。

更开放，平台在弹性分布式训练框架、提升分配效率等方面都借鉴了大量的开源代码和设计思路。因此我们也希望通过开源平台来回馈社区，一起分享、一起学习，一起进步，和社区一起将平台做得更好。