AI Runtime上的快速容错PyTorch训练

在大规模训练中，你的训练效率由一个指标决定："有效吞吐"，即GPU花在生产性计算上的时间比例，而非等待或从故障中恢复的时间。由于大规模下GPU故障是常态，快速自动恢复的能力是维持高有效吞吐和管理总GPU开销的唯一方法。

两个子系统决定恢复的成败，但它们经常被忽视：为加速器提供数据的数据管道，以及保存状态以便作业恢复的检查点机制。任何一个出错，每次故障都会让你付出远超应有的空闲GPU时间。即使在无故障情况下，数据管道跟不上加速器也会悄然饿死GPU，像故障一样侵蚀有效吞吐。我们将探讨两者的机制和权衡，以及每个如何影响你的有效吞吐和总GPU开销。参见配套的训练性能和弹性指南获取代码指针和示例。

关于同一问题基础设施方面——集群如何在任务宕机前检测和隔离不健康的GPU，参见配套文章：如何让GPU在Databricks AI中保持可靠。

为什么大规模下故障是预期情况

随着任务中GPU数量增加，任务在无中断情况下完成全时长的概率急剧下降。一个有用的粗略模型来自Databricks配套文章，假设每个GPU年化故障率约为1%。在此假设下，文章指出"一个256-GPU任务运行30天约有19%的概率遇到故障。在1024 GPU时，概率升至57%。"这些仅指基础设施层面问题。

为了将这一估计与现实对应，608个H100 GPU的delta超级计算机每1.9小时出现一次故障，这意味着对于32 GPU任务，平均故障时间为36小时。主要结论是：你的训练任务很可能在某个时刻失败，做出正确的决策能使你的模型具有弹性，并减少失败发生时损失的总时间。

影响1：检查点格式决定你能承受的保存频率

检查点是决定恢复能力的关键，而选择的机制对你保存的频率有直接影响。这是有效吞吐的最大杠杆：如果每天检查点一次，那么一次故障平均需要重跑12小时的重复工作，才能恢复到故障发生时的状态。

单一torch.save的瓶颈

许多团队第一个检查点实现是在rank 0上简单的torch.save。根据模型训练方式，可能有两个问题：

1. 对于分布式训练，它会将所有状态收集到rank 0，并写入一个单一文件。
2. 单个进程同步写入整个检查点。在保存到远程对象存储（如Unity Catalog，UC）时，可能因网络传输而阻塞。

这种阻塞行为让GPU闲置，降低有效吞吐。但有一种方法可以减少GPU花在检查点上的时间：Torch的分布式检查点API。

分布式检查点（DCP）：每个rank写入自己的分片

PyTorch的分布式检查点反转了设计。每个rank并行写入自己独特的分片，同时附带一个小.metadata文件描述分片如何组合成完整张量。

保存时间大约与rank数成1/N比例减少，并且因为.metadata文件记录了全局布局，同一检查点可以重新加载到不同数量的GPU上。DCP重新规划每个新rank需要哪些字节，因此丢失节点后恢复到容量减小的集群也能正常工作。

即使对于纯数据并行任务，DCP也值得使用

一个常见假设是DCP仅适用于分片模型，数据并行（DDP）任务中每个rank持有相同的权重副本，没有什么好处。事实并非如此，DCP对模型状态分片并在每个worker上并行写入，即使对于DDP训练任务也是如此。

这也是你转向FSDP或张量并行时需要的API，所以尽早采用意味着你永远不必在最糟糕的时候重写弹性代码。

异步保存使保存频率几乎免费

即使并行写入，同步保存也会阻塞训练直到数据持久化到存储，对于大型检查点保存到远程卷，会浪费数十秒的加速器空闲时间。async_save将操作拆分：快速复制到暂存缓冲区，然后后台上传与训练重叠。

训练循环仅为暂存复制付费，而非上传。原本花费数十秒空闲时间的检查点现在几乎免费，这正是下一节中频繁检查点变得可行的原因。

在AI Runtime上，UCVolumeWriter和UCVolumeReader实现了针对UC卷的DCP，通过本地NVMe暂存I/O，并在数据完全落盘后才标记检查点完成。详见性能和弹性指南获取完整细节和代码示例。

上述不包括torch.save的网络存储时间。

影响2：检查点频率决定恢复成本

这里各种因素相互叠加。当任务失败时，它会丢失自上次有效检查点以来的所有内容，必须重新计算。所以每次故障预期浪费的工作量约为检查点间隔的一半，而便宜的异步保存让你可以缩小这个间隔。

将间隔缩短为原来的十分之一，预期恢复时间也会缩短为原来的十分之一。回想 Llama 3 图中每天约 8.6 次中断的数据：在该故障率下，每 2 小时检查点一次意味着你预计每天浪费 8.6 小时在重新训练上，有效吞吐量为 64%。每 30 分钟检查点一次，你只需花费 2.15 小时，有效吞吐量为 91%。

恢复还必须是自动的。重启时，作业应找到最近已完成写入的检查点，跳过因崩溃而半写的部分，并在无需人工干预的情况下从中恢复。DCP 使其可靠：.metadata 文件仅在所有分片落盘后写入，因此其存在是选择该“保存完整”标记的可信依据。

影响 3：数据加载决定你的 GPU 是否闲置

