tiny-llm 🤏🧠📚

A hands-on, notebook-by-notebook journey through how LLMs actually work — built entirely from scratch, one small concept at a time.

This is a learning project. Every stage is deliberately small enough to read top-to-bottom in one sitting, deliberately trivial enough to train on a laptop CPU in seconds, and deliberately unglamorous — no framework doing the heavy lifting, no imported "transformer" class. Just raw PyTorch, a training loop, and enough curiosity to keep asking "but why does this work?" 🎓

If you've ever wanted to understand GPT from the tokenizer up, this repo is the scenic route.

🧭 Why this exists

Most "build an LLM" tutorials jump straight into multi-head attention and a few thousand lines of code. This repo takes the opposite approach: start with the smallest possible model that can still learn something, and grow it — piece by piece, notebook by notebook — until it's a real transformer capable of loading actual GPT-2 weights.

Each step introduces exactly one new idea, so you can watch the model gain capability in slow motion:

flowchart LR
    A["dummy 🌱\n1 linear weight"] --> B["v0 🔤\n+ embeddings"]
    B --> C["v1 👀\n+ self-attention"]
    C --> D["v2 💾\n+ save model"]
    D --> E["v3 📂\n+ load model"]
    E --> F["v4 🎯\n+ fine-tuning"]
    F --> G["final 🚀\nfull GPT-2 (355M)"]

    style A fill:#e8f5e9,stroke:#2e7d32
    style B fill:#e3f2fd,stroke:#1565c0
    style C fill:#fff3e0,stroke:#ef6c00
    style D fill:#f3e5f5,stroke:#6a1b9a
    style E fill:#f3e5f5,stroke:#6a1b9a
    style F fill:#fce4ec,stroke:#ad1457
    style G fill:#fffde7,stroke:#f9a825

Loading

📖 The roadmap

Stage	Notebook	What it adds	Trainable params
🌱 dummy	llm-from-scratch-dummy/dummy.ipynb	Word-level tokenizer + a bare y = w·x + b model. Proves the pipeline: tokenize → dataloader → forward → loss → backprop → generate.	2
🔤 v0	llm-from-scratch-v0/v0.ipynb	Real token + positional embeddings feeding a linear output head, trained with cross-entropy over a 70-word vocab.	3,504
👀 v1	llm-from-scratch-v1/v1.ipynb	A single causal self-attention layer — queries, keys, values, a causal mask, and softmax attention weights (plus a peek at how attention entropy shifts across blocks).	5K+
💾 v2	llm-from-scratch-v2-v3-v4-final/v2.ipynb	Stacks multiple transformer blocks and saves trained weights to disk (state_dict).	—
📂 v3	llm-from-scratch-v2-v3-v4-final/v3.ipynb	Loads the saved weights back into a fresh model and generates from them — no retraining needed.	—
🎯 v4	llm-from-scratch-v2-v3-v4-final/v4.ipynb	Fine-tunes the saved model on new data, showing how a pretrained checkpoint adapts to a new corpus.	—
🚀 final	llm-from-scratch-v2-v3-v4-final/model.py	A complete, from-scratch GPT-2 architecture (multi-head attention, LayerNorm, GELU feed-forward, residual streams) that downloads and loads real OpenAI GPT-2 (355M) weights and generates text.	355M
🧩 bonus	embedding-model-main/ (separate branch)	A standalone NumPy-only word-embedding demo — no PyTorch, no autograd, just Jaccard similarity and gradient-free updates pulling similar-context words together in 2D space.	—

🏗️ The architecture, visualized

Everything from v1 onward converges toward the same shape: the architecture from "Attention Is All You Need", applied to decoder-only (GPT-style) language modeling.

The full model (v2 → final GPT-2)

flowchart TD
    IN["Input token ids"] --> TOK["Token Embedding\n(nn.Embedding)"]
    IN --> POS["Positional Embedding\n(nn.Embedding)"]
    TOK --> SUM(("+"))
    POS --> SUM
    SUM --> DROP["Dropout"]
    DROP --> TB

    subgraph TB["Transformer Block × N"]
        direction TB
        LN1["LayerNorm"] --> MHA["Masked Multi-Head\nSelf-Attention"]
        MHA --> ADD1(("+ residual"))
        ADD1 --> LN2["LayerNorm"]
        LN2 --> FF["Feed Forward\n(Linear → GELU → Linear)"]
        FF --> ADD2(("+ residual"))
    end

    TB --> FN["Final LayerNorm"]
    FN --> HEAD["Linear Output Head\n(tied to token embedding)"]
    HEAD --> OUT["Next-token logits"]

Loading

Inside one attention head

flowchart LR
    X["Input embeddings"] --> Q["Query = X · Wq"]
    X --> K["Key = X · Wk"]
    X --> V["Value = X · Wv"]
    Q --> SC["Scores = Q · Kᵀ / √d"]
    K --> SC
    SC --> MASK["Causal mask\n(block future tokens)"]
    MASK --> SM["Softmax"]
    SM --> CTX["Context = weights · V"]
    V --> CTX

Loading

v1 implements this with a single head; the v2-v3-v4-final model generalizes it to multi-head attention by splitting d_out across NUM_HEAD parallel heads, exactly as described in the paper included at Attention Is All You Need.pdf.

📂 Repo structure

tiny-llm/
├── llm-from-scratch-dummy/         # Stage 0 — tokenize + 1 linear weight
├── llm-from-scratch-v0/            # Stage 1 — embeddings
├── llm-from-scratch-v1/            # Stage 2 — self-attention
└── llm-from-scratch-v2-v3-v4-final/
    ├── v2.ipynb                    # save
    ├── v3.ipynb                    # load
    ├── v4.ipynb                    # fine-tune
    ├── model.py                    # full GPT-2 architecture
    ├── download.py                 # fetches real GPT-2 (355M) weights
    └── helper.py                   # maps OpenAI checkpoint → our model

Each folder has its own README with a deeper dive into that stage.

🚀 Getting started

# pick a stage
cd llm-from-scratch-v1

# set up an isolated environment
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

# open the notebook and run it top to bottom
jupyter notebook v1.ipynb

To play with a real, pretrained GPT-2:

cd llm-from-scratch-v2-v3-v4-final
pip install -r requirements.txt
python download.py   # downloads GPT-2 355M from OpenAI
python model.py       # loads weights and generates text

🌱 Contributing — this is a learning project, so dig in!

This repo is meant to be read, run, broken, and extended — not just cloned. If you're learning LLMs too, there's a lot of room to grow it further:

• 🧪 Add a v5 — KV-caching, rotary embeddings, or grouped-query attention
• 📊 Add visualizations (attention maps, loss curves, embedding projections)
• 🧵 Port a stage to a different framework (JAX, pure NumPy, MLX) to compare
• 🐛 Spot a bug or a confusing explanation? Open an issue or a PR
• 📝 Improve a stage's README with more intuition or diagrams

No contribution is too small — clarifying a comment or fixing a typo in a README is just as welcome as a new notebook. Fork it, branch off, and send a PR. Let's learn this together. 🙌

📚 References

• Vaswani et al., Attention Is All You Need (2017)
• OpenAI GPT-2 — source of the pretrained weights used in the final stage

📝 Notes

• Diagrams (*.excalidraw) and local virtual environments (.venv/) are intentionally excluded from version control — they're scratch/local artifacts, not source of truth.
• Each stage lives on its own branch during development; check the open PRs/branches for work in progress.