Skip to content

PyTorch DataLoader

Create a loader with StatefulDataLoader.from_storage. The loader owns its dataset reader while open. Consume an iterator completely or close it before starting another on the same loader. See save and resume for checkpoint lifecycle details.

Install with uv pip install "guideloom[torch]" and import the loader from guideloom or guideloom.pytorch. Stream storage requires sequence_length; document storage does not. local_batch_size is per rank, and global_batch_size is local_batch_size * world_size. len(loader) reports the number of batches per epoch, including the final partial global batch when drop_last=False.

guideloom.pytorch.StatefulDataLoader dataclass

Bases: Iterable[BatchT], Sized, SyncContextManager

Stateful iterable loader that partitions a dataset across ranks.

Use the loader as a context manager so its dataset reader is opened and closed. Exhaust an iterator, or close it, before starting another iterator on the same loader.

Parameters:

Name Type Description Default
dataset DatasetAbstract[SourceT]

Dataset to load from.

 required
collate_fn Callable[[Sequence[SourceT]], BatchT]

Function that collates samples into a batch.

 required
device device

Device batches are copied to. Defaults to the current CUDA device when CUDA is available, otherwise CPU. Copies onto CUDA use non_blocking=True.

 _get_default_device()
local_batch_size int

Samples per batch for this rank. Defaults to 1.

 1
rank int

This rank's index in the distributed job. Defaults to 0.

 0
world_size int

Number of ranks. Defaults to 1.

 1
num_workers int

Worker threads or processes. Defaults to 0, which does not load batches concurrently.

 0
parallel_method (thread, process)

How collate_fn is parallelized. Defaults to thread.

 'thread'
shuffle bool

Whether to shuffle sample indexes. Defaults to False.

 False
shuffle_seed int | None

Seed for shuffling. None makes shuffling nondeterministic. Defaults to None. Ignored when shuffle is False.

 None
pin_memory bool

Whether to pin CPU memory before a CUDA copy. Defaults to False.

 False
drop_last bool

Whether to drop the last global batch when it is smaller than local_batch_size * world_size. Defaults to True. When False and world_size > 1, a rank can receive an empty sample list, collate_fn([]), so every rank takes the same number of steps and can call a collective such as fabric.barrier(). collate_fn must accept that empty list. When world_size == 1, collate_fn never receives an empty list.

 True
prefetch int

Batches to prefetch. Defaults to 0.

 0

global_batch_size: int property

from_storage(storage: StreamStorageAbstract[InT, OutT] | DocumentStorageAbstract[InT, OutT], collate_fn: Callable[[Sequence[OutT]], cBatchT], *, sequence_length: int | None = None, device: torch.device | None = None, local_batch_size: int = 1, rank: int = 0, world_size: int = 1, num_workers: int = 0, parallel_method: Literal['thread', 'process'] = 'thread', shuffle: bool = False, shuffle_seed: int | None = None, pin_memory: bool = False, drop_last: bool = True, prefetch: int = 0) -> StatefulDataLoader[OutT, cBatchT] classmethod 

from_storage(storage: StreamStorageAbstract[InT, OutT], collate_fn: Callable[[Sequence[OutT]], cBatchT], *, sequence_length: int, device: torch.device | None = None, local_batch_size: int = 1, rank: int = 0, world_size: int = 1, num_workers: int = 0, parallel_method: Literal['thread', 'process'] = 'thread', shuffle: bool = False, shuffle_seed: int | None = None, pin_memory: bool = False, drop_last: bool = True, prefetch: int = 0) -> StatefulDataLoader[OutT, cBatchT]

from_storage(storage: StreamStorageAbstract[InT, OutT], collate_fn: Callable[[Sequence[OutT]], cBatchT], *, sequence_length: None = None, device: torch.device | None = None, local_batch_size: int = 1, rank: int = 0, world_size: int = 1, num_workers: int = 0, parallel_method: Literal['thread', 'process'] = 'thread', shuffle: bool = False, shuffle_seed: int | None = None, pin_memory: bool = False, drop_last: bool = True, prefetch: int = 0) -> NoReturn

from_storage(storage: DocumentStorageAbstract[InT, OutT], collate_fn: Callable[[Sequence[OutT]], cBatchT], *, sequence_length: None = None, device: torch.device | None = None, local_batch_size: int = 1, rank: int = 0, world_size: int = 1, num_workers: int = 0, parallel_method: Literal['thread', 'process'] = 'thread', shuffle: bool = False, shuffle_seed: int | None = None, pin_memory: bool = False, drop_last: bool = True, prefetch: int = 0) -> StatefulDataLoader[OutT, cBatchT]

