kempnerforge.data

Data pipeline for KempnerForge.

Public API:

• MemoryMappedDataset / HuggingFaceDataset: Dataset implementations

• DistributedSampler: Rank-aware, deterministic sampling

• StatefulDataLoader: Checkpoint-resumable data loading

class kempnerforge.data.DistributedSampler

Bases: ~torch.utils.data.Sampler.

Deterministic distributed sampler with skip-ahead support.

Partitions dataset indices across data-parallel ranks. Each rank sees a unique, non-overlapping subset of the data.

Parameters:

• dataset – Dataset to sample from.

• num_replicas – Number of data-parallel ranks (default: world_size).

• rank – Current rank (default: from dist).

• shuffle – Whether to shuffle indices.

• seed – Base random seed for deterministic shuffling.

• drop_last – Drop samples that don’t divide evenly across ranks.

__init__(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=True)

Parameters:

• dataset (torch.utils.data.Dataset)

• num_replicas (int | None)

• rank (int | None)

• shuffle (bool)

• seed (int)

• drop_last (bool)

Return type:

None

set_epoch(epoch)

Set the epoch for deterministic re-shuffling.

Parameters:

epoch (int)

Return type:

None

set_skip(skip)

Set number of samples to skip (for resumption after checkpoint).

Parameters:

skip (int)

Return type:

None

state_dict()

Return checkpoint state. Keys: epoch, seed, num_replicas, rank.

Return type:

dict

load_state_dict(state)

Restore from checkpoint. Only epoch is restored; rank info is local.

Parameters:

state (dict)

Return type:

None

class kempnerforge.data.HuggingFaceDataset

Bases: Dataset

HuggingFace dataset with on-the-fly tokenization and sequence packing.

Loads a HuggingFace dataset, tokenizes text on the fly, and packs multiple documents into fixed-length sequences (separated by EOS tokens).

Parameters:

• dataset_name – HuggingFace dataset name (e.g., “allenai/c4”).

• dataset_config – Optional config name (e.g., “wikitext-2-raw-v1”).

• split – Dataset split (“train”, “validation”, etc.).

• text_field – Name of the text column.

• seq_len – Sequence length for packing.

• tokenizer_path – Path or name for HuggingFace tokenizer.

__init__(dataset_name, split, text_field, seq_len, tokenizer_path, dataset_config=None, pack_sequences=False)

Parameters:

• dataset_name (str)

• split (str)

• text_field (str)

• seq_len (int)

• tokenizer_path (str)

• dataset_config (str | None)

• pack_sequences (bool)

Return type:

None

state_dict()

Return checkpoint state. Keys: epoch, sample_idx, total_sequences.

Return type:

dict

load_state_dict(state)

Restore from checkpoint. Restores epoch and sample_idx.

Parameters:

state (dict)

Return type:

None

class kempnerforge.data.MemoryMappedDataset

Bases: Dataset

Pre-tokenized dataset backed by memory-mapped token files.

Expects 1D arrays of uint16/uint32 token IDs that have been pre-packed into fixed-length sequences, as .npy files or raw .bin shards. Each matched file is read by its own suffix, so a pattern may match both formats; a matched file with any other suffix is rejected.

A .npy file carries its dtype in its header. A .bin file does not: its dtype is token_dtype, else the tokenized_info.dtype recorded in the metadata.yaml in its directory (tatm writes one). The two must agree when both are given, and a shard whose size is not a whole number of tokens of that dtype is rejected. When neither declares a dtype, the shards are read as uint32 with a warning, provided each is a whole number of uint32 tokens; otherwise the dataset refuses to guess and raises.

File layout: each file stores a flat array of tokens. The dataset splits them into non-overlapping chunks of seq_len tokens. Multiple files are concatenated logically.

Parameters:

• data_dir – Directory containing .npy or .bin token files.

• seq_len – Sequence length (number of tokens per sample).

• file_pattern – Glob pattern for data files.

• token_dtype – "uint16" or "uint32" for the .bin files, or "" to take it from their metadata.yaml as described above.

__init__(data_dir, seq_len, file_pattern='*.npy', pack_sequences=False, eos_token_id=None, token_dtype='')

Parameters:

• data_dir (str)

• seq_len (int)

• file_pattern (str)

• pack_sequences (bool)

• eos_token_id (int | None)

• token_dtype (str)

Return type:

None

state_dict()

Return checkpoint state. Keys: epoch, total_samples.

Return type:

dict

load_state_dict(state)

Restore from checkpoint. Only epoch is restored; sample count is derived.

Parameters:

state (dict)

Return type:

None

close()

Release the underlying mmaps. Preferred path; do not rely on __del__.

Return type:

None

class kempnerforge.data.StatefulDataLoader

Bases: object

Stateful wrapper around PyTorch DataLoader.

Tracks iteration progress so training can resume from the exact position after a checkpoint load.

Parameters:

• dataset – Dataset to load from.

• batch_size – Per-device micro-batch size.

• sampler – Distributed sampler (created automatically if None).

• config – Data pipeline configuration.

• collate_fn – Optional custom batch collator. When None, uses PyTorch’s default collation. VLM training passes VLMCollator so that the fixed-length padding and image_positions slot reach the batch.

__init__(dataset, batch_size, sampler=None, config=None, collate_fn=None)

Parameters:

• dataset (torch.utils.data.Dataset)

• batch_size (int)

• sampler (DistributedSampler | MixtureSampler | None)

• config (DataConfig | None)

• collate_fn (Callable | None)

Return type:

None

state_dict()

Return checkpoint state. Keys: epoch, batches_yielded, sampler.

Return type:

dict

load_state_dict(state)

Restore from checkpoint. __iter__ re-applies the sampler skip from _batches_yielded, so double-resume within the same epoch stays aligned.

Parameters:

state (dict)

Return type:

None

Modules

dataloader	Distributed, stateful DataLoader for KempnerForge.
dataset	Dataset implementations for KempnerForge.
sampler	Distributed sampler with deterministic shuffling and skip-ahead.
video_dataset	Video dataset and collator for the VLM video path (WebVid-style layout).
video_io	Video frame sampling and decoding for the VLM video path.
vlm_dataset	VLM dataset and collator (Joint-Decoder).