[Paper Review] 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 Learning25 references4 citations

TL;DR

This paper evaluates checkpointing in large-scale deep learning training across PyTorch, TensorFlow, and Chainer, measuring performance overhead, file size, scalability, and determinism. It reveals significant checkpointing bottlenecks—especially single-node checkpointing and non-deterministic behavior—highlighting the need for HPC-optimized, scalable checkpointing solutions in deep learning frameworks.

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.

Motivation & Objective

• To evaluate the performance and efficiency of checkpointing mechanisms in major deep learning frameworks used in HPC environments.
• To analyze the computational overhead, file size, and I/O characteristics of checkpointing at scale across multiple GPUs and nodes.
• To investigate deterministic training behavior and its reproducibility across framework restarts.
• To identify limitations in existing checkpointing support, especially for multi-node and model-parallel training.
• To provide actionable insights for users selecting fault-tolerant frameworks and for developers improving checkpointing in deep learning systems.

Proposed method

• Benchmarking three state-of-the-art deep learning frameworks—Chainer, PyTorch, and TensorFlow—on the CIFAR-10 dataset using data parallelism across multiple GPUs.
• Measuring checkpointing overhead, including time taken, I/O bandwidth, and GPU utilization during checkpoint operations.
• Analyzing checkpoint file formats, sizes, and structure across frameworks under varying model architectures and scale.
• Evaluating deterministic training and restart behavior by comparing model weights and gradients across multiple runs with identical seeds.
• Assessing multi-node checkpointing capabilities and scalability, particularly the bottleneck caused by single-node checkpointing in distributed setups.
• Using a combination of controlled experiments and profiling to isolate performance bottlenecks in checkpointing pipelines.

Experimental results

Research questions

• RQ1 How does checkpointing overhead scale with the number of GPUs and nodes in distributed deep learning training?
• RQ2 What are the differences in file size, format, and I/O patterns across PyTorch, TensorFlow, and Chainer?
• RQ3 To what extent do current frameworks support deterministic training and reproducible restarts?
• RQ4 What are the scalability limitations of existing checkpointing mechanisms in multi-node HPC environments?
• RQ5 How do model size and architecture affect checkpointing performance and memory usage?

Key findings

• Checkpointing overhead increases significantly with scale, particularly when multiple GPUs are used, due to idle GPU time during checkpoint operations.
• Chainer exhibits highly variable checkpoint file sizes depending on model architecture, unlike PyTorch and TensorFlow, which show more consistent sizing.
• Only Chainer supports multi-node checkpointing, but it produces redundant copies without improving checkpointing time, indicating a lack of true parallelization.
• Despite offering deterministic training options, all frameworks fail to guarantee 100% reproducibility across restarts, undermining reproducibility in critical applications.
• The single-node checkpointing bottleneck limits scalability in multi-node HPC clusters, as only one node performs the checkpoint, creating a performance hotspot.
• Existing checkpointing mechanisms do not support partial or model-parallel checkpointing, making them unsuitable for large-scale, model-parallel training scenarios.