HDF5 Tokenized Corpus

A tokenized corpus on disk is a flat stream of integer token ids, and the trainer needs O(1) random access into it. A modern language-model training run reads tokens at hundreds of thousands of samples per second across dozens of workers. JSONL on disk dies at the first cold-cache page fault: the JSON parser is slow, document boundaries are not addressable, and seeking to "sample 4,217,884" means scanning the file. Parquet compresses well but is still a poor fit, because the trainer does not want columns. It wants the flat stream.

HDF5 fits because it offers a chunked, resizable, integer-only dataset whose chunks are page-cache friendly at read time. The trainer asks for tokens[3200000 : 3208192] and HDF5 copies that hyperslab from the page cache into a freshly allocated NumPy array. The cost is one open file handle and a chunk-sized page-cache footprint per worker, which is small next to the cost of decoding JSONL.

The hard part is the write side. Resizable datasets are easy to misuse. Write one document at a time and the file fragments until it is unusable. Write all documents in one resize and a process death loses the whole shard. The discipline that works is buffer-then-extend, with a buffer size that matches the chunk size, plus a sharded write that splits the workload across files so a crash loses at most one shard.

Resizable HDF5 done right

A resizable HDF5 dataset is one created with maxshape=(None,), so it can grow after creation, and a fixed chunks=(chunk_size,). Writing buffers tokens in a NumPy array of length chunk_size. When the buffer fills, the dataset is resized by exactly chunk_size and the buffer is written into the new range. At end of shard the residual buffer goes into a final partial range. Every write is contiguous and chunk-aligned except that last one. The final resize uses the exact count, so the dataset has no padding, but the reader still truncates at the token_count recorded in the shard's HDF5 attributes. Tokens are stored as np.uint16, so token ids are limited to 65535 and a larger vocabulary needs a wider dtype.

Sharded write

A shard here is one HDF5 file holding the tokens of one input shard. A single HDF5 file is a single point of failure, so the pipeline splits the work: each input shard from Phase 19 lesson 42 produces one HDF5 output shard. A shards.json index records, per shard, the file path, the token count, the document count, and a sha256 over the tokens. The trainer reads shards.json to compute global offsets and to validate the corpus. The demo's write_corpus loops over the shards one after another, but because each shard is an independent file, nothing stops you from writing them in parallel.

Chunked read

A chunked read is a request for tokens[start:stop] that HDF5 serves from the page cache once the chunk is hot. MmapTokenStore (the name is historical: it opens each shard through h5py in swmr=True mode and does not use np.memmap) maps a global range onto shard files, and get_slice copies the requested range into a freshly allocated NumPy array. The worker does not materialise the whole file. The slice is copied into the dataloader's batch buffer, and the dataloader copies that into a pinned-memory training tensor at batch time. The hot path has one syscall per chunk transition, and everything else is RAM access.

Sliding-window dataloader

A sliding-window dataloader samples fixed-length training sequences from the global token stream. This is the stage that knows about training-sequence length. It picks a random start index, reads window_size + 1 tokens, and returns (input, target) = (tokens[:-1], tokens[1:]). Document boundaries are not enforced, so a window may straddle two documents, with the boundary token (id 0) that the writer placed after each document, so the model learns to use the separator. The dataloader itself injects nothing; it only slices. This is the standard packing rule, and the one a beginner forgets, ending up with a corpus that is 8 percent training boundary tokens and 92 percent natural text.

code/main.py implements five main pieces:

• Tokenizer: a byte-level deterministic tokenizer, good enough for the demo. The interface is encode(text) -> list[int] and vocab_size.
• HDF5ShardWriter: opens a resizable integer dataset, buffers tokens to chunk size, resizes and writes in fixed-size strides, and records token_count and sha256 as HDF5 attributes on close. add_boundary() writes the boundary token (id 0), and write_shard calls it after every document.
• ShardedTokenizationPipeline: iterates input documents, routes them to a writer, and emits a shards.json index.
• MmapTokenStore: opens shard files with h5py in swmr=True mode, computes global offsets, and exposes a single get_slice(start, stop) API that copies the range into a NumPy array.
• SlidingWindowDataloader: picks random windows from the global stream and yields (input_ids, target_ids) NumPy arrays.

