commit a8262fc01e4d84a744c7693c047460b2a770c0a9 Author: wehub-resource-sync Date: Mon Jul 13 13:10:22 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..dc4672f --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,51 @@ +name: docs + +# Build the MkDocs Material site and publish it via GitHub Pages. +# Pages "Source" must be set to "GitHub Actions" (Settings -> Pages) — which it already is. +# No gh-pages branch needed; this uses the official Pages deployment actions. + +on: + push: + branches: [main] + paths: + - "docs/**" + - "mkdocs.yml" + - "POST_TRAINING.md" + - ".github/workflows/docs.yml" + workflow_dispatch: {} + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: pages + cancel-in-progress: false + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install docs dependencies + run: pip install mkdocs mkdocs-material pymdown-extensions + - name: Build the site (strict) + run: mkdocs build --strict --site-dir ./site + - uses: actions/configure-pages@v5 + - uses: actions/upload-pages-artifact@v3 + with: + path: ./site + + deploy: + needs: build + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - id: deployment + uses: actions/deploy-pages@v4 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8420144 --- /dev/null +++ b/.gitignore @@ -0,0 +1,22 @@ +__pycache__/ +*.pyc +*.pyo +.ipynb_checkpoints/ + +# build / packaging artifacts (generated by `pip install -e .`) +*.egg-info/ +build/ +dist/ + +# generated docs site +site/ +.superpowers/ + +# data / model artifacts (kept on the /ephemeral disk, not committed) +data/ +models/ +*.h5 +*.pt +*.jsonl +wandb/ +/ephemeral/ diff --git a/.streamlit/config.toml b/.streamlit/config.toml new file mode 100644 index 0000000..3055b11 --- /dev/null +++ b/.streamlit/config.toml @@ -0,0 +1,13 @@ +[theme] +primaryColor = "#16a085" +backgroundColor = "#ffffff" +secondaryBackgroundColor = "#f3faf8" +textColor = "#0a3d33" +font = "sans serif" + +[server] +headless = true +runOnSave = true + +[browser] +gatherUsageStats = false diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..4f1b584 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Fareed Khan + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/POST_TRAINING.md b/POST_TRAINING.md new file mode 100644 index 0000000..5f2c10e --- /dev/null +++ b/POST_TRAINING.md @@ -0,0 +1,196 @@ +# Post-Training from Scratch: SFT · Reward Model · PPO · DPO · GRPO + +This adds a complete, **from-scratch** (pure PyTorch — no `trl`/`peft`/`transformers`) +post-training suite on top of the repo's own `Transformer`, covering the full modern +pipeline used to turn a base LM into an aligned, reasoning model: + +``` +Base (pretrained) ──► SFT ──► Reward Model ──► PPO ┐ + │ ├─► GRPO / RLVR (GSM8K) + └────────► DPO / ORPO / KTO ─┘ +``` + +Everything is self-contained on the repo's custom model and trained/evaluated on **real +public datasets** (Alpaca, Dolly, Anthropic HH-RLHF, UltraFeedback, GSM8K). The headline +metric is **greedy GSM8K accuracy across stages**. + +> **Expectation:** a ~400M model pretrained from scratch on 2×H100 is coherent and +> instruction-followable and shows real before/after gains at each stage, but its +> *absolute* GSM8K score stays modest — frontier numbers need far more pretraining +> compute. The value here is the authentic end-to-end pipeline on real data with real eval. + +--- + +## 0. Environment + +H100s need CUDA-12 wheels (the original `requirements.txt` pins cu118 for the legacy +pretraining path and is left untouched). Use a separate venv + `requirements-post.txt`, +and keep all large artifacts on the big `/ephemeral` disk. + +```bash +python3 -m venv /ephemeral/venv && source /ephemeral/venv/bin/activate +pip install -r requirements-post.txt # torch cu121 + datasets/wandb/tiktoken/h5py +export HF_HOME=/ephemeral/hf_cache +``` + +Run everything from the repo root with `PYTHONPATH=.`. Single GPU: `python scripts/X.py`. +Both GPUs: `torchrun --standalone --nproc_per_node=2 scripts/X.py` (DDP + bf16, one code path). + +Config lives in [config/post_training_config.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/post_training_config.py) as +per-stage dataclasses; **any field is a CLI override**, e.g. `--lr 2e-5 --batch_size 16`. +The default base model is ~400M (`n_embed=1024, n_head=16, n_blocks=24, context_length=1024`). + +--- + +## 1. Pretrain the base (the long pole) + +```bash +# one-time data prep (Pile -> flat-token HDF5 on /ephemeral) +PYTHONPATH=. python scripts/prepare_pretrain_data.py --split val --out /ephemeral/data/pile_dev.h5 +PYTHONPATH=. python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out /ephemeral/data/pile_train.h5 + +# pretrain (multi-day for full quality; checkpoints continuously to /ephemeral/ckpts/base_pretrained.pt) +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py +``` + +[scripts/pretrain_base.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/pretrain_base.py) upgrades the original +`train_transformer.py` recipe with DDP, bf16 autocast, gradient accumulation, a cosine LR +schedule with warmup, and periodic checkpointing — everything needed to train a mid-size +model on 2×H100. The original training script is untouched. + +--- + +## 2. SFT (instruction tuning) + +```bash +PYTHONPATH=. python scripts/prepare_sft_data.py --context_length 1024 # Alpaca+Dolly+GSM8K -> packed HDF5 +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_sft.py # -> /ephemeral/ckpts/sft.pt +``` + +- **From-scratch loss:** prompt-masked next-token CE in + [src/post_training/sft.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/sft.py) (`sft_loss`). Only assistant tokens + are trained (mask from the chat template); sequence **packing** fills each row to the + context length. +- **Chat format:** [src/post_training/chat_template.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py). + The r50k_base tokenizer has only `<|endoftext|>` as a special token, so role markers + (`<|user|>`, `<|assistant|>`, ``, ``) are ordinary tokens the model learns. + GSM8K is reformatted into `…N` so the model learns the + exact output structure the RL verifier rewards. +- **Eval:** masked dev perplexity + greedy GSM8K dev accuracy. + +## 3. Reward Model + +```bash +PYTHONPATH=. python scripts/prepare_preference_data.py --source both # HH-RLHF + UltraFeedback -> JSONL +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_reward.py # -> reward.pt +``` + +- **From-scratch:** a scalar reward head on the SFT backbone + ([src/post_training/reward_model.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/reward_model.py)), trained with the + **Bradley-Terry** pairwise loss in + [src/post_training/reward_train.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/reward_train.py). The reward is read + off the last real token (causal attention makes right-padding safe — no attention mask + needed). **Eval:** held-out preference accuracy (expect ~0.65–0.75 on noisy real data). + +## 4. DPO / ORPO / KTO + +```bash +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo --beta 0.1 +# --loss_type orpo (reference-free, folds SFT+alignment into one stage) +# --loss_type kto (unpaired, reference-KL baseline) +``` + +- **From-scratch:** all three objectives in [src/post_training/dpo.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/dpo.py). + Policy initialized from SFT; a frozen deep copy is the reference (ORPO needs none). + Operates on summed response log-probs from + [src/post_training/rollout.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rollout.py) (`sequence_logprobs`). + **Eval:** implicit-reward accuracy/margins + GSM8K dev. + +## 5. PPO (classic RLHF) + +```bash +PYTHONPATH=. python scripts/prepare_rl_prompts.py # GSM8K + arithmetic warm-up -> JSONL +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py --reward_source verifier +# --reward_source rm to use the trained reward model instead of the GSM8K checker +``` + +- **From-scratch:** GAE, clipped policy/value losses in + [src/post_training/ppo.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/ppo.py); the actor-critic shares the backbone + via [src/post_training/value_head.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/value_head.py). Each iteration: + roll out → score (verifier or RM) → add per-token **KL-to-reference** penalty → GAE → + several clipped-surrogate epochs. **Eval:** reward / KL / clip-fraction / value-loss curves + and greedy GSM8K test accuracy. + +## 6. GRPO / RLVR (the 2025 frontier; DeepSeek-R1 style) + +```bash +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py --group_size 8 +``` + +- **From-scratch:** group-relative advantages + token-level clipped surrogate with a k3 KL + penalty in [src/post_training/grpo.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/grpo.py). **No critic** — the + baseline is each prompt's own group of G samples. An **arithmetic curriculum** runs for + the first `curriculum_iters` iterations so the policy has non-zero reward variance before + full GSM8K. **Eval:** mean group reward, informative-group fraction, KL, GSM8K test accuracy. + +--- + +## 7. Inference / chat (any stage checkpoint) + +[scripts/chat.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/chat.py) loads **any** checkpoint (base/sft/dpo/ppo/grpo) — +reading the model dims from the checkpoint itself — and generates with the chat template +(instruction models) or as raw continuation (the base model): + +```bash +# instruction-tuned models (chat template applied automatically) +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt --prompt "What is 13 + 29?" +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/grpo.pt --prompt "..." --greedy +# base model continuation +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/base_pretrained.pt --raw --prompt "Once upon a time" +# interactive REPL (omit --prompt); sampling via --temperature/--top_p/--top_k or --greedy +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt +``` + +Generation reuses the same tested core as training/eval +([src/post_training/rollout.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rollout.py), +[src/post_training/inference.py](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py)). + +## 8. The across-stages results table + +```bash +for s in base_pretrained sft dpo ppo grpo; do + PYTHONPATH=. python scripts/eval_post_training.py --ckpt /ephemeral/ckpts/$s.pt \ + --label $s --limit 200 --append /ephemeral/logs/stage_table.jsonl +done +PYTHONPATH=. python scripts/eval_post_training.py --table /ephemeral/logs/stage_table.jsonl +``` + +Every trainer also writes a JSONL metrics file under `/ephemeral/logs/` (plottable without +any external service); pass `--use_wandb true` to also mirror to Weights & Biases. + +--- + +## Design notes (why it's built this way) + +- **Wrap, don't rewrite.** The educational `Transformer`/`Block`/`Head`/`MLP` are unchanged + except one additive method, `forward_hidden` (returns post-final-LN hidden states the + heads consume). Value head, reward head, and all RL log-prob math compose around it. +- **Causal attention ⇒ right-padding is safe.** The last real token never attends to + padding after it, so RM (last-token reward) and DPO (masked response) need no attention + mask; the response mask zeros padded positions in the loss. +- **fp32 log-probs.** PPO/GRPO/DPO subtract log-probs, so they're always computed in fp32 + even under bf16 autocast. +- **Context cap.** Learned absolute positions cap any sequence at `context_length`; rollouts + enforce `prompt + generation ≤ context_length`. +- **Reward hacking / KL control.** Verifier rewards are correctness-dominant with a small, + bounded format bonus; a KL-to-reference penalty anchors RL to the SFT policy. + +## Tests + +```bash +PYTHONPATH=. python tests/test_post_training_smoke.py # core math: log-probs, heads, parsing, masking +``` + +Each trainer also runs end-to-end on a tiny model in seconds (see the smoke commands used +during development), and the PPO/GRPO math has standalone unit checks (GAE, clipped losses, +group advantages, k3 KL). diff --git a/README.md b/README.md new file mode 100644 index 0000000..5eb6adb --- /dev/null +++ b/README.md @@ -0,0 +1,762 @@ +![main image](https://cdn-images-1.medium.com/max/5200/1*r99Hq3YBd5FTTWLNYKKvPw.png) + +
+ + +# Train LLM From Scratch + +![Python](https://img.shields.io/badge/Python-3.9%2B-blue) ![License](https://img.shields.io/badge/License-MIT-green) ![Contributions](https://img.shields.io/badge/Contributions-Welcome-blue) [![Docs](https://img.shields.io/badge/Docs-Available-success)](https://fareedkhan-dev.github.io/train-llm-from-scratch/) + +**I am Looking for a PhD position in AI**. [GitHub](https://github.com/FareedKhan-dev) + +
+ +I implemented a transformer model from scratch using PyTorch, based on the paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). You can use my scripts to train your own **billion** or **million** parameter LLM using a single GPU. + +This started as a pretraining tutorial. It now goes all the way from raw text to an aligned, reasoning style model, with every algorithm hand written in plain PyTorch (no `trl`, no `peft`, no `transformers`). The whole journey is one idea repeated: turn text into numbers, predict the next token, then keep changing the data and the loss until the model does what we want. + +![From raw text to an aligned reasoning model](images/00_pipeline.png) + +Here is the path we will walk, end to end: + +``` +raw text -> tokens -> a Transformer -> next-token loss -> a base model +base model -> SFT -> Reward Model -> {PPO, DPO} -> GRPO -> evaluation and chat +``` + +Below is the output of a trained 13 million parameter LLM, just so you can see where the small end of this starts: + +``` +In ***1978, The park was returned to the factory-plate that +the public share to the lower of the electronic fence that +follow from the Station's cities. The Canal of ancient Western +nations were confined to the city spot. The villages were directly +linked to cities in China that revolt that the US budget and in +Odambinais is uncertain and fortune established in rural areas. +``` + + +## Table of Contents +- [Who this is for](#who-this-is-for) +- [Prerequisites and Training Time](#prerequisites-and-training-time) +- [Setup](#setup) +- [Code Structure](#code-structure) +- [Step 1: Preparing the Data](#step-1-preparing-the-data) +- [Step 2: The Model, Built From Small Pieces](#step-2-the-model-built-from-small-pieces) + - [Multi Layer Perceptron (MLP)](#multi-layer-perceptron-mlp) + - [Single Head Attention](#single-head-attention) + - [Multi Head Attention](#multi-head-attention) + - [The Transformer Block](#the-transformer-block) + - [The Full Transformer](#the-full-transformer) +- [Step 3: Pretraining the Base Model](#step-3-pretraining-the-base-model) +- [Step 4: Generating Text](#step-4-generating-text) +- [Step 5: Post-Training, Turning a Base Model Into an Assistant](#step-5-post-training-turning-a-base-model-into-an-assistant) + - [SFT (Supervised Fine-Tuning)](#sft-supervised-fine-tuning) + - [The Reward Model](#the-reward-model) + - [DPO, ORPO and KTO](#dpo-orpo-and-kto) + - [PPO](#ppo) + - [GRPO / RLVR](#grpo--rlvr) +- [Step 6: Evaluation](#step-6-evaluation) +- [Step 7: Talking to the Model](#step-7-talking-to-the-model) +- [The Streamlit Control Panel](#the-streamlit-control-panel) +- [The Documentation Site](#the-documentation-site) +- [Run the Whole Thing](#run-the-whole-thing) +- [What's Next](#whats-next) + +## Who this is for + +I tried to write this so one page works for very different readers: + +- If you are a **student**, read top to bottom. Every block of code comes after a plain explanation of what it does and why, and most blocks are followed by the output you should expect. +- If you are a **developer**, the commands and file paths are all here. You can copy, run, and read the referenced source files directly. +- If you are a **researcher**, the post-training half is the interesting part: SFT, a Bradley-Terry reward model, PPO with GAE, DPO/ORPO/KTO, and GRPO, all from scratch on the same small Transformer, trained on real public datasets. + +Every diagram in this README is colored the same way, so the colors mean something: + +- green is raw data +- teal is stored, tokenized data on disk +- blue is a plain processing step +- yellow is the model or a training step +- orange is the reinforcement learning and reward parts +- red is a loss +- grey is a saved checkpoint +- purple is the final output or evaluation + +## Prerequisites and Training Time + +You need a basic understanding of object oriented programming, neural networks, and PyTorch. Below are some resources to help you get started: + +| Topic | Video Link | +|---------------------|-----------------------------------------------------------| +| OOP | [OOP Video](https://www.youtube.com/watch?v=Ej_02ICOIgs) | +| Neural Network | [Neural Network Video](https://www.youtube.com/watch?v=Jy4wM2X21u0) | +| Pytorch | [Pytorch Video](https://www.youtube.com/watch?v=V_xro1bcAuA) | + +You will need a GPU to train. A free Colab or Kaggle T4 is enough for the 13 million parameter model, but it will not fit a billion parameter model. Here is a rough guide: + +| GPU Name | Memory | 2B LLM Training | 13M LLM Training | Max Practical LLM Size (Training) | +|--------------------------|--------|-----------------|------------------|-----------------------------------| +| NVIDIA A100 | 40 GB | ✔ | ✔ | ~6B to 8B | +| NVIDIA V100 | 16 GB | ✘ | ✔ | ~2B | +| NVIDIA RTX 4090 | 24 GB | ✔ | ✔ | ~4B | +| NVIDIA RTX 5090 | 32 GB | ✔ | ✔ | 13M verified, larger configs TBD | +| NVIDIA RTX 3090 | 24 GB | ✔ | ✔ | ~3.5B to 4B | +| NVIDIA RTX 4080 | 16 GB | ✘ | ✔ | ~2B | +| NVIDIA RTX 4060 | 8 GB | ✘ | ✔ | ~1B | +| Tesla T4 | 16 GB | ✘ | ✔ | ~1.5B to 2B | + +If a large config runs out of memory, the pretraining script has opt-in flags (`--amp`, `--grad-checkpointing`, `--grad-accum`) that bring the memory down a lot. More on those later. + +## Setup + +Clone the repository and install it in editable mode. The editable install puts `config`, `src`, `data_loader`, and `ui` on your import path, so you do not need to set `PYTHONPATH` by hand anymore: + +```bash +git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git +cd train-llm-from-scratch +pip install -e . +``` + +There are optional extras for the parts you want: + +```bash +pip install -e ".[train]" # datasets + wandb, for downloading data and logging +pip install -e ".[ui]" # streamlit + pandas + altair, for the control panel +pip install -e ".[docs]" # mkdocs, for the documentation site +pip install -e ".[all]" # everything +``` + +There are two config systems, and it helps to know which is which from the start: + +- `config/config.py` is the original, simple config for the legacy pretraining script `scripts/train_transformer.py`. It is plain Python constants. +- `config/post_training_config.py` plus the JSON files in `configs/` drive everything else (pretraining the bigger base, SFT, reward, DPO, PPO, GRPO). You edit a small JSON file per stage, and any field can also be overridden on the command line, for example `--lr 2e-5 --batch_size 16`. + +For fast checks there is a tiny `configs/smoke/` variant of every stage that shrinks the model so a full run finishes in seconds on a CPU or a single GPU. + +## Code Structure + +```bash +train-llm-from-scratch/ +├── src/ +│ ├── models/ # the Transformer, built from small pieces +│ │ ├── mlp.py # the feed-forward block +│ │ ├── attention.py # single head and multi head attention +│ │ ├── transformer_block.py # one block: attention + MLP + residuals +│ │ └── transformer.py # the full model: embeddings + blocks + lm_head +│ └── post_training/ # SFT, reward model, PPO, DPO, GRPO, eval, inference +├── config/ +│ ├── config.py # legacy pretraining config (plain constants) +│ ├── post_training_config.py # dataclasses for every post-training stage +│ └── loader.py # merges defaults < base.json < stage.json < CLI +├── configs/ # editable JSON, one file per stage (+ smoke/) +├── data_loader/ # batch iterators for each kind of data +├── scripts/ # every runnable step lives here +├── ui/ # the Streamlit control panel +├── docs/ # the MkDocs site (theory + diagrams) +├── images/ # the diagrams in this README (+ the generator) +└── pyproject.toml # pip install -e . +``` + +## Step 1: Preparing the Data + +A model only ever sees integers. So the first job is always the same: take text, turn it into token ids, and store those ids on disk in a format that is fast to read during training. We do this four times, once for each kind of training we will do later. + +![The data pipeline](images/01_data.png) + +The four streams are: + +1. **Pretraining text** from [The Pile](https://huggingface.co/datasets/monology/pile-uncopyrighted), stored as a flat array of token ids in an HDF5 file. +2. **Instruction data** (Alpaca, Dolly, GSM8K) for SFT, packed into fixed length rows with a mask that says which tokens are the assistant's answer. +3. **Preference pairs** (Anthropic HH-RLHF, UltraFeedback) for the reward model and DPO, stored as `{prompt, chosen, rejected}`. +4. **RL prompts** (GSM8K and a small arithmetic warm-up) for PPO and GRPO, stored as `{prompt, gold}`. + +### Tokenization + +We use the `r50k_base` tokenizer from OpenAI's `tiktoken`, the same one GPT-3 used. Text becomes a list of integers, and we append a special `<|endoftext|>` token (id 50256) at the end of every document so the model learns where one piece of text stops and the next begins. + +![Tokenization](images/02_tokenization.png) + +For the legacy path, download a slice of The Pile and tokenize it into HDF5: + +```bash +python scripts/data_download.py # downloads the validation file + 1 training shard +python scripts/data_preprocess.py # tokenizes to data/train/pile_train.h5 and data/val/pile_dev.h5 +``` + +The newer, faster path streams and batch-encodes the same data straight into a flat token array: + +```bash +python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5 +python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5 +``` + +Once tokenized, the data is just a long line of integers. Here is a real peek at the validation file I prepared for this README (8.76 million tokens), the first ten ids, and what they decode back to: + +```python +#### OUTPUT #### +dtype: int32 | shape: (8762951,) | total tokens: 8762951 +first 10 token ids: [18610, 286, 3993, 3081, 319, 4088, 11, 4640, 2163, 11] +decoded back to text: +'Effect of sleep quality on memory, executive function, and language + performance in patients with refractory focal epilepsy ...' +``` + +That is the whole idea of tokenization in one output: text in, a flat array of integers out, and the integers decode straight back to the original words. + +### The chat format and loss mask + +For everything after pretraining the model has to know who is talking. The `r50k_base` tokenizer has only one special token, so instead of inventing new ones we use plain text role markers that the model simply learns during SFT. A single turn looks like this (see `src/post_training/chat_template.py`): + +``` +<|user|> +{user content}<|endoftext|><|assistant|> +{assistant content}<|endoftext|> +``` + +For math and reasoning we ask the assistant to show its work in a fixed structure, because the reinforcement learning reward later checks the number inside the answer tags: + +``` +step by step reasoning ...42 +``` + +The important trick is the **loss mask**. When we encode a conversation we also build a 0/1 mask that is 1 only on the assistant tokens (and the `<|endoftext|>` that ends the turn). That way SFT trains the model to write answers, not to parrot the prompt back. Here is the exact code that builds the ids and the aligned mask: + +```python +def encode_chat(messages, add_generation_prompt=False): + ids, mask = [], [] + for m in messages: + role = m["role"] + # Role header is always masked out (we never train the model to emit it). + header_ids = _encode_ordinary(_header_for(role)) + ids.extend(header_ids) + mask.extend([0] * len(header_ids)) + + content_ids = _encode_ordinary(m["content"]) + is_completion = role == "assistant" + ids.extend(content_ids) + mask.extend([1 if is_completion else 0] * len(content_ids)) # train on assistant only + + ids.append(EOT_ID) # turn terminator + mask.append(1 if is_completion else 0) # learn to stop + return ids, mask +``` + +Here is a real rendered conversation and the verifier reward in action, printed from this repo: + +```python +#### OUTPUT #### +rendered chat: +<|user|> +What is 13 + 29?<|endoftext|><|assistant|> +13 + 29 = 4242<|endoftext|> + +extract_answer("42") -> 42.0 +reward_gsm8k("42", 42.0) -> 1.2 # correct AND well formatted +reward_gsm8k("7", 42.0) -> 0.2 # wrong, but it used the format +``` + +And here is one real packed SFT row, showing how only the assistant tokens are trained (the mask is 1 on 48 of the 512 tokens in this row): + +```python +#### OUTPUT #### +tokens shape: (2131, 512) | loss_mask shape: (2131, 512) +row 0: trained (mask=1) tokens = 48 / 512 +row 0 decoded: + <|user|> + What is the world's oldest annual marathon based on the reference text below? ... + <|assistant|> + The Boston Marathon is the world's oldest annual marathon, beginning on April 19th 1897. +``` + +The data prep scripts for these stages are: + +```bash +python scripts/prepare_sft_data.py # Alpaca + Dolly + GSM8K -> sft_packed.h5 +python scripts/prepare_preference_data.py # HH-RLHF + UltraFeedback -> preferences.jsonl +python scripts/prepare_rl_prompts.py # GSM8K + arithmetic -> rl_prompts.jsonl +``` + +## Step 2: The Model, Built From Small Pieces + +A Transformer looks scary as one block of code, so we build it from four small pieces and then stack them. Each piece is a tiny `nn.Module`. We start at the bottom. + +### Multi Layer Perceptron (MLP) + +The MLP is the part of each block that does the per-token "thinking". It takes each token vector, expands it to four times its size, applies a `ReLU`, and squeezes it back down. The expansion gives the layer room to mix features before projecting back. + +![MLP](images/03_mlp.png) + +```python +class MLP(nn.Module): + """A simple Multi-Layer Perceptron with one hidden layer.""" + def __init__(self, n_embed): + super().__init__() + self.hidden = nn.Linear(n_embed, 4 * n_embed) # expand to 4x + self.relu = nn.ReLU() # non-linearity + self.proj = nn.Linear(4 * n_embed, n_embed) # project back down + + def forward(self, x): + x = self.relu(self.hidden(x)) + x = self.proj(x) + return x +``` + +The `__init__` sets up the two linear layers and the activation. The `forward` runs them in order. Input and output shapes are the same, `(B, T, n_embed)`, so blocks can be stacked without any reshaping. The code is in `src/models/mlp.py`. + +### Single Head Attention + +Attention is the part that lets a token look at other tokens. Each head builds three views of the input: a query (what am I looking for), a key (what do I contain), and a value (what I will pass on if chosen). We score every query against every key, scale the scores, hide the future with a causal mask, turn the scores into weights with a softmax, and take a weighted sum of the values. + +![Single Head Attention](images/04_attention_head.png) + +```python +class Head(nn.Module): + """A single attention head with causal masking.""" + def __init__(self, head_size, n_embed, context_length): + super().__init__() + self.key = nn.Linear(n_embed, head_size, bias=False) + self.query = nn.Linear(n_embed, head_size, bias=False) + self.value = nn.Linear(n_embed, head_size, bias=False) + # a lower-triangular matrix used to mask out future positions + self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length))) + + def forward(self, x): + B, T, C = x.shape + k = self.key(x) + q = self.query(x) + scale_factor = 1 / math.sqrt(C) + attn_weights = q @ k.transpose(-2, -1) * scale_factor # (B, T, T) scores + attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # no peeking ahead + attn_weights = F.softmax(attn_weights, dim=-1) + v = self.value(x) + out = attn_weights @ v # weighted sum of values + return out +``` + +The causal mask is what makes this a language model: position `t` can only attend to positions `0..t`, never to the future it is trying to predict. The code is in `src/models/attention.py`. + +### Multi Head Attention + +One head learns one kind of relationship. We want many, running in parallel, so the model can track several patterns at once (for example, a pronoun and the noun it refers to). We run `n_head` heads, concatenate their outputs, and pass the result through one more linear layer. + +![Multi Head Attention](images/05_multi_head_attention.png) + +```python +class MultiHeadAttention(nn.Module): + def __init__(self, n_head, n_embed, context_length): + super().__init__() + self.heads = nn.ModuleList( + [Head(n_embed // n_head, n_embed, context_length) for _ in range(n_head)] + ) + self.proj = nn.Linear(n_embed, n_embed) # mixes the heads back together + + def forward(self, x): + x = torch.cat([h(x) for h in self.heads], dim=-1) # concat along the feature dim + x = self.proj(x) + return x +``` + +Each head works in a smaller subspace of size `n_embed // n_head`, so the concatenation lands right back at `n_embed`. The final projection lets the heads talk to each other. + +### The Transformer Block + +Now we combine attention and the MLP into one block. The block uses pre-norm residual connections: we normalize, run a sub-layer, and add the result back to the input. The "add back" (the residual) is what lets gradients flow through a deep stack without vanishing. + +![The Transformer Block](images/06_transformer_block.png) + +```python +class Block(nn.Module): + def __init__(self, n_head, n_embed, context_length): + super().__init__() + self.ln1 = nn.LayerNorm(n_embed) + self.attn = MultiHeadAttention(n_head, n_embed, context_length) + self.ln2 = nn.LayerNorm(n_embed) + self.mlp = MLP(n_embed) + + def forward(self, x): + x = x + self.attn(self.ln1(x)) # attention sub-layer + residual + x = x + self.mlp(self.ln2(x)) # MLP sub-layer + residual + return x +``` + +Read `x = x + self.attn(self.ln1(x))` as "look at the other tokens, then add what you learned back onto yourself". The MLP line is the same idea for the per-token thinking. The code is in `src/models/transformer_block.py`. + +### The Full Transformer + +Finally we wrap everything. Token ids become vectors through an embedding table, we add a position embedding so the model knows token order, we run the stack of blocks, normalize one last time, and project to vocabulary-sized scores called logits. If we pass targets, the model also returns the cross-entropy loss. + +![The Full Transformer](images/07_transformer.png) + +```python +class Transformer(nn.Module): + def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS): + super().__init__() + self.token_embed = nn.Embedding(vocab_size, n_embed) + self.position_embed = nn.Embedding(context_length, n_embed) + self.attn_blocks = nn.ModuleList( + [Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)] + ) + self.layer_norm = nn.LayerNorm(n_embed) + self.lm_head = nn.Linear(n_embed, vocab_size) + self.register_buffer('pos_idxs', torch.arange(context_length)) + + def forward(self, idx, targets=None): + x = self.forward_hidden(idx) # token + position embeddings, then the blocks + final norm + logits = self.lm_head(x) # (B, T, vocab_size) + loss = None + if targets is not None: + B, T, C = logits.shape + # reshape (not view): the target slice is not contiguous, so .view() fails on CPU + flat_logits = logits.reshape(B * T, C) + targets = targets.reshape(B * T).long() + loss = F.cross_entropy(flat_logits, targets) + return logits, loss +``` + +One small detail worth pointing out: we use `.reshape` and not `.view` on the targets. The target batch is a non-contiguous slice of the data, and `.view` refuses to work on that on CPU. `.reshape` handles both cases and is identical in every other way. The full model, including `forward_hidden` (which the reward and value heads reuse later) and `generate`, lives in `src/models/transformer.py`. + +When we build the model it prints its parameter count. Here are the three sizes used in this repo: + +```python +#### OUTPUT #### +13M small config (n_embed=128, n_head=8, n_blocks=1): 13,142,656 params +this tutorial's base (n_embed=512, n_head=8, n_blocks=8): 77,031,552 params +post-training default (n_embed=1024, n_head=16, n_blocks=24): 406,359,168 params +``` + +## Step 3: Pretraining the Base Model + +Pretraining is the long pole. We read random windows of tokens, ask the model to predict the next token at every position, measure how wrong it was with cross-entropy, and nudge the weights. We repeat that a few thousand times. + +![The pretraining loop](images/08_training_loop.png) + +The simplest version is the original `scripts/train_transformer.py`, which reads `config/config.py` and trains on one GPU. To train the 13 million parameter model, set these values in `config/config.py`: + +```python +VOCAB_SIZE = 50304 +CONTEXT_LENGTH = 128 +N_EMBED = 128 +N_HEAD = 8 +N_BLOCKS = 1 +``` + +then run: + +```bash +python scripts/train_transformer.py +``` + +For long runs you can save periodic checkpoints and resume after an interruption: + +```bash +python scripts/train_transformer.py --checkpoint-every 1000 --keep-last 3 +python scripts/train_transformer.py --resume latest +``` + +If a bigger config does not fit in memory, turn on the opt-in memory savers (all off by default, so default behavior never changes): + +```bash +python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8 +``` + +The bigger, modern path is `scripts/pretrain_base.py`. It is the same recipe with the things you need to train a mid-size base: DistributedDataParallel across GPUs, bf16 autocast, gradient accumulation, a cosine learning-rate schedule with warmup, and periodic checkpoints. One GPU or many, same command shape: + +```bash +# one GPU +python scripts/pretrain_base.py +# both GPUs +torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py +``` + +The core of the loop is small. Each step pulls a batch, runs the forward pass under bf16, scales the loss for gradient accumulation, backpropagates, clips the gradient, and steps the optimizer: + +```python +for micro in range(cfg.grad_accum): + xb, yb = next(batch_iter) + with amp_autocast(cfg.amp_dtype, ctx.device): + _, loss = model(xb, yb) + loss = loss / cfg.grad_accum # so the accumulated gradient is the full-batch mean + loss.backward() +torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) +optimizer.step() +``` + +While it trains it prints the step, the loss, the learning rate, and the throughput, and at eval time it also prints the dev loss and peak GPU memory: + +``` +#### OUTPUT #### +Model parameters: 77,031,552 (~77M) | world_size=2 +Effective batch = 24*4*2 = 192 seqs/step +step 0 | loss 11.1393 | lr 7.50e-06 | 0 tok/s +step 20 | loss 8.6159 | lr 1.57e-04 | 148,936 tok/s +step 100 | loss 6.3108 | lr 6.00e-04 | 150,609 tok/s + [eval] step 100 | train 6.2501 | dev 6.1745 +step 500 | loss 4.5317 | lr 5.39e-04 | 137,334 tok/s + [eval] step 500 | train 4.7066 | dev 4.6499 +step 1000 | loss 4.0100 | lr 3.48e-04 | 131,348 tok/s + [eval] step 1000 | train 4.1200 | dev 4.1419 +step 1500 | loss 3.6483 | lr 1.45e-04 | 123,781 tok/s + [eval] step 1500 | train 3.8393 | dev 3.8985 +step 1900 | loss 3.7725 | lr 6.36e-05 | 151,488 tok/s + [eval] step 1900 | train 3.7345 | dev 3.7607 +Done. Final checkpoint -> /ephemeral/ckpts/base_pretrained.pt +``` + +### The loss curve + +Here is the real training and dev loss for the 77 million parameter base I trained for this README on 2x L40 GPUs. The loss starts near `ln(vocab_size)`, which is about 10.8 (the loss of a model that guesses uniformly), and drops as the model learns the statistics of the text: + +![Training and validation loss](images/loss_curve.png) + +This run started at a loss of **11.14** (just above the uniform-guess line at `ln(50304) = 10.83`) and came down to about **3.73 train / 3.76 dev** after 2000 steps, training at roughly **130,000 to 150,000 tokens per second** across the two L40s. The curve going down is the whole story of pretraining: the model is slowly compressing the patterns of language into its weights. + +## Step 4: Generating Text + +A trained model predicts a distribution over the next token. To generate, we sample one token from that distribution, append it, and feed the longer sequence back in. We repeat until we have enough tokens. + +```python +def generate(self, idx, max_new_tokens): + for _ in range(max_new_tokens): + idx_cond = idx[:, -self.context_length:] # never look back further than the context window + logits, _ = self(idx_cond) + logits = logits[:, -1, :] # only the last position matters for the next token + probs = F.softmax(logits, dim=-1) + idx_next = torch.multinomial(probs, num_samples=1) + idx = torch.cat((idx, idx_next), dim=1) + return idx +``` + +Run it from a saved checkpoint: + +```bash +python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "The" --max_new_tokens 100 +``` + +The 13 million parameter model already produces real words and roughly correct grammar, which is the encouraging part of starting small. + +## Step 5: Post-Training, Turning a Base Model Into an Assistant + +A base model can continue text, but it cannot follow instructions or reason on purpose. That takes post-training. The good news is that the model never changes. We reuse the exact same Transformer backbone and only change two things at each stage: the data, and the loss. + +![The post-training pipeline](images/00_pipeline.png) + +The one design idea that makes all of this fit in a small repo is **wrap, do not rewrite**. The educational `Transformer` gains a single extra method, `forward_hidden`, which returns the hidden states right before `lm_head`. The reward head, the value head, and all the log-probability math compose around that one method. Nothing in `src/models/` had to be rewritten. + +### SFT (Supervised Fine-Tuning) + +SFT teaches the base model to answer in the chat format. It is still next-token prediction, with one change: the loss is only counted on the assistant tokens, using the mask we built back in Step 1. + +![SFT](images/09_sft.png) + +```python +def sft_loss(logits, tokens, loss_mask): + logits = logits[:, :-1, :] # predict token t+1 from position t + targets = tokens[:, 1:] + mask = loss_mask[:, 1:].to(logits.dtype) + + V = logits.size(-1) + ce = F.cross_entropy(logits.reshape(-1, V).float(), targets.reshape(-1).long(), reduction="none") + ce = ce.view(targets.shape) * mask # zero out the prompt positions + return ce.sum() / mask.sum().clamp(min=1.0) # average over assistant tokens only +``` + +Prepare the data and train: + +```bash +python scripts/prepare_sft_data.py --context_length 1024 +torchrun --standalone --nproc_per_node=2 scripts/train_sft.py +``` + +The loss code is in `src/post_training/sft.py`, the trainer in `scripts/train_sft.py`. + +### The Reward Model + +To do reinforcement learning we need a number that says how good an answer is. One way to get it is to train a reward model on human preference pairs. We put a small linear head on top of the SFT backbone that reads one scalar off the last real token, and we train it with the Bradley-Terry loss so the chosen answer always scores higher than the rejected one. + +![The reward model](images/10_reward_model.png) + +```python +def bradley_terry_loss(chosen_rewards, rejected_rewards): + """Mean -log sigmoid(chosen - rejected) over a batch of preference pairs.""" + return -F.logsigmoid(chosen_rewards - rejected_rewards).mean() +``` + +```bash +python scripts/prepare_preference_data.py --source both +torchrun --standalone --nproc_per_node=2 scripts/train_reward.py +``` + +The headline metric is preference accuracy, the fraction of held-out pairs where the model scores the chosen answer higher. In this run on 7974 real preference pairs it reached **0.574** (above the 0.5 chance line); with a larger model and more data this climbs toward 0.65 to 0.75. + +``` +#### OUTPUT #### +Reward model from sft.pt | 7974 pairs | total_steps=996 + [eval] step 250 | test_acc 0.539 | margin 0.006 + [eval] step 750 | test_acc 0.576 | margin 0.063 +Done RM. test_acc 0.574 margin 0.063 -> reward.pt +``` + +The code is in `src/post_training/reward_model.py` and `src/post_training/reward_train.py`. + +### DPO, ORPO and KTO + +DPO skips the reward model and the RL loop entirely. It works directly on preference pairs by comparing how much more likely the policy makes the chosen answer (relative to a frozen reference copy of the SFT model) than the rejected one. + +![DPO](images/11_dpo.png) + +```python +def dpo_loss(policy_chosen_logps, policy_rejected_logps, + ref_chosen_logps, ref_rejected_logps, beta=0.1): + pi_logratios = policy_chosen_logps - policy_rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logits = pi_logratios - ref_logratios + loss = -F.logsigmoid(beta * logits).mean() + return loss, ... +``` + +```bash +torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo +# --loss_type orpo reference free, folds SFT and alignment into one stage +# --loss_type kto works from an unpaired desirable / undesirable signal +``` + +In this run, DPO reached an implicit-reward accuracy of **0.574** on the held-out pairs (the fraction where the policy prefers the chosen response more than the frozen reference does). All three objectives are in `src/post_training/dpo.py`. + +### PPO + +PPO is the classic RLHF loop. For each prompt the model writes an answer (a rollout), we score it (with the reward model or with a GSM8K answer checker), add a small per-token penalty for drifting away from the reference model, estimate how good each token was with GAE, and take a few clipped gradient steps. + +![PPO](images/12_ppo.png) + +The two load-bearing pieces, the advantage estimate and the clipped policy loss, are short: + +```python +def ppo_policy_loss(new_logp, old_logp, advantages, mask, clip=0.2): + ratio = torch.exp(new_logp - old_logp) + surr1 = ratio * advantages + surr2 = torch.clamp(ratio, 1.0 - clip, 1.0 + clip) * advantages # the clip keeps the step small + loss = -masked_mean(torch.min(surr1, surr2), mask) + return loss, ... +``` + +```bash +python scripts/prepare_rl_prompts.py +torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py --reward_source verifier +# --reward_source rm to use the trained reward model instead of the answer checker +``` + +The actor-critic shares the backbone through a small value head (`src/post_training/value_head.py`), and the GAE, clipped policy loss, and clipped value loss are in `src/post_training/ppo.py`. + +### GRPO / RLVR + +GRPO is the 2025, DeepSeek-R1 style method. It throws away the value network. For each prompt it samples a whole group of answers, scores them with a verifiable reward (did the final number match the gold answer), and uses the group's own mean and standard deviation as the baseline. The advantage of an answer is just how much better it did than its group. + +![GRPO](images/13_grpo.png) + +```python +def group_advantages(rewards, group_size, eps=1e-4): + r = rewards.view(-1, group_size) + mean = r.mean(dim=1, keepdim=True) + std = r.std(dim=1, keepdim=True) + adv = (r - mean) / (std + eps) # how much better than the rest of my group + return adv.reshape(-1) +``` + +```bash +torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py --group_size 8 +``` + +A short arithmetic curriculum runs first, so the model gets some non-zero reward to learn from before it faces full GSM8K. The group-relative advantage, the clipped surrogate, and the k3 KL penalty are in `src/post_training/grpo.py`. + +## Step 6: Evaluation + +The single number that ties all the stages together is greedy GSM8K accuracy. We give the model a math question, let it generate, pull the number out of the `` tags, and check it against the gold answer. + +![Evaluation](images/14_evaluation.png) + +```bash +for s in base_pretrained sft dpo ppo grpo; do + python scripts/eval_post_training.py --ckpt models/$s.pt --label $s --limit 200 --append logs/table.jsonl +done +python scripts/eval_post_training.py --table logs/table.jsonl +``` + +This builds the headline comparison so you can track every stage on one axis, Base, SFT, DPO, PPO, and GRPO, all with the same greedy decoding and the same verifiable reward: parse the number inside the `` tags and check it against the gold answer. The same command runs on any checkpoint, so you fill in the table for your own run, and the bigger the base model and the more pretraining compute you give it, the higher these scores go. + +### What changes after SFT + +The clearest effect of SFT is a change in behavior. The base model only knows how to continue text, so when you give it a question it just writes more text. After SFT the model has learned the chat format: it answers inside the `......` structure it was trained on, which is exactly the structure the reward model and the RL stages then optimize against. That learned format is the foundation every stage after it builds on. You can talk to any stage's checkpoint with `scripts/chat.py` and watch this directly, which is the next section. + +## Step 7: Talking to the Model + +`scripts/chat.py` loads any checkpoint, reads the model dimensions from the checkpoint itself, and lets you talk to it. It applies the chat template for instruction models, or treats the base model as raw continuation. + +![Inference and chat](images/15_inference.png) + +```bash +# instruction-tuned models +python scripts/chat.py --ckpt models/sft.pt --prompt "What is 13 + 29?" +python scripts/chat.py --ckpt models/grpo.pt --prompt "..." --greedy +# base model, raw continuation +python scripts/chat.py --ckpt models/base_pretrained.pt --raw --prompt "Once upon a time" +# interactive, no --prompt +python scripts/chat.py --ckpt models/sft.pt +``` + +Generation reuses the same tested core as training and eval, so what you see in chat is exactly what the RL stages optimized. Sampling is controlled by `--temperature`, `--top_p`, `--top_k`, or `--greedy`. The code is in `src/post_training/inference.py`. + +## The Streamlit Control Panel + +If you would rather click than type, there is a small control panel that can launch every stage, watch the loss live, run evaluation, and chat with a checkpoint: + +```bash +pip install -e ".[ui]" +streamlit run ui/app.py +``` + +It has one page per stage (Data, Pretrain, SFT, Reward, DPO, PPO, GRPO, Evaluate, Chat), each one a form over the same JSON configs you would edit by hand. + +## The Documentation Site + +Every stage has a deeper write-up, with the theory, the diagram, the real code, and what each metric means, on the documentation site: + +**https://fareedkhan-dev.github.io/train-llm-from-scratch/** + +To run it locally: + +```bash +pip install -e ".[docs]" +mkdocs serve +``` + +There is also a Foundations section that explains the ideas this code assumes you know (tokenization, the decoder-only Transformer, attention, objectives, optimization, and generation). + +## Run the Whole Thing + +Once the base is pretrained and the data is prepared, one script runs the entire post-training chain and prints the across-stages table: + +```bash +bash scripts/run_posttraining.sh # uses both GPUs via torchrun +NPROC=1 bash scripts/run_posttraining.sh # single GPU +``` + +For a fast end-to-end check on a tiny model that finishes in seconds, every stage has a smoke config: + +```bash +python tests/test_post_training_smoke.py # core math, on CPU +python scripts/train_sft.py --config configs/smoke/sft.json # a real (tiny) training run +``` + +## What's Next + +I recommend you start by training the 13 million parameter model, see it produce sensible words, then scale `n_embed` and `n_blocks` up with the memory flags until you hit your GPU limit. After that, walk the post-training chain one stage at a time and watch the GSM8K number move. Every stage is small enough to read in one sitting, and they all share the same model. + +If you want to go deeper on any single stage, the documentation site has a focused page for each one. + +
+ +Wanna chat on something? [My Linkedin](https://www.linkedin.com/in/fareed-khan-dev/) + +## Star History + +[![](https://api.star-history.com/svg?repos=FareedKhan-dev/train-llm-from-scratch&type=Date)](https://star-history.com/#FareedKhan-dev/train-llm-from-scratch&Date) diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..1ea4031 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`FareedKhan-dev/train-llm-from-scratch` +- 原始仓库：https://github.com/FareedKhan-dev/train-llm-from-scratch +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/config/__init__.py b/config/__init__.py new file mode 100644 index 0000000..77a0021 --- /dev/null +++ b/config/__init__.py @@ -0,0 +1,2 @@ +"""Configuration package: legacy ``config.config`` (pretraining), the post-training +dataclasses (``config.post_training_config``), and the JSON ``config.loader``.""" diff --git a/config/config.py b/config/config.py new file mode 100644 index 0000000..5e69eca --- /dev/null +++ b/config/config.py @@ -0,0 +1,68 @@ +# --- Configuration --- + +import torch + +# Define vocabulary size and transformer configuration (3 Billion) +VOCAB_SIZE = 50304 # Number of unique tokens in the vocabulary +CONTEXT_LENGTH = 512 # Maximum sequence length for the model +N_EMBED = 2048 # Dimension of the embedding space +N_HEAD = 16 # Number of attention heads in each transformer block +N_BLOCKS = 64 # Number of transformer blocks in the model + +# Paths to training and development datasets +TRAIN_PATH = "data/train/pile_train.h5" # File path for the training dataset +DEV_PATH = "data/val/pile_dev.h5" # File path for the validation dataset + +# Transformer training parameters +T_BATCH_SIZE = 32 # Number of samples per training batch +T_CONTEXT_LENGTH = 16 # Context length for training batches +T_TRAIN_STEPS = 200000 # Total number of training steps +T_EVAL_STEPS = 1000 # Frequency (in steps) to perform evaluation +T_EVAL_ITERS = 250 # Number of iterations to evaluate the model +T_LR_DECAY_STEP = 50000 # Step at which to decay the learning rate +T_LR = 5e-4 # Initial learning rate for training +T_LR_DECAYED = 5e-5 # Learning rate after decay +T_OUT_PATH = "models/transformer_B.pt" # Path to save the trained model +T_CHECKPOINT_STEPS = 0 # Save periodic checkpoints every N steps (0 disables) +T_KEEP_LAST_CHECKPOINTS = 3 # Number of periodic checkpoints to keep (0 keeps all) +T_CHECKPOINT_DIR = None # Optional checkpoint directory override + +# Memory-optimisation knobs (all OFF by default => unchanged behaviour/numerics). +# These let large configs fit in less VRAM; enable them on the CLI or here. See issue #5. +USE_AMP = False # bf16/fp16 autocast (CUDA only; ignored on CPU) +AMP_DTYPE = "bf16" # "bf16" (no GradScaler) or "fp16" (GradScaler) +USE_GRADIENT_CHECKPOINTING = False # recompute block activations in backward to save VRAM +GRAD_ACCUM_STEPS = 1 # micro-batches per optimizer step (effective batch x N) +REPORT_MEMORY_BUDGET = False # print a rough VRAM budget before training (CUDA only) + +# Device configuration +DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' + +# Store all configurations in a dictionary for easy access and modification +default_config = { + 'vocab_size': VOCAB_SIZE, + 'context_length': CONTEXT_LENGTH, + 'n_embed': N_EMBED, + 'n_head': N_HEAD, + 'n_blocks': N_BLOCKS, + 'train_path': TRAIN_PATH, + 'dev_path': DEV_PATH, + 't_batch_size': T_BATCH_SIZE, + 't_context_length': T_CONTEXT_LENGTH, + 't_train_steps': T_TRAIN_STEPS, + 't_eval_steps': T_EVAL_STEPS, + 't_eval_iters': T_EVAL_ITERS, + 't_lr_decay_step': T_LR_DECAY_STEP, + 't_lr': T_LR, + 't_lr_decayed': T_LR_DECAYED, + 't_out_path': T_OUT_PATH, + 't_checkpoint_steps': T_CHECKPOINT_STEPS, + 't_keep_last_checkpoints': T_KEEP_LAST_CHECKPOINTS, + 't_checkpoint_dir': T_CHECKPOINT_DIR, + 'use_amp': USE_AMP, + 'amp_dtype': AMP_DTYPE, + 'use_gradient_checkpointing': USE_GRADIENT_CHECKPOINTING, + 'grad_accum_steps': GRAD_ACCUM_STEPS, + 'report_memory_budget': REPORT_MEMORY_BUDGET, + 'device': DEVICE, +} diff --git a/config/loader.py b/config/loader.py new file mode 100644 index 0000000..4c1c347 --- /dev/null +++ b/config/loader.py @@ -0,0 +1,85 @@ +""" +JSON config loader for the post-training stages. + +Industry-standard, learning-friendly: each stage's knobs live in a small editable JSON file +under ``configs/`` (typed by the dataclasses in :mod:`config.post_training_config`). This +loader resolves a final dataclass instance by merging four layers, lowest precedence first: + + 1. dataclass field defaults (config.post_training_config.Config) + 2. configs/base.json (shared model + runtime fields) + 3. the stage JSON (configs/sft.json, ...) (that stage's hyperparameters) + 4. CLI --field overrides (highest precedence) + +JSON ``null`` maps to Python ``None`` cleanly -- the right way to set ``str | None`` fields +like ``amp_dtype``. Unknown keys are warned about, not fatal. + +When ``json_path`` lives in a sub-dir with its own ``base.json`` (e.g. ``configs/smoke/sft.json``), +that sibling ``base.json`` is used automatically -- so the smoke configs shrink the model too. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import fields +from typing import Any + + +def _deep_merge(dst: dict, src: dict) -> dict: + """Recursively merge ``src`` into ``dst`` (nested-dict aware; future-proof).""" + for k, v in src.items(): + if isinstance(v, dict) and isinstance(dst.get(k), dict): + _deep_merge(dst[k], v) + else: + dst[k] = v + return dst + + +def _resolve_base(json_path: str | None, base_path: str | None) -> str: + if base_path is not None: + return base_path + if json_path: + sibling = os.path.join(os.path.dirname(json_path), "base.json") + if os.path.exists(sibling): + return sibling + return "configs/base.json" + + +def load_config( + cfg_cls, + json_path: str | None = None, + overrides: dict[str, Any] | None = None, + *, + base_path: str | None = None, +): + """ + Resolve ``cfg_cls`` from ``base.json`` + the stage JSON + CLI overrides. + + Args: + cfg_cls: the stage dataclass (e.g. ``SFTConfig``). + json_path: path to the stage JSON (e.g. ``configs/sft.json``); None = base + defaults. + overrides: parsed CLI ``--field`` values (None values are ignored). + base_path: shared base JSON; if None, uses the sibling ``base.json`` of ``json_path`` + (so ``configs/smoke/sft.json`` picks up ``configs/smoke/base.json``), else + ``configs/base.json``. + + Returns: + an instance of ``cfg_cls`` with the resolved values. + """ + base = _resolve_base(json_path, base_path) + merged: dict[str, Any] = {} + for path in (base, json_path): + if path and os.path.exists(path): + with open(path) as fh: + _deep_merge(merged, json.load(fh)) + + field_names = {f.name for f in fields(cfg_cls)} + for key in list(merged): + if key not in field_names: + print(f"[config] ignoring unknown key '{key}' for {cfg_cls.__name__}") + merged.pop(key) + + if overrides: + merged.update({k: v for k, v in overrides.items() if k in field_names and v is not None}) + + return cfg_cls(**merged) diff --git a/config/post_training_config.py b/config/post_training_config.py new file mode 100644 index 0000000..a0f83f0 --- /dev/null +++ b/config/post_training_config.py @@ -0,0 +1,178 @@ +""" +Configuration for the from-scratch post-training suite. + +Kept entirely separate from ``config/config.py`` (which import-executes for the original +pretraining path and must stay untouched). Each stage is a frozen-ish dataclass that +inherits the shared :class:`BaseModelConfig` model/runtime fields and adds its own +hyperparameters. Construct with overrides, e.g. ``SFTConfig(lr=2e-5, batch_size=16)``. + +The default base model is ~400M parameters (n_embed=1024, n_head=16, n_blocks=24, +context_length=1024) -- the "mid" size chosen so real datasets (Alpaca, HH-RLHF, GSM8K) +give meaningful results while still fitting comfortably on one H100 and training in a +reasonable time on 2x H100. A tiny ``SMOKE`` variant is provided for fast CPU/1-GPU tests. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace + + +# Shared paths (all heavy artifacts live on the 1.5TB /ephemeral disk). +EPHEMERAL = "/ephemeral" +CKPT_DIR = f"{EPHEMERAL}/ckpts" +DATA_DIR = f"{EPHEMERAL}/data" +LOG_DIR = f"{EPHEMERAL}/logs" + + +@dataclass +class BaseModelConfig: + # --- model architecture (must match across all stages + the pretrained ckpt) --- + vocab_size: int = 50304 + context_length: int = 1024 + n_embed: int = 1024 + n_head: int = 16 + n_blocks: int = 24 + + # --- runtime --- + device: str = "cuda" + amp_dtype: str | None = "bf16" # None | "bf16"; bf16 needs no GradScaler on H100 + seed: int = 1337 + compile: bool = False # torch.compile the model (big speedup, slow 1st step) + ckpt_dir: str = CKPT_DIR + log_dir: str = LOG_DIR + use_wandb: bool = False + wandb_project: str = "train-llm-from-scratch-posttrain" + + +@dataclass +class PretrainConfig(BaseModelConfig): + """Pretrain the mid base model from scratch on the Pile HDF5 (mix in task text late).""" + train_path: str = "data/train/pile_train.h5" + dev_path: str = "data/val/pile_dev.h5" + batch_size: int = 24 # per-GPU micro-batch + grad_accum: int = 8 # effective batch = batch_size * grad_accum * world + train_steps: int = 200_000 + eval_steps: int = 1_000 + eval_iters: int = 100 + warmup_steps: int = 2_000 + lr: float = 3e-4 + min_lr: float = 3e-5 + weight_decay: float = 0.1 + grad_clip: float = 1.0 + out_ckpt: str = f"{CKPT_DIR}/base_pretrained.pt" + save_every: int = 2_000 + + +@dataclass +class SFTConfig(BaseModelConfig): + pretrained_ckpt: str = f"{CKPT_DIR}/base_pretrained.pt" + data_path: str = f"{DATA_DIR}/sft_packed.h5" + out_ckpt: str = f"{CKPT_DIR}/sft.pt" + batch_size: int = 16 + grad_accum: int = 2 + epochs: int = 3 + max_steps: int = -1 # -1 = run full epochs + eval_steps: int = 200 + warmup_steps: int = 100 + lr: float = 1e-5 + min_lr: float = 1e-6 + weight_decay: float = 0.0 + grad_clip: float = 1.0 + save_every: int = 500 + + +@dataclass +class RewardConfig(BaseModelConfig): + sft_ckpt: str = f"{CKPT_DIR}/sft.pt" + pref_path: str = f"{DATA_DIR}/preferences.jsonl" + out_ckpt: str = f"{CKPT_DIR}/reward.pt" + batch_size: int = 8 # pairs per step (2x sequences through the model) + epochs: int = 1 + eval_steps: int = 200 + warmup_steps: int = 50 + lr: float = 1e-5 + weight_decay: float = 0.0 + grad_clip: float = 1.0 + max_len: int = 768 + save_every: int = 500 + + +@dataclass +class DPOConfig(BaseModelConfig): + sft_ckpt: str = f"{CKPT_DIR}/sft.pt" # init policy + frozen reference + pref_path: str = f"{DATA_DIR}/preferences.jsonl" + out_ckpt: str = f"{CKPT_DIR}/dpo.pt" + loss_type: str = "dpo" # "dpo" | "orpo" | "kto" + beta: float = 0.1 + orpo_lambda: float = 1.0 # ORPO odds-ratio weight (loss_type="orpo") + batch_size: int = 8 + epochs: int = 1 + eval_steps: int = 200 + warmup_steps: int = 50 + lr: float = 5e-7 + weight_decay: float = 0.0 + grad_clip: float = 1.0 + max_len: int = 768 + save_every: int = 500 + + +@dataclass +class PPOConfig(BaseModelConfig): + sft_ckpt: str = f"{CKPT_DIR}/sft.pt" + reward_ckpt: str = f"{CKPT_DIR}/reward.pt" # used when reward_source="rm" + prompt_path: str = f"{DATA_DIR}/rl_prompts_train.jsonl" + eval_prompt_path: str = f"{DATA_DIR}/rl_prompts_test.jsonl" + out_ckpt: str = f"{CKPT_DIR}/ppo.pt" + reward_source: str = "verifier" # "verifier" (GSM8K checker) | "rm" (reward model) + iterations: int = 1_000 + prompts_per_iter: int = 32 # prompts sampled per PPO iteration (per rank) + rollout_len: int = 300 + temperature: float = 1.0 + top_p: float = 1.0 + ppo_epochs: int = 4 + minibatch_size: int = 16 + clip: float = 0.2 + vf_clip: float = 0.2 + vf_coef: float = 0.5 + ent_coef: float = 0.0 + gamma: float = 1.0 + gae_lambda: float = 0.95 + kl_coef: float = 0.05 # penalty on KL(policy || ref) added to reward + lr: float = 1e-6 + grad_clip: float = 1.0 + eval_every: int = 50 + save_every: int = 100 + + +@dataclass +class GRPOConfig(BaseModelConfig): + sft_ckpt: str = f"{CKPT_DIR}/sft.pt" + prompt_path: str = f"{DATA_DIR}/rl_prompts_train.jsonl" + eval_prompt_path: str = f"{DATA_DIR}/rl_prompts_test.jsonl" + curriculum_path: str = f"{DATA_DIR}/arithmetic_prompts.jsonl" # warm-up before GSM8K + curriculum_iters: int = 100 # iterations on the arithmetic warm-up before GSM8K + out_ckpt: str = f"{CKPT_DIR}/grpo.pt" + iterations: int = 1_000 + prompts_per_iter: int = 8 # distinct prompts per iter (per rank) + group_size: int = 8 # samples per prompt (group) + rollout_len: int = 300 + temperature: float = 1.0 + top_p: float = 1.0 + grpo_epochs: int = 1 + clip: float = 0.2 + kl_coef: float = 0.04 # KL(policy || ref) penalty term in the loss + lr: float = 1e-6 + grad_clip: float = 1.0 + eval_every: int = 50 + save_every: int = 100 + + +# Tiny config for fast smoke tests (CPU or a single GPU, seconds not hours). +SMOKE = dict( + vocab_size=256, context_length=64, n_embed=64, n_head=4, n_blocks=2, device="cpu", amp_dtype=None +) + + +def smoke(cfg_cls): + """Return an instance of ``cfg_cls`` shrunk to the tiny SMOKE model dims.""" + return replace(cfg_cls(), **SMOKE) diff --git a/configs/base.json b/configs/base.json new file mode 100644 index 0000000..b9730b9 --- /dev/null +++ b/configs/base.json @@ -0,0 +1,15 @@ +{ + "vocab_size": 50304, + "context_length": 1024, + "n_embed": 1024, + "n_head": 16, + "n_blocks": 24, + "device": "cuda", + "amp_dtype": "bf16", + "seed": 1337, + "compile": false, + "ckpt_dir": "/ephemeral/ckpts", + "log_dir": "/ephemeral/logs", + "use_wandb": false, + "wandb_project": "train-llm-from-scratch-posttrain" +} diff --git a/configs/dpo.json b/configs/dpo.json new file mode 100644 index 0000000..3b31b16 --- /dev/null +++ b/configs/dpo.json @@ -0,0 +1,17 @@ +{ + "sft_ckpt": "/ephemeral/ckpts/sft.pt", + "pref_path": "/ephemeral/data/preferences.jsonl", + "out_ckpt": "/ephemeral/ckpts/dpo.pt", + "loss_type": "dpo", + "beta": 0.1, + "orpo_lambda": 1.0, + "batch_size": 8, + "epochs": 1, + "eval_steps": 200, + "warmup_steps": 50, + "lr": 5e-07, + "weight_decay": 0.0, + "grad_clip": 1.0, + "max_len": 768, + "save_every": 500 +} diff --git a/configs/grpo.json b/configs/grpo.json new file mode 100644 index 0000000..6da1a66 --- /dev/null +++ b/configs/grpo.json @@ -0,0 +1,21 @@ +{ + "sft_ckpt": "/ephemeral/ckpts/sft.pt", + "prompt_path": "/ephemeral/data/rl_prompts_train.jsonl", + "eval_prompt_path": "/ephemeral/data/rl_prompts_test.jsonl", + "curriculum_path": "/ephemeral/data/arithmetic_prompts.jsonl", + "curriculum_iters": 100, + "out_ckpt": "/ephemeral/ckpts/grpo.pt", + "iterations": 1000, + "prompts_per_iter": 8, + "group_size": 8, + "rollout_len": 300, + "temperature": 1.0, + "top_p": 1.0, + "grpo_epochs": 1, + "clip": 0.2, + "kl_coef": 0.04, + "lr": 1e-06, + "grad_clip": 1.0, + "eval_every": 50, + "save_every": 100 +} diff --git a/configs/ppo.json b/configs/ppo.json new file mode 100644 index 0000000..5c45b4f --- /dev/null +++ b/configs/ppo.json @@ -0,0 +1,26 @@ +{ + "sft_ckpt": "/ephemeral/ckpts/sft.pt", + "reward_ckpt": "/ephemeral/ckpts/reward.pt", + "prompt_path": "/ephemeral/data/rl_prompts_train.jsonl", + "eval_prompt_path": "/ephemeral/data/rl_prompts_test.jsonl", + "out_ckpt": "/ephemeral/ckpts/ppo.pt", + "reward_source": "verifier", + "iterations": 1000, + "prompts_per_iter": 32, + "rollout_len": 300, + "temperature": 1.0, + "top_p": 1.0, + "ppo_epochs": 4, + "minibatch_size": 16, + "clip": 0.2, + "vf_clip": 0.2, + "vf_coef": 0.5, + "ent_coef": 0.0, + "gamma": 1.0, + "gae_lambda": 0.95, + "kl_coef": 0.05, + "lr": 1e-06, + "grad_clip": 1.0, + "eval_every": 50, + "save_every": 100 +} diff --git a/configs/pretrain.json b/configs/pretrain.json new file mode 100644 index 0000000..6b7f077 --- /dev/null +++ b/configs/pretrain.json @@ -0,0 +1,16 @@ +{ + "train_path": "/ephemeral/data/pile_train.h5", + "dev_path": "/ephemeral/data/pile_dev.h5", + "batch_size": 8, + "grad_accum": 12, + "train_steps": 50000, + "eval_steps": 1000, + "eval_iters": 100, + "warmup_steps": 2000, + "lr": 0.0003, + "min_lr": 3e-05, + "weight_decay": 0.1, + "grad_clip": 1.0, + "out_ckpt": "/ephemeral/ckpts/base_pretrained.pt", + "save_every": 1000 +} diff --git a/configs/reward.json b/configs/reward.json new file mode 100644 index 0000000..4037732 --- /dev/null +++ b/configs/reward.json @@ -0,0 +1,14 @@ +{ + "sft_ckpt": "/ephemeral/ckpts/sft.pt", + "pref_path": "/ephemeral/data/preferences.jsonl", + "out_ckpt": "/ephemeral/ckpts/reward.pt", + "batch_size": 8, + "epochs": 1, + "eval_steps": 200, + "warmup_steps": 50, + "lr": 1e-05, + "weight_decay": 0.0, + "grad_clip": 1.0, + "max_len": 768, + "save_every": 500 +} diff --git a/configs/sft.json b/configs/sft.json new file mode 100644 index 0000000..b59a839 --- /dev/null +++ b/configs/sft.json @@ -0,0 +1,16 @@ +{ + "pretrained_ckpt": "/ephemeral/ckpts/base_pretrained.pt", + "data_path": "/ephemeral/data/sft_packed.h5", + "out_ckpt": "/ephemeral/ckpts/sft.pt", + "batch_size": 16, + "grad_accum": 2, + "epochs": 3, + "max_steps": -1, + "eval_steps": 200, + "warmup_steps": 100, + "lr": 1e-05, + "min_lr": 1e-06, + "weight_decay": 0.0, + "grad_clip": 1.0, + "save_every": 500 +} diff --git a/configs/smoke/base.json b/configs/smoke/base.json new file mode 100644 index 0000000..19f90c8 --- /dev/null +++ b/configs/smoke/base.json @@ -0,0 +1,9 @@ +{ + "vocab_size": 50304, + "context_length": 256, + "n_embed": 128, + "n_head": 4, + "n_blocks": 2, + "device": "cpu", + "amp_dtype": null +} diff --git a/configs/smoke/dpo.json b/configs/smoke/dpo.json new file mode 100644 index 0000000..02a571c --- /dev/null +++ b/configs/smoke/dpo.json @@ -0,0 +1,8 @@ +{ + "batch_size": 4, + "epochs": 1, + "eval_steps": 10000, + "warmup_steps": 2, + "max_len": 256, + "save_every": 10000 +} diff --git a/configs/smoke/grpo.json b/configs/smoke/grpo.json new file mode 100644 index 0000000..8cce3d6 --- /dev/null +++ b/configs/smoke/grpo.json @@ -0,0 +1,10 @@ +{ + "iterations": 2, + "prompts_per_iter": 2, + "group_size": 4, + "rollout_len": 32, + "curriculum_iters": 1, + "grpo_epochs": 1, + "eval_every": 10000, + "save_every": 10000 +} diff --git a/configs/smoke/ppo.json b/configs/smoke/ppo.json new file mode 100644 index 0000000..5d07eb8 --- /dev/null +++ b/configs/smoke/ppo.json @@ -0,0 +1,9 @@ +{ + "iterations": 2, + "prompts_per_iter": 4, + "rollout_len": 32, + "ppo_epochs": 2, + "minibatch_size": 2, + "eval_every": 10000, + "save_every": 10000 +} diff --git a/configs/smoke/pretrain.json b/configs/smoke/pretrain.json new file mode 100644 index 0000000..ad122b0 --- /dev/null +++ b/configs/smoke/pretrain.json @@ -0,0 +1,10 @@ +{ + "batch_size": 4, + "grad_accum": 1, + "train_steps": 20, + "eval_steps": 10000, + "warmup_steps": 2, + "save_every": 10000, + "train_path": "/ephemeral/data/pile_train.h5", + "dev_path": "/ephemeral/data/pile_dev.h5" +} diff --git a/configs/smoke/reward.json b/configs/smoke/reward.json new file mode 100644 index 0000000..02a571c --- /dev/null +++ b/configs/smoke/reward.json @@ -0,0 +1,8 @@ +{ + "batch_size": 4, + "epochs": 1, + "eval_steps": 10000, + "warmup_steps": 2, + "max_len": 256, + "save_every": 10000 +} diff --git a/configs/smoke/sft.json b/configs/smoke/sft.json new file mode 100644 index 0000000..832fc5d --- /dev/null +++ b/configs/smoke/sft.json @@ -0,0 +1,7 @@ +{ + "max_steps": 10, + "batch_size": 4, + "eval_steps": 10000, + "warmup_steps": 2, + "save_every": 10000 +} diff --git a/data_loader/__init__.py b/data_loader/__init__.py new file mode 100644 index 0000000..d766b6f --- /dev/null +++ b/data_loader/__init__.py @@ -0,0 +1 @@ +"""Data loaders: pretraining HDF5 iterator + SFT / preference / RL-prompt iterators.""" diff --git a/data_loader/data_loader.py b/data_loader/data_loader.py new file mode 100644 index 0000000..4375928 --- /dev/null +++ b/data_loader/data_loader.py @@ -0,0 +1,78 @@ +import torch +import numpy as np +import h5py +from typing import Iterator, Tuple + +def get_batch_iterator(data_path: str, batch_size: int, context_length: int, device: str = "cpu") -> Iterator[Tuple[torch.Tensor, torch.Tensor]]: + """ + Creates an iterator for generating batches of data from an HDF5 file. + + Args: + data_path (str): Path to the HDF5 file containing tokenized data. + batch_size (int): Number of sequences in each batch. + context_length (int): Length of each sequence. + device (str, optional): Device to load the data onto ('cpu' or 'cuda'). Defaults to "cpu". + + Yields: + tuple: A tuple containing input sequences (xb) and target sequences (yb). + """ + # Open the HDF5 file in read mode + with h5py.File(data_path, 'r') as hdf5_file: + + # Extract the dataset of tokenized sequences + dataset = hdf5_file['tokens'] + + # Get the total size of the dataset + dataset_size = dataset.shape[0] + + # Calculate the number of examples (sequences) that can be made from the data + n_examples = (dataset_size - 1) // context_length + + # Create an array of indices for examples and shuffle them for randomness + example_idxs = np.arange(n_examples) + np.random.shuffle(example_idxs) + + # Initialize epoch counter and example counter + epochs = 0 + counter = 0 + + while True: + # Check if the current batch exceeds the number of available examples + if counter + batch_size > n_examples: + # Shuffle the indices again and reset the counter to 0 + np.random.shuffle(example_idxs) + counter = 0 + print(f"Finished epoch {epochs}") # Print epoch number when an epoch finishes + epochs += 1 # Increment the epoch counter + + # Select a batch of random indices to generate sequences + random_indices = example_idxs[counter:counter+batch_size] * context_length + + # Retrieve sequences from the dataset based on the random indices + random_samples = torch.tensor(np.array([dataset[idx:idx+context_length+1] for idx in random_indices])) + + # Separate the input sequences (xb) and target sequences (yb) + xb = random_samples[:, :context_length].to(device) # Input sequence (first half of the random sample) + yb = random_samples[:, 1:context_length+1].to(device) # Target sequence (second half of the random sample) + + # Increment the counter to move to the next batch + counter += batch_size + + # Yield the input and target sequences as a tuple for the current batch + yield xb, yb + +if __name__ == '__main__': + # Example Usage (requires a dummy HDF5 file for testing) + # Create a dummy HDF5 file + import os + dummy_data_path = "dummy_data.h5" + if not os.path.exists(dummy_data_path): + with h5py.File(dummy_data_path, 'w') as f: + f.create_dataset('tokens', data=np.arange(1000)) + + batch_size = 4 + context_length = 10 + for xb, yb in get_batch_iterator(dummy_data_path, batch_size, context_length): + print("Input Batch Shape:", xb.shape) + print("Target Batch Shape:", yb.shape) + break \ No newline at end of file diff --git a/data_loader/preference_dataset.py b/data_loader/preference_dataset.py new file mode 100644 index 0000000..50b939a --- /dev/null +++ b/data_loader/preference_dataset.py @@ -0,0 +1,79 @@ +""" +Batch iterator over preference pairs for reward-model and DPO training. + +Reads a JSONL file of ``{"prompt", "chosen", "rejected"}`` (produced by +``scripts/prepare_preference_data.py``). Each side is rendered through the chat template +so we get, for the chosen and rejected responses to the same prompt: + - token ids of ``prompt + response + EOT`` + - a response mask (1 over the completion, used by DPO) + - the true sequence length (used by the reward model to read the last-token reward) + +Right-padding is safe here because the model's attention is causal: the last real token +never attends to padding that comes after it, and the response mask zeros padded +positions in the loss. +""" + +from __future__ import annotations + +import json +from typing import Iterator + +import numpy as np +import torch + +from src.post_training.chat_template import EOT_ID, encode_chat + + +def _encode_side(prompt: str, response: str, max_len: int) -> tuple[list[int], list[int]]: + ids, mask = encode_chat([{"role": "user", "content": prompt}, + {"role": "assistant", "content": response}]) + return ids[:max_len], mask[:max_len] + + +def _collate(rows: list[dict], max_len: int, device: str) -> dict: + enc = [( _encode_side(r["prompt"], r["chosen"], max_len), + _encode_side(r["prompt"], r["rejected"], max_len)) for r in rows] + # Pad chosen and rejected to a single common length so they can share one forward. + L = max(max(len(c[0]), len(j[0])) for c, j in enc) + + def pad(seq, fill): + return seq + [fill] * (L - len(seq)) + + ch_ids, ch_mask, ch_len, rj_ids, rj_mask, rj_len = [], [], [], [], [], [] + for (cids, cmask), (jids, jmask) in enc: + ch_len.append(len(cids)); rj_len.append(len(jids)) + ch_ids.append(pad(cids, EOT_ID)); ch_mask.append(pad(cmask, 0)) + rj_ids.append(pad(jids, EOT_ID)); rj_mask.append(pad(jmask, 0)) + + t = lambda a, dt: torch.tensor(a, dtype=dt, device=device) + return { + "chosen_ids": t(ch_ids, torch.long), "chosen_mask": t(ch_mask, torch.long), "chosen_len": t(ch_len, torch.long), + "rejected_ids": t(rj_ids, torch.long), "rejected_mask": t(rj_mask, torch.long), "rejected_len": t(rj_len, torch.long), + } + + +def get_preference_iterator( + path: str, + batch_size: int, + max_len: int, + device: str = "cpu", + *, + rank: int = 0, + world_size: int = 1, + shuffle: bool = True, + infinite: bool = True, +) -> Iterator[dict]: + """Yield collated preference batches (dict of tensors). Rows are sharded across ranks.""" + with open(path) as f: + rows = [json.loads(line) for line in f if line.strip()] + rows = rows[rank::world_size] + rng = np.random.default_rng(7 + rank) + while True: + order = np.arange(len(rows)) + if shuffle: + rng.shuffle(order) + for s in range(0, len(order) - batch_size + 1, batch_size): + batch = [rows[i] for i in order[s:s + batch_size]] + yield _collate(batch, max_len, device) + if not infinite: + return diff --git a/data_loader/prompt_dataset.py b/data_loader/prompt_dataset.py new file mode 100644 index 0000000..426d62a --- /dev/null +++ b/data_loader/prompt_dataset.py @@ -0,0 +1,36 @@ +""" +Prompt iterator for RL (PPO / GRPO). Reads a JSONL of ``{"prompt", "gold"}`` (gold is the +verifiable numeric answer, or null) and yields lists of sampled rows. Rows are sharded +across ranks so each GPU optimizes on its own prompts. +""" + +from __future__ import annotations + +import json +from typing import Iterator + +import numpy as np + + +def load_prompt_rows(path: str) -> list[dict]: + with open(path) as f: + return [json.loads(line) for line in f if line.strip()] + + +def get_prompt_iterator( + path: str, + prompts_per_iter: int, + *, + rank: int = 0, + world_size: int = 1, + seed: int = 0, + shuffle: bool = True, +) -> Iterator[list[dict]]: + """Infinitely yield lists of ``prompts_per_iter`` rows (this rank's shard).""" + rows = load_prompt_rows(path)[rank::world_size] + rng = np.random.default_rng(seed + rank) + n = len(rows) + while True: + idx = rng.permutation(n) if shuffle else np.arange(n) + for s in range(0, n - prompts_per_iter + 1, prompts_per_iter): + yield [rows[i] for i in idx[s:s + prompts_per_iter]] diff --git a/data_loader/sft_dataset.py b/data_loader/sft_dataset.py new file mode 100644 index 0000000..dabdd25 --- /dev/null +++ b/data_loader/sft_dataset.py @@ -0,0 +1,52 @@ +""" +Batch iterator for packed SFT data (mirrors ``data_loader.data_loader.get_batch_iterator`` +but carries the per-token loss mask alongside the tokens). + +The HDF5 file has two aligned datasets, ``tokens`` and ``loss_mask``, both shape +(N, context_length), produced by ``scripts/prepare_sft_data.py``. +""" + +from __future__ import annotations + +from typing import Iterator + +import h5py +import numpy as np +import torch + + +def get_sft_batch_iterator( + data_path: str, + batch_size: int, + device: str = "cpu", + *, + rank: int = 0, + world_size: int = 1, + shuffle: bool = True, + infinite: bool = True, +) -> Iterator[tuple[torch.Tensor, torch.Tensor, int]]: + """ + Yield ``(tokens, loss_mask, epoch)`` batches from a packed SFT HDF5 file. + + Rows are sharded across ranks (each rank sees a disjoint stride) so DDP data-parallel + training covers the dataset once per epoch. Set ``infinite=False`` to stop after one + pass (used for evaluation). + """ + with h5py.File(data_path, "r") as f: + tokens = f["tokens"] + masks = f["loss_mask"] + n = tokens.shape[0] + idxs = np.arange(rank, n, world_size) + epoch = 0 + rng = np.random.default_rng(1234 + rank) + while True: + if shuffle: + rng.shuffle(idxs) + for start in range(0, len(idxs) - batch_size + 1, batch_size): + batch = np.sort(idxs[start: start + batch_size]) # sorted for h5py fancy-index + tk = torch.tensor(np.asarray(tokens[batch]), dtype=torch.long, device=device) + mk = torch.tensor(np.asarray(masks[batch]), dtype=torch.long, device=device) + yield tk, mk, epoch + epoch += 1 + if not infinite: + return diff --git a/docs/01_data_pipeline.md b/docs/01_data_pipeline.md new file mode 100644 index 0000000..403acc2 --- /dev/null +++ b/docs/01_data_pipeline.md @@ -0,0 +1,162 @@ + +# Data Handling & Preprocessing + +Every stage needs data in a different shape, and getting these shapes right is honestly half the +battle — a misaligned loss mask or a mis-parsed gold answer will silently wreck training. So before +any model code, here is exactly how I download and preprocess each dataset, and the format each +trainer expects. + +For the first-principles background behind these shapes, read +[Tokenization & Data Shapes](foundations/tokenization.md). That page explains why pretraining uses +flat token streams, why SFT needs `loss_mask`, and why preference data must preserve a shared prompt. + +There are **four** data pipelines, all feeding off real public datasets: + +![Data preprocessing pipelines](diagrams/01_data_pipeline.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart TD + subgraph PT[" 1 · Pretraining "] + P1([Pile .jsonl.zst]):::data --> P2[stream-decompress
+ tiktoken r50k_base]:::proc --> P3[(pile_train.h5
flat tokens)]:::store + end + subgraph SF[" 2 · SFT "] + S1([Alpaca · Dolly · GSM8K]):::data --> S2[render chat template
mask prompt tokens]:::proc --> S3[pack to 1024]:::proc --> S4[(sft_packed.h5
tokens + loss_mask)]:::store + end + subgraph PF[" 3 · Preferences "] + F1([HH-RLHF · UltraFeedback]):::data --> F2[split prompt /
chosen / rejected]:::proc --> F3[(preferences.jsonl)]:::store + end + subgraph RL[" 4 · RL prompts "] + R1([GSM8K · arithmetic]):::data --> R2[extract numeric
gold answer]:::proc --> R3[(rl_prompts.jsonl
prompt + gold)]:::store + end + P3 --> M1{{pretrain_base.py}}:::model + S4 --> M2{{train_sft.py}}:::model + F3 --> M3{{train_reward.py · train_dpo.py}}:::rl + R3 --> M4{{train_ppo.py · train_grpo.py}}:::rl + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef store fill:#cdece8,stroke:#16a085,stroke-width:2px,color:#0a3d33; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; +``` + +
+ +Everything lands on the big `/ephemeral` disk and uses the OpenAI **`r50k_base`** tokenizer +(`vocab_size = 50304`, the only special token is `<|endoftext|>` = id `50256`). + +## 1 · Pretraining data (Pile → flat-token HDF5) + +[`scripts/prepare_pretrain_data.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/prepare_pretrain_data.py) streams the compressed +Pile shards, batch-tokenizes with tiktoken, and writes one flat `int32` token array to HDF5 (far +faster than the original per-document resize). Each document is terminated with `<|endoftext|>`: + +```python +for ids in enc.encode_ordinary_batch(docs): + buf.extend(ids) + buf.append(EOT_ID) # 50256 separates documents + if len(buf) >= WRITE_CHUNK: + flush() # append ~8M tokens to the HDF5 dataset at once +``` + +```bash +PYTHONPATH=. python scripts/prepare_pretrain_data.py --split val --out /ephemeral/data/pile_dev.h5 +PYTHONPATH=. python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out /ephemeral/data/pile_train.h5 +``` + +The base [`get_batch_iterator`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/data_loader/data_loader.py) then slices random +`context_length + 1` windows out of this flat array for next-token training. + +## 2 · SFT data (instructions → packed tokens **+ loss mask**) + +This is the subtle one. We only want to train the model to produce the **assistant** tokens, not to +parrot the prompt. The chat format ([`chat_template.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py)) uses +plain-text role markers (since `r50k_base` has no spare special tokens) and `<|endoftext|>` as the +turn terminator: + +``` +<|user|> +{question}<|endoftext|><|assistant|> +{answer}<|endoftext|> +``` + +[`encode_chat`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py#L95) builds the token ids **and** an aligned +`loss_mask` that is `1` only over the assistant completion (and its terminating EOT): + +```python +content_ids = _encode_ordinary(m["content"]) +is_completion = role == "assistant" +ids.extend(content_ids) +mask.extend([1 if is_completion else 0] * len(content_ids)) # train ONLY assistant tokens +ids.append(EOT_ID) +mask.append(1 if is_completion else 0) # ...and teach it to stop +``` + +[`prepare_sft_data.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/prepare_sft_data.py) renders Alpaca + Dolly + GSM8K through this, +reformatting GSM8K into the `…N` structure (so the model learns the +exact shape the RL verifier later rewards), then [`pack_examples`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/sft.py#L41) +concatenates everything and slices it into fixed `1024`-token rows, writing two aligned HDF5 datasets, +`tokens` and `loss_mask`. + +```bash +PYTHONPATH=. python scripts/prepare_sft_data.py --context_length 1024 --out_dir /ephemeral/data +``` + +I verified on the real file that the mask covers exactly `4` and excludes the user +question — that alignment is what makes SFT work. + +## 3 · Preference data (→ `{prompt, chosen, rejected}` JSONL) + +[`prepare_preference_data.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/prepare_preference_data.py) pulls **Anthropic/hh-rlhf** and +**HuggingFaceH4/ultrafeedback_binarized** and normalizes both to one schema. For HH-RLHF I split each +dialogue at the last `Assistant:` turn so the chosen/rejected share a prompt and differ only in the +final response: + +```python +def _split_hh(text): + idx = text.rfind("\n\nAssistant:") + return text[:idx].strip(), text[idx + len("\n\nAssistant:"):].strip() +``` + +Output is `preferences.jsonl` (train) + `preferences_test.jsonl` (held-out, for measuring reward-model +accuracy). [`preference_dataset.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/data_loader/preference_dataset.py) tokenizes each side through +the same chat template and right-pads a batch — which is safe because the model's attention is +**causal**, so the last real token never attends to padding after it (no attention mask needed). + +```bash +PYTHONPATH=. python scripts/prepare_preference_data.py --source both --max_per_source 40000 +``` + +## 4 · RL prompt data (→ `{prompt, gold}` JSONL) + +[`prepare_rl_prompts.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/prepare_rl_prompts.py) turns GSM8K into prompts with a **verifiable +numeric gold answer** (parsed from the dataset's `#### N`), plus a programmatic **arithmetic curriculum** +that even a weak policy can partly solve — so RL has non-zero reward signal to bootstrap from: + +```python +gold = gsm8k_gold_answer(ex["answer"]) # the number after '####' +rows.append({"prompt": ex["question"].strip(), "gold": gold}) +``` + +```bash +PYTHONPATH=. python scripts/prepare_rl_prompts.py --out_dir /ephemeral/data +``` + +I cross-checked the emitted gold answers 50/50 against the live GSM8K dataset — they match exactly, +which matters because the verifier reward ([08_evaluation.md](08_evaluation.md)) is only as trustworthy +as the gold it compares against. + +## What you end up with + +| File | Shape | Used by | +|---|---|---| +| `pile_train.h5` / `pile_dev.h5` | flat `int32` tokens | pretraining | +| `sft_packed.h5` | `tokens` + `loss_mask`, `(N, 1024)` | SFT | +| `preferences.jsonl` (+ `_test`) | `{prompt, chosen, rejected}` | Reward Model, DPO | +| `rl_prompts_train.jsonl` / `_test` | `{prompt, gold}` | PPO, GRPO | +| `arithmetic_prompts.jsonl` | `{prompt, gold}` | GRPO curriculum warm-up | +
+ +➡️ Next: [Stage 1 — Pretraining](02_pretraining.md). diff --git a/docs/02_pretraining.md b/docs/02_pretraining.md new file mode 100644 index 0000000..b8cd08a --- /dev/null +++ b/docs/02_pretraining.md @@ -0,0 +1,102 @@ + +# Stage 1 — Pretraining the base model + +Everything downstream is only as good as the base model, so the first thing I do is pretrain a +**~400M-parameter** version of this repo's own `Transformer` from scratch on the Pile. The original +[`train_transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_transformer.py) is a clean single-GPU loop; for a mid-size +model on 2×H100 I wrote [`pretrain_base.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/pretrain_base.py), which adds the few things +that actually matter at this scale — DistributedDataParallel, bf16 autocast, gradient accumulation, a +cosine LR schedule with warmup, and periodic checkpointing — without touching the model itself. + +If the architecture or training terms are unfamiliar, read the foundations chapters first: +[Decoder-Only Transformer](foundations/transformer.md), +[Attention, Masks & Heads](foundations/attention.md), +[Objectives, Losses & Perplexity](foundations/objectives.md), and +[Optimization & Training Systems](foundations/optimization.md). + +![Pretraining loop](diagrams/02_pretraining.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + H[(pile_train.h5)]:::store --> IT[get_batch_iterator
random windows]:::proc + IT --> FWD{{forward
bf16 autocast}}:::model + FWD --> L[cross-entropy loss]:::loss + L --> BWD[backward
x grad_accum]:::model + BWD --> CL[clip grad norm 1.0]:::proc + CL --> ST[AdamW step
cosine LR + warmup]:::model + ST -->|every 1000 steps| CK[(base_pretrained.pt)]:::ckpt + ST -->|next step| IT + classDef store fill:#cdece8,stroke:#16a085,stroke-width:2px,color:#0a3d33; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; +``` + +
+ +## The model + +The base config lives in [`config/post_training_config.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/post_training_config.py) +(`BaseModelConfig`): `n_embed=1024, n_head=16, n_blocks=24, context_length=1024` → ~406M params. The +context length is bumped to 1024 (vs the original 512) so GSM8K reasoning chains fit later. + +## The training step + +The heart of [`pretrain_base.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/pretrain_base.py) is a gradient-accumulation loop under +bf16 autocast, syncing gradients across GPUs only on the last micro-step: + +```python +for micro in range(cfg.grad_accum): + xb, yb = next(batch_iter) + sync = (micro == cfg.grad_accum - 1) or not ctx.enabled + cm = model.no_sync() if (ctx.enabled and not sync) else _nullcm() + with cm, amp_autocast(cfg.amp_dtype, ctx.device): # bf16 on H100, no GradScaler needed + _, loss = model(xb, yb) + loss = loss / cfg.grad_accum + loss.backward() + +torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) # stability +optimizer.step() +``` + +A few choices worth calling out: +- **bf16 autocast** ([`amp_autocast`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/utils.py)) needs no `GradScaler` (unlike + fp16), which keeps the loop clean. Master weights stay fp32. +- **AdamW with a weight-decay split** ([`configure_optimizer`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/optim.py)) — decay + the 2-D weight matrices, never the biases / norms / embeddings (the standard GPT recipe). +- **Cosine LR with warmup** ([`cosine_lr`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/optim.py)) — linear ramp for + `warmup_steps`, then cosine decay to `min_lr`. +- **DDP** ([`distributed.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/distributed.py)) — each rank seeds its data shuffle + differently so the two GPUs see different windows; only rank 0 logs and checkpoints. + +## Run it + +```bash +# single GPU +PYTHONPATH=. python scripts/pretrain_base.py +# both H100s (effective batch = batch_size * grad_accum * num_gpus) +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True PYTHONPATH=. \ + torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py \ + --batch_size 8 --grad_accum 12 --train_steps 50000 +``` + +> **Why `batch_size 8`?** The repo's educational attention materializes a `(B, n_head, T, T)` tensor +> per block, so memory is dominated by the sequence-length term. At context 1024, batch 8 fits an 80GB +> H100 comfortably under DDP; we recover the effective batch with `grad_accum`. + +## What the numbers mean + +- **train_loss** — running cross-entropy; starts near `ln(vocab) ≈ 10.8` and should fall steadily + (mine went `11.06 → 8.6 → 6.0 → …`). This is the single best health signal. +- **tok/s** — throughput (~32k/s combined on 2×H100 here). +- **eval train/dev** — averaged loss on held-out windows, printed every `eval_steps`; watch the dev + loss to spot overfitting. + +Checkpoints are written to `/ephemeral/ckpts/base_pretrained.pt` every `save_every` steps and carry the +config, so every later stage can rebuild the exact model with [`load_backbone_from_ckpt`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/utils.py). + +➡️ Next: [Stage 2 — SFT](03_sft.md). diff --git a/docs/03_sft.md b/docs/03_sft.md new file mode 100644 index 0000000..fb15a7a --- /dev/null +++ b/docs/03_sft.md @@ -0,0 +1,95 @@ + +# Stage 2 — Supervised Fine-Tuning (SFT) + +The base model can continue text but it doesn't know it's supposed to *answer* you. SFT fixes that by +showing it thousands of `(instruction, response)` pairs and training it to produce the **response**. +The only real difference from pretraining is a per-token **loss mask**: we compute the loss on the +assistant tokens and ignore the prompt. + +The exact token/mask mechanics are explained in +[Tokenization & Data Shapes](foundations/tokenization.md), and the masked objective is derived in +[Objectives, Losses & Perplexity](foundations/objectives.md). + +![SFT masked-loss flow](diagrams/03_sft.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + H[(sft_packed.h5
tokens + loss_mask)]:::store --> B[batch:
tokens, mask]:::proc + B --> M{{Transformer
logits}}:::model + M --> SH[shift: predict t+1]:::proc + SH --> CE[token cross-entropy]:::loss + MASK([loss_mask = 1 on
assistant tokens]):::data --> CE + CE --> AVG[average over
masked tokens only]:::loss --> UPD[AdamW step]:::model + classDef store fill:#cdece8,stroke:#16a085,stroke-width:2px,color:#0a3d33; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; +``` + +
+ +## The masked loss + +The whole stage hinges on [`sft_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/sft.py#L18). It's ordinary next-token +cross-entropy, except every target position is weighted by the mask so only completion tokens count: + +```python +def sft_loss(logits, tokens, loss_mask): + logits = logits[:, :-1, :] # predict token t+1 from position t (same shift as pretraining) + targets = tokens[:, 1:] + mask = loss_mask[:, 1:].to(logits.dtype) + V = logits.size(-1) + ce = F.cross_entropy(logits.reshape(-1, V).float(), targets.reshape(-1).long(), reduction="none") + ce = ce.view(targets.shape) * mask + return ce.sum() / mask.sum().clamp(min=1.0) # mean over ASSISTANT tokens only +``` + +The mask itself was produced at data-prep time by +[`encode_chat`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py#L95) (see +[01_data_pipeline.md](01_data_pipeline.md)) and packed alongside the tokens. The `.float()` on the +logits keeps the cross-entropy numerically clean under bf16. + +## The trainer + +[`train_sft.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_sft.py) loads the pretrained base with +[`load_backbone_from_ckpt`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/utils.py), then runs a compact loop — autocast forward, +masked loss, clip, step, cosine LR — with periodic dev evaluation: + +```python +tokens, mask, epoch = next(train_it) +with amp_autocast(cfg.amp_dtype, ctx.device): + logits, _ = model(tokens) + loss = sft_loss(logits, tokens, mask) +loss.backward() +torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) +optimizer.step() +``` + +Batches come from [`get_sft_batch_iterator`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/data_loader/sft_dataset.py), which shards the packed +rows across DDP ranks and yields `(tokens, loss_mask, epoch)`. + +## Run it + +```bash +PYTHONPATH=. python scripts/train_sft.py # single GPU +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_sft.py # both GPUs +# tune: --lr 1e-5 --epochs 3 --batch_size 16 +``` + +## What the numbers mean + +- **train_loss / ppl** — masked cross-entropy (and its perplexity) over assistant tokens; should drop + well below the base model's loss. To sanity-check the mechanics I ran an *overfit* test on 8 rows and + watched the loss collapse `11.0 → 4.7`, confirming the gradient path learns. +- **dev_loss** — the same masked loss on a held-out split (`sft_dev_packed.h5`); the honest signal. +- **GSM8K dev accuracy** — after SFT the model both follows instructions *and* emits the + `…` format, so this should rise above the base model (see [08_evaluation.md](08_evaluation.md)). + +The result is saved to `/ephemeral/ckpts/sft.pt` and becomes the starting point for the reward model, +DPO, PPO and GRPO. + +➡️ Next: [Stage 3 — Reward Model](04_reward_model.md) or jump to [DPO](05_dpo.md). diff --git a/docs/04_reward_model.md b/docs/04_reward_model.md new file mode 100644 index 0000000..892f2d5 --- /dev/null +++ b/docs/04_reward_model.md @@ -0,0 +1,105 @@ + +# Stage 3 — Reward Model + +To do classic RLHF (PPO) we need something that scores a response with a single number: higher = more +preferred. That's the reward model. I build it by putting a tiny scalar head on top of the SFT +backbone and training it on human preference pairs with the **Bradley-Terry** loss — the same recipe +as InstructGPT. + +This page assumes you already know how the backbone produces hidden states. If not, start with +[Decoder-Only Transformer](foundations/transformer.md). The preference-data shape is covered in +[Tokenization & Data Shapes](foundations/tokenization.md). + +![Reward model (Bradley-Terry)](diagrams/04_reward_model.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + P([preference pair
prompt + chosen / rejected]):::data --> BB{{SFT backbone
forward_hidden}}:::model + BB --> LT[take last real token]:::proc + LT --> RH[reward head
Linear→1]:::model + RH --> RC([r_chosen]):::rl + RH --> RR([r_rejected]):::rl + RC --> BT[Bradley-Terry
-log σ r_chosen - r_rejected]:::loss + RR --> BT + BT --> UPD[AdamW step]:::model + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; +``` + +
+ +## The model: a scalar head on the backbone + +[`RewardModel`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/reward_model.py#L37) wraps a `Transformer`, drops the `lm_head`, +and reads the reward off the **last real token's** hidden state (the InstructGPT convention). Because +attention is causal, that last token has seen the whole sequence and never attends to the right-padding +after it — so we need no attention mask: + +```python +class RewardModel(nn.Module): + def __init__(self, transformer): + self.transformer = transformer + self.reward_head = nn.Linear(transformer.lm_head.in_features, 1, bias=False) + + def forward(self, idx, seq_lengths=None): + rewards = self.reward_head(self.transformer.forward_hidden(idx)).squeeze(-1) # (B, T) + return gather_last(rewards, seq_lengths) # reward at the last real token -> (B,) +``` + +[`gather_last`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/utils.py) just indexes `rewards[i, seq_lengths[i]-1]`. + +## The objective: Bradley-Terry + +[`bradley_terry_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/reward_train.py#L18) pushes the chosen reward above the +rejected one. That's the entire training signal: + +```python +def bradley_terry_loss(chosen_rewards, rejected_rewards): + return -F.logsigmoid(chosen_rewards - rejected_rewards).mean() +``` + +[`preference_accuracy`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/reward_train.py#L23) — the fraction of pairs where +`r_chosen > r_rejected` — is the metric I actually watch. + +## The trainer + +[`train_reward.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_reward.py) initializes the backbone from `sft.pt`, then for each +batch runs the **chosen and rejected sequences through the model in a single forward** (concatenated to +`2B`), splits the rewards, and applies the loss: + +```python +ids = torch.cat([batch["chosen_ids"], batch["rejected_ids"]], dim=0) +lens = torch.cat([batch["chosen_len"], batch["rejected_len"]], dim=0) +rewards = rm(ids, seq_lengths=lens).float() +chosen_r, rejected_r = rewards[:B], rewards[B:] +loss = bradley_terry_loss(chosen_r, rejected_r) +``` + +Pairs come from [`get_preference_iterator`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/data_loader/preference_dataset.py), which right-pads each +batch (safe under causal attention) and tracks the true length of each side. + +## Run it + +```bash +PYTHONPATH=. python scripts/train_reward.py +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_reward.py +# tune: --lr 1e-5 --max_len 768 +``` + +## What the numbers mean + +- **loss** — Bradley-Terry; starts at `-log σ(0) = 0.693` (chance) and drops as the gap widens. +- **train_acc / test_acc** — preference accuracy. On clean fixtures it goes to `1.0`; on **real, noisy** + HH-RLHF / UltraFeedback expect roughly **0.65–0.75** — that's normal, human preferences are noisy. +- **margin** — mean `r_chosen − r_rejected`; a useful "is it still separating them" signal. + +Saved to `/ephemeral/ckpts/reward.pt`; PPO loads it with +[`load_reward_model`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/reward_model.py) when `--reward_source rm`. + +➡️ Next: [Stage 5 — PPO](06_ppo.md) (which consumes this), or the RM-free path: [DPO](05_dpo.md). diff --git a/docs/05_dpo.md b/docs/05_dpo.md new file mode 100644 index 0000000..52b9425 --- /dev/null +++ b/docs/05_dpo.md @@ -0,0 +1,107 @@ + +# Stage 4 — DPO (and ORPO / KTO) + +Direct Preference Optimization is the shortcut around RLHF: instead of training a reward model and then +running an RL loop, DPO optimizes the policy *directly* on preference pairs, using a frozen copy of the +SFT model as a reference anchor. No reward model, no rollouts, no value function — just one clean loss. +I also implemented two popular variants behind a `--loss_type` flag: **ORPO** (reference-free) and +**KTO** (works from a desirable/undesirable signal). + +For the sequence log-probability notation used here, see +[Objectives, Losses & Perplexity](foundations/objectives.md). + +![DPO / ORPO / KTO](diagrams/05_dpo.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + P([chosen / rejected]):::data --> POL{{policy
trainable}}:::model + P --> REF{{reference
frozen SFT copy}}:::ckpt + POL --> LPP[seq log-probs
π_chosen, π_rejected]:::proc + REF --> LPR[seq log-probs
ref_chosen, ref_rejected]:::proc + LPP --> D[DPO loss
-log σ β·Δlogratios]:::loss + LPR --> D + D --> UPD[AdamW step]:::model + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; +``` + +
+ +## Sequence log-probs (the shared ingredient) + +DPO compares how much *more* likely the policy makes the chosen response vs the rejected one, relative +to the reference. So I need the **summed log-prob of each response** under both models. +[`sequence_logprobs`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rollout.py#L268) does exactly that (and is reused by PPO/GRPO): + +```python +def sequence_logprobs(model, sequences, response_mask, *, temperature=1.0, requires_grad=True): + lp, mask = compute_logprobs(model, sequences, response_mask, temperature=temperature, requires_grad=requires_grad) + m = mask.to(lp.dtype) + return (lp * m).sum(dim=-1), m.sum(dim=-1) # (summed logprob, #tokens) per sequence +``` + +## The DPO loss + +[`dpo_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/dpo.py#L21) is the canonical objective. The β temperature controls how +hard it pushes away from the reference: + +```python +def dpo_loss(policy_chosen_logps, policy_rejected_logps, ref_chosen_logps, ref_rejected_logps, beta=0.1): + pi_logratios = policy_chosen_logps - policy_rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logits = pi_logratios - ref_logratios + loss = -F.logsigmoid(beta * logits).mean() + chosen_reward = beta * (policy_chosen_logps - ref_chosen_logps).detach() + rejected_reward = beta * (policy_rejected_logps - ref_rejected_logps).detach() + return loss, chosen_reward, rejected_reward +``` + +The two **variants** ([`orpo_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/dpo.py#L48), +[`kto_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/dpo.py#L71)) live in the same file: +- **ORPO** — reference-free: combines the SFT negative-log-likelihood on the chosen response with an + odds-ratio preference term, folding SFT + alignment into one stage (no frozen reference needed). +- **KTO** — treats chosen as *desirable* and rejected as *undesirable* against a reference-KL baseline + estimated from the batch; useful when you only have thumbs-up/down rather than pairs. + +## The trainer + +[`train_dpo.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_dpo.py) loads the policy from `sft.pt`, makes a frozen reference with +[`make_frozen_copy`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/utils.py) (skipped for ORPO), and computes policy + reference +log-probs each step: + +```python +policy = load_backbone_from_ckpt(cfg, cfg.sft_ckpt, ctx.device) +ref = make_frozen_copy(policy, device=ctx.device) if cfg.loss_type != "orpo" else None +... +loss, cr, rr = _compute_losses(policy, ref, batch, cfg, ctx) # picks dpo/orpo/kto by cfg.loss_type +loss.backward() +``` + +## Run it + +```bash +PYTHONPATH=. python scripts/train_dpo.py --loss_type dpo --beta 0.1 +PYTHONPATH=. python scripts/train_dpo.py --loss_type orpo --orpo_lambda 1.0 +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py +``` + +> DPO uses a **small** learning rate (`5e-7` by default) — it's easy to over-push away from the +> reference and degrade the model, so go gentle. + +## What the numbers mean + +- **loss** — DPO/KTO start near `0.693`; ORPO starts higher (it includes the NLL term). +- **acc** — implicit-reward accuracy: fraction of pairs where the policy's implicit reward prefers the + chosen response. Should climb above 0.5. +- **r_chosen / r_rejected** — the implicit rewards `β·(logπ − logref)`; the gap (margin) should widen. +- **GSM8K dev accuracy** — the real downstream check. + +Saved to `/ephemeral/ckpts/dpo.pt`. + +➡️ Next: the RL path — [PPO](06_ppo.md) and [GRPO](07_grpo.md). diff --git a/docs/06_ppo.md b/docs/06_ppo.md new file mode 100644 index 0000000..103c74c --- /dev/null +++ b/docs/06_ppo.md @@ -0,0 +1,116 @@ + +# Stage 5 — PPO (classic RLHF) + +This is the original ChatGPT recipe: let the model generate, score the generations with a reward, and +nudge the policy toward higher-reward behaviour using Proximal Policy Optimization — with a value +network (critic) for variance reduction and a KL penalty to keep it from drifting too far from the SFT +model. I wrote the whole loop from scratch: rollout → reward → GAE advantages → clipped update. + +The clipped objective and policy-ratio notation are introduced in +[Objectives, Losses & Perplexity](foundations/objectives.md). The optimizer and stability pieces are +covered in [Optimization & Training Systems](foundations/optimization.md). + +![PPO loop](diagrams/06_ppo.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + PR([GSM8K prompts]):::data --> RO[rollout
generate_with_logprobs]:::proc + RO --> SC{score: verifier
or reward model}:::rl + SC --> KL[+ per-token
KL-to-ref penalty]:::rl + KL --> GAE[compute_gae
advantages + returns]:::proc + GAE --> UP{{clipped update
policy + value, K epochs}}:::model + UP -->|sync old policy| RO + REF{{frozen ref}}:::ckpt + REF -. KL .-> KL + VH{{value head}}:::model + VH -. value .-> GAE + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; +``` + +
+ +## The actor-critic + +PPO needs a per-token value estimate `V(s_t)` next to the policy logits. I get both from one backbone +with [`TransformerWithValueHead`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/value_head.py#L19) — it reuses +`forward_hidden` + `lm_head` for the policy and adds a small scalar value head (initialized to ~0 so the +critic doesn't destabilize early training): + +```python +def forward(self, idx): + hidden = self.transformer.forward_hidden(idx) + logits = self.transformer.lm_head(hidden) # policy + values = self.value_head(hidden).squeeze(-1) # critic, (B, T) + return logits, values +``` + +## Rollout + log-probs + +[`rollout_prompts`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rollout.py#L180) length-buckets the prompts and samples a +completion for each, and [`generate_with_logprobs`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rollout.py#L94) records the +sampling log-probs. Log-probs are always taken in **fp32** ([`compute_logprobs`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rollout.py#L233)) +because PPO subtracts them and bf16 rounding there is harmful. + +## GAE — Generalized Advantage Estimation + +[`compute_gae`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/ppo.py#L24) works in the "action frame" (index `t` = producing +token `t+1`), bootstrapping only while the next action is still a response token: + +```python +for t in reversed(range(L)): + nonterminal = m[:, t + 1] if t + 1 < L else 0.0 # episode ends after the last response token + delta = rewards[:, t] + gamma * values_next[:, t] * nonterminal - values[:, t] + lastgae = delta + gamma * lam * nonterminal * lastgae + adv[:, t] = lastgae +returns = adv + values +``` + +The per-token reward is the **KL-to-reference penalty** at every response token, plus the scalar task +reward added at the **last** response token. Advantages are then normalized with +[`whiten`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/ppo.py#L60). + +## The clipped objective + +[`ppo_policy_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/ppo.py#L68) is the standard clipped surrogate; +[`ppo_value_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/ppo.py#L84) clips the value update too: + +```python +ratio = torch.exp(new_logp - old_logp) +surr1 = ratio * advantages +surr2 = torch.clamp(ratio, 1.0 - clip, 1.0 + clip) * advantages +loss = -masked_mean(torch.min(surr1, surr2), mask) +``` + +[`train_ppo.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_ppo.py) ties it together: rollout once, compute old log-probs / ref +log-probs / values, build rewards, GAE, then run `ppo_epochs` of minibatched clipped updates. + +## Run it + +```bash +PYTHONPATH=. python scripts/train_ppo.py --reward_source verifier # GSM8K checker as reward +PYTHONPATH=. python scripts/train_ppo.py --reward_source rm # use the trained reward.pt +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py +``` + +## What the numbers mean + +- **reward** — mean task reward per iteration; the headline curve, should trend up. +- **KL_ref** — mean KL of the policy from the SFT reference; must stay **bounded**. If it blows up the + model is degenerating — lower the LR or raise `--kl_coef`. +- **clipfrac** — fraction of tokens hitting the PPO clip; a health/▒step-size signal. +- **value_loss** — critic regression error. +- **GSM8K test accuracy** — the real outcome, evaluated every `--eval_every`. + +> PPO is the touchy one: small LR (`1e-6`), `clip 0.2`, grad-clip 1.0, and watch KL. I verified the loop +> truly *optimizes* by giving it a learnable synthetic reward — reward climbed `0.10 → 1.00`. + +Saved to `/ephemeral/ckpts/ppo.pt`. + +➡️ Next: [Stage 6 — GRPO](07_grpo.md), which drops the critic entirely. diff --git a/docs/07_grpo.md b/docs/07_grpo.md new file mode 100644 index 0000000..796aa5a --- /dev/null +++ b/docs/07_grpo.md @@ -0,0 +1,102 @@ + +# Stage 6 — GRPO / RLVR (the reasoning frontier) + +GRPO (Group Relative Policy Optimization) is the algorithm behind DeepSeek-R1, and it's beautifully +simple: **throw away PPO's value network**. For each prompt, sample a whole *group* of answers, score +them with a verifiable reward, and use the group's own mean/std as the baseline. The advantage is just +"how much better than your groupmates was this answer?" — no critic to train, no value loss. + +For the group-relative advantage formula and how it relates to PPO-style policy ratios, see +[Objectives, Losses & Perplexity](foundations/objectives.md). + +![GRPO loop](diagrams/07_grpo.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + PR([prompt]):::data --> G[sample a GROUP
of G answers]:::proc + G --> V{verifier reward
per answer}:::rl + V --> A[group advantage
r - mean / std]:::proc + A --> L[clipped surrogate
+ k3 KL to ref]:::loss + L --> UPD{{policy update}}:::model + UPD -->|next prompt| PR + REF{{frozen ref}}:::ckpt + REF -. KL .-> L + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; +``` + +
+ +## Group-relative advantage + +[`group_advantages`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/grpo.py#L17) is the whole idea — standardize rewards *within +each group*, so a good answer is one that beats its siblings on the same prompt: + +```python +def group_advantages(rewards, group_size, eps=1e-4): + r = rewards.view(-1, group_size) # rewards laid out group-contiguously + adv = (r - r.mean(1, keepdim=True)) / (r.std(1, keepdim=True) + eps) + return adv.reshape(-1) +``` + +A nice property: if every answer in a group gets the same reward (all right or all wrong), the std-based +advantage is ~0 and that group simply contributes no gradient — so I log the fraction of *informative* +groups as a health metric. + +## The loss: clipped surrogate + KL + +[`grpo_loss`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/grpo.py#L37) applies the same PPO-style token-level clipped surrogate +(advantage broadcast across a completion's tokens) plus a per-token KL penalty to the reference, using +Schulman's non-negative **k3** estimator ([`k3_kl`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/grpo.py#L31)): + +```python +ratio = torch.exp(new_logp - old_logp) +surrogate = torch.min(ratio * adv, torch.clamp(ratio, 1 - clip, 1 + clip) * adv) +kl = k3_kl(new_logp, ref_logp) # exp(Δ) - Δ - 1, always ≥ 0 +loss = -masked_mean(surrogate - kl_coef * kl, resp_mask) +``` + +## The trainer + curriculum + +[`train_grpo.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/train_grpo.py) loads the policy from `sft.pt`, replicates each prompt `G` +times (group-contiguously), rolls them out, scores with the GSM8K verifier, and updates. It runs an +**arithmetic curriculum** for the first `--curriculum_iters` iterations so the policy earns some reward +*before* facing full GSM8K (otherwise every group is all-wrong and there's no signal): + +```python +rows = next(warm_it if it < cfg.curriculum_iters else main_it) +prompts = [p for p in base_prompts for _ in range(G)] # group-contiguous +rewards = torch.tensor([reward_gsm8k(responses[i], golds[i]) for i in range(len(prompts))]) +adv = group_advantages(rewards, G) +``` + +## Run it + +```bash +PYTHONPATH=. python scripts/train_grpo.py --group_size 8 +PYTHONPATH=. torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py +# tune: --curriculum_iters 100 --kl_coef 0.04 --temperature 1.0 +``` + +## What the numbers mean + +- **reward** — mean verifier reward across the group samples; the curve you want climbing. +- **informative** — fraction of groups with non-zero reward spread (groups that actually teach + something). If this collapses to 0, raise temperature / group size or stay longer on the curriculum. +- **KL** — KL to the reference; keep it bounded. +- **GSM8K test accuracy** — the headline reasoning metric, evaluated every `--eval_every`. + +> I verified the GRPO path genuinely optimizes: with a learnable reward the mean reward climbed +> **0.10 → 0.69 → 1.00** in ~15 iterations and saturated. PPO and GRPO share the same rollout/log-prob +> core, so this also exercises the common machinery. + +Saved to `/ephemeral/ckpts/grpo.pt`. + +➡️ Next: [measure all stages on GSM8K](08_evaluation.md) and [chat with the result](09_inference.md). diff --git a/docs/08_evaluation.md b/docs/08_evaluation.md new file mode 100644 index 0000000..a5c2941 --- /dev/null +++ b/docs/08_evaluation.md @@ -0,0 +1,99 @@ + +# Evaluation + +A pipeline is only believable if you can measure it, so I evaluate every stage on the **same** held-out +GSM8K test set with greedy decoding. The headline deliverable is a single table: GSM8K accuracy as it +moves Base → SFT → DPO → PPO → GRPO. The reward is *verifiable* — I parse the model's final number and +compare it to the gold answer — so the score is objective, not a judgment call. + +![Evaluation flow](diagrams/08_evaluation.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + CK[(stage checkpoint)]:::ckpt --> GEN[batched_generate
greedy, length-bucketed]:::proc + Q([GSM8K question]):::data --> GEN + GEN --> EX[extract_answer
answer / #### / last number]:::proc + EX --> CMP{== gold?}:::eval + GOLD([gold number]):::data --> CMP + CMP --> ACC[accuracy across
Base→SFT→DPO→PPO→GRPO]:::eval + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; +``` + +
+ +## Generation: length-bucketed, greedy + +The educational model has no padding-aware attention mask, so +[`batched_generate`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py#L24) groups prompts of equal length and decodes +each bucket together; `greedy=True` forces argmax (`top_k=1`) for comparable, deterministic numbers. + +## Scoring: a verifiable reward + +[`gsm8k_accuracy`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py#L77) generates an answer per question and checks +it with the verifier: + +```python +prompts = [encode_prompt([{"role": "user", "content": q}]) for q, _ in qa_pairs] +responses = batched_generate(model, prompts, max_new_tokens, device=device, greedy=greedy) +correct = sum(is_correct(resp, gsm8k_gold_answer(ans)) for (q, ans), resp in zip(qa_pairs, responses)) +``` + +The reward/checker lives in [`rewards/`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rewards/). `extract_answer` is tolerant — +it prefers an `…` tag, then a GSM8K-style `#### N`, then falls back to the last number +in the text — and [`reward_gsm8k`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/rewards/verifiers.py#L35) is +**correctness-dominant** with only a small, bounded format bonus, to discourage reward hacking: + +```python +r = 0.0 +if _answers_match(extract_answer(text), gold): r += 1.0 # the reward that matters +if has_well_formed_answer(text): r += 0.2 # small format nudge +return min(r, 1.2) # clipped +``` + +I sanity-checked this scorer independently of any model: feeding it correct answers scores **100/100** +and wrong answers **0/100** false positives, with gold cross-verified against the live GSM8K dataset. + +## The across-stages table + +[`eval_post_training.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/eval_post_training.py) loads any checkpoint (reading its dims from +the stored `cfg`), scores it, and appends a row to a JSONL you can render as a table: + +```bash +for s in base_pretrained sft dpo ppo grpo; do + PYTHONPATH=. python scripts/eval_post_training.py --ckpt /ephemeral/ckpts/$s.pt \ + --label $s --limit 200 --append /ephemeral/logs/stage_table.jsonl +done +PYTHONPATH=. python scripts/eval_post_training.py --table /ephemeral/logs/stage_table.jsonl +``` + +``` +stage GSM8K acc n +------------------------------------ +base_pretrained ... 200 +sft ... 200 +dpo ... 200 +ppo ... 200 +grpo ... 200 +``` + +## In-training metrics + +Each trainer also writes a metrics JSONL under `/ephemeral/logs/` (via +[`MetricsLogger`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/logging_utils.py)) — train/dev loss for SFT, preference accuracy +for the reward model, implicit-reward accuracy for DPO, and reward/KL/clip-fraction + GSM8K accuracy for +PPO/GRPO. Pass `--use_wandb true` to also mirror to Weights & Biases; the JSONL is always written so you +can plot offline. + +## What "good" looks like at this scale + +A ~400M from-scratch model won't top the GSM8K leaderboard — the point is the **relative** climb across +stages and bounded KL during RL. Expect modest absolute numbers but a clear, real before/after gain at +each step. + +➡️ Next: [talk to any checkpoint](09_inference.md). diff --git a/docs/09_inference.md b/docs/09_inference.md new file mode 100644 index 0000000..8e9df26 --- /dev/null +++ b/docs/09_inference.md @@ -0,0 +1,92 @@ + +# Inference & Chat + +Training is only satisfying if you can actually *talk* to the result. The original +[`generate_text.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/generate_text.py) does raw continuation for the base model, but it's +hard-wired to the legacy config and has no chat template — so I added a small inference layer that loads +**any** stage checkpoint (base / SFT / DPO / PPO / GRPO) and talks to it correctly. + +For the underlying decoding loop, context cropping, temperature, and stop-token behavior, read +[Generation & Sampling](foundations/generation.md). + +![Inference / chat flow](diagrams/09_inference.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart LR + CK[(any checkpoint)]:::ckpt --> LD[load_model_from_ckpt
dims from stored cfg]:::proc + LD --> MODE{chat or raw?}:::proc + MODE -->|instruction model| CT[wrap in chat template]:::proc + MODE -->|base model| RAW[raw prefix]:::proc + CT --> GEN{{generate
temperature / top-p / greedy}}:::model + RAW --> GEN + GEN --> DEC([decode → reply]):::eval + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; +``` + +
+ +## Load any checkpoint by its stored config + +[`load_model_from_ckpt`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py) reads the model dimensions from the +checkpoint's saved `cfg`, so you never re-specify `n_embed`/`n_blocks`, and it tolerates DDP / +reward-head key prefixes: + +```python +ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) +cfg = {**(ck.get("cfg") or {}), **(overrides or {})} +model = Transformer(n_head=cfg["n_head"], n_embed=cfg["n_embed"], ...) +state = {k.removeprefix("module.").removeprefix("transformer."): v for k, v in state.items()} +``` + +## Chat vs raw + +[`generate_reply`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/inference.py#L37) has two modes, reusing the same tested +generation core as training/eval ([`batched_generate`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/evaluation.py#L24)): + +- **chat** (default) — wraps your text in the chat template (optionally with a `system` message) and + returns the decoded assistant turn. Use this for SFT/DPO/PPO/GRPO checkpoints. +- **raw** (`--raw`) — treats your text as a prefix and returns the base model's continuation (no + template). Use this for `base_pretrained.pt`. + +```python +if raw: + ids = get_tokenizer().encode_ordinary(user_text) +else: + ids = encode_prompt([{"role": "user", "content": user_text}]) # ...ends at <|assistant|> +out = batched_generate(model, [ids], max_new_tokens, device=device, + temperature=temperature, top_k=top_k, top_p=top_p, greedy=greedy) +``` + +Decoding is defensive — [`decode`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/chat_template.py) drops the EOT terminator and +any padding-vocab ids (the model's vocab is padded to 50304 but r50k_base only decodes 0–50255). + +## The CLI + +[`scripts/chat.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/scripts/chat.py) is one-shot or an interactive REPL: + +```bash +# instruction-tuned models (chat template applied automatically) +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt --prompt "What is 13 + 29?" +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/grpo.pt --prompt "..." --greedy +# base-model continuation +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/base_pretrained.pt --raw --prompt "Once upon a time" +# interactive REPL (omit --prompt) +PYTHONPATH=. python scripts/chat.py --ckpt /ephemeral/ckpts/sft.pt +``` + +Sampling controls: `--temperature`, `--top_p`, `--top_k`, or `--greedy` for deterministic argmax. Runs +on `--device cuda` or `cpu` (both verified). + +## Sampling knobs, briefly + +- **greedy** — reproducible, best for eval / math (`--greedy`). +- **temperature** — higher = more random; ~`0.7–1.0` for open-ended chat. +- **top_p / top_k** — nucleus / top-k truncation to cut the long tail of unlikely tokens. + +That's the full loop: pretrain → align → reason → measure → chat. Back to the [overview](README.md). diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..04db457 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,108 @@ + +# Post-Training & Alignment — Overview + +When I first trained this transformer from scratch, it could *continue* text but it couldn't +*follow instructions* or *reason*. That's what post-training fixes. This `docs/` folder walks +through the whole journey I built on top of the base model — every stage written from scratch +in plain PyTorch (no `trl`, no `peft`, no `transformers`), trained on real public datasets, and +runnable on a single GPU or scaled across multiple GPUs with DDP. + +If you are new to LLM training internals, start with the new +**[LLM Foundations](foundations/README.md)** section before reading the stage pages. It explains the +token shapes, decoder-only Transformer, attention masks, objectives, optimization loop, and generation +mechanics that every later page relies on. + +## Recommended reading order + +1. **Foundations first**: + [Tokenization](foundations/tokenization.md) -> + [Transformer](foundations/transformer.md) -> + [Attention](foundations/attention.md) -> + [Objectives](foundations/objectives.md) -> + [Optimization](foundations/optimization.md) -> + [Generation](foundations/generation.md). +2. **Then the full pipeline**: + [Data](01_data_pipeline.md) -> + [Pretraining](02_pretraining.md) -> + [SFT](03_sft.md) -> + [Reward Model](04_reward_model.md) -> + [DPO](05_dpo.md) -> + [PPO](06_ppo.md) -> + [GRPO](07_grpo.md). +3. **Finally run and inspect**: + [Evaluation](08_evaluation.md), [Inference / Chat](09_inference.md), and the + [command cheatsheet](howto/commands.md). + +The pipeline mirrors how modern aligned/reasoning models are actually built: + +![Post-training pipeline](diagrams/00_overview.png) + +
+Mermaid source (live, editable) + +```mermaid +flowchart TD + PILE([The Pile
9.8B tokens]):::data --> PRE{{Pretrain
~400M base}}:::model + PRE --> BASE[(base_pretrained.pt)]:::ckpt + BASE --> SFT{{SFT
Alpaca · Dolly · GSM8K}}:::model + SFT --> SFTCK[(sft.pt)]:::ckpt + SFTCK --> RM{{Reward Model
Bradley-Terry}}:::rl + SFTCK --> DPO{{DPO / ORPO / KTO
preference}}:::rl + RM --> RMCK[(reward.pt)]:::ckpt + RMCK -->|reward signal| PPO{{PPO
GAE + clip + KL}}:::rl + SFTCK --> PPO + SFTCK --> GRPO{{GRPO / RLVR
group-relative}}:::rl + PPO --> EVAL([GSM8K eval
+ chat / inference]):::eval + DPO --> EVAL + GRPO --> EVAL + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef ckpt fill:#eeeeee,stroke:#555555,stroke-width:2px,color:#222; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; +``` + +
+ +## The stages, in order + +| # | Stage | What it teaches the model | Doc | +|---|---|---|---| +| 1 | **Pretraining** | language itself (next-token prediction on the Pile) | [02_pretraining.md](02_pretraining.md) | +| 2 | **SFT** | to follow instructions & produce the `/` format | [03_sft.md](03_sft.md) | +| 3 | **Reward Model** | to score which answer humans prefer | [04_reward_model.md](04_reward_model.md) | +| 4 | **DPO / ORPO / KTO** | to prefer better answers *without* an RL loop | [05_dpo.md](05_dpo.md) | +| 5 | **PPO** | to maximize a reward (RM or verifier) with the classic RLHF loop | [06_ppo.md](06_ppo.md) | +| 6 | **GRPO / RLVR** | to reason, using verifiable rewards (DeepSeek-R1 style) | [07_grpo.md](07_grpo.md) | +| — | **Data pipeline** | how every dataset above is downloaded & preprocessed | [01_data_pipeline.md](01_data_pipeline.md) | +| — | **Evaluation** | how I measure GSM8K accuracy across all stages | [08_evaluation.md](08_evaluation.md) | +| — | **Inference / chat** | how to actually talk to any checkpoint | [09_inference.md](09_inference.md) | + +## The one design rule: *wrap, don't rewrite* + +Everything here sits on top of the original [`Transformer`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py). I changed the +educational model in exactly **one** place — I added a [`forward_hidden`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py#L56) +method that returns the final hidden states the `lm_head` consumes. Every post-training head (a value +head for PPO, a scalar reward head for the reward model) and every RL log-prob computation composes +*around* that one method, so the from-scratch model you already understand stays intact. + +## Colour legend (used in every diagram in these docs) + +🟩 data / corpus · 🟦 preprocessing · 🟦‍⬛ storage (HDF5 / JSONL) · 🟨 model / training loop +· 🟧 RL / reward · 🟥 loss / objective · 🟪 evaluation · ⬜ checkpoint + +> Each diagram is a hand-drawn, colour-coded Mermaid sketch, **pre-rendered to a PNG and embedded as +> an image** (GitHub's live Mermaid doesn't reliably do `look: handDrawn`, and some viewers — e.g. the +> VS Code preview — block SVGs, so an embedded PNG shows everywhere). The editable Mermaid source sits +> in a collapsible *"Mermaid source"* block under each image. To regenerate the images after editing, +> see [diagrams/README.md](diagrams/README.md). + +## Run the whole thing + +Once the base model has pretrained ([02_pretraining.md](02_pretraining.md)), the entire chain is one script: + +```bash +bash scripts/run_posttraining.sh # SFT -> RM -> DPO -> PPO -> GRPO -> eval table +``` + +See [POST_TRAINING.md](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/POST_TRAINING.md) for the condensed command reference. diff --git a/docs/assets/foundations-tutorial-screenshot.png b/docs/assets/foundations-tutorial-screenshot.png new file mode 100644 index 0000000..220c6eb Binary files /dev/null and b/docs/assets/foundations-tutorial-screenshot.png differ diff --git a/docs/diagrams/00_overview.png b/docs/diagrams/00_overview.png new file mode 100644 index 0000000..289b15a Binary files /dev/null and b/docs/diagrams/00_overview.png differ diff --git a/docs/diagrams/01_data_pipeline.png b/docs/diagrams/01_data_pipeline.png new file mode 100644 index 0000000..8c0cca4 Binary files /dev/null and b/docs/diagrams/01_data_pipeline.png differ diff --git a/docs/diagrams/02_pretraining.png b/docs/diagrams/02_pretraining.png new file mode 100644 index 0000000..53b5f1b Binary files /dev/null and b/docs/diagrams/02_pretraining.png differ diff --git a/docs/diagrams/03_sft.png b/docs/diagrams/03_sft.png new file mode 100644 index 0000000..7bd2cfc Binary files /dev/null and b/docs/diagrams/03_sft.png differ diff --git a/docs/diagrams/04_reward_model.png b/docs/diagrams/04_reward_model.png new file mode 100644 index 0000000..df92740 Binary files /dev/null and b/docs/diagrams/04_reward_model.png differ diff --git a/docs/diagrams/05_dpo.png b/docs/diagrams/05_dpo.png new file mode 100644 index 0000000..f6ad7e5 Binary files /dev/null and b/docs/diagrams/05_dpo.png differ diff --git a/docs/diagrams/06_ppo.png b/docs/diagrams/06_ppo.png new file mode 100644 index 0000000..8fa51f4 Binary files /dev/null and b/docs/diagrams/06_ppo.png differ diff --git a/docs/diagrams/07_grpo.png b/docs/diagrams/07_grpo.png new file mode 100644 index 0000000..2703ab4 Binary files /dev/null and b/docs/diagrams/07_grpo.png differ diff --git a/docs/diagrams/08_evaluation.png b/docs/diagrams/08_evaluation.png new file mode 100644 index 0000000..74bb66f Binary files /dev/null and b/docs/diagrams/08_evaluation.png differ diff --git a/docs/diagrams/09_inference.png b/docs/diagrams/09_inference.png new file mode 100644 index 0000000..6b91092 Binary files /dev/null and b/docs/diagrams/09_inference.png differ diff --git a/docs/diagrams/README.md b/docs/diagrams/README.md new file mode 100644 index 0000000..f449e11 --- /dev/null +++ b/docs/diagrams/README.md @@ -0,0 +1,30 @@ +# Diagrams + +The documentation diagrams are **hand-drawn, colour-coded Mermaid sketches**, pre-rendered to PNG and +embedded as images in the docs. We pre-render (rather than rely on live ```` ```mermaid ```` blocks) +because GitHub's live Mermaid renderer does not reliably support the `look: handDrawn` style, and some +markdown viewers (e.g. the VS Code preview) block SVGs — an embedded **PNG** shows the hand-drawn look +everywhere. Each doc also keeps the editable Mermaid source in a collapsible *"Mermaid source"* block +under its image. + +## Files + +- `src/*.mmd` — the canonical hand-drawn Mermaid sources (with `look: handDrawn` + the colour palette). +- `*.png` — the rendered images embedded by the docs (and `README.png` for the top-level README). + +## Regenerate after editing + +Edit the relevant `src/.mmd`, then from the repo root: + +```bash +bash scripts/render_diagrams.sh +``` + +That re-renders every `src/*.mmd` to `docs/diagrams/.png`. Requires the Mermaid CLI +(`npm i -g @mermaid-js/mermaid-cli`) and a Chrome/Chromium for headless rendering +(set `CHROME=/path/to/chrome` if it isn't at `/usr/bin/google-chrome-stable`). + +## Colour legend + +🟩 data / corpus · 🟦 preprocessing · teal storage (HDF5 / JSONL) · 🟨 model / training loop · +🟧 RL / reward · 🟥 loss / objective · 🟪 evaluation · ⬜ checkpoint diff --git a/docs/diagrams/README.png b/docs/diagrams/README.png new file mode 100644 index 0000000..6fd7449 Binary files /dev/null and b/docs/diagrams/README.png differ diff --git a/docs/diagrams/src/00_overview.mmd b/docs/diagrams/src/00_overview.mmd new file mode 100644 index 0000000..b13b547 --- /dev/null +++ b/docs/diagrams/src/00_overview.mmd @@ -0,0 +1,20 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart TD + PILE([The Pile
9.8B tokens]):::data --> PRE{{Pretrain
~400M base}}:::model + PRE --> BASE[(base_pretrained.pt)]:::ckpt + BASE --> SFT{{SFT
Alpaca · Dolly · GSM8K}}:::model + SFT --> SFTCK[(sft.pt)]:::ckpt + SFTCK --> RM{{Reward Model
Bradley-Terry}}:::rl + SFTCK --> DPO{{DPO / ORPO / KTO
preference}}:::rl + RM --> RMCK[(reward.pt)]:::ckpt + RMCK -->|reward signal| PPO{{PPO
GAE + clip + KL}}:::rl + SFTCK --> PPO + SFTCK --> GRPO{{GRPO / RLVR
group-relative}}:::rl + PPO --> EVAL([GSM8K eval
+ chat / inference]):::eval + DPO --> EVAL + GRPO --> EVAL + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef ckpt fill:#eeeeee,stroke:#555555,stroke-width:2px,color:#222; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; diff --git a/docs/diagrams/src/01_data_pipeline.mmd b/docs/diagrams/src/01_data_pipeline.mmd new file mode 100644 index 0000000..6f7d254 --- /dev/null +++ b/docs/diagrams/src/01_data_pipeline.mmd @@ -0,0 +1,23 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart TD + subgraph PT[" 1 · Pretraining "] + P1([Pile .jsonl.zst]):::data --> P2[stream-decompress
+ tiktoken r50k_base]:::proc --> P3[(pile_train.h5
flat tokens)]:::store + end + subgraph SF[" 2 · SFT "] + S1([Alpaca · Dolly · GSM8K]):::data --> S2[render chat template
mask prompt tokens]:::proc --> S3[pack to 1024]:::proc --> S4[(sft_packed.h5
tokens + loss_mask)]:::store + end + subgraph PF[" 3 · Preferences "] + F1([HH-RLHF · UltraFeedback]):::data --> F2[split prompt /
chosen / rejected]:::proc --> F3[(preferences.jsonl)]:::store + end + subgraph RL[" 4 · RL prompts "] + R1([GSM8K · arithmetic]):::data --> R2[extract numeric
gold answer]:::proc --> R3[(rl_prompts.jsonl
prompt + gold)]:::store + end + P3 --> M1{{pretrain_base.py}}:::model + S4 --> M2{{train_sft.py}}:::model + F3 --> M3{{train_reward.py · train_dpo.py}}:::rl + R3 --> M4{{train_ppo.py · train_grpo.py}}:::rl + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef store fill:#cdece8,stroke:#16a085,stroke-width:2px,color:#0a3d33; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; diff --git a/docs/diagrams/src/02_pretraining.mmd b/docs/diagrams/src/02_pretraining.mmd new file mode 100644 index 0000000..6b05c1b --- /dev/null +++ b/docs/diagrams/src/02_pretraining.mmd @@ -0,0 +1,15 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + H[(pile_train.h5)]:::store --> IT[get_batch_iterator
random windows]:::proc + IT --> FWD{{forward
bf16 autocast}}:::model + FWD --> L[cross-entropy loss]:::loss + L --> BWD[backward
x grad_accum]:::model + BWD --> CL[clip grad norm 1.0]:::proc + CL --> ST[AdamW step
cosine LR + warmup]:::model + ST -->|every 1000 steps| CK[(base_pretrained.pt)]:::ckpt + ST -->|next step| IT + classDef store fill:#cdece8,stroke:#16a085,stroke-width:2px,color:#0a3d33; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; diff --git a/docs/diagrams/src/03_sft.mmd b/docs/diagrams/src/03_sft.mmd new file mode 100644 index 0000000..3d47086 --- /dev/null +++ b/docs/diagrams/src/03_sft.mmd @@ -0,0 +1,13 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + H[(sft_packed.h5
tokens + loss_mask)]:::store --> B[batch:
tokens, mask]:::proc + B --> M{{Transformer
logits}}:::model + M --> SH[shift: predict t+1]:::proc + SH --> CE[token cross-entropy]:::loss + MASK([loss_mask = 1 on
assistant tokens]):::data --> CE + CE --> AVG[average over
masked tokens only]:::loss --> UPD[AdamW step]:::model + classDef store fill:#cdece8,stroke:#16a085,stroke-width:2px,color:#0a3d33; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; diff --git a/docs/diagrams/src/04_reward_model.mmd b/docs/diagrams/src/04_reward_model.mmd new file mode 100644 index 0000000..844c141 --- /dev/null +++ b/docs/diagrams/src/04_reward_model.mmd @@ -0,0 +1,15 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + P([preference pair
prompt + chosen / rejected]):::data --> BB{{SFT backbone
forward_hidden}}:::model + BB --> LT[take last real token]:::proc + LT --> RH[reward head
Linear→1]:::model + RH --> RC([r_chosen]):::rl + RH --> RR([r_rejected]):::rl + RC --> BT[Bradley-Terry
-log σ r_chosen - r_rejected]:::loss + RR --> BT + BT --> UPD[AdamW step]:::model + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; diff --git a/docs/diagrams/src/05_dpo.mmd b/docs/diagrams/src/05_dpo.mmd new file mode 100644 index 0000000..3ca5f59 --- /dev/null +++ b/docs/diagrams/src/05_dpo.mmd @@ -0,0 +1,14 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + P([chosen / rejected]):::data --> POL{{policy
trainable}}:::model + P --> REF{{reference
frozen SFT copy}}:::ckpt + POL --> LPP[seq log-probs
π_chosen, π_rejected]:::proc + REF --> LPR[seq log-probs
ref_chosen, ref_rejected]:::proc + LPP --> D[DPO loss
-log σ β·Δlogratios]:::loss + LPR --> D + D --> UPD[AdamW step]:::model + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; diff --git a/docs/diagrams/src/06_ppo.mmd b/docs/diagrams/src/06_ppo.mmd new file mode 100644 index 0000000..1b8aa79 --- /dev/null +++ b/docs/diagrams/src/06_ppo.mmd @@ -0,0 +1,17 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + PR([GSM8K prompts]):::data --> RO[rollout
generate_with_logprobs]:::proc + RO --> SC{score: verifier
or reward model}:::rl + SC --> KL[+ per-token
KL-to-ref penalty]:::rl + KL --> GAE[compute_gae
advantages + returns]:::proc + GAE --> UP{{clipped update
policy + value, K epochs}}:::model + UP -->|sync old policy| RO + REF{{frozen ref}}:::ckpt + REF -. KL .-> KL + VH{{value head}}:::model + VH -. value .-> GAE + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; diff --git a/docs/diagrams/src/07_grpo.mmd b/docs/diagrams/src/07_grpo.mmd new file mode 100644 index 0000000..080f872 --- /dev/null +++ b/docs/diagrams/src/07_grpo.mmd @@ -0,0 +1,16 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + PR([prompt]):::data --> G[sample a GROUP
of G answers]:::proc + G --> V{verifier reward
per answer}:::rl + V --> A[group advantage
r - mean / std]:::proc + A --> L[clipped surrogate
+ k3 KL to ref]:::loss + L --> UPD{{policy update}}:::model + UPD -->|next prompt| PR + REF{{frozen ref}}:::ckpt + REF -. KL .-> L + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; diff --git a/docs/diagrams/src/08_evaluation.mmd b/docs/diagrams/src/08_evaluation.mmd new file mode 100644 index 0000000..1862d30 --- /dev/null +++ b/docs/diagrams/src/08_evaluation.mmd @@ -0,0 +1,12 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + CK[(stage checkpoint)]:::ckpt --> GEN[batched_generate
greedy, length-bucketed]:::proc + Q([GSM8K question]):::data --> GEN + GEN --> EX[extract_answer
answer / #### / last number]:::proc + EX --> CMP{== gold?}:::eval + GOLD([gold number]):::data --> CMP + CMP --> ACC[accuracy across
Base→SFT→DPO→PPO→GRPO]:::eval + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; diff --git a/docs/diagrams/src/09_inference.mmd b/docs/diagrams/src/09_inference.mmd new file mode 100644 index 0000000..d16651e --- /dev/null +++ b/docs/diagrams/src/09_inference.mmd @@ -0,0 +1,13 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart LR + CK[(any checkpoint)]:::ckpt --> LD[load_model_from_ckpt
dims from stored cfg]:::proc + LD --> MODE{chat or raw?}:::proc + MODE -->|instruction model| CT[wrap in chat template]:::proc + MODE -->|base model| RAW[raw prefix]:::proc + CT --> GEN{{generate
temperature / top-p / greedy}}:::model + RAW --> GEN + GEN --> DEC([decode → reply]):::eval + classDef ckpt fill:#eeeeee,stroke:#555,stroke-width:2px,color:#222; + classDef proc fill:#d6e8ff,stroke:#2c6fbb,stroke-width:2px,color:#0d2c52; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; diff --git a/docs/diagrams/src/README.mmd b/docs/diagrams/src/README.mmd new file mode 100644 index 0000000..442e337 --- /dev/null +++ b/docs/diagrams/src/README.mmd @@ -0,0 +1,20 @@ +%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%% +flowchart TD + PILE([The Pile
9.8B tokens]):::data --> PRE{{Pretrain
~400M base}}:::model + PRE --> BASE[(base_pretrained.pt)]:::ckpt + BASE --> SFT{{SFT
Alpaca · Dolly · GSM8K}}:::model + SFT --> SFTCK[(sft.pt)]:::ckpt + SFTCK --> RM{{Reward Model
Bradley-Terry}}:::rl + SFTCK --> DPO{{DPO / ORPO / KTO}}:::rl + RM --> RMCK[(reward.pt)]:::ckpt + RMCK -->|reward signal| PPO{{PPO
GAE + clip + KL}}:::rl + SFTCK --> PPO + SFTCK --> GRPO{{GRPO / RLVR
group-relative}}:::rl + PPO --> EVAL([GSM8K eval
+ chat]):::eval + DPO --> EVAL + GRPO --> EVAL + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef rl fill:#ffd9b3,stroke:#e67e22,stroke-width:2px,color:#6b3500; + classDef ckpt fill:#eeeeee,stroke:#555555,stroke-width:2px,color:#222; + classDef eval fill:#e8d6ff,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a; diff --git a/docs/foundations/README.md b/docs/foundations/README.md new file mode 100644 index 0000000..e82d288 --- /dev/null +++ b/docs/foundations/README.md @@ -0,0 +1,108 @@ +# LLM Foundations + +This section explains the ideas that the training code assumes you know. It is not a PyTorch primer. +It is a bridge between the source files in this repository and the core concepts behind modern +decoder-only language models. + +The priority is: + +1. understand the base model and pretraining mechanics; +2. connect those mechanics to every existing stage in the site; +3. use the same language later for SFT, reward models, DPO, PPO, and GRPO. + +## The whole training story + +```mermaid +flowchart LR + RAW[raw text] --> TOK[tokenizer] + TOK --> IDS[token ids] + IDS --> BASE[decoder-only Transformer] + BASE --> CE[next-token cross-entropy] + CE --> CKPT[base checkpoint] + CKPT --> SFT[SFT: instruction following] + SFT --> PREF[preference optimization] + PREF --> RL[RL / verifier optimization] + RL --> CHAT[inference and chat] + + classDef data fill:#d6ffd9,stroke:#27ae60,stroke-width:2px,color:#143d1a; + classDef model fill:#ffe8a3,stroke:#d48806,stroke-width:2px,color:#5a3d00; + classDef loss fill:#ffd6d6,stroke:#c0392b,stroke-width:2px,color:#5c1212; + class RAW,TOK,IDS data; + class BASE,CKPT,SFT,PREF,RL,CHAT model; + class CE loss; +``` + +At the base level, an LLM is a conditional probability model: + +\[ +p_\theta(x_1, x_2, \ldots, x_T) += \prod_{t=1}^{T} p_\theta(x_t \mid x_{