A0 — Data Pipeline and Random Number Generation

Files: llmc/rand.h, llmc/tokenizer.h, llmc/dataloader.h, llmc/sampler.h
Needs: no GPU
Builds on: —
Check with: ./grade A0

Overview

Before any neural network can train, you need three things: a way to turn text into numbers, a way to feed those numbers to the model in batches, and a source of randomness. A0 builds all three.

You will implement:

Mersenne Twister (MT19937) — the same PRNG used by PyTorch’s CPU backend, so your results can be numerically verified against reference values.
BPE Tokenizer — a byte-pair encoding decoder that maps integer token IDs back to text strings.
DataLoader — a sharded, optionally-shuffled data pipeline that reads pre-tokenised .bin files and yields (input, target) batches.
Softmax sampler — used at inference time to sample the next token from a probability distribution.

None of these functions touch GPU memory. Everything here runs on the CPU and provides the scaffolding that everything later depends on. No GPU is required for this assignment.

Every function you must write is marked in the source with a // TODO (A0) comment followed by a // HINT: line. To list them all:

grep -rn "TODO (A0)" llmc/

When these notes and the in-file HINT disagree, the code is authoritative — please report the discrepancy so it can be fixed.

Where this fits

flowchart LR
    A0["A0
tokens, batches, PRNG"] --> A1["A1
forward pass (CPU)"]
    A1 --> A2["A2
backward + AdamW (CPU)"]
    A2 --> A3["A3
port to CUDA (1 GPU)"]
    A3 --> A4["A4
scale to N GPUs"]
    A3 --> A5["A5
schedules, clipping, checkpoints"]
    style A0 fill:#2563eb,stroke:#1e40af,color:#fff

You are building the layer that turns a text file into (inputs, targets) tensors, plus the random number generator everything else leans on. Nothing here does any maths on the model — but everything later reads its batches through your DataLoader, and A4’s bit-exact multi-GPU requirement and A5’s resumable checkpoints both ultimately depend on your PRNG being reproducible.

flowchart LR
    txt["tinyshakespeare.txt"] -->|"tiktoken (Python)"| shards["shard_*.bin
256-int header
+ uint16 token stream"]
    shards --> loader["DataLoader
Part 3"]
    rng["MT19937
Part 1"] -->|"shard + sample
permutations"| loader
    loader --> batch["inputs[B,T]
targets[B,T]
targets = inputs shifted by 1"]
    batch --> model["GPT-2
(A1-A4)"]
    model --> logits["logits[B,T,Vp]"]
    logits --> samp["sample_softmax
Part 4"]
    samp --> tokid["token id"]
    tokid --> tok["Tokenizer
Part 2"]
    tok --> out["generated text"]

Why a Mersenne Twister and not rand()? Because everything here is checked by numerical comparison against PyTorch. PyTorch’s CPU backend uses MT19937, so a faithful port means your weight initialisation and shuffling can be checked against a reference bit for bit. That property is what makes A4’s “same loss on 1 GPU and 4 GPUs” requirement checkable at all.

Why is the data pre-tokenised? Tokenisation is a string-processing problem, not a numerical one, and it is done once. Keeping it in Python leaves the C side as a pure uint16 stream, which is what lets the DataLoader be a thin layer over fread. You implement the decoder half in Part 2 because generation needs it.

Background Reading

Read these in order; the first two are the only ones you need before starting.

Source	Read	Why it matters here
The header comment of `llmc/rand.h`	all of it	It contains the exact PyTorch snippet your implementation must match, plus the golden vector the grader checks. Start here.
`llmc/dataloader.h` struct definitions	`DataLoader`, `EvalLoader`	Every field you must populate, with a comment saying what it means. The notes below are a gloss on these.
Sennrich et al. 2016, Neural Machine Translation of Rare Words with Subword Units	§3.2	The BPE algorithm. You only implement the decoder, but this explains why a token can be an arbitrary byte string and why the vocab file is length-prefixed rather than newline-delimited.
Karpathy, Let’s build the GPT Tokenizer	first 45 min	The most direct explanation of why tokenisation is fiddly — merges, byte fallback, and why `" the"` and `"the"` are different tokens.
Matsumoto & Nishimura 1998, Mersenne Twister	§2 (the recurrence), §4 (tempering)	The source of the magic constants. If you would rather not read the paper, the Wikipedia pseudocode is equivalent — but note that PyTorch’s `left_`/`next_` bookkeeping differs from the textbook presentation, and the code is what you must match.
Box & Muller 1958, A Note on the Generation of Random Normal Deviates	2 pages	Where `sqrt(-2 ln u1) * cos(2 pi u2)` comes from. Explains why the transform naturally produces samples in pairs, which is why `normal_fill_16` writes `data[t]` and `data[t+8]` together.

Part 1 — Mersenne Twister (`llmc/rand.h`)

This is a port of PyTorch’s MT19937, not the textbook version. The state struct is:

typedef struct {
    unsigned long long seed_;
    int left_;                              // words remaining before the next twist
    unsigned int next_;                     // read cursor into state_
    unsigned int state_[MERSENNE_STATE_N];  // N = 624
    unsigned int MATRIX_A[2];               // {0x0, 0x9908b0df}, indexed by the low bit of y
} mt19937_state;

MERSENNE_STATE_N = 624, MERSENNE_STATE_M = 397, UMASK = 0x80000000, LMASK = 0x7fffffff.

Functions to implement (all tagged // TODO (A0)):

Function	Return	Notes
`manual_seed`	`void`	Seed the state array; see below
`next_state`	`void`	Regenerate all 624 words (the “twist”); reset `left_` and `next_`
`randint32`	`unsigned int`	Extract and temper one word; twist first if the state is exhausted
`randint64`	`unsigned long long`	Two `randint32` calls: `(hi << 32) \\| lo`
`randfloat32`	`float`	`(randint32(state) & ((1ull<<24)-1)) * (1.0f / (1ull<<24))` — 24 mantissa bits
`randfloat64`	`double`	Same idea from `randint64`, using 53 bits
`uniform_`	`void`	Fill `numel` floats with `randfloat32(state) * (to - from) + from`
`normal_fill_16`	`void`	Box-Muller over a 16-element block that already holds uniforms
`normal_fill`	`void`	Fill `numel` uniforms, then apply `normal_fill_16` in blocks of 16
`normal_`	`void`	Dispatch: `numel >= 16` uses `normal_fill`, otherwise the float64 path
`init_identity_permutation`	`void`	Fill `data[i] = i`
`random_permutation`	`void`	Fisher-Yates shuffle using `randint32`

Seeding (manual_seed):

MATRIX_A[0] = 0x0
MATRIX_A[1] = 0x9908b0df
state_[0] = seed & 0xffffffff
for j in 1..N-1:
    state_[j] = (1812433253 * (state_[j-1] ^ (state_[j-1] >> 30)) + j) & 0xffffffff
left_ = 1
next_ = 0

Twist (next_state) — note PyTorch splits this into three loops and uses the MATRIX_A lookup instead of a branch, which is what makes it constant-time:

left_ = N,  next_ = 0
for j in 0 .. N-M-1:
    y = (state_[j] & UMASK) | (state_[j+1] & LMASK)
    state_[j] = state_[j+M] ^ (y >> 1) ^ MATRIX_A[y & 1]
for j in N-M .. N-2:
    y = (state_[j] & UMASK) | (state_[j+1] & LMASK)
    state_[j] = state_[j + (M-N)] ^ (y >> 1) ^ MATRIX_A[y & 1]
y = (state_[N-1] & UMASK) | (state_[0] & LMASK)
state_[N-1] = state_[M-1] ^ (y >> 1) ^ MATRIX_A[y & 1]

Extraction and tempering (randint32):

if state is NULL: return 0
if MATRIX_A is not initialised: manual_seed(state, 5489)   // auto-init
if (--left_ <= 0): next_state(state)
y = state_[next_++]
y ^= (y >> 11)
y ^= (y << 7)  & 0x9d2c5680
y ^= (y << 15) & 0xefc60000
y ^= (y >> 18)
return y

Box-Muller (normal_fill_16) operates on 16 uniforms in place, pairing element t with element t+8:

for t in 0..7:
    u1 = 1 - data[t]
    u2 = data[t + 8]
    radius = sqrt(-2 * log(u1 + 1e-12))
    theta  = 2 * pi * u2
    data[t]     = radius * cos(theta) * std + mean
    data[t + 8] = radius * sin(theta) * std + mean

normal_fill then handles the numel % 16 != 0 tail by recomputing the last 16 values — re-drawing 16 fresh uniforms over data + numel - 16 and calling normal_fill_16 again. This overlapping-tail behaviour is what PyTorch does; matching it is required for the golden values below.

Verification. llmc/rand.h documents a golden vector in its file header, generated from torch.manual_seed(137). After seeding with 137, the first five randint32 calls must produce:

Seeding with 42 instead, the first randint32 must be 1608637542. Both are checked by ./grade A0.

Part 2 — Tokenizer (`llmc/tokenizer.h`)

The tokenizer reads a .bin vocabulary file produced by train_gpt2.py and stores token strings. At inference, tokenizer_decode maps a token ID to its string. Tokens can be arbitrary bytes (including non-UTF-8 sequences), so the file format stores each token as a length-prefixed byte sequence.

Functions to implement:

Function	Return	Notes
`safe_printf`	`void`	Print a token string, skipping bytes that are neither `isprint` nor `isspace`
`tokenizer_init`	`void`	Open the vocab file, validate the header, read the token table
`tokenizer_decode`	`const char*`	Return `token_table[token_id]`, or NULL if uninitialised or out of range
`tokenizer_free`	`void`	Free each token string and the table itself

Header format (first 256 uint32_t = 1024 bytes):

header[0] = 20240328  // magic
header[1] = version   // 1 or 2
header[2] = vocab_size
header[3] = eot_token // version 2 only

Version handling is part of the exercise: for version == 1 the EOT id was not stored, so assert vocab_size == 50257 and hard-code eot_token = 50256. For version == 2, read it from header[3]. Any other version is a fatal error.

Each token then follows as (unsigned char length, char bytes[length]). Allocate length + 1 and null-terminate so the string can be printed. Assert length > 0.

If the file cannot be opened, do not exit — set init_ok = 0 and return, so that training can proceed without a tokenizer. tokenizer_decode and tokenizer_free must both respect init_ok.

safe_printf detail: only single-byte tokens are filtered. If piece[1] == '\0' the token is one byte, and it is printed only if isprint or isspace. Multi-byte tokens are printed as-is.

Part 3 — DataLoader (`llmc/dataloader.h`)

Training data is stored as one or more .bin shards. Each shard is a flat array of uint16_t token IDs preceded by a 256-int header. The DataLoader iterates over shards, loading one at a time, and yields overlapping (input[B,T], target[B,T]) windows where target[b,t] = input[b,t+1].

Train shard header:

header[0] = 20240520  // magic
header[1] = 1         // version
header[2] = ntok      // number of tokens in this shard

Eval shard header uses magic 20240522.

Functions to implement:

Function	Return	Notes
`dataloader_load_shard_`	`int64_t`	Open shard `loader->glob_result.gl_pathv[shard_index]`, validate magic/version, return the token count
`prepare_intra_shard_indices_`	`void`	Allocate and fill `intra_shard_indices` with a fresh permutation for the current shard
`dataloader_reset`	`void`	Reset to shard 0, sample 0; reshuffle shard order when `should_shuffle`
`dataloader_advance_`	`void`	Move to the next shard (wrapping), reload it, reset the intra-shard position
`dataloader_init`	`void`	glob the pattern, size the buffers, seed `shuffle_rng`, then call `dataloader_reset`
`dataloader_load_batch`	`void`	Read `B*T+1` tokens at the current position into `inputs`/`targets`
`dataloader_next_batch`	`void`	Advance the shard if the current one is exhausted, then `dataloader_load_batch`
`dataloader_resume`	`void`	Jump to a given `(shard_idx, sample_idx)` without reshuffling
`dataloader_free`	`void`	Free buffers, close the file, `globfree`
`evalloader_init`	`void`	Like `dataloader_init` but for eval shards; reads groups of `ASSUMED_NUM_COMPLETIONS`
`evalloader_reset`	`void`	Reset shard and sample counters
`evalloader_next_example_`	`void`	Read one example (context + all completions) into the batch at a given offset
`evalloader_next_batch`	`void`	Fill a whole batch by repeatedly calling `evalloader_next_example_`
`evalloader_stat_losses`	`int`	Return the number of correct examples in this batch
`evalloader_free`	`void`	Free all resources

Note the return type of evalloader_stat_losses: it is an int — a count of correct examples that the caller accumulates into an accuracy — not a mean loss. An example is “correct” when the completion with the lowest mask-averaged loss is the one named by loader->label[i]. Only tokens with mask == 1 participate in the average, and a fully-empty (padding) example must not be counted.

Sharding across processes. Each of the num_processes ranks reads a disjoint stride of the shard: rank r starts at local_batch_offset_bytes = r * B * T * sizeof(uint16_t) and the loader advances by total_batch_size_bytes = num_processes * B * T * sizeof(uint16_t) per batch. dataloader_init sets both, validates that every shard holds at least num_processes * B * T + 1 tokens, and sizes shard_num_samples accordingly.

Shuffle mode. When should_shuffle = 1, the shard order is a permutation held in shard_indices, and the sample order within a shard is a permutation held in intra_shard_indices. Both are produced with random_permutation from Part 1 driven by loader->shuffle_rng, which dataloader_init seeds with manual_seed(&rng, 42 + process_rank). This is why A0 builds the PRNG first: dataloader_resume in A5 has to restore that exact RNG state to make a run bit-exactly resumable.

Part 4 — Softmax Sampler (`llmc/sampler.h`)

At inference time, after the forward pass produces logits, we convert them to probabilities and sample.

Functions to implement:

Function	Return	Notes
`random_u32`	`unsigned int`	xorshift*, see below
`random_f32`	`float`	`(random_u32(state) >> 8) / 16777216.0f`
`sample_softmax`	`int`	Softmax over `n` logits, then walk the CDF until it exceeds `coin`

random_u32 is xorshift* — a different algorithm from the plain xorshift you may find first on Wikipedia. Shift amounts are 12/25/27 and there is a final multiply:

*state ^= *state >> 12;
*state ^= *state << 25;
*state ^= *state >> 27;
return (*state * 0x2545F4914F6CDD1Dull) >> 32;

Note the state is unsigned long long (64-bit) and the return is the high 32 bits of the product.

sample_softmax receives raw (unnormalised) logits and a coin in [0,1). Accumulate exp(logit) into a double normaliser, scale coin by it, then walk the array accumulating exp(logit) until the running total exceeds the scaled coin. Return the last index as a fallback for rounding error.

This PRNG is deliberately not the MT19937 from rand.h: sampling happens once per generated token and does not need to match PyTorch, so the cheaper generator is used.

Testing

A0 has two test harnesses.

1. The DataLoader test (this is the main one for A0). It synthesises its own shards, so it needs no downloaded data:

cd dev/test && make test_dataloader && ./test_dataloader

Build it at -O0 when debugging — at -O3 an unimplemented function leaves the buffers uninitialised, the reads are undefined behaviour, and the compiler may delete the checks entirely, which looks like a pass:

cd dev/test && make TEST_CFLAGS="-O0 -g -fsanitize=address" test_dataloader && ./test_dataloader

It exercises single-shard iteration, multi-process striding, shuffled iteration, and resumption. It prints a line per subtest and exits non-zero on failure.

2. The MT19937 golden vector. Compile a scratch file against your rand.h:

printf '#include <stdio.h>\n#include "llmc/rand.h"\nint main(){mt19937_state s;manual_seed(&s,137);for(int i=0;i<5;i++)printf("%%u\\n",randint32(&s));return 0;}\n' > /tmp/mt.c && cc -O2 -I. -o /tmp/mt /tmp/mt.c -lm && /tmp/mt

This must print the five values listed in Part 1.

What A0 does not test. make test_gpt2 builds the full model test, which needs the A1/A2 forward and backward passes. It will fail until those are done — that is expected. For A0 all we require of it is that it compiles and links:

make test_gpt2

Checking your work

./grade A0

This builds everything, runs the MT19937 / sampler / tokenizer checks and dev/test/test_dataloader, and prints your score against the table below. Anything it cannot verify on your machine (a missing gpt2_tokenizer.bin, say) is reported as SKIP and left out of the total rather than counted as zero.

The grader builds the test harnesses at -O0 with ASan/UBSan. This matters: at -O3 an unimplemented function leaves buffers uninitialised, the reads become undefined behaviour, and the compiler is entitled to delete the very checks being graded – which looks like a pass.

Scoring

What ./grade A0 checks, and what each part is worth.

rand.h is scored one row per function, in the order the functions appear in the header. They build on each other — everything below randint32 consumes it, and randint32 consumes manual_seed and next_state — so when several rows fail at once, the topmost failing row is where the bug is.

Component	Points
`manual_seed` — seeded state array, plus `left_`/`next_`/`MATRIX_A`	2
`next_state` — the twist, including the `j + M - N` wraparound and the last word	3
`randint32` — golden vectors (seeds 137 and 42) and the draw-624 twist boundary	4
`randint64` — two draws combined high half first, stream advanced by exactly two	2
`randfloat32` — golden vector (24 mantissa bits), range, chi-squared uniformity	3
`randfloat64` — golden vector (53 mantissa bits), range, mean	2
`uniform_` — scale and shift, exactly one draw per element, values in `[from, to)`	3
`normal_fill_16` — Box-Muller on a fixed input (half credit if `mean`/`std` are ignored)	2
`normal_fill` — aligned blocks of 16 (1), and the ragged tail when `numel % 16 != 0` (1)	2
`normal_` — both dispatch paths, RNG state afterwards, and sample mean/std over 100k draws	4
`init_identity_permutation`	1
`random_permutation` — Fisher-Yates golden vectors (1), still a valid permutation of `0..n-1` (1)	2
`safe_printf` — drops unprintable single-byte tokens, passes multi-byte tokens through verbatim	4
Tokenizer initialises, handles both header versions and a missing file, and decodes a known token	11
DataLoader yields correct `(input, target)` pairs on a small shard	20
DataLoader shard rollover, multi-process striding, and shuffled resume	15
`sample_softmax` returns a valid index in `[0, n)` with the right distribution	10
Code compiles cleanly; no memory leaks (valgrind on Linux, or `-fsanitize=address`)	10
Total	100

A0 — Data Pipeline and Random Number Generation

Overview

Where this fits

Background Reading

Part 1 — Mersenne Twister (llmc/rand.h)

Part 2 — Tokenizer (llmc/tokenizer.h)

Part 3 — DataLoader (llmc/dataloader.h)

Part 4 — Softmax Sampler (llmc/sampler.h)

Testing

Checking your work

Scoring

Part 1 — Mersenne Twister (`llmc/rand.h`)

Part 2 — Tokenizer (`llmc/tokenizer.h`)

Part 3 — DataLoader (`llmc/dataloader.h`)

Part 4 — Softmax Sampler (`llmc/sampler.h`)