Checkpointing

Applies to: Embedded ✅, Enterprise ✅ StreamingDataset provides checkpointing in case your model training fails partway through a run. Use StreamingDataset.state_dict() to capture the current number of samples consumed from each split in a Python dictionary, save it alongside your model weights, and load_state_dict() to restore the last saved position — see Resuming for the load side. This page assumes the rank, world_size, num_splits and shuffle_seed terminology from Distributed Training. A stored StreamingDataset state reflects only round-robin cycles that have been completed across all splits; samples completed partway through a cycle at the time of saving can be replayed. To ensure efficient saving, make sure global_batch_size is a multiple of num_splits, as in the example below.

import torch
from lancedb.streaming import StreamingDataset

global_batch_size = 48  # one sample per split in each global step
batch_size = global_batch_size // world_size

dataset = StreamingDataset(
    table,
    num_splits=48,
    shuffle_seed=42,
    epoch=current_epoch,
    rank=rank,
    world_size=world_size,
)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=batch_size,
    num_workers=0,
)

for step, batch in enumerate(dataloader):
    train_step(batch)

    if (step + 1) % checkpoint_interval == 0:
        torch.save(
            {"model": model.state_dict(), "dataset": dataset.state_dict()},
            f"checkpoint_{step + 1}.pt",
        )

Checkpointing with Multiple DataLoader Workers

Python’s built-in torch.utils.data.DataLoader produces a safe checkpoint when num_workers = 0. With num_workers > 0, calling state_dict() often raises RuntimeError. In this case, use StreamingDataLoader for multi-worker prefetch and exact resumability. It carries a state snapshot alongside every internal batch and commits it to the parent StreamingDataset only when that batch is returned to the trainer. The trainer receives the same collated batch it would receive from a standard DataLoader.

from lancedb.streaming import StreamingDataset, StreamingDataLoader

dataset = StreamingDataset(
    table,
    num_splits=48,
    shuffle_seed=42,
    epoch=current_epoch,
    rank=rank,
    world_size=world_size,
)
dataloader = StreamingDataLoader(
    dataset,
    batch_size=batch_size,
    num_workers=4,
)

for step, batch in enumerate(dataloader):
    train_step(batch)

    if (step + 1) % checkpoint_interval == 0:
        torch.save(
            {"model": model.state_dict(), "dataset": dataset.state_dict()},
            f"checkpoint_{step + 1}.pt",
        )

StreamingDataLoader accepts the same arguments as torch.utils.data.DataLoader, with a few restrictions:

dataset must be a StreamingDataset (subclasses that override __iter__ are not supported).
in_order = True is required so that consumer-committed checkpoints stay deterministic.
persistent_workers = True is not supported, because prefetched worker copies cannot be restored from parent-committed state.
drop_last = True is not supported, because incomplete tails discarded by worker replicas cannot be checkpointed topology-independently.

With more than one worker, call state_dict() at a complete logical step boundary where every split assigned to the rank has the same consumed-sample count. Calling it mid-step raises RuntimeError asking you to consume more batches first.

Introduction

Getting Started

Concepts

Data Ingestion

Search

Data Curation

Feature Engineering

Training

Management

Maintenance

Security

Deployment

SQL Reference

Checkpointing with Multiple DataLoader Workers

Next steps

Resuming

Distributed Training

​Checkpointing with Multiple DataLoader Workers

​Next steps

Resuming

Distributed Training

Checkpointing with Multiple DataLoader Workers

Next steps