阿里云 ACK 云原生 AI 套件中的分布式弹性训练实践

众所周知，随着时间的推移，算力成为了 AI 行业演进一个不可或缺的因素。在数据量日益庞大、模型体量不断增加的今天，企业对分布式算力和模型训练效率的需求成为了首要的任务。如何更好的、更高效率的以及更具性价比的利用算力，使用更低的成本来训练不断的迭代 AI 模型，变成了迫切需要解决的问题。而分布式训练的演进很好的体现了 AI 模型发展的过程。

Distributed Training

分布式训练一般分为两种类型，数据并行和模型并行。数据并行是指每个训练 Worker 都保存有一份模型的副本，然后将大规模的数据进行切分，分布到每个训练 Worker 上进行计算，最后再进行集合通信统一计算结果的过程。在相当一段的时间内，该种模式是分布式训练的主流模式，因为模型的规模足以放进单个训练 Worker 之中，而数据的规模才是整体训练效率的瓶颈。利用分布式数据并行可以充分利用集群中的算力资源，并行处理庞大的数据集，以达到加速训练的效果。

而模型并行则是在模型非常庞大的情况下将模型进行切分，分布到不同的训练 Worker 中，然后训练数据按照模型的结构分布经过不同的训练 Worker，以达到用分布式的算力来训练大模型的效果。现在的大语言模型由于其体量的庞大，所以一般都是使用模型并行的模式来进行训练。

基于数据并行的分布式训练又分为两种不同的架构。

• Parameter Server 架构：
  有一个中心化的 Parameter Server 用来存储分布式训练的梯度，每一个训练 Worker 在进行每个 Step 的训练前都需要先从 Parameter Server 中先 pull 到最新的梯度信息，在这个 step 训练结束后再将训练的结果梯度 push 到 Parameter Server。在 Tensorflow 中，PS 模式的训练一般为异步的分布式训练，该情况下对于任务的全部 Worker 来说，其无需等待其他 Worker 的训练流程去同步梯度，只需要根据流程完成自身的训练即可。这种模式多用于基于 Tensorflow 的搜推广场景。
• AllReduce 架构：一个去中心化的同步的分布式训练模式，在分布式训练中一般采用 Ring-All Reduce，每个训练 Worker 只与自身左右 Rank 的 Worker 进行通信，这样就能形成一个通信环，经过环形通信可以使得每一个 Worker 中的梯度都同步到了其他的 Worker 中并完成计算。这种模式多用于 CV、NLP 的应用场景。

Elastic Training

上面我们介绍了分布式训练，现在来了解一下弹性分布式训练。什么是弹性训练？具体可以总结为三大块的能力：

• 训练规模弹性改变：这里主要指的是弹性改变训练的 Worker 数目，扩容增加 Worker 数量以提升训练速度，缩容减少 Worker 数量以腾出部分集群资源；
• 训练过程弹性容错：由于部分因素导致任务异常或可预见问题如 Spot 回收事件预示下的整体任务的容错，避免因为少部分 Worker 失败而直接导致的整个任务的失败；