$ show rasbt/LLMs-from-scratch
This repository is the official companion code for the book *Build a Large Language Model (From Scratch)* by Sebastian Raschka (105k+ stars). It provides a step-by-step, educational implementation of GPT-2 and its modern variants (Llama 3, Qwen3, Gemma 3/4, Olmo3) entirely in PyTorch, organized by book chapters (2–7 plus appendices). The codebase is a mix of Jupyter notebooks for interactive learning and Python modules that form a reusable pip package (`llms-from-scratch`), accompanied by extensive CI/CD pipelines, unit tests, and performance analysis scripts. Architecturally, the core is a modular transformer library that includes tokenization (BPE, tiktoken, custom), multi-head attention with multiple variants (GQA, MLA, SWA, DSA, cross-layer KV sharing), feedforward blocks, normalization (LayerNorm, RMSNorm), and LoRA. The code supports pretraining from scratch, classification finetuning (spam, IMDB), instruction finetuning (with/without LoRA), and preference tuning (DPO). It can load pretrained weights from OpenAI, Hugging Face, and local checkpoints. The package exposes clean interfaces (`GPTModel`, `generate`, `train_model_simple`) and includes bonus materials for KV cache implementations, mixture-of-experts, memory estimators, and user interfaces (Chainlit). While primarily educational, the code exhibits production-quality patterns: isolated attention modules, configurable model architecture, well-structured training loops, and robust weight loading with shape validation. The repository is actively maintained with recent commits, cross-platform support (Linux, macOS, Windows), and a strong emphasis on reproducibility (fixed seeds, deterministic settings, environment definitions via pixi/uv/conda). Duplication exists across bonus notebooks and a few hardcoded paths, but the core package and test suite are well-coordinated.
The LLM identified 4 files worth highlighting for reuse:
Implements a GPT-2 model augmented with a key-value (KV) cache for efficient autoregressive generation. The KV cache is stored as a list of tuples (keys, values) and reused across decode steps, avoiding recomputation of past attention states. This pattern is isolated (no deep internal imports beyond torch), demonstrates non-obvious understanding of transformer inference optimization, and can be extracted verbatim to speed up any decoder-only model. It also includes a streaming generation function (`generate_text_simple_cached`) that returns tokens incrementally, showcasing a real-time inference pattern.
Contains a self-contained implementation of Low-Rank Adaptation (LoRA) for PyTorch linear layers. The `LinearWithLoRA` class wraps an existing `nn.Linear` with a low-rank decomposition (A, B matrices) and provides methods to merge/unmerge weights for inference. The `replace_linear_with_lora` function recursively swaps linear layers in any model, supporting configurable rank, alpha, and target modules. This pattern is isolated, highly reusable for parameter-efficient fine-tuning of any transformer, and demonstrates clean modular design without requiring external LoRA libraries.
A Jupyter notebook that implements Byte Pair Encoding (BPE) tokenization entirely from scratch, including training (`get_merges`), encoding, and decoding with a `BPETokenizerSimple` class. It handles GPT-2-style merge order, special tokens, and regex-based word splitting. The notebook compares its output against OpenAI's tiktoken and the original GPT-2 BPE files for correctness. This component is a reusable teaching tool for subword tokenization and demonstrates a clean, isolated implementation that can be extracted into a standalone library for educational or custom tokenizer needs.
Exports the core `GPTModel` and `GPTModelFast` classes with configurable architecture: embedding dimension, number of layers, attention heads, feedforward expansion, dropout, and context length. It assembles transformer blocks (`TransformerBlock`) from reusable submodules (`MultiHeadAttention`, `FeedForward`, `LayerNorm`). The class structure is clean, isolated, and parameterized via a config dictionary, making it easy to instantiate different GPT variants (e.g., GPT-2 small, medium). This component is the central architectural pattern of the repository and demonstrates high craftsmanship in modular transformer design.