CheckFreq: Frequent, Fine-Grained DNN Checkpointing

文章标题：CheckFreq：频繁、细粒度的 DNN 检查点
作者/机构：Jayashree Mohan (UT Austin), Amar Phanishayee (Microsoft Research), Vijay Chidambaram (UT Austin and VMware research)

A1 主要贡献

深度神经网络（DNN）的训练是一项资源密集且耗时的任务。训练过程中，模型在 GPU 上进行计算以学习权重，这个过程会重复多个轮次（epoch）。学习到的权重驻留在 GPU 内存中，并偶尔被检查点化（写入持久存储）以实现容错。传统上，模型参数在轮次边界进行检查点，但对于现代深度网络，一个轮次可能运行数小时。因此，由于抢占、节点故障或进程失败导致的训练中断，会在恢复时造成数小时的 GPU 工作损失。

核心问题：传统的 DNN 训练检查点策略（在轮次边界）过于粗粒度，导致在频繁中断（例如，在可抢占虚拟机或大规模集群中，故障间隔时间可能很短）的环境下，会浪费大量的 GPU 计算时间。同时，简单地增加检查点频率会引入高昂的运行时开销（检查点停顿），并可能因数据加载器状态无法恢复而影响模型准确性。

研究目标：本文旨在设计并实现一个自动化的、细粒度的检查点框架，能够在不显著增加运行时开销、不影响模型准确性的前提下，实现迭代级别的频繁检查点，从而将故障恢复时间从数小时缩短至秒级。

创新点：
本文提出了 CheckFreq，一个用于 DNN 训练的细粒度检查点框架，通过以下几点创新来解决上述问题：

1. 算法化的检查点频率确定：CheckFreq 使用系统的在线分析（online profiling）方法，在迭代粒度上算法化地确定检查点频率。它会根据模型大小、存储带宽和迭代时间等因素来计算最佳频率。
2. 动态自适应速率调整：为了将检查点开销限制在用户设定的范围内，CheckFreq 在运行时使用自适应速率调整技术，动态监控作业运行时间，并相应地增加或减少检查点频率。
3. 保证数据不变性：为了维持“每个轮次中，数据集中的每个项目只使用一次”的数据不变性，CheckFreq 引入了一个轻量级的可恢复迭代器（resumable iterator），它能够对数据加载器的状态进行检查点。
4. 低成本的两阶段检查点机制：CheckFreq 通过引入两阶段检查点（two-phase checkpointing）来减少检查点成本，该机制将检查点操作与计算流水线化。第一阶段是内存快照（snapshot()），第二阶段是异步持久化（persist()）。

实验表明，在各种模型、存储后端和 GPU 上，CheckFreq 能够将恢复时间从数小时减少到几秒钟，同时将运行时开销限制在 3.5% 以内。

A3 背景知识与关键观察

2 背景知识

DNN 计算模型。训练深度神经网络（DNN）是确定网络中权重和偏置（统称为可学习参数）集合的过程。训练完成后，DNN 使用学习到的权重来计算输出。DNN 训练从一组随机选择的可学习参数开始，并以称为迭代的步骤迭代进行。每次迭代处理数据集的一个称为小批量（minibatch）的小的不相交子集。当整个数据集被精确处理一次后，一个轮次（epoch）就完成了。训练的每个迭代按顺序执行以下步骤：
* 数据增强：从存储中获取一个小批量数据，并应用随机的预处理操作。例如，在流行的图像分类模型（如 ResNets）中，预处理包括随机裁剪输入图像、调整大小、旋转和翻转。
* 前向传播：将模型函数应用于小批量数据以获得预测结果。
* 反向传播：使用损失函数来确定预测与正确答案的偏差程度；DNN 中的每一层都会计算损失的梯度。
* 权重更新：使用反向传播中计算出的梯度来更新可学习的模型参数。
在训练结束时（通常在固定数量的轮次后），最终学习到的参数被保存到持久存储中。为了对模型进行推理，用学习到的参数初始化 DNN 并预测输出。

检查点。训练 DNN 是一项非常耗时的任务。例如，最先进的语言建模网络 BERT-large 在 16 个 V100 GPU 上并行训练需要 2.5 天【8, NVIDIA AI. BERT Meets GPUs】。由于可学习参数在训练期间保存在 GPU 内存中，任何由于进程崩溃、服务器崩溃、作业或虚拟机抢占或作业迁移导致的训练中断，都会导致已学习的模型状态丢失。这个状态的大小通常从几百 MB 到几百 GB 不等【36, OpenAI. GPT-3 Checkpoint】。因此，花费在训练上的数小时 GPU 时间将会丢失。为了解决这个问题，模型状态通常在轮次边界进行检查点；即为了容错而写入持久存储。当训练作业恢复时，可以加载此检查点，以确保进度不会完全丢失。

