Start with a clear goal: durable, secure, versioned checkpoints that allow fast reload (possibly partial), cross-language inference, and reproducible experiments.Pickle pitfalls- Security: never unpickle data from untrusted sources — it can execute arbitrary code.- Versioning: pickle ties to Python class/module paths and library versions; loading after refactor or different package versions often fails.- Use pickle only for trusted, short-term artifacts (e.g., ephemeral experiment artifacts). Prefer safer, versioned formats for long-term storage.Joblib and memory-mapping- joblib.dump/load is convenient for large numpy arrays and sklearn pipelines; supports compression.- For huge arrays, use memory-mapping to avoid copying into RAM:python
import joblib
joblib.dump(arr, 'arr.pkl') # save
arr = joblib.load('arr.pkl', mmap_mode='r') # memory-mapped readonly
- mmap_mode='r+' allows in-place updates on disk-backed arrays. Good for large embedding tables.Atomic checkpointing- Always write checkpoints atomically to avoid corrupt files on crash: write to a temp file and atomically rename/replace.python
import tempfile, os
def atomic_save(path, data_bytes):
dirn = os.path.dirname(path)
fd, tmp = tempfile.mkstemp(dir=dirn)
with os.fdopen(fd, 'wb') as f:
f.write(data_bytes)
os.replace(tmp, path) # atomic on POSIX/Windows
- Keep multiple recent checkpoints (e.g., last N), and a manifest with checksums and metadata (timestamp, training step, git commit, library versions).Storing optimizer and training state- For full reproducibility, checkpoint model weights, optimizer state (momentum, learning rate schedule), RNG states (python, numpy, framework), and training step.- PyTorch example:python
torch.save({
'epoch': epoch,
'model_state': model.state_dict(),
'optim_state': optimizer.state_dict(),
'rng_python': random.getstate(),
'rng_numpy': np.random.get_state(),
'rng_cuda': torch.random.get_rng_state_all()
}, tmp_path)
os.replace(tmp_path, final_path)
- Beware that optimizer state can be large (e.g., Adam has two moment tensors per parameter). Consider sharded checkpoints or excluding/resetting optimizer state for cheap fine-tuning.Sharding and streaming for huge models- Split large tensors into shards (per-layer or per-parameter-file) to parallelize upload/download and reduce memory during load.- Use formats that support streaming reads (HDF5, Zarr, custom shard+manifest).Cross-language interoperability- ONNX: good for exporting model structure and weights for inference across Python/C++/.NET; not ideal for optimizer or training state. Use for production inference paths.- HDF5: hierarchical, supports named datasets (weights, metadata). Many languages and libs support HDF5; good for large arrays and partial reads.- FlatBuffers/Protobuf: schema-enforced, versionable; good for metadata and small tensors.- Recommend: export model for inference in ONNX (or framework-specific runtime) and keep training checkpoints (optimizer, RNG) in framework-native format. Provide a conversion pipeline and document versions.Best practices summary- Use atomic saves + manifest/checksums + rotation (keep last N).- Use joblib/HDF5/mmap/Zarr for large arrays to enable partial/memory-mapped loads.- Save optimizer and RNG state for exact resumes; consider sharding to handle size.- Avoid untrusted pickle; prefer standardized, versioned formats for long-term storage and cross-language use (ONNX for inference, HDF5/Zarr for large arrays).- Record environment: framework versions, CUDA, hardware, git commit — embed in metadata for reproducibility.