Resuming

Applies to: Embedded ✅, Enterprise ✅ Resuming loads a StreamingDataset checkpoint saved by Checkpointing back into a fresh dataset and DataLoader, so a failed or restarted run picks up where it left off rather than rereading from the start. This page assumes the rank, world_size, num_splits and shuffle_seed terminology from Distributed Training.

Resuming on the same topology

import torch
from lancedb.streaming import StreamingDataset

checkpoint = torch.load("checkpoint_100.pt")
model.load_state_dict(checkpoint["model"])

dataset_state = checkpoint["dataset"]
dataset = StreamingDataset(
    table,
    num_splits=dataset_state["num_splits"],
    shuffle_seed=dataset_state["shuffle_seed"],
    epoch=dataset_state["epoch"],
    rank=rank,
    world_size=world_size,  # same as the run that saved the checkpoint
)
dataset.load_state_dict(dataset_state)

# The combined batch size across every rank, the same value used when the
# checkpoint was saved; each rank's DataLoader gets its world_size share of it.
global_batch_size = 48
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=global_batch_size // world_size,
    num_workers=0,
)

for batch in dataloader:
    train_step(batch)

load_state_dict() rejects a checkpoint whose num_splits or shuffle_seed differs from the new dataset. For an exact resume, also use the same table snapshot, epoch, shuffle, filter, and other data-selection settings, and keep world_size the same — changing it is a topology change, covered next. The number of DataLoader workers may still change as long as the batch-size divisibility constraints from Distributed Training hold.

Resuming Across Different Topologies

When training across ranks, each rank owns its own subset of splits and only its own splits have exact progress. To resume on a different world_size, collect the state_dict() from every rank of the previous run and merge them with StreamingDataset.merge_state_dicts before calling load_state_dict() on the new run. The merge is topology-agnostic: pass the full list of per-rank states in, and hand the identical merged dict to every rank of the resumed job, regardless of whether the topology grew, shrank, or stayed the same.

# On each rank of the previous run:
state = dataset.state_dict()
# ...gather `state` from every rank into a single list `states` (for example
# via torch.distributed all-gather or by saving one file per rank).

merged = StreamingDataset.merge_state_dicts(states)

# On each rank of the resumed run (world_size may differ):
dataset = StreamingDataset(
    table,
    num_splits=merged["num_splits"],
    shuffle_seed=merged["shuffle_seed"],
    epoch=merged["epoch"],
    rank=rank,
    world_size=new_world_size,
)
dataset.load_state_dict(merged)

Introduction

Getting Started

Concepts

Data Ingestion

Search

Data Curation

Feature Engineering

Training

Management

Maintenance

Security

Deployment

SQL Reference

Resuming on the same topology

Resuming Across Different Topologies

Next steps

Checkpointing

Distributed Training

​Resuming on the same topology

​Resuming Across Different Topologies

​Next steps

Checkpointing

Distributed Training

Resuming on the same topology

Resuming Across Different Topologies

Next steps