恢复时间。当 DNN 训练作业中断时，它会回滚到最后一个已检查点的已完成轮次，如图 1 所示。请注意，在最后一个检查点和中断点之间执行的所有 GPU 工作都会丢失，并且在训练恢复时必须重做。由于中断而丢失的 GPU 时间量被称为恢复时间。换句话说，这是将模型恢复到中断前状态所花费的时间。

3 当前检查点技术的现状

我们分析了 PyTorch【5, PyTorch. https://github.com/pytorch/pytorch, 2019】、TensorFlow【6, Martín Abadi et al., Tensorflow: A system for large-scale machine learning, OSDI 16】和 MxNet【11, Tianqi Chen et al., Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems, CoRR, 2015】等流行的开源机器学习训练框架中检查点的现状。我们发现开源 ML 训练框架中的检查点存在不正确和效率低下的问题。

• 正确性：训练脚本中使用的检查点机制可能在作业失败或中断时导致检查点文件的丢失或损坏。
• 效率：检查点效率低下。检查点频率以临时方式确定，通常在轮次边界，这导致恢复时损失数小时的 GPU 时间。此外，缺乏对细粒度检查点的支持；现有的数据迭代器不支持在迭代边界恢复训练状态，并导致高昂的检查点停顿。

3.1 检查点机制不正确

覆盖写入导致的损坏。一些由 PyTorch【38, PyTorch. PyTorch Training Examples. https://github.com/pytorch/examples/tree/master/imagenet】维护的官方训练工作负载，在每个轮次结束时覆盖同一个检查点文件以减少存储使用。然而，这在检查点操作期间发生崩溃时，会暴露检查点文件被损坏的风险。先前的工作 【37, Thanumalayan Sankaranarayana Pillai et al., All file systems are not created equal: On the complexity of crafting crash-consistent applications, OSDI ’14】已经表明，不同的文件系统对覆盖写入的处理方式不同；在 ext3 的写回模式下，崩溃可能导致非原子数据更新从而造成数据损坏，而在 ext4 上则可能截断文件导致数据丢失。无论哪种情况，检查点文件都变得不可用；训练必须从第一个轮次重新开始。

检查点文件可能不会持久化。分析训练框架用于检查点的原语，例如 torch.save，发现它们不会对检查点文件执行 fsync()。我们验证了这可能导致数据丢失。此外，简单地执行频繁的同步 fsync() 会显著影响训练性能（§5.3.1）。

3.2 检查点机制效率低下

检查点执行稀疏且方式随意。训练作业中没有系统的检查点策略；检查点间隔的选择是随意的。例如，一些作业在训练期间不进行检查点，而另一些作业仅在大量轮次（训练的 60%）过后才开始进行检查点。总的来说，我们观察到检查点通常在轮次边界执行，仅提供有限的容错能力；在作业中断的情况下，训练将从上一个完成的轮次恢复，这可能损失数小时的 GPU 训练时间，这些时间必须重做。例如，当使用 ImageNet 在 V100 GPU 上训练 ResNext101 时，如果作业中断，平均会损失两个小时的 GPU 时间（§5.5）。

简单的频繁检查点调度会导致检查点停顿。提供更高的容错能力需要在比轮次边界更频繁的级别上执行检查点，即在迭代边界。然而，简单地增加检查点频率会在训练中引入大量的检查点停顿。由于模型权重在迭代之间不断更新，检查点需要训练短暂暂停以准确捕获模型权重。我们将此开销（即 GPU 空闲等待检查点完成的时间）称为检查点停顿。因此，给定一个 DNN（因为检查点的大小在不同 DNN 之间从 100MB 到 100GB 不等）和存储带宽，找到正确的检查点频率以最小化检查点停顿至关重要。

训练期间违反数据不变性会影响模型准确性。每个轮次以随机顺序对数据集进行完整遍历，并保持每个数据项在每个轮次中仅被看到一次的不变性。在轮次边界进行检查点的好处之一是，数据迭代器状态无需持久化，因为它在轮次结束时被重置。在更细的粒度（即迭代级别）进行检查点，需要基础设施支持来恢复数据迭代器的状态。我们注意到，在一些不为每个批次执行随机预处理操作的 NLP 模型的自定义数据加载器中，存在持久化迭代器状态的支持。然而，对于每批次对输入数据应用随机变换的图像和视频模型，PyTorch、MxNet 中的现有数据加载器以及像 NVIDIA 的 DALI 这样的最先进数据管道，在迭代边界是不可恢复的。结果是，在中断情况下，它们违反了数据不变性，导致像流行的 ResNet18 模型准确性下降高达 13%（图 6）。

3.3 总结

