Struct AsyncCheckpointer 

pub struct AsyncCheckpointer { /* private fields */ }

Asynchronous checkpointer that writes shards to disk in a background thread.

Before writing, tensor data is staged (copied) into CPU memory so that training can continue on GPU without blocking. The actual file I/O happens in a separate thread.

Usage

let ckpt = AsyncCheckpointer::new(
    checkpoint_dir.clone(),
    rank,
    world_size,
    shard_metadata.clone(),
);

// Non-blocking: training continues immediately.
let mut future = ckpt.save_async(&state_dict)?;

// ... continue training ...

// When you need to be sure it finished:
future.wait()?;

Implementations

impl AsyncCheckpointer

pub fn new( dir: PathBuf, rank: usize, world_size: usize, shard_spec: ShardMetadata, ) -> Self

Create a new async checkpointer.

• dir: directory where shard files will be written.
• rank: this process’s rank.
• world_size: total number of ranks.
• shard_spec: metadata describing how tensors are sharded.

pub fn dir(&self) -> &Path

The checkpoint directory.

pub fn rank(&self) -> usize

This process’s rank.

pub fn world_size(&self) -> usize

Total number of ranks.

pub fn save_async( &self, state_dict: &HashMap<String, Tensor<f32>>, ) -> Result<CheckpointFuture, DistCheckpointError>

Start an asynchronous checkpoint.

This immediately copies all tensor data to CPU-owned Vecs (staging), then spawns a background thread that serializes and writes the shard file. For GPU tensors, the staging step transfers data to host memory.

Returns a CheckpointFuture that can be polled or waited on.

Errors

Returns an error immediately if another async save is already in flight, or if staging (GPU-to-CPU copy) fails.