提高 AI 训练算力效率：蚂蚁 DLRover 故障自愈技术的创新实践

2024-04-09
浙江
本文字数：2780 字
阅读完需：约 9 分钟

本文来自蚂蚁 DLRover 开源负责人王勤龙（花名长凡）在 2024 全球开发者先锋大会(GDC)的分享——《DLRover 训练故障自愈：大幅提升大规模 AI 训练的算力效率》。

王勤龙，长期在蚂蚁从事 AI 基础设施的研发，主导了蚂蚁分布式训练的弹性容错与自动扩缩容项目的建设。先后参与多个开源项目，如 ElasticDL 和 DLRover，开放原子基金会 2023 活力开源贡献者，蚂蚁集团 2022 T-Star 优秀工程师。目前为蚂蚁 AI Infra 开源项目 DLRover 的架构师，专注于打造稳定、可扩展和高效的大规模分布式训练系统。

当前大规模语言模型训练需要大量的加速卡来训练，如 GPU 等。由于 GPU 机器的故障率较高，频繁的故障会导致训练中断、计算浪费和集群空转，从而造成大量的时间和算力浪费。为此，DLRover 开源了训练故障自愈技术，长凡在演讲中介绍了 DLRover 如何通过快速的节点状态检测、弹性扩缩容、动态组网和 Flash Checkpoint 等技术，最大程度地降低故障导致的算力浪费。

大规模 AI 分布式训练面临的挑战

随着大模型从十亿参数迈向万亿参数，训练规模的增长导致了集群成本的飙升，同时也影响了系统的稳定性。如此规模的系统带来的高额运维成本，成为了大模型训练过程中亟需解决的问题。

参数量大：大模型参数迈向万亿，如何在庞大的计算资源上实现高效的模型训练？
成本高：大规模集群需要更多的硬件资源，导致成本显著增加，如何降低 AI 训练成本？
稳定性：随着集群规模的扩大，单节点故障率不变的情况下，整体故障概率增加，如何提高分布式训练的稳定性？

蚂蚁 AI 训练工程实践

上图展示了蚂蚁 AI 训练的工程技术栈，分布式训练引擎 DLRover 支持了蚂蚁的对话、代码、视频和多模态等模型的多样训练任务。以下是 DLRover 提供的主要功能：

训练故障自愈：将千卡分布式训练有效时间占比 >97%，降低大规模训练故障的算力成本；
训练优化 ATorch：自动根据模型和硬件选择最优的分布式训练策略。将千卡(A100)集群硬件算力利用率 >60% ；
训练优化器：相比 AdamW 提升 1.5x 的收敛加速。相关成果分别发表在 ECML PKDD ’21 ，KDD’23，NeurIPS ’23；
显存与传输优化 GLake：显存+传输一体化优化和全局显存优化，将训练显存需求降低 2-10 倍。成果发布在 ASPLOS’24。

机器故障大幅增加训练成本

蚂蚁之所以会特别关注训练故障的问题，主要因为训练过程中的机器故障大幅增加了训练成本。例如，Meta 训练 OPT-175B 模型时，使用了 992 80GB A100 GPU，共 124 台 8 卡机器，按照 AWS 的 GPU 价格，每天需要耗费约 70 万。因故障导致训练周期延长 20 多天，从而增加了上千万元的算力成本。

为什么训练故障会造成这么大的影响呢？

首先，分布式训练是多个节点协同工作的，任一节点发生故障（无论是软件、硬件、网卡或 GPU 问题），整个训练流程均需暂停。
其次，训练故障发生后，排查费时费力，比如现在大家常用的手动检查方式二分法，查一次至少需要 1-2 个小时。
最后，训练是有状态的，训练重启需要从之前的训练状态里面恢复再继续，隔一段时间就要保存训练状态。而保存的过程耗时很长，而且故障回滚也会造成计算的浪费。上方的右图展示了我们上线故障自愈之前的训练耗时分布，可以看到有效训练时间大概只有 72%。