总而言之，我们观察到当今的检查点机制是不正确的，可能导致检查点数据丢失或损坏。此外，检查点策略是临时的；没有系统的方法来确定应该多久进行一次检查点，以同时最小化恢复时间和产生低的检查点停顿。最小化恢复时间的解决方案是执行频繁的、迭代级别的检查点。然而，执行正确且高效的细粒度检查点是具有挑战性的。我们需要（1）低成本的检查点机制，（2）轻量级、可恢复的数据迭代器以保持模型准确性，以及（3）一种系统地确定检查点频率的方法。

A2 方法细节

我们提出了 CheckFreq 的目标及其提供的恢复保证。然后，我们概述了 CheckFreq 的整体架构，并讨论了 CheckFreq 为实现所列目标而使用的技术。

4.1 目标

• 正确性。CheckFreq 旨在提供频繁、迭代级别的检查点，这些检查点是一致且持久的。
• 对模型准确性无影响。CheckFreq 旨在通过确保在中断后恢复训练时数据不变性成立，从而不影响模型的统计效率。
• 自动频率选择。CheckFreq 旨在根据正在训练的模型和训练环境（GPU 型号、存储类型、迭代时间）自动确定和调整检查点频率。检查点频率影响恢复时间，即将模型状态恢复到中断前状态所需的时间。
• 低检查点停顿。CheckFreq 旨在减少训练期间的检查点停顿，以便频繁检查点的运行时开销较低（例如，< 5%）。
• 最小化代码更改。CheckFreq 旨在以最小的训练代码更改来实现检查点管理和恢复的自动化。

4.2 CheckFreq 恢复保证

中断作业的恢复。一个被中断的作业会从磁盘上最新的可用检查点恢复训练。在传统的基于轮次的检查点中，无论作业何时被中断，训练都会从上一个轮次边界恢复，如图 1 所示。如果一个作业每个轮次执行 $n$ 次迭代，每次迭代耗时 $t_i$，那么该作业的平均恢复时间 $R_{avg}$ 为：

这是因为，当在一个轮次中途被中断时，该轮次中已完成的工作在恢复时必须重做，因为状态被重置到上一个轮次结束时。因此，基于轮次的检查点的恢复时间 R 的界限为：

注意，$n * t_i$ 是一个轮次的持续时间；它可以长达几个小时。

CheckFreq 的恢复时间保证。CheckFreq 旨在为恢复时间提供一个严格的界限，并采用更细粒度的方法在迭代边界进行检查点。CheckFreq 保证在任何时间点，系统中最多只有一个正在进行的检查点操作。当被中断时，它最多回滚一个检查点——要么是最后启动的检查点（如果它完成了），要么是前一个检查点，如图 2 所示。如果 CheckFreq 自动确定的频率是 k 次迭代，那么 CheckFreq 保证恢复时间 R 的界限为：

我们稍后在评估中（§5.4）将展示，所选择的检查点频率 $k$ 比 $n$ 小 100-300 倍，因此与基于轮次的检查点相比，恢复时间减少了几个数量级。

4.3 设计

架构概览。CheckFreq 的架构如图 3 所示。CheckFreq 有三个主要组成部分：一个可恢复的数据迭代器，它向训练作业返回一个小批量数据；一个反馈驱动的检查点策略，它决定何时触发检查点；以及一个低成本的检查点机制，它分为 snapshot() 和 persist() 两个阶段。CheckFreq 监控每个检查点间隔内产生的运行时开销；这被用作反馈来动态调整检查点频率，以确保运行时开销不超过用户给定的限制 $p$（例如 5%）。当中断发生时，CheckFreq 恢复最新的可用检查点并继续训练。我们下面详细描述每个组件。

4.3.1 检查点机制

同步检查点的问题。如今的 DNN 检查点是同步执行的；训练会暂停直到检查点操作完成。然而，同步检查点会引入大量的检查点停顿，如果频繁执行，会导致巨大的运行时开销。换句话说，同步检查点的成本（$T_c$）很高。例如，考虑一个每三次迭代进行一次检查点的策略。模型状态在权重更新阶段之后写入磁盘，该阶段根据反向传播中计算的梯度更新权重。如图 4a 所示，检查点成本发生在关键路径上，导致高昂的检查点停顿，这会显著减慢端到端的训练时间。为了在这种开销 $p$ 内掩盖如此高的检查点成本，需要不频繁地执行检查点，这反过来又会导致高昂的恢复成本。

两阶段检查点。CheckFreq 旨在通过减少检查点停顿来降低中断事件中的恢复成本。为实现低检查点成本，CheckFreq 引入了一种 DNN 感知的两阶段检查点机制。CheckFreq 将检查点分为两个阶段：snapshot() 和 persist()，并将每个阶段与计算流水线化。CheckFreq 两阶段检查点背后的主要洞见是，它利用了 DNN 计算模型（§2）来在现代加速器（如 GPU）上流水线化检查点操作。

