llm-library is a minimal language-model library I built from first principles in PyTorch. It covers the path from raw text and BPE tokenization to training a decoder-only Transformer, saving checkpoints, and generating text.
The goal was not to hide the implementation behind higher-level abstractions, but to write the core components explicitly and understand what happens inside the model and training loop.
The source code is available in the public llm-library repository. The current public milestone is v0.1.0.
From Text to Tokens
I implemented BPE training, encoding, and decoding. The tokenizer handles special-token boundaries and can encode an iterable lazily, so a large input file does not have to be loaded entirely into memory.
The complete data path includes training a vocabulary and merge rules, encoding the training and validation corpora into token IDs, and loading random batches from the resulting NumPy arrays.
Transformer Architecture
The model is a decoder-only Transformer composed of token embeddings, repeated pre-normalization Transformer blocks, a final RMS normalization, and a language-model head.
Each block combines:
- RMSNorm;
- packed query, key, and value projections;
- rotary positional embeddings applied to queries and keys;
- causal multi-head self-attention;
- a SwiGLU feed-forward network;
- residual connections around attention and the feed-forward network.
I also implemented the underlying linear and embedding layers, softmax, and cross-entropy loss. PyTorch still provides tensors, automatic differentiation, and the low-level kernels; the purpose of the library is to make the model and training logic explicit rather than to replace those foundations yet.
Optimization, Training, and Generation
The training pipeline includes a custom AdamW optimizer, cosine learning-rate scheduling with warmup, gradient clipping, validation, metric recording, and checkpoint save and resume support.
The generation path loads a trained checkpoint and performs autoregressive sampling with temperature and top-p controls. The input is truncated to the configured context length as generation progresses.
Verification
The implementations are checked through correctness tests adapted from Stanford CS336. They cover the tokenizer, BPE training, model components, loss, optimizer, learning-rate schedule, gradient clipping, batching, and checkpoint serialization.
GitHub Actions reproduces the locked development environment, runs Ruff, and executes the complete test suite on every push and pull request. The public v0.1.0 state passes both the local checks and CI.
TinyStories Training
I used TinyStories to verify that the individual components also work together in an actual training run. The repository includes the vocabulary, merge rules, reproducible commands, and checkpoints obtained after training on approximately 16 million and 80 million tokens.
The documented configuration uses a context length of 256, four Transformer layers, a model dimension of 512, 16 attention heads, and a batch size of 32. A 5,000-step run with this configuration processes approximately 40 million tokens.
To make the progression inspectable, I generated text from the same prompt using checkpoints after approximately 16, 40, and 80 million training tokens. The later samples are more coherent within this specific run, but they are qualitative examples rather than a general model-quality benchmark.
Scope and Next Steps
Version 0.1.0 focuses on the core model, pretraining, and inference components of a Transformer language model. It does not include distributed training, supervised fine-tuning, RLHF or RLVR, large-scale evaluation, or production serving.
The current implementation relies on PyTorch tensors and autograd, while einops is used in some places to keep tensor operations readable. The next directions I want to explore are a custom autograd engine and efficient implementations of performance-critical components in C++.