DLRover 训练故障自愈流程

DLRover 引入了一套完善的故障自愈机制，包括自动节点检测、自动扩缩容及快速 checkpoint 功能。

自动节点检测

自动节点检测是指在训练中断时，系统会自动进行硬件和通信检测，识别并隔离故障节点，随后启动新的节点替代。

自动扩缩容

针对训练过程中可能出现的节点变化，DLRover 运用自动扩缩容技术，通过动态组网和弹性调度确保训练在可控节点数量下持续进行，全程无需人工介入。

以右图所示为例，初始状态下系统中有 6 个运行正常的节点，这些节点在 K8s 环境中分别对应着 6 个 Pod。运行一段时间后，Pod 5 由于硬件故障宕机。系统会立即启动响应机制，尝试从集群中调度一个新的 Pod 6 替代失效的 Pod 5。但由于集群内已有 1 台机器出现故障，导致可用计算资源受限，新的 Pod 6 无法成功调度。这时，系统会智能地进行自动化调整，将该作业的节点规模缩减至 4 个，以确保作业能在现有可用资源下继续运行。这里选择保留 4 个而非 5 个节点的原因在于，考虑到大模型算法的实际需求，需要对节点数量加以合理控制，避免资源过度分配而影响训练效率。当先前故障的第 6 个节点修复并重新接入集群，系统将恢复至 6 个节点的状态，从而加快整个训练进程。

Flash Checkpoint

训练故障恢复过程中，关键在于模型状态的保存与恢复。传统的 Checkpoint 方法由于保存耗时较长，往往导致训练效率低下。为解决这一问题，DLRover 创新性地提出了 Flash Checkpoint 方案，能够在训练过程中近乎实时地将模型状态从 GPU 显存导出至内存，同时辅以内存间的备份机制，确保即使节点故障，也能迅速从备份节点内存中恢复训练状态，极大地缩短了故障恢复时间。

如上图所示，实验证明，采用 Flash Checkpoint 方案后，无论是单机多卡环境还是多机多卡环境，模型状态的保存和恢复速度均有显著提升。

训练自愈在国产卡训练的实践

除了支持 GPU，DLRover 故障自愈还支持国产加速卡的分布式训练。例如，我们在华为 Ascend 910 平台上运行 LLama-7B 模型时，采用了 256 张卡进行大规模训练。起初，我们采用的是 KubeFlow 的 PyTorchJob，但该工具不具备容错功能，导致训练进程在持续约十几个小时后会自动终止，一旦发生这种情况，用户必须手动重新提交任务；否则，集群资源将处于闲置状态。第二个图表描绘了在启用训练故障自愈功能后的整个训练过程。在训练进行到 20 小时时，出现了一次通信 timeout 故障，此时系统自动重启了训练进程并恢复训练。大约四十多个小时后，遇到一次机器硬件故障，系统迅速隔离故障机器并重启了一个 pod 继续训练。除了对华为 Ascend 910 的支持，我们还兼容了阿里的含光 PPU，并与沐曦科技合作，在其自主研发的 GPU 上运用 DLRover 进行大模型训练。

训练自愈在千卡训练的实践

上图展示了 DLRover 训练故障自愈在千卡训练上的实践效果：使用 1000 多张 H800 卡运行大型模型训练，在故障频率为每天一次的情况下，引入训练故障自愈功能后，有效训练时间占比超过了 97%。 右侧对比表格显示，在采用阿里云高性能存储 FSDP 的情况下，单次保存仍需大约五分钟的时间，而我们的 Flash Checkpoint 技术仅需 0.3 秒即可完成。此外，通过优化，节点效率提升了近一分钟左右。在导出间隔方面，原先每 2 小时执行一次导出操作，而在上线 Flash Checkpoint 功能后，可以实现每 10 分钟一次的高频导出。一周内 save 操作的累计耗时几乎可以忽略不计。同时，回滚时间相较于之前降低了 5 倍左右。