阶段 1：snapshot()。第一阶段是 snapshot() 阶段，在迭代的权重更新步骤之后执行。在这里，模型状态的副本在内存中被捕获，以便可以异步地将其写出到存储。由于模型状态驻留在 GPU 内存中，snapshot() 涉及将模型参数从 GPU 复制到 CPU 内存。在关键路径上同步执行此操作会导致不可忽略的 snapshot() 开销，如图 4b 所示。因此，CheckFreq 小心地将 snapshot() 与计算流水线化。
* 流水线化 snapshot()。将 snapshot() 与计算流水线化必须谨慎进行，以确保模型参数的一致性并保持随机梯度下降（SGD）的正确性，SGD 是学习算法常用的一种流行优化技术。简单地将它们流水线化可能导致不一致的快照，其中包含一个迭代的部分权重更新和另一个迭代的其余部分。CheckFreq 利用 DNN 学习结构来实现正确、流水线化的快照。我们观察到，可学习的模型参数在一次迭代的反向传播之后，在一个称为权重更新的步骤中在 GPU 内存中更新。因此，我们可以将迭代 $i$ 的 snapshot() 与计算流水线化，直到迭代 $i+1$ 的权重更新。如果 snapshot() 到那时还没有完成，那么迭代 $i+1$ 会等待，直到正在进行的 snapshot() 成功完成，如图 4c 所示。这种紧密耦合是确保一致快照所必需的；否则我们可能会捕获一个被后续迭代部分更新的状态，这反过来会影响学习算法的正确性【28, Qi Meng et al., Convergence analysis of distributed stochastic gradient descent with shuffling, Neurocomputing, 2019】。

基于 GPU 的 snapshot()。尽管 snapshot() 与后续迭代的计算流水线化，但在无法完全隐藏将模型状态从 GPU 复制到 CPU 的成本的情况下，它仍可能导致检查点停顿。因此，在可行时，CheckFreq 使用基于 GPU 的 snapshot() 进一步优化此操作。我们观察到，在 GPU 内存中执行 snapshot() 的成本比将其执行到 CPU 内存要便宜一个数量级，因为后者在关键路径上涉及 GPU 到 CPU 的复制。因此，CheckFreq 采用以下方法：
(a) 当训练环境中有足够的空闲 GPU 内存来容纳快照副本时，我们在 GPU 上将快照保存在 GPU 内存中。然后，persist() 阶段异步地将快照复制到 CPU 内存，然后再到磁盘。
(b) 如果没有，CheckFreq 直接将快照保存到 CPU 内存中。这可能会在关键路径上引入停顿。
(c) CheckFreq 适当地调整检查点频率，以最小化 snapshot() 的开销（在(b)中可能特别大）和 persist() 中的停顿。

阶段 2：persist()。检查点的第二阶段是 persist() 阶段，它异步地将快照写入持久存储，这类似于已充分探索的异步检查点技术【33, Bogdan Nicolae et al., Deepfreeze: Towards scalable asynchronous checkpointing of deep learning models, CCGRID 2020】【34, Bogdan Nicolae et al., Veloc: Towards high performance adaptive asynchronous checkpointing at large scale, IPDPS 2019】【40, Faisal Shahzad et al., An evaluation of different I/O techniques for checkpoint/restart, IPDPSW 2013】【45, Devesh Tiwari et al., Lazy checkpointing: Exploiting temporal locality in failures to mitigate checkpointing overheads on extreme-scale systems, DSN 2014】。然而，为了提供 §4.2 中讨论的有界回滚保证，persist() 与计算紧密耦合。CheckFreq 将 persist() 操作作为后台进程执行，并监控其进度。当策略确定触发下一个检查点时，会检查正在进行的 persist() 操作的进度。如果 persist() 尚未完成，则计算进程会等待，直到正在进行的检查点操作完成。这确保了在任何时间点最多只有一个正在进行的检查点操作，如果作业被中断，它最多回滚到前一个检查点。虽然在下一个检查点被触发时放弃一个正在进行的检查点可能很诱人，但这是一种棘手且有风险的操作。假设我们放弃当前的检查点并开始写入下一个，此时的失败可能最终导致两个检查点都丢失。这可能是一系列连锁反应；如果最近的所有检查点都被放弃，失败可能导致回滚到一个非常旧的检查点，从而导致高昂的恢复时间。由于 CheckFreq 旨在保证我们最多回滚到前一个检查点，因此它不会放弃任何正在运行的检查点。

