29.6x Less Checkpoint Work for PyTorch Model Saving — Exact Reload

data
1

29.6x Less Checkpoint Work for PyTorch Model Saving — Exact Reload

I’ve been experimenting with a different way to handle persistent model checkpoints when only a small fraction of model state changes between versions.

Instead of writing the entire state again for every checkpoint, HKD Checkpoint stores the initial state and then records the active changes required for deterministic continuation.

A filesystem benchmark using actual file writes, fsync(), and exact reload:

  • 2,000,000 float32 elements
  • 30 checkpoint versions
  • 1,000 changed elements per update
  • exact reconstruction: True

Results:

Method Total time Total bytes
torch.save 0.621 s ~240 MB
Safetensors 0.388 s ~240 MB
torch.distributed.checkpoint.save 0.451 s ~240 MB
torch.distributed.checkpoint.async_save 0.438 s ~240 MB
HKD Checkpoint 0.0134 s 8.35 MB

That is approximately:

28.75x fewer bytes written

and, in this benchmark:

29.05x–46.50x faster actual checkpoint writing

with exact reconstruction of the final state.

The important distinction is that this isn’t intended as a faster implementation of writing the same full checkpoint repeatedly. The idea is to avoid doing most of that work when state changes are sparse.

I’ve also separated the more general idea into HKD Incremental, which applies the same active-change principle to Python computations where a small part of a large state changes between versions.

HKD Checkpoint:

I’d particularly like feedback from people working with large embeddings, sparse training, frequent checkpoints, distributed training, or workloads where checkpoint I/O is significant.

Does this match a checkpointing bottleneck you’ve encountered in real PyTorch workloads?