Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

9 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

tinytransformer

tinytransformer is a GPT-2 small (124M) inference engine written from scratch in C++17. It loads the official OpenAI weights in safetensors format, tokenizes text with a from-scratch byte-level BPE tokenizer, runs the full transformer forward pass with a KV cache for autoregressive decoding, and produces logits that match HuggingFace's transformers to within a few parts in 10^5. There are no machine learning frameworks involved; the only third-party dependency is the C++ standard library.

Why

Modern inference stacks hide the actual computation behind many layers of abstraction. Writing the whole path yourself, from the byte-to-unicode table in the tokenizer to the cache-blocked matmul to the KV cache indexing, is the most direct way to understand what a transformer forward pass really costs and where the time and memory actually go. The validation step against HuggingFace keeps that understanding honest: if a single sign or transpose is wrong, the logits diverge and the test fails.

Building

Requirements are a C++17 compiler and CMake 3.16+. There are no external library dependencies. (libz is not actually required by the current code; safetensors and the GPT-2 vocab files are uncompressed, so nothing links against zlib.)

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build

This produces build/tinytransformer (the CLI), build/bench (the profiler), and the test binaries build/test_ops, build/test_tokenizer, and build/test_weights. The build uses -O2 -Wall -Wextra -Werror.

Weights and vocab

The weights and vocabulary files are not checked in (the weights/ directory is gitignored). Download them once:

mkdir -p weights && cd weights
curl -L -o vocab.bpe    https://openaipublic.blob.core.windows.net/gpt-2/models/117M/vocab.bpe
curl -L -o encoder.json https://openaipublic.blob.core.windows.net/gpt-2/models/117M/encoder.json
curl -L -o model.safetensors "https://huggingface.co/openai-community/gpt2/resolve/main/model.safetensors?download=true"
cd ..

Running

Generate text (temperature sampling with top-k):

./build/tinytransformer generate "The quick brown fox" --max-tokens 50 --temp 0.8 --top-k 40 --seed 42

Generate deterministically (greedy / argmax):

./build/tinytransformer generate "The meaning of life is" --max-tokens 30 --greedy

Print the top-5 next-token logits for a prompt (used by validation):

./build/tinytransformer logits "The meaning of life is"

Run the tests and the benchmark:

ctest --test-dir build --output-on-failure
./build/bench weights "The quick brown fox jumps over the" 100

Validating against HuggingFace

tools/validate.py compares the engine's last-token logits against transformers using the same weights and the same token IDs, so the comparison isolates the model arithmetic. Set up a virtualenv with torch and transformers, point a small directory at the local weights, and run it:

python3 -m venv .venv && .venv/bin/pip install torch transformers
mkdir -p weights/hf_gpt2
# config.json comes from the gpt2 repo; the safetensors file is the one above.
ln -s ../model.safetensors weights/hf_gpt2/model.safetensors
.venv/bin/python tools/validate.py "The meaning of life is" ./build/tinytransformer

A passing run reports a top-1 match and a maximum absolute error over the top-5 logits below 0.01. In practice the error is on the order of 1e-4 (single-precision accumulation order differences), and greedy decoding matches HuggingFace token-for-token.

Architecture walkthrough

src/tensor.{h,cpp} is a flat row-major float32 tensor with precomputed strides and rank 1 through 4 indexing. src/ops.{h,cpp} holds the numerics: a cache-blocked matmul, a batched matmul, bias broadcast, row-wise softmax, the GPT-2 tanh-approximation GELU, and layer norm over the last dimension. src/tokenizer.{h,cpp} is the byte-level BPE tokenizer: the GPT-2 regex pre-tokenizer, the byte-to-unicode table, the merge ranks from vocab.bpe, and a minimal JSON reader for encoder.json. src/weights.{h,cpp} parses the safetensors header (an 8-byte little-endian length followed by a JSON map) and loads every F32 tensor, then arranges them into a typed GPT2Weights struct. src/transformer.{h,cpp} is the model: token and position embeddings, twelve pre-norm transformer blocks with multi-head causal self-attention and a 4x GELU MLP, a final layer norm, and the tied-embedding output projection, plus the KV cache, the samplers, and the generation loop. src/main.cpp is the CLI and src/bench.cpp is the profiler.