可恢复的轻量级数据迭代器。DNN 训练工作负载通过数据迭代器提供的薄 API 与 CheckFreq 交互。DNN 训练中数据迭代器的功能是向 GPU 返回一批预处理过的数据项，以使数据不变性成立——每个轮次以随机顺序精确处理所有数据项一次。虽然 PyTorch 中的原生迭代器和像 DALI【4, NVIDIA DALI. https://github.com/NVIDIA/DALI, 2018】这样的最先进数据管道在通常情况下支持这一点，但如果训练被中断，它们缺乏可恢复性。例如，考虑一个包含从 1 到 8 的八个数据项的数据集。在一个轮次中，处理数据项的顺序可能如图 5a 所示。假设我们在每次迭代处理一个数据项后对模型状态进行检查点。如果训练在该轮次中途被中断，数据迭代器会丢失状态，并以数据集的随机洗牌顺序恢复，如图 5b 所示，导致在一个轮次中数据项被重复和遗漏，违反了数据不变性。CheckFreq 的数据迭代器使用以下技术来支持恢复：
* 它在每个轮次使用一个依赖于轮次编号的种子来洗牌数据项。因此，要重新创建相同的洗牌顺序，只需持久化当前的轮次 ID 和到目前为止已处理的数据项数量（这使得迭代器检查点变得轻量级）。当训练恢复时，迭代器会重建洗牌顺序，并确定性地从上次检查点中断的地方重新开始，如图 5c 所示。

总结。两阶段检查点机制与可恢复的数据迭代器一起，提供了正确、低成本的检查点。下一个要回答的重要问题是，我们应该多久对模型进行一次检查点？

4.3.2 检查点策略

确定检查点频率。为了执行自动的、迭代级别的检查点，我们必须确定执行检查点的频率。一方面，我们可以在每次迭代后都进行检查点，提供低恢复成本但可能高的运行时开销。另一方面，我们可以在轮次边界执行粗粒度的检查点，导致高恢复成本但低运行时开销。一个有效的检查点策略必须在恢复成本和运行时开销之间找到正确的平衡，将两者都最小化。CheckFreq 检查点策略背后的主要思想是每 $k$ 次迭代（称为检查点频率）启动一次检查点，这样一次检查点操作的开销可以被分摊到 $k$ 次迭代中。虽然 HPC 中的先前工作已经探索了基于集群中故障分布来确定检查点频率的方法【12, John T. Daly, A higher order estimate of the optimum checkpoint interval for restart dumps, Future Gener. Comput. Syst., 2006】【14, Sheng Di et al., Optimization of multi-level checkpoint model for large scale HPC applications, IPDPS 2014】【15, Sheng Di et al., Optimization of multi-level checkpoint model for large scale HPC applications, IPDPS 2014】，但 CheckFreq 根据 DNN 和硬件特性找到能够掩盖检查点开销的最短间隔。

系统化的在线分析。CheckFreq 采用基于系统性分析的方法来确定检查点频率。选择的频率应使得由于检查点引入的运行时开销在实际计算时间的百分比 $p$ 以内，其中 $p$ 是用户决定的允许开销（例如 5%）。CheckFreq 按如下方式确定初始检查点频率。当训练作业开始时，CheckFreq 的数据迭代器（§4.3.1）会自动分析几个影响检查点频率的迭代级别和检查点特定指标——迭代时间（$T_i$）、执行权重更新的时间（$T_w$）、创建内存中 GPU 副本的时间（$T_g$）、创建内存中 CPU 副本的时间（$T_c$）、写入存储的时间（$T_s$）、检查点大小（$m$）、峰值 GPU 内存利用率（$M$）和总 GPU 内存（$M_{max}$）。基于 CheckFreq 的两阶段检查点机制，频率确定算法如算法 1 所示。
该算法提供两个输出：1）检查点频率 $k$，即每次检查点之间经过的迭代次数，以及 2）snapshot() 模式（基于 CPU 或 GPU）。算法首先根据可用的空闲 GPU 内存确定快照模式；如果有足够的空间在 GPU 内存中快照模型状态，则模式设置为 GPU，否则首选模式设置为基于 CPU 的快照。根据所选模式，算法估计在如前所述（§4.3.1）以紧密耦合方式流水线化检查点和计算后，在关键路径上产生的开销。然后，它确定摊销此开销所需的迭代次数，以使产生的总运行时开销低于阈值 $p$。例如，假设检查点操作的成本和一次迭代的持续时间都是 1 个时间单位。如果运行时开销的阈值设置为 5%，那么 CheckFreq 选择每 20 次迭代进行一次检查点。

自适应速率调整。当模型的训练环境在作业的整个运行时间内保持不变时，基于静态分析的频率确定效果很好。然而，在实践中，在线分析器估计的检查点成本可能会发生偏差，导致运行时开销高于预期。例如，一个作业可能因与其他并发运行的作业共享存储读/写而面临写干扰，这会影响写入检查点的时间。因此，CheckFreq 使用自适应速率调整技术来执行反馈驱动的频率更改。CheckFreq 的迭代器在运行时（在初始频率确定之后）监控作业的运行时间和检查点的实际成本。如果观察到的运行时超过了期望的开销，则使用这些值重新计算检查点频率。其思想是确保整体运行时开销不超过阈值 $p$。

