> ## Documentation Index
> Fetch the complete documentation index at: https://docs.lancedb.com/llms.txt
> Use this file to discover all available pages before exploring further.

# Checkpointing

> Save StreamingDataset's read position alongside your model, so a failed run can pick back up.

**Applies to:** Embedded ✅, Enterprise ✅

[`StreamingDataset`](https://docs.lancedb.com/training/data-loading) 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](https://docs.lancedb.com/training/resuming.md) for the load side.

This page assumes the `rank`, `world_size`, `num_splits` and `shuffle_seed`
terminology from [Distributed Training](https://docs.lancedb.com/training/distributed-training.md#parallelism-and-elasticity).

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.

```py Python icon=Python theme={"theme":{"light":"vitesse-light","dark":"catppuccin-mocha"}}
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`.

```py Python icon=Python theme={"theme":{"light":"vitesse-light","dark":"catppuccin-mocha"}}
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.

## Next steps

> [card] [Resuming](https://docs.lancedb.com/training/resuming.md)
>
> Load a saved checkpoint back, on the same or a different topology.

> [card] [Distributed Training](https://docs.lancedb.com/training/distributed-training.md)
>
> The rank, world size and split terminology checkpointing builds on.


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.