[论文解读] A Study of Checkpointing in Large Scale Training of Deep Neural Networks

Elvis Rojas, Albert Njoroge Kahira+3 more|arXiv (Cornell University)|Dec 1, 2020

Adversarial Robustness in Machine Learning参考文献 25被引用 4

一句话总结

本文评估了在 PyTorch、TensorFlow 和 Chainer 中大规模深度学习训练的检查点机制，测量了性能开销、文件大小、可扩展性及确定性。研究揭示了显著的检查点瓶颈，尤其是单节点检查点和非确定性行为，凸显了在深度学习框架中需要开发面向高性能计算（HPC）的、可扩展的检查点解决方案。

ABSTRACT

Deep learning (DL) applications are increasingly being deployed on HPC systems, to leverage the massive parallelism and computing power of those systems for DL model training. While significant effort has been put to facilitate distributed training by DL frameworks, fault tolerance has been largely ignored. In this work, we evaluate checkpoint-restart, a common fault tolerance technique in HPC workloads. We perform experiments with three state-of-the-art DL frameworks common in HPC Chainer, PyTorch, and TensorFlow). We evaluate the computational cost of checkpointing, file formats and file sizes, the impact of scale, and deterministic checkpointing. Our evaluation shows some critical differences in checkpoint mechanisms and exposes several bottlenecks in existing checkpointing implementations. We provide discussion points that can aid users in selecting a fault-tolerant framework to use in HPC. We also provide takeaway points that framework developers can use to facilitate better checkpointing of DL workloads in HPC.

研究动机与目标

• 评估在 HPC 环境中使用的主流深度学习框架的检查点机制的性能与效率。
• 分析在多个 GPU 和节点上大规模训练时，检查点的计算开销、文件大小及 I/O 特性。
• 研究训练行为的确定性及其在框架重启后的可重现性。
• 识别现有检查点支持的局限性，特别是针对多节点和模型并行训练的场景。
• 为用户选择容错框架以及开发者改进深度学习系统中的检查点机制提供可操作的见解。

提出的方法

• 在 CIFAR-10 数据集上，使用多 GPU 数据并行方式对三种最先进的深度学习框架——Chainer、PyTorch 和 TensorFlow 进行基准测试。
• 测量检查点开销，包括检查点操作期间的时间消耗、I/O 带宽及 GPU 利用率。
• 分析在不同模型架构和规模下，各框架的检查点文件格式、大小及结构。
• 通过比较使用相同随机种子的多次运行中的模型权重和梯度，评估训练的确定性及重启行为。
• 评估多节点检查点能力及可扩展性，特别关注分布式设置中单节点检查点造成的瓶颈。
• 结合受控实验与性能分析，隔离检查点流水线中的性能瓶颈。

实验结果

研究问题

• RQ1 在分布式深度学习训练中，检查点开销如何随 GPU 和节点数量的增加而变化？
• RQ2 PyTorch、TensorFlow 和 Chainer 在文件大小、格式及 I/O 模式方面有何差异？
• RQ3 当前框架在多轮训练中支持确定性训练和可重现重启的程度如何？
• RQ4 现有检查点机制在多节点 HPC 环境中的可扩展性限制是什么？
• RQ5 模型大小和架构如何影响检查点性能及内存使用？

主要发现

• 随着规模扩大，检查点开销显著增加，尤其是在使用多个 GPU 时，这是由于检查点操作期间 GPU 处于空闲状态所致。
• Chainer 的检查点文件大小因模型架构而异，变化极大；而 PyTorch 和 TensorFlow 的文件大小则更为一致。
• 仅 Chainer 支持多节点检查点，但其生成冗余副本且未缩短检查点时间，表明缺乏真正的并行化。
• 尽管所有框架都提供确定性训练选项，但均无法在重启后保证 100% 的可重现性，这在关键应用中严重损害了可重现性。
• 单节点检查点瓶颈限制了多节点 HPC 集群的可扩展性，因为仅一个节点执行检查点操作，形成性能热点。
• 现有检查点机制不支持部分检查点或模型并行检查点，因此不适用于大规模、模型并行训练场景。