训练作业以其最慢输入的速度进行。当加速器等待下一批次时，有效吞吐量会降低，因为 GPU 只是闲置。解决此问题的唯一方法是确保你的输入流水线将下一步的数据准备与当前步骤的计算重叠，如下所示：

我们经常看到客户转向将数据加载与计算重叠后，挂钟时间减少 20-50%。

直接从远程存储读取的成本

在受管平台上，训练数据位于远程对象存储中。在 AI Runtime 上，Unity Catalog (UC) 卷以网络挂载的形式呈现。

每次访问都直接从该挂载读取文件，会将你的步长时间绑定到网络延迟，并在每个 epoch 重新下载相同的文件。解决方案是使用数据加载器，在首次访问时将每个文件复制到本地快速存储，后续读取从该本地缓存提供，并在 GPU 计算时并行获取后续文件。

使用 AI Runtime，UCVolumeDataset 和 DataLoader 正是这样做的（参见指南中的代码示例） 。UCVolumeDataset 从 UC 卷流式传输文件，首次访问时将每个文件缓存到本地 NVMe，并在各 rank 和工作进程之间对文件进行分区，以便每个加速器获得不相交、不重叠的切片。我们的DataLoader 是 PyTorch DataLoader 的直接子类，其默认设置针对此路径进行了调整，因此文件在 GPU 计算时并发获取和缓存，而不是在训练线程上逐个获取。

示例：从 UC 文件训练图像模型

考虑一个简单的图像分类工作负载：从 UC 卷解码 JPEG、增强并训练视觉模型。让我们看看在同一 GPU、模型和批次大小下执行此操作的两种方式：标准的 PyTorch Dataset 从 UC 卷读取，与 UCVolumeDataset 加上 Databricks DataLoader 默认设置。

你不必猜测时间去了哪里

作为工程 DataLoader 的一部分，我们确保其将指标记录到 MLFlow，使你可以一眼看出数据管道是否阻塞训练。

指标 fetch_seconds 明确测量数据加载器产生一个批次所需的时间，在这段时间内你的 GPU 处于闲置状态。

影响 4：忘记数据管道会悄悄损坏你的模型

最后一个弹性 bug 不会产生错误消息、崩溃或失败的作业，只会产生一个比预期稍差的模型。当你对模型、优化器和步骤进行检查点，但没有对数据管道在数据集中的位置进行检查点时，就会发生这种情况。

考虑一个在 epoch 中途被中断的作业。它正确恢复了模型并继续训练循环，但数据加载器从数据集开头重新开始。

恢复的作业会重新训练本 epoch 已经见过的示例，并可能跳过尚未到达的示例。在规模带来的许多重启中，这会悄悄使数据分布产生偏差。模型仍然训练；只是训练了错误的数据采样，这正是那种静默故障，代价最高，因为作业完成且没有人看到问题，直到指标令人失望。

解决方案是将数据位置视为检查点的一部分。根据你的管道，这意味着跟踪样本或分片偏移并在恢复时跳过，自定义数据集序列化其自身位置，或在 epoch 边界设置检查点。所有这些都建立在一个先决条件上：确定性。 洗牌和增强从随机数生成器获取，因此这些种子和 RNG 状态也必须包含在检查点中，否则重启后的数据顺序将与之前不匹配，保存的位置会指向错误的样本。

种子、可复现顺序和可恢复的数据管道是同一个想法的三种表达。指南 涵盖了每种策略的代码。

总结

快速、容错的训练来自一系列叠加的决策：

1. 使用 Distributed Checkpoint 而不是 torch.save，即使对于 DDP 也是如此，这样保存是并行且廉价的，而不是串行瓶颈。
2. 异步保存使检查点几乎免费，这让你可以频繁保存。
3. 自动恢复到最近的有效检查点，这样故障仅损失几分钟的重算，而不是几小时。
4. 将数据加载与计算重叠，通过从远程存储缓存和预取，使加速器永远不会因等待输入而闲置。这是每一步都节省的周期性 GPU 小时数。
5. 对数据管道和 RNG 状态进行检查点，以便恢复的作业继续使用正确的数据，而不是悄悄损坏你的模型。

统一原则：频繁、廉价、完整的检查点将硬件故障从结束作业的事件变成四舍五入的误差，而重叠的输入管道让加速器在此期间保持忙碌。廉价（异步）保存使频繁成为可能；完整保存（模型、数据和 RNG）使恢复正确。两者到位，且机群检测并隔离故障硬件，无论底层集群多么不稳定，你的有效训练时间都会接近硬件允许的上限。

参考资料

• Meta，《Llama 3 模型群》(2024)：§3.3.1，关于在多达16,384个H100 GPU上的训练可靠性和中断分析。
• 《表征Hopper H100和Ampere A100 GPU的韧性》(2025)：Delta系统为期2.5年的实地研究；按作业规模划分的MTBE和故障概率。
• Databricks，《我们如何在Databricks AI中保持GPU可靠》：集群级GPU健康检查和概率故障模型。

准备好试用了吗？请参阅《训练性能和弹性指南》，位于Databricks AI Runtime文档中，获取完整代码，并阅读《我们如何在Databricks AI中保持GPU可靠》，了解基础设施方面的内容。