4.4 实现

实现细节。我们将 CheckFreq 实现为 PyTorch 的一个可插拔模块。CheckFreq 的数据迭代器是在 PyTorch 的最先进数据管道 DALI 之上实现的。CheckFreq 可以用作 PyTorch 中现有数据加载器的直接替代品。CheckFreq 通过分析第一个轮次中前 1% 的迭代或前 50 次迭代（取较小者）来确定初始检查点频率。因此，在这个初始阶段不执行检查点，这只是总运行时间的一小部分。此外，我们将分析的指标和确定的策略缓存在持久存储上，以便在作业崩溃后恢复时可以跳过分析。CheckFreq 内部使用 torch.save()，然后是 fsync() 来执行持久化，从而保证持久性。为消除数据损坏的可能性，CheckFreq 总是将检查点写入一个新文件。然而，为了保持空间利用率有界，CheckFreq 在任何给定时间只在磁盘上维护两个检查点；一个已完成的检查点和另一个正在进行的检查点。此外，在轮次边界执行的检查点会被保留（用户可以关闭此功能）。CheckFreq 用一个信号量包装优化器中的权重更新步骤，该信号量等待正在进行的 snapshot()，以确保在模型状态被下一次迭代更新之前完成其副本的创建。

A4 实验

5.1 实验环境

对比基线。我们使用 PyTorch 中基于轮次的检查点（epoch-based checkpointing）作为基准，并结合了最先进的数据流水线 DALI【4, NVIDIA DALI. https://github.com/NVIDIA/DALI, 2018】。对于所有模型（BERT除外），基线为在轮次边界进行检查点。BERT 以迭代为单位进行训练，因此我们使用默认的 200 次迭代作为检查点间隔的基线。为了确保持久和正确的检查点，我们在检查点操作返回后显式地刷新（flush）检查点文件。

硬件配置。我们评估了 CheckFreq 在两代 GPU 上的性能，配置如表 2 所示：
* Conf-Volta: 8x NVIDIA Volta V100 GPU，1.8TB SSD 用于持久存储。
* Conf-Pascal: 8x NVIDIA Pascal 1080Ti GPU，1.8TB HDD 用于持久存储。
两台服务器均配备 24 个 CPU 核心和 500GB DRAM。

软件配置。
* 操作系统：64 位 Ubuntu 16.04
* 依赖库：CUDA toolkit 10.0，PyTorch 1.1.0

模型与数据集。
我们使用了 7 个 DNN 模型进行评估：
* ResNet18, ResNet50, ResNext101, DenseNet121, VGG16, InceptionV3：均在 ImageNet-1k 数据集上训练。
* Bert-Large：在 Wikipedia 和 BookCorpus 数据集上进行预训练。
每个模型都使用文献中报告的默认小批量（minibatch）大小。

5.2 实验结果

准确性影响

实验内容：为了验证可恢复数据迭代器的必要性，我们训练 ResNet18 模型达到 69.5% 的目标准确率或 70 个轮次。实验设置了三种场景：(1) 无中断正常训练；(2) 基线中断，使用 DALI 迭代器并在每次中断前进行迭代级检查点，中断频率为每 7 分钟一次（约每两个轮次）；(3) CheckFreq 中断，使用 CheckFreq 的可恢复迭代器，中断和检查点设置同(2)。
实验结果：如图 6 所示，使用基线迭代器进行迭代级检查点会导致模型最终准确率下降高达 13%，因为它破坏了数据不变性（数据项在一个轮次中被重复或遗漏）。相比之下，使用 CheckFreq 的可恢复迭代器，模型能够达到与无中断训练几乎相同的目标准确率。
结论：可恢复的数据迭代器对于实现正确、细粒度的检查点至关重要。CheckFreq 的迭代器状态检查点开销极小（仅需持久化两个整数），因此是轻量级的。

检查点机制的性能

检查点停顿：
实验内容：我们在 Conf-Pascal 服务器上，以 CheckFreq 为各模型选择的检查点频率，比较了 CheckFreq 的两阶段检查点机制与基线同步检查点机制所产生的运行时开销。
实验结果：如图 7 所示，基线同步检查点在频繁操作下会产生 17% 至 73% 的高昂运行时开销。而 CheckFreq 通过将检查点与计算流水线化，成功将开销控制在 3.5% 以内。
结论：CheckFreq 的两阶段流水线机制显著减少了检查点停顿。

收益分解：
实验内容：为了分析 snapshot() 和 persist() 流水线化的各自贡献，我们在两台服务器上训练 VGG16，并评估三种设置：(1) 基线同步模式；(2) 仅 persist() 流水线化；(3) snapshot() 和 persist() 均流水线化。
实验结果：如表 3 所示，与仅流水线化 persist() 相比，同时流水线化两个阶段可以将检查点成本降低 5-18 倍。在存储较慢的 Conf-Pascal 上，persist() 流水线化的收益更明显。在存储较快的 Conf-Volta 上，snapshot() 和 persist() 的成本相当，因此流水线化 snapshot() 带来了显著的加速。
结论：将 snapshot() 和 persist() 两个阶段都与计算流水线化是实现低检查点开销的关键。