Key implementation decisions

Matmul loop order. The core matmul iterates (i, k, j) rather than the textbook (i, j, k). With row-major storage, fixing i and k lets the inner loop over j stream contiguously through a row of B and a row of C, scaling by the scalar A[i,k]. This keeps the hot inner loop sequential in memory and auto-vectorizable, which matters because the two MLP matmuls (768x3072 and 3072x768) dominate decode time. There is no BLAS and no hand-written SIMD.

KV cache design. Each layer keeps two growable std::vector<float> buffers, one for keys and one for values, each laid out as [total_len, n_embd] to mirror the QKV projection output. On the prompt pass (n_past == 0) all rows are appended at once; on each subsequent step a single new token's K and V rows are appended, and attention for the new query runs against the full cached history. This makes the per-step cost linear in context length instead of quadratic, and a query at absolute position p simply attends to cached positions 0..p, so the causal mask is implicit in the loop bound rather than an explicit -1e10 add.

BPE from scratch instead of loaded. The tokenizer is reimplemented rather than shelling out to a library because tokenization is where silent correctness bugs hide: the byte-to-unicode mapping, the ordered merge application, and the regex pre-tokenization all have to match OpenAI's encoder.py exactly or the token IDs drift in ways that are invisible until the logits are wrong. The encode path is checked against the canonical encoder for several inputs, including contractions and mixed whitespace, and decode(encode(s)) == s round-trips for arbitrary bytes. The one documented approximation is Unicode letter/number classification for the regex stage: ASCII is exact, and non-ASCII codepoints are treated as letters, which is correct for Latin text but can merge exotic Unicode punctuation into letter runs. The byte-level encoding itself is exact for all 256 byte values, so round-tripping always holds.

Tensor memory layout. Tensors are a single contiguous std::vector<float> in row-major order with explicit strides, never a nested or block-sparse structure. This keeps every operation a straightforward pointer walk, lets weights be memcpy'd directly out of the safetensors blob (which is already row-major F32), and means the GPT-2 Conv1D weights, stored as [in, out], are consumed as-is by x @ W with no transpose at load time.

Benchmark results

Measured on Apple Silicon (arm64), single-threaded, compiled with clang at -O2, generating 100 tokens after a 7-token prompt:

Benchmark Results
-----------------
Prompt tokens:     7
Generated tokens:  100
Total time:        2.20s
Prompt forward:    56.1 ms
Throughput:        46.6 tok/s (gen)
Per-token latency: mean 21.5 ms, p95 22.5 ms
Peak memory:       1.627 GB

Per-layer breakdown (avg over 12 layers, 100 gen steps):
  attention:    0.172 ms (35.0%)
  mlp:          0.317 ms (64.5%)
  layer_norm:   0.003 ms ( 0.5%)

At single-token decode with short context, the MLP matmuls dominate (about two thirds of per-layer time) and attention is comparatively cheap; the attention share grows with context length as the per-step scores-and-values work scales with the cache size. Peak resident memory is dominated by the 124M parameters held as float32 (roughly 0.5 GB) plus the embedding and logit working set.

What is missing versus production engines

This is a reference implementation, not a serving stack. There is no quantization: weights are kept in float32, so an int8 or int4 build would cut memory roughly 4x to 8x and speed up the memory-bound matmuls. There is no explicit SIMD or threading; the matmul relies on the compiler's auto-vectorization and runs on a single core, where a real engine would use blocked, multithreaded, NEON/AVX kernels or a tuned BLAS. There is no batching, so multiple sequences cannot share a forward pass, and the KV cache is a plain growable buffer with no paging, eviction, or preallocation to a maximum context. The sampler covers greedy and temperature/top-k only, with no top-p, repetition penalty, or beam search, and the logit projection always computes the full vocabulary. These are deliberate omissions in favor of code that is short enough to read end to end and that validates exactly against the reference.

About

From-scratch GPT-2 (124M) inference engine in C++17: byte-level BPE tokenizer, safetensors loader, cache-blocked matmul, and a KV-cached transformer. No ML frameworks; logits match HuggingFace to ~1e-4.

Topics

Resources

Stars

Watchers

Forks

Contributors

Languages