ByteCheckpoint: A Unified Checkpointing System for Large Foundation Model Development

TL;DR

ByteCheckpoint is a high-performance, flexible checkpointing system designed for large-scale large foundation model training, enabling efficient resharding, multi-framework support, and scalable I/O operations.

Contribution

It introduces a parallelism-agnostic checkpoint representation and a generic workflow, significantly improving efficiency and flexibility over existing systems.

Findings

Reduces checkpoint runtime stalls by 54.20x on average.

Achieves up to 9.96x faster checkpoint saving.

Achieves up to 8.80x faster checkpoint loading.

Abstract

Checkpointing to preserve training states is crucial during the development of Large Foundation Models (LFMs), for training resumption upon various failures or changes in GPU resources and parallelism configurations. In addition, saved checkpoints are dispatched to evaluation tasks or transferred across different training stages (e.g., from pre-training to post-training). All these scenarios require resharding distributed checkpoints from one parallelism to another. In production environments, different LFMs are trained with various frameworks and storage backends, depending on model sizes and training scales. A high-performance checkpointing system is needed to enable efficient checkpoint management at scale throughout the lifecycle of LFM development. We introduce ByteCheckpoint, an industrial-grade checkpointing system for large-scale LFM training. ByteCheckpoint features: a parallelism-agnostic checkpoint representation that enables efficient load-time checkpoint resharding; a generic checkpoint saving/loading workflow to accommodate multiple training frameworks and support different storage backends; full-stack optimizations to ensure high I/O efficiency and scalability; a suite of monitoring tools to streamline large-scale performance analysis and bottleneck detection. Compared to existing open-source checkpointing systems [52, 58], ByteCheckpoint significantly reduces runtime checkpoint stalls, achieving an average reduction of 54.20x. For saving and loading times, ByteCheckpoint achieves improvements of up to 9.96x and 8.80x, respectively.