检查点策略

检查点频率：
实验内容：我们展示了在 3.5% 的开销阈值下，CheckFreq 为不同模型确定的检查点频率。
实验结果：如表 4 所示，在 Conf-Pascal 上，CheckFreq 执行的检查点频率比基于轮次的方法高 83-278 倍。在 Conf-Volta 上，则高 25-100 倍。这些结果都将运行时开销控制在 3.5% 以内。
结论：CheckFreq 能够在低开销下实现远高于传统方法的检查点频率，为快速恢复奠定了基础。

自适应频率调整：
实验内容：我们运行一个 VGG16 训练作业（Job-A），在 100 次迭代后启动第二个 VGG16 作业（Job-B）以产生存储带宽竞争。我们比较了有无自适应频率调整功能时 Job-A 的运行时开销。
实验结果：如表 5 所示，在没有竞争时，Job-A 的开销为 5%。引入 Job-B 后，静态频率策略导致开销飙升至 35%。而 CheckFreq 的自适应调整机制能够动态改变检查点频率，将开销始终控制在 5%。
结论：自适应频率调整对于在共享资源环境中保持低且可预测的检查点开销至关重要。

恢复时间

实验内容：我们比较了在发生中断时，使用 CheckFreq 和基于轮次的检查点策略的平均恢复时间。
实验结果：如表 6 所示，CheckFreq 极大地缩短了恢复时间。例如，在 V100 GPU 上训练 ResNext101 时，平均恢复时间从 2 小时减少到 32 秒。在多 GPU 场景下，恢复时间同样从数十分钟减少到几秒钟。
结论：通过频繁的细粒度检查点，CheckFreq 将因中断而损失的计算时间从小时/分钟级别降低到了秒级。

端到端训练

实验内容：我们模拟了一个可抢占的集群环境，设置平均抢占间隔为 5 小时。在此环境下，我们使用 CheckFreq 和基线策略训练 ResNet50（在 Conf-Pascal 上）和 ResNext101（在 Conf-Volta 上）以达到目标准确率。
实验结果：如图 8 所示，对于 ResNet50，CheckFreq 将端到端训练时间缩短了 2 倍，因为它将每次中断的恢复时间从 1.9 小时减少到不足一分钟。对于 ResNext101，训练时间也加快了 1.6 倍。
结论：在存在中断的真实训练环境中，CheckFreq 通过最小化恢复时间，显著提升了端到端的训练效率。

A7 补充细节

6 讨论

在分布式集群训练中的适用性。CheckFreq 目前适用于分布式数据并行（DDP）模式，其中每个节点只有一个 GPU（rank 0）负责检查点。虽然我们的结果展示了单 GPU 和多 GPU 训练，但将其扩展到多节点设置是直接的。在多 GPU 和多节点设置中，DDP 的检查点机制是相同的。模型权重通常在每次迭代或累积几十次迭代后在不同工作节点（同一节点或分布式集群中）之间同步；因此每个节点在这些同步点看到相同版本的权重。因此，CheckFreq 的一个实例在每个节点上运行，并在同步边界持久化一个相同的检查点以供本地恢复。由于每个节点独立且并行地持久化检查点，因此检查点没有额外的同步开销。

通用性。CheckFreq 专注于优化检查点，这是目前 DNN 训练作业从故障中恢复的主要方式。虽然我们的论文关注于数据并行训练，但先前在模型或流水线并行方面的研究也依赖于检查点。使用 CheckFreq，在小批量边界（每 n 次迭代）进行检查点时，每个流水线阶段仅持久化该工作节点托管的参数和优化器状态的子集。CheckFreq 还支持在流水线并行训练中在小批量边界内（每 m 个微批次）进行检查点，因为 CheckFreq 的迭代器控制每个微批次进入流水线。在微批次粒度上进行检查点需要存储额外的模型状态——特别是在每个阶段累积的权重梯度，以及参数和优化器状态。我们将把 CheckFreq 的实现集成到支持流水线并行的框架中作为未来的工作。虽然我们在 PyTorch 中实现了 CheckFreq，但我们可以通过将特定框架的 API 包装到 CheckFreq 暴露的 API 中，将其扩展到其他框架，如 TF 和 MxNet。

7 相关工作