from_storage(storage: DocumentStorageAbstract[InT, OutT], collate_fn: Callable[[Sequence[OutT]], cBatchT], *, sequence_length: int, device: torch.device | None = None, local_batch_size: int = 1, rank: int = 0, world_size: int = 1, num_workers: int = 0, parallel_method: Literal['thread', 'process'] = 'thread', shuffle: bool = False, shuffle_seed: int | None = None, pin_memory: bool = False, drop_last: bool = True, prefetch: int = 0) -> NoReturn

Create a loader by wrapping storage in the matching dataset.

Stream storage, including streams with metadata, is wrapped in StreamDataset and requires sequence_length. Document storage is wrapped in DocumentDataset. The remaining arguments match the constructor.

Parameters:

Name Type Description Default
storage StreamStorageAbstract[InT, OutT] | DocumentStorageAbstract[InT, OutT]

Storage to read from.

 required
collate_fn Callable[[Sequence[OutT]], cBatchT]

Function to collate samples into a batch.

 required
sequence_length int | None

Stream elements per stream sample. Required for stream storage and omitted for document storage. Defaults to None.

 None
device device | None

Target device for data placement. When omitted, uses CPU or the first CUDA device. Defaults to None.

 None
local_batch_size int

Number of samples per batch for this rank. Defaults to 1.

 1
rank int

Local rank, used for data partitioning. Defaults to 0.

 0
world_size int

Total number of ranks. Defaults to 1.

 1
num_workers int

Number of worker threads or processes. Defaults to 0.

 0
parallel_method (thread, process)

Parallelization strategy. Defaults to thread.

 'thread'
shuffle bool

Whether to shuffle the dataset. Defaults to False.

 False
shuffle_seed int | None

Seed for shuffling. Defaults to None.

 None
pin_memory bool

Whether to use pinned memory for CUDA. Defaults to False.

 False
drop_last bool

Whether to drop the last incomplete global batch. Defaults to True.

 True
prefetch int

Number of batches to prefetch. Defaults to 0.

 0

Returns:

Type Description
StatefulDataLoader[OutT, cBatchT]

StatefulDataLoader[OutT, cBatchT]: Loader over the dataset derived from storage.

open() -> None

Open the dataset reader.

Raises:

Type Description
ValueError

If the dataset contains fewer observations than world_size * local_batch_size.

close() -> None

Close the dataset reader.

__iter__() -> _SyncIteratorState[BatchT]

__len__() -> int

state_dict() -> StatefulDataLoaderStateDict

Return the checkpoint for batches already delivered to the consumer.

Prefetched batches that have not been yielded do not advance observation_offset. Each yielded batch advances the offset by global_batch_size, including padding in an incomplete final batch. The result contains epoch, observation_offset, shuffle, shuffle_seed, and drop_last.

Returns:

Name Type Description
StatefulDataLoaderStateDict StatefulDataLoaderStateDict

Loader checkpoint.

load_state_dict(state_dict: StatefulDataLoaderStateDict) -> None

Restore epoch and observation_offset from a checkpoint.

Call this before starting an iterator. A warning is emitted when shuffle, shuffle_seed, or drop_last differ from this loader. Those saved values are not applied. The global observation offset allows resuming with a different batch size or world size.

Parameters:

Name Type Description Default
state_dict StatefulDataLoaderStateDict

Checkpoint from state_dict.

 required

Raises:

Type Description
RuntimeError

If an iterator on this loader is still active.

get_default_device() -> torch.device staticmethod

Return the current CUDA device, or CPU when CUDA is unavailable.

Returns:

Type Description
device

torch.device: cuda on torch.cuda.current_device() when CUDA is available, otherwise CPU.

guideloom.types.StatefulDataLoaderStateDict

Bases: TypedDict

Checkpoint returned by StatefulDataLoader.state_dict.

epoch and observation_offset identify the global sample position delivered to the consumer. Prefetched batches that have not been yielded are not included. shuffle, shuffle_seed, and drop_last record how the loader was constructed. They are not applied by load_state_dict.

Attributes:

Name Type Description
epoch int

Epoch index. It increments when an iterator is exhausted.
observation_offset int

Global sample position in the current epoch. Advances by the global batch size for each delivered batch, including padding in an incomplete final batch.
shuffle bool

Whether that loader was constructed with shuffling enabled.
shuffle_seed int | None

Seed used for shuffling. None is not reproducible.
drop_last bool

Whether that loader drops an incomplete global batch.

epoch: int instance-attribute

observation_offset: int instance-attribute

shuffle: bool instance-attribute

shuffle_seed: int | None instance-attribute

drop_last: bool instance-attribute