作者 | 王勤龙
责编 | 夏萌
出品丨AI 科技大本营（ID：rgznai100）
4 月 25、26 日，全球机器学习技术大会在上海环球港凯悦大酒店举行！蚂蚁集团 DLRover 开源负责人王勤龙在大会上发表了《DLRover 训练故障自愈：大幅提升大规模 AI 训练的算力效率》主题演讲，分享如何在千卡规模大模型训练作业下，快速故障自愈，王勤龙介绍了 DLRover 背后的技术原理和使用案例，以及 DLRover 在社区大模型的实战效果。
王勤龙，长期在蚂蚁从事 AI 基础设施的研发，主导了蚂蚁分布式训练的弹性容错与自动扩缩容项目的建设。先后参与多个开源项目，如 ElasticDL 和 DLRover，开放原子基金会 2023 活力开源贡献者，蚂蚁集团 2022 T-Star 优秀工程师。目前为蚂蚁 AI Infra 开源项目 DLRover 的架构师，专注于打造稳定、可扩展和高效的大规模分布式训练系统。
大模型训练及挑战
大模型训练基本流程如上图所示，需要经过训练样本数据集的准备，Transformer 模型的构造，预训练，模型微调，最终构建成为用户 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 在 2022 年公布了其大模型训练的实际数据，在训练 OPT-175B 模型时，使用了 992 80GB A100 GPU，共 124 台 8 卡机器，按照 AWS 的 GPU 价格，每天需要耗费约 70 万。因故障导致训练周期延长 20 多天，从而增加了上千万元的算力成本。
下面的图片展示了在阿里云集群进行大模型训练时遇到的故障分布情况，这些故障有的可以通过重启解决，而另一些则无法通过重启修复。例如，掉卡问题，因为重启后故障卡仍然是损坏的。必须要先更换损坏的机器，才能重新启动并恢复系统。
为什么训练故障会造成这么大的影响呢？首先，分布式训练是多个节点协同工作的，任一节点发生故障（无论是软件、硬件、网卡或 GPU 问题），整个训练流程均需暂停。其次，训练故障发生后，排查费时费力，比如现在大家常用的手动检查方式二分法，查一次至少需要 1-2 个小时。最后，训练是有状态的，训练重启需要从之前的训练状态里面恢复再继续，隔一段时间就要保存训练状态。而保存的过程耗时很长，而且故障回滚也会造成计算的浪费。
上方的右图展示了我们上线故障自愈之前的训练耗时分布，可以看到 Checkpoint 的相关时间占比大概是 8%，手动节点检测的时间占比大概是 5%，网络故障导致停机时间约 7%， 硬件故障导致停机时间约 8%，最终有效训练时间大概只有 72%。
DLRover 训练故障自愈功能总览
上图是 DLRover 在故障自愈技术上的两大核心功能。首先，通过 Flash Checkpoint 能够在几乎不停止训练流程的前提下快速保存状态，实现高频备份。这意味着一旦遇到故障，系统能够立即从最近的检查点恢复，减少数据丢失和训练时间损耗。
其次，DLRover 利用 Kubernetes 实现了智能化的弹性调度机制。该机制能够自动化应对节点故障，例如在 100 台机器的集群中若有一台失效，系统会自动调整至 99 台机器继续训练，无需人工介入。此外，它兼容Kubeflow 和 PyTorchJob，强化了节点健康监测能力，确保在 10 分钟内迅速识别并响应任何故障，维持训练作业的连续性和稳定性。
DLRover 弹性容错训练
DLRover 采用了 master-worker 的架构，这在机器学习的初期并不常见。在此设计中，master 作为控制中心，负责关键任务如节点调度、状态监控、网络配置管理和故障日志分析，而不会运行训练代码。通常部署于 CPU 节点上。worker承担实际训练负荷，每个节点会运行多个子进程来利用节点的多块 GPU，以加速计算任务。
另外，为了加强系统的健壮性，我们对 worker 上的 Elastic Agent 进行了定制化增强，使其能更有效地进行故障检测与定位，确保训练过程中的稳定性和效率。
接下来是故障检测流程。当训练过程中遭遇故障导致任务中断时，直观表现是训练暂停，但具体原因及故障源头并不直接显现，因为一旦发生故障，所有相关机器会同步停止。为解决这一问题，故障发生后，我们立即在所有机器上执行检测脚本。一旦检测到某节点未通过检验，会立即通知 Kubernetes 集群移除该故障节点，并重新部署一个新的替代节点。新节点与现有节点共同完成进一步的健康检查，一切无误后，自动重启训练任务。
值得注意的是，若故障节点隔离后导致资源不足，我们将执行缩容策略（后面会具体介绍）。而当原先的故障机器恢复正常，系统将自动执行扩容操作，确保训练高效持续进行。
接下来是故障诊断流程，通过以下综合手段，实现快速准确地故障定位与处理：
首先，Agent 收集各训练进程的错误信息，将这些错误栈汇总至 master 节点。随后，master 节点分析汇总的错误数据，以准确定位问题所在机器。例如，若某机器日志显示 ECC 错误，则直接判定该机器故障并将其排除。
此外，Kubernetes 的退出码也可以用来辅助诊断，如退出码 137 通常指示底层计算平台因检测到问题而终止该机器运行；退出码 128 则意味着设备不被识别，可能是 GPU 驱动存在故障。还有大量的故障是不能通过退出码来检测的，常见的比如说网络抖动的 timeout。
还有许多故障，如网络波动导致的超时，无法仅凭退出码识别。我们会采用更通用的策略：无论故障具体性质，首要目标是迅速识别并踢掉故障节点，然后通知 master 去具体检测问题出在哪里。
首先，在所有节点上执行矩阵乘法运算。随后，将节点配对分组，例如在含 6 个节点的 Pod 中，节点被分为(0,1)，(2,3)，(4,5) 三组，并进行 AllGather 通信检测。若 4 和 5 之间有通信故障，而其他组通信正常，则可断定故障存在于节点 4 或 5。
接下来，将疑似故障节点与已知正常节点重新配对以进一步测试，例如将 0 与 5 组合检测。通过对比结果，确切识别出故障节点。此自动化检测流程能在十分钟内准确诊断出故障机器。
前面讨论了系统中断与故障检测的情况，但还需解决机器卡死的识别问题。NCCL 设置的默认超时时间为 30 分钟，允许数据重传以减少误报，然而这可能导致实际卡顿时，每张卡白白等待 30 分钟，累计起来损失巨大。
为精确诊断卡死，建议采用一种精细化的 profiling 工具。当监测到程序暂停推进，例如一分钟内程序栈无变化，即记录各卡的栈信息，对比分析差异。例如，若发现 4 台 rank 中有 3 台执行 Sync 操作，而 1 台执行 wait 操作，即可定位该设备存在问题。
进一步，我们将关键的 CUDA 通信 kernel 和计算 kernel 进行劫持，在其执行前后插入 event 监控，通过计算事件间隔来判断运算是否正常进行。例如，若某运算超过预期的 30 秒仍未完成，可视为卡死，并自动输出相关日志及调用栈，提交至 master 进行对比，迅速定位故障机器。
故障机器确定后，考虑到成本与效率，以往训练中虽有备份机制，但数量有限。此时，引入弹性扩缩容策略显得尤为重要。假设原集群有 100 个节点，一旦某节点故障，剩余 99 节点能继续训练任务；待故障节点修复后，系统能自动恢复至 100 节点运行，且此过程无需人工介入，确保高效稳定的训练环境。
DLRover Flash Checkpoint
训练故障恢复过程中，关键在于模型状态的保存与恢复。传统的 Checkpoint 方法由于保存耗时较长，往往导致训练效率低下。为解决这一问题，DLRover 创新性地提出了 Flash Checkpoint 方案，能够在训练过程中近乎实时地将模型状态从 GPU 显存导出至内存，同时辅以内存间的备份机制，确保即使节点故障，也能迅速从备份节点内存中恢复训练状态，极大地缩短了故障恢复时间。
针对现在大家常用的 Megatron-LM，Checkpoint 的导出过程需要一个集中式的进程来协调完成，这不仅引入了额外的通信负担和内存消耗，还导致了较高的时间成本。而 DLRover 经过优化采取了一种创新方法，利用分布式导出策略，使得每个计算节点（rank）能够独立保存和加载自己的 Checkpoint，从而有效避免了额外的通信和内存需求，大大提升了效率。
在进行模型 Checkpoint 的创建时，还有一个细节值得留意。模型的训练基于数据，假设我们在训练进程的第 1000 步保存了 Checkpoint。如果之后重新启动训练但未考虑数据进度，直接从头开始重新消费数据将导致两个问题：后续新数据可能被遗漏，同时前期数据可能会被重复使用。
为解决这一问题，我们引入了 Distributed Sampler 策略。该策略在保存 Checkpoint 时，不仅记录模型状态，还会一并保存数据读取的偏移位置。这样一来，当加载 Checkpoint 恢复训练时，数据集会从之前保存的偏移点继续加载，继而推进训练，从而确保模型训练数据的连续性和一致性，避免了数据的错漏或重复处理。
在上述图表中，我们展示了一项单机多 GPU（A100）环境下的实验结果，旨在比较不同存储方案对训练过程中 Checkpoint 保存所造成阻塞时间的影响。实验表明，存储系统的性能直接影响效率：采用较低效的存储方式直接将 Checkpoint 写入磁盘时，训练会被显著阻塞，时间延长。具体而言，针对约 20GB 大小的 1.5B 模型 Checkpoint，若使用 NAS 存储，写入时间大约在 2-3 分钟；而采取一种优化策略，即异步先将数据暂存于内存中，能大幅缩短该过程，平均只需约 1 秒钟，显著提升了训练的连续性和效率。
DLRover 的 Flash Checkpoint 特性广泛兼容各大主流大模型训练框架，包括 DDP、FSDP、DeepSpeed、Megatron-LM、transformers.Trainer 以及 Ascend-DDP，为每种框架定制了专属 API，确保了极高的易用性——用户几乎无需调整现有的训练代码，开箱即用。具体而言，DeepSpeed 框架的用户仅需在执行 Checkpoint 时，通过调用 DLRover 的 save 接口来进行，而 Megatron-LM 的集成更为简便，仅需将原生的 Checkpoint 导入语句替换为 DLRover 提供的导入方式即可。
DLRover 分布式训练实践
我们对每个故障场景进行了一系列实验，以评估系统的容错能力、处理慢节点的能力以及扩缩容的灵活性。具体实验如下：
节点容错实验：通过手动关闭一些节点，测试集群是否能够快速恢复；
慢节点实验：使用 chaosblade 工具将节点的 CPU 负载提高至 90%，从而模拟出耗时较高的慢节点情况；
扩缩容实验：模拟了机器资源紧张的场景。例如，一个作业配置了 4 个节点，但实际只启动了 3 个，这 3 个节点仍然可以正常训练。过一段时间后，我们模拟将一个节点隔离，此时可用于训练的 Pod 数减少为 3 个。当这台机器重新回到调度队列中，可用的 Pods 数量又能增加至 4 个，此时 Dataloader 将从上一次的 Checkpoint 继续训练，而不是重新开始。
DLRover 在国产卡的实践
除了支持 GPU，DLRover 故障自愈还支持国产加速卡的分布式训练。例如，我们在华为 Ascend 910 平台上运行 LLama-7B 模型时，采用了 256 张卡进行大规模训练。起初，我们采用的是 KubeFlow 的 PyTorchJob，但该工具不具备容错功能，导致训练进程在持续约十几个小时后会自动终止，一旦发生这种情况，用户必须手动重新提交任务；否则，集群资源将处于闲置状态。
第二个图表描绘了在启用训练故障自愈功能后的整个训练过程。在训练进行到 20 小时的时候，出现了一次通信 timeout 故障，此时系统自动重启了训练进程并恢复训练。大约四十多个小时后，遇到一次机器硬件故障，系统迅速隔离故障机器并重启了一个 pod 继续训练。
除了对华为 Ascend 910 的支持，我们还兼容了阿里的含光 PPU，并与沐曦科技合作，在其自主研发的 千卡 GPU 上运用 DLRover 进行训练 LLAMA2-65B 模型。
DLRover 千卡千亿模型训练实践
上图展示了 DLRover 训练故障自愈在千卡训练上的实践效果：使用 1000 多张 H800 卡运行大型模型训练，在故障频率为每天一次的情况下，引入训练故障自愈功能后，有效训练时间占比超过了 97%。
右侧对比表格显示，在采用阿里云高性能存储 FSDP 的情况下，单次保存仍需大约五分钟的时间，而我们的 Flash Checkpoint 技术仅需 0.3 秒即可完成。此外，通过优化，节点效率提升了近一分钟左右。
在导出间隔方面，原先每 2 小时执行一次导出操作，而在上线 Flash Checkpoint 功能后，可以实现每 10 分钟一次的高频导出。一周内 save 操作的累计耗时几乎可以忽略不计。同时，回滚时间相较于之前降低了 5 倍左右。
DLRover 计划&社区共建
DLRover 目前已经对外发布了 3 个大的版本，预计在 6 月份会发布 V0.4.0，发布基于 CUDA Event 的运行时故障检测。
V0.1.0(2023/07)：k8s/ray 上 TensorFlow PS 的弹性容错和自动扩缩容；
V0.2.0(2023/09) ：k8s 上 PyTorch 同步训练的节点检测和弹性容错；
V0.3.5(2024/03)：Flash Checkpoint 与国产卡故障检测；
未来规划方面，DLRover 会持续在节点调度与管理、编译优化 lynx、训练优化框架 ATorch 和国产卡训练几个方向持续优化和改进 DLRover 的功能：
节点调度与管理：通信拓扑感知调度，降低 Ring-based AllReduce 通信时顶层交换机的通信流量；硬件 Metric 收集与故障预测；
编译优化 lynx：计算图调度优化，达到最优的通信与计算重叠， 隐藏通信时延；SPMD 自动分布式训练；
训练优化框架 ATorch：RLHF 训练优化；分布式训练的初始化加速；自动训练加速配置 auto_accelerate；
国产卡训练：将故障自愈、训练加速等功能扩展到国产卡上；提供国产卡训练的最佳实践；
DLRover 开源&共建
DLRover：https://github.com/intelligent-machine-learning/dlrover
GLake：https://github.com/intelligent-machine-learning/glake
技术进步始于开放合作，欢迎大家到 GitHub 上关注和参与我们的开源项目。
3万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