异步 DNN 检查点。虽然最近的工作如 DeepFreeze【33, Bogdan Nicolae et al., Deepfreeze: Towards scalable asynchronous checkpointing of deep learning models, CCGRID 2020】在执行异步 DNN 检查点时采用了与 CheckFreq 类似的 IO 流水线技术，但它只考虑了 CPU 集群。它没有考虑到在使用最先进的 GPU 进行训练时，在内存中快照模型状态的成本。我们的工作表明，在现代 ML 优化服务器上，快照模型状态（从 GPU 复制到 CPU）的成本是显著的，并展示了如何将此传输与计算流水线化，以及利用空闲的 GPU 能力来实现快速快照。此外，DeepFreeze 需要手动干预来为给定的模型、硬件和训练环境调整检查点频率，而 CheckFreq 为用户屏蔽了这些复杂性，并分析性地确定了检查点的最佳参数。与 DeepFreeze 使用静态检查点频率不同，CheckFreq 在共享集群设置中也很有益，因为它根据其他作业引起的内存和存储干扰来调整检查点频率，以最小化检查点停顿。

HPC 中的异步检查点。HPC 的先前工作【34, Bogdan Nicolae et al., Veloc: Towards high performance adaptive asynchronous checkpointing at large scale, IPDPS 2019】【40, Faisal Shahzad et al., An evaluation of different I/O techniques for checkpoint/restart, IPDPSW 2013】【45, Devesh Tiwari et al., Lazy checkpointing: Exploiting temporal locality in failures to mitigate checkpointing overheads on extreme-scale systems, DSN 2014】使用异步检查点来掩盖 IO 延迟。区分 DNN 检查点与传统 HPC 检查点的一个关键挑战是，由于 GPU 日益增长的计算能力，从 GPU 到 CPU 同步执行内存中模型状态副本的成本很高。CheckFreq 利用 DNN 学习结构，甚至小心地将内存中快照与计算流水线化，以执行正确、一致的检查点。此外，CheckFreq 通过利用空闲的 GPU 内存和计算能力（如果可能）来执行快速快照，进一步降低了检查点延迟。

HPC 中的检查点间隔估计。先前的工作【12, John T. Daly, A higher order estimate of the optimum checkpoint interval for restart dumps, Future Gener. Comput. Syst., 2006】【14, Sheng Di et al., Optimization of multi-level checkpoint model for large scale HPC applications, IPDPS 2014】【15, Sheng Di et al., Optimization of multi-level checkpoint model for large scale HPC applications, IPDPS 2014】根据系统中观察到的故障分布来确定大规模 HPC 应用的检查点间隔。CheckFreq 以 DNN 感知的方式执行此操作，通过利用 DNN 训练的确定性、重复性结构来系统地分析运行时的资源利用情况。

自适应检查点。在 HPC 应用中使用自适应进行故障管理【25, Zhiling Lan and Yawei Li. Adaptive fault management of parallel applications for high-performance computing. IEEE Trans. Computers, 2008】的思想已被用于决定何时进行检查点，这基于一个故障预测模块。CheckFreq 在 DNN 检查点频率中引入了自适应性。它根据正在训练的模型的特性、系统硬件以及其他作业的干扰来识别并动态调整检查点频率。

TensorFlow Checkpoint Manager。TF 检查点管理器【43, TensorFlow. Tensorflow checkpoint manager. https://www.tensorflow.org/api_docs/python/tf/train/CheckpointManager】允许在用户给定的时间间隔进行检查点，并支持持久化迭代器状态。然而，它有三个缺点。首先，检查点频率由用户临时决定；如果选择不当，会引入大的检查点停顿。其次，如果涉及随机数据转换，它无法对迭代器状态进行检查点；这在大多数基于图像的模型中很常见 【44, TensorFlow. Tensorflow iterator checkpointing. https://www.tensorflow.org/guide/data#iterator_checkpointing】。最后，即使在可以持久化迭代器状态的情况下 ，TF 也会将整个操作图与预取项一起写入存储，导致检查点大小很大。CheckFreq 通过自动调整检查点频率和使用轻量级、可恢复的数据迭代器来解决这些挑战。

框架透明的检查点。像 CRIU【1, CRIU checkpointing. https://criu.org/Main_Page】这样的透明检查点技术可以备份整个虚拟机状态以实现容错；然而，它们不 对 GPU 或加速器状态进行检查点。即使它们要捕获整个设备状态，仅设备状态就比在迭代边界捕获的模型状态大一个数量级，使得频繁的 CRIU 检查点不切实际。因此，在这项工作中，我们专注于 DNN 容错的主导方法——框架辅助的模型状态检查点。

A5 结论

本文介绍了 CheckFreq，一个用于 DNN 训练的自动化、细粒度检查点框架。CheckFreq 通过使用可恢复的数据迭代器、流水线化的两阶段检查点机制以及自动确定和调整检查点频率，实现了迭代级别的一致、低成本检查点。当作业中断时，CheckFreq 将流行 DNN 的恢复时间从数小时减少到几秒钟，同时产生的运行时开销很低。