Implementation:Microsoft Onnxruntime Checkpoint API
| Knowledge Sources | |
|---|---|
| Domains | Training, API, Checkpoint |
| Last Updated | 2026-02-10 04:00 GMT |
Overview
Defines the CheckpointState data class and functions for saving and loading ORT training checkpoints as flatbuffer files.
Description
The Checkpoint API provides the core checkpoint persistence layer for ONNX Runtime Training. The `CheckpointState` struct aggregates three categories of training state: module checkpoint state (trainable and non-trainable parameters), optimizer checkpoint state (learning rate, step count, momentum values), and a user-defined property bag (e.g., epoch number, best score). The API exposes free functions to save these states to a flatbuffer-formatted checkpoint file and to load them back, supporting both file paths and in-memory byte buffers. An additional overload saves ONNX TensorProto initializers directly as a checkpoint, with an external data threshold (defaulting to 1.8 GB) to handle models that exceed the 2 GB protobuf limit. A `LoadCheckpointToModel` function is available in non-minimal builds to populate an ONNX ModelProto from a checkpoint file.
Usage
Use this header when implementing checkpoint save/load workflows in the ORT Training API. It is used by the `TrainingSession` and higher-level C API wrappers to persist and restore training state between sessions.
Code Reference
Source Location
- Repository: Microsoft_Onnxruntime
- File: orttraining/orttraining/training_api/checkpoint.h
- Lines: 1-105
Signature
struct CheckpointState {
ModuleCheckpointState module_checkpoint_state;
OptimizerCheckpointState optimizer_checkpoint_state;
PropertyBag property_bag;
bool has_external_data = false;
};
PathString ExternalCheckpointDataPath(const PathString& checkpoint_path);
Status SaveCheckpoint(const CheckpointState& state, const PathString& checkpoint_path,
const bool include_optimizer_state);
Status SaveCheckpoint(gsl::span<const ONNX_NAMESPACE::TensorProto> trainable_tensor_protos,
gsl::span<const ONNX_NAMESPACE::TensorProto> non_trainable_tensor_protos,
const PathString& checkpoint_path, const bool nominal_checkpoint,
const size_t external_data_threshold = 1800 * 1024 * 1024);
Status LoadCheckpoint(const PathString& checkpoint_path, CheckpointState& checkpoint_state);
Status LoadCheckpointFromBuffer(gsl::span<const uint8_t> checkpoint_bytes,
CheckpointState& checkpoint_state);
Status LoadCheckpointToModel(const PathString& checkpoint_path,
ONNX_NAMESPACE::ModelProto& model_proto);
Import
#include "orttraining/training_api/checkpoint.h"
I/O Contract
| Function | Inputs | Outputs | Description |
|---|---|---|---|
| SaveCheckpoint (state) | CheckpointState, PathString, bool | Status | Serializes module, optimizer, and property states to a flatbuffer checkpoint file |
| SaveCheckpoint (protos) | TensorProto spans, PathString, bool, size_t | Status | Serializes ONNX initializers directly to a checkpoint file with optional external data |
| LoadCheckpoint | PathString | CheckpointState (out), Status | Loads checkpoint from a file path into a CheckpointState struct |
| LoadCheckpointFromBuffer | byte span | CheckpointState (out), Status | Loads checkpoint from an in-memory buffer |
| LoadCheckpointToModel | PathString | ModelProto (out), Status | Loads checkpoint state into an ONNX ModelProto (non-minimal build only) |
| ExternalCheckpointDataPath | PathString | PathString | Returns the external data file path for a given checkpoint path |
Usage Examples
#include "orttraining/training_api/checkpoint.h"
using namespace onnxruntime::training::api;
// Save training state
CheckpointState state;
// ... populate state with module/optimizer states ...
Status save_status = SaveCheckpoint(state, ORT_TSTR("model_checkpoint.ckpt"), true);
// Load training state
CheckpointState loaded_state;
Status load_status = LoadCheckpoint(ORT_TSTR("model_checkpoint.ckpt"), loaded_state);
// Load from buffer
std::vector<uint8_t> buffer = ReadFileToBuffer("model_checkpoint.ckpt");
CheckpointState buffer_state;
Status buf_status = LoadCheckpointFromBuffer(gsl::make_span(buffer), buffer_state);