# Universal Checkpointing: A Flexible and Efficient Distributed Checkpointing System for Large-Scale DNN Training with Reconfigurable Parallelis

arXiv: https://arxiv.org/abs/2406.18820 · HTML: https://mlsanity.com/paper/2406.18820/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 55 references that are in the map.
- First posted on arXiv 27 Jun 2024.

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [BLOOM: A 176B-Parameter Open-Access Multilingual Language Model](https://mlsanity.com/paper/2211.05100/) (2022)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [A Study of BFLOAT16 for Deep Learning Training](https://mlsanity.com/paper/1905.12322/) (2019)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs](https://mlsanity.com/paper/2402.15627/) (2024)
- [NVIDIA Tensor Core Programmability, Performance & Precision](https://mlsanity.com/paper/1803.04014/) (2018)
- [Bamboo: Making Preemptible Instances Resilient for Affordable Training of Large DNNs](https://mlsanity.com/paper/2204.12013/) (2022)
- [Parcae: Proactive, Liveput-Optimized DNN Training on Preemptible Instances](https://mlsanity.com/paper/2403.14097/) (2024)
- [Tenplex: Dynamic Parallelism for Deep Learning using Parallelizable Tensor Collections](https://mlsanity.com/paper/2312.05181/) (2023)
- [Singularity: Planet-Scale, Preemptive and Elastic Scheduling of AI Workloads](https://mlsanity.com/paper/2202.07848/) (2022)
- [FastPersist: Accelerating Model Checkpointing in Deep Learning](https://mlsanity.com/paper/2406.13768/) (2024)
- [VirtualFlow: Decoupling Deep Learning Models from the Underlying Hardware](https://mlsanity.com/paper/2009.09523/) (2020)

## Built on by (most cited first, 9 of 9)

- [YuLan-Mini: An Open Data-efficient Language Model](https://mlsanity.com/paper/2412.17743/) (2024)
- [Story of Two GPUs: Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs](https://mlsanity.com/paper/2503.11901/) (2025)
- [TTrace: Lightweight Error Checking and Diagnosis for Distributed Training](https://mlsanity.com/paper/2506.09280/) (2025)
- [Nereus: Adaptive Parallelism for LLM Post-Training](https://mlsanity.com/paper/2609.34645/) (2026)
- [Learning In Chaos: Efficient Autoscaling and Self-Healing for Multi-Party Distributed Training](https://mlsanity.com/paper/2505.12815/) (2025)
- [TierCheck: Tiered Checkpointing for Fault Tolerance in Large Language Model Training](https://mlsanity.com/paper/2605.17821/) (2026)
- [PHOENIX: Resilient LLM Training with Hot-Swapping via Zero-Overhead Checkpoint](https://mlsanity.com/paper/2607.01646/) (2026)
- [ZeroLock: Concurrent Memory-Efficient LLM Training via Modular Update Decoupling](https://mlsanity.com/paper/2608.07974/) (2026)
- [Parallelism Strategy Chaining for Fast Training Convergence](https://mlsanity.com/paper/2609.07236/) (2026)

## Abstract (arXiv)

Deep neural network (DNN) training continues to scale rapidly in terms of model size, data volume, and sequence length, to the point where multiple machines are required to fit large models for training. Different distributed and parallel training strategies have been developed to support large-scale DNN training by partitioning the training state across GPUs. However, existing DNN training systems provide very limited support for reconfiguring parallelism strategies in the middle of the training via checkpointing. This limitation arises because distributed checkpoints are tightly coupled to specific model parallelism and hardware configurations, preventing large-scale training jobs from efficiently adapting to hardware failures or resource elasticity. This paper presents Universal Checkpointing (UCP), a novel checkpointing system that enables flexible and efficient DNN training with reconfigurable parallelism. UCP overcomes challenges in existing systems by decoupling checkpoint structure from parallel training strategies and hardware configurations. In addition, we present a pattern-based reconfiguration pipeline that enables automatic, flexible, and efficient mapping of checkpoint state to various parallelism strategies. Evaluation on a range of DNN models, including state-of-the-art dense and sparse LLMs, shows that UCP enables reconfiguration for a broader set of widely used parallelism strategies than existing solutions while adding negligible reconfiguration cost. UCP has been successfully employed in real LLM training workloads, greatly enhancing their flexibility and resilience to dynamic hardware environments.