字节跳动 Seed 开源大模型检查点系统 ByteCheckpoint

👋 Hi, everyone!

We are ByteDance Seed team.

You can get to know us better through the following channels👇

ByteCheckpoint is a unified, efficient and production-grade checkpointing system for large foundation model development.

ByteCheckpoint is the open-source implementation of our research paper: ByteCheckpoint: A Unified Checkpointing System for Large Foundation Model Development.

ByteCheckpoint is easy to use and efficient with:

✔ **Framework-Agnostic API**: Provides a unified checkpointing entrypoint, i.e., `bytecheckpoint.save` and `bytecheckpoint.load`, to support various parallelism configurations across different frameworks.

✔ **Load-time Checkpoint Resharding**: Enables seamless checkpoint reloading with arbitrary new parallelism configurations, eliminating the need for manual resharding scripts.

✔ **Optimized I/O Performance**: Integrates advanced techniques such as asynchronous and parallel I/O, D2H tensor copying with pinned memory, load-balanced checkpointing, decomposed tensor representation.

✔ **Comprehensive Toolset**: Provides utilities for checkpoint merging/conversion/modification and metadata/tensor file inspection. Enables flexible checkpoint transfer and management. 📰 News [2025/04] We officially released ByteCheckpoint! 🔥

[2024/12] ByteCheckpoint is accepted to NSDI 2025. 🚀 Getting started Installation

Install ByteCheckpoint from source.

Install ByteCheckpoint from PyPI. Basic Usage

We introduce how to use Bytecheckpoint to save, load, and merge checkpoint.

In ByteCheckpoint, a checkpoint consists of three parts (folders): `model`: It contains model checkpoint, including one ``.metadata`` checkpoint metadata file and multiple `.distcp` tensor data files. `optimizer`: It contains optimizer checkpoint, including one ``.metadata`` checkpoint metadata file and multiple `.distcp` tensor data files. `extra_state`: It contains user-saved pickable objects, e.g., the dataloader state dictionary and RNG states. Save and Load Checkpoint

Get model, optimizer, and extra states (RNG states, learning rate scheduler) from training code.

Save them with ByteCheckpoint `save` API.

Load them with ByteCheckpoint `load` API. The model and optimizer will be loaded in an in-place manner. The extra state will be loaded in `checkpoint_state["extra_state"]`. Training Code Example (FSDP)

A simple single-machine FSDP training demo with ByteCheckpoint is on demo/fsdp_save_reshard.py

Start training and save checkpoint at each step:

Load checkpoint and resume training:

For multi-machine training, we recommend operating checkpoint in a shared file system that supports POSIX semantics, such as NFS. Merge Model checkpoint

To merge model checkpoint, you can use `scripts/merge_bcp.py`

Merge saved checkpoint in the demo training code with safetensors format:

🔧 Advanced Usage Guide API Arguments Enable `fast_saving` and `fast_loading` to use asynchronous and parallel I/O techniques. Enable `save_decomposed_model_optimizer` and `load_decomposed_model_optimizer` for FSDP (`use_orig_params=True` is required) to obtain model/optimizer state dict without additional communication and GPU-CPU synchronization. Pass the `role` keyword (e.g., actor, critic) to support checkpointing in multi-role training scenarios, such as PPO training. Enable `strict` in `load` API to check whether the fqns in a given state_dict are strictly the same as those recorded in the .metadata file. Configuration Enable `BYTECHECKPOINT_ENABLE_TREE_TOPO` to improve the stability of large-scale planning for model/optimizer planning. Enable `BYTECHECKPOINT_ENABLE_PINNED_MEM_D2H` to use the pinned CPU memory pool to accelerate D2H tensor copying. Adjust `BYTECHECKPOINT_STORE_WORKER_COUNT` and `BYTECHECKPOINT_LOAD_WORKER_COUNT` to tune the I/O performance.

Please refer to config.py for more details. 🤝 Contribution Guide

Community contributions are welcome. Please checkout Contribution Guidance. Code Formatting

We use `ruff` to enforce strict code formatting when reviewing PRs. To reformat your code locally, make sure you have installed the latest version of `ruff`.

Then you can format code with: Testing Run local tests with: 📄 License This project is licensed under Apache License 2.0. See the `LICENSE` file for details. 😊 Citation and Acknowledgement If you find this project helpful, please give us a star ⭐ and cite our paper:

ByteCheckpoint is inspired by the design of PyTorch Distributed Checkpoint (DCP). 🌱 About ByteDance Seed Team

Founded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society.