The demo at the bottom of the file builds a tiny in-memory corpus and tokenizes it into two shards. It then opens them through MmapTokenStore, runs the dataloader for 10 batches, and prints the per-batch shape and a checksum.

Around those sit a few helpers: JSONLSource streams text from a JSONL file, pack_documents packs tokenized documents into fixed-length groups, tokenize_jsonl_path tokenizes one JSONL file into one shard, load_index reads shards.json back, and validate_corpus recomputes each shard's sha256 and returns the failing shard ids.

Run it:

python3 code/main.py

The script exits zero and prints batch checksums.

Four patterns carry this design into a real training run.

Chunk size equals the typical read. The trainer reads window_size + 1 tokens per sample. Set the HDF5 chunk to a multiple of window_size and reads are page-cache aligned. Mismatched chunks halve the throughput because every sample touches two chunks.

Token count in attributes, not in the dataset. The last range is shorter than a chunk, because the chunk size does not divide the shard's token total. This writer resizes by the exact count, so the dataset ends at the last real token and nothing is padded. Store the real token_count as an HDF5 attribute on the dataset anyway and have the reader truncate at that value. A writer that pre-sizes whole chunks leaves a zero-filled tail, and a reader that trusts the dataset length walks off the end into zero-padded tokens, and the model learns to predict zero.

Sharded sha256 with parallel verification. Each shard has its own sha256 over the token bytes, so the trainer can verify all shards in parallel before training starts. A wrong sha256 fails the run early, not on epoch three after sixteen hours.

swmr=True on both sides, with libver="latest" on the writer. Single-Writer-Multiple-Reader mode requires the writer to open with libver="latest", create every dataset up front, then set file.swmr_mode = True. After that the writer must call dataset.flush() after each resize so reader workers (opened with swmr=True) see consistent data. Skipping libver="latest" or enabling SWMR after structural changes is a common source of "file is locked" failures.

Three habits keep the pipeline recoverable once it runs against real data:

• One HDF5 per source shard. The downloader (lesson 42) emits one shard per URL, and tokenization (this lesson) emits one HDF5 per source shard. The 1:1 mapping makes resume and partial-failure recovery trivial.
• Boundary token id. The boundary token (id 0) is part of the tokenizer vocab and is the only token the pipeline injects: write_shard calls add_boundary() after every document, and the dataloader injects nothing. The training loss masks the boundary token if the model is supposed to ignore it. Otherwise the model learns to use it as a sequence separator.
• shards.json as the source of truth. Adding a shard means writing the HDF5, computing its sha256, and appending an entry. The trainer reads the file once at startup and never touches the directory listing.

On a real project, outputs/skill-hdf5-tokenized-corpus.md would describe which tokenizer feeds the pipeline, what chunk size matches the trainer's window, where shards.json lives in version control, and how dataloader workers are sharded across files. This lesson ships the engine.

1. Add a --compression gzip flag to the HDF5 writer and measure the throughput cost on the demo corpus. Defend the default you pick.
2. The dataloader already takes a seed. Extend it so each of N workers derives a distinct seed from it, then verify that two runs with the same seed produce identical batches per worker and that workers do not repeat each other's windows.
3. validate_corpus already recomputes each shard's sha256. Extend it to also compare the token_count attribute against shards.json and check that global_start offsets are contiguous, then expose it as a --validate mode that CI runs before training starts.
4. Compare dataloader throughput at chunk sizes equal to, half of, and twice the window size. Report the page-cache effect.
5. Add a --max-document-tokens flag that truncates very long documents at write time. Defend the trade-off against deciding at read time.

• HDF5 chunking documentation - the chunked, resizable dataset layout this lesson uses
• h5py user guide - Python bindings for HDF5
• NumPy memory mapping - the true memory-mapped alternative; this lesson's store reads through h5py instead
• Phase 19 · 42 - the downloader whose output this lesson tokenizes
• Phase 19 · 44 - the cosine schedule that consumes this dataloader
• Phase 19 · 45 - the AMP loop that wraps the training step