chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
+22
@@ -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/
|
||||
@@ -0,0 +1,13 @@
|
||||
[theme]
|
||||
primaryColor = "#16a085"
|
||||
backgroundColor = "#ffffff"
|
||||
secondaryBackgroundColor = "#f3faf8"
|
||||
textColor = "#0a3d33"
|
||||
font = "sans serif"
|
||||
|
||||
[server]
|
||||
headless = true
|
||||
runOnSave = true
|
||||
|
||||
[browser]
|
||||
gatherUsageStats = false
|
||||
@@ -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.
|
||||
@@ -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|>`, `<think>`, `<answer>`) are ordinary tokens the model learns.
|
||||
GSM8K is reformatted into `<think>…</think><answer>N</answer>` 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).
|
||||
@@ -0,0 +1,762 @@
|
||||
|
||||
|
||||
<div align="center">
|
||||
|
||||
<!-- omit in toc -->
|
||||
# Train LLM From Scratch
|
||||
|
||||
   [](https://fareedkhan-dev.github.io/train-llm-from-scratch/)
|
||||
|
||||
**I am Looking for a PhD position in AI**. [GitHub](https://github.com/FareedKhan-dev)
|
||||
|
||||
</div>
|
||||
|
||||
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.
|
||||
|
||||
|
||||
|
||||
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.
|
||||
```
|
||||
|
||||
<!-- omit in toc -->
|
||||
## 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 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.
|
||||
|
||||
|
||||
|
||||
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:
|
||||
|
||||
```
|
||||
<think>step by step reasoning ...</think><answer>42</answer>
|
||||
```
|
||||
|
||||
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|>
|
||||
<think>13 + 29 = 42</think><answer>42</answer><|endoftext|>
|
||||
|
||||
extract_answer("<answer>42</answer>") -> 42.0
|
||||
reward_gsm8k("<answer>42</answer>", 42.0) -> 1.2 # correct AND well formatted
|
||||
reward_gsm8k("<answer>7</answer>", 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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
```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 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:
|
||||
|
||||
|
||||
|
||||
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 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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
|
||||
|
||||
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.
|
||||
|
||||
|
||||
|
||||
```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 `<answer>` tags, and check it against the gold answer.
|
||||
|
||||
|
||||
|
||||
```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 `<answer>` 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 `<think>...</think><answer>...</answer>` 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.
|
||||
|
||||
|
||||
|
||||
```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.
|
||||
|
||||
<hr>
|
||||
|
||||
Wanna chat on something? [My Linkedin](https://www.linkedin.com/in/fareed-khan-dev/)
|
||||
|
||||
## Star History
|
||||
|
||||
[](https://star-history.com/#FareedKhan-dev/train-llm-from-scratch&Date)
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`FareedKhan-dev/train-llm-from-scratch`
|
||||
- 原始仓库：https://github.com/FareedKhan-dev/train-llm-from-scratch
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
@@ -0,0 +1,2 @@
|
||||
"""Configuration package: legacy ``config.config`` (pretraining), the post-training
|
||||
dataclasses (``config.post_training_config``), and the JSON ``config.loader``."""
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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.<Stage>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)
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"vocab_size": 50304,
|
||||
"context_length": 256,
|
||||
"n_embed": 128,
|
||||
"n_head": 4,
|
||||
"n_blocks": 2,
|
||||
"device": "cpu",
|
||||
"amp_dtype": null
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"batch_size": 4,
|
||||
"epochs": 1,
|
||||
"eval_steps": 10000,
|
||||
"warmup_steps": 2,
|
||||
"max_len": 256,
|
||||
"save_every": 10000
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"batch_size": 4,
|
||||
"epochs": 1,
|
||||
"eval_steps": 10000,
|
||||
"warmup_steps": 2,
|
||||
"max_len": 256,
|
||||
"save_every": 10000
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"max_steps": 10,
|
||||
"batch_size": 4,
|
||||
"eval_steps": 10000,
|
||||
"warmup_steps": 2,
|
||||
"save_every": 10000
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
"""Data loaders: pretraining HDF5 iterator + SFT / preference / RL-prompt iterators."""
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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]]
|
||||
@@ -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
|
||||
@@ -0,0 +1,162 @@
|
||||
<!-- omit in toc -->
|
||||
# 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:
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
subgraph PT[" 1 · Pretraining "]
|
||||
P1([Pile .jsonl.zst]):::data --> P2[stream-decompress<br/>+ tiktoken r50k_base]:::proc --> P3[(pile_train.h5<br/>flat tokens)]:::store
|
||||
end
|
||||
subgraph SF[" 2 · SFT "]
|
||||
S1([Alpaca · Dolly · GSM8K]):::data --> S2[render chat template<br/>mask prompt tokens]:::proc --> S3[pack to 1024]:::proc --> S4[(sft_packed.h5<br/>tokens + loss_mask)]:::store
|
||||
end
|
||||
subgraph PF[" 3 · Preferences "]
|
||||
F1([HH-RLHF · UltraFeedback]):::data --> F2[split prompt /<br/>chosen / rejected]:::proc --> F3[(preferences.jsonl)]:::store
|
||||
end
|
||||
subgraph RL[" 4 · RL prompts "]
|
||||
R1([GSM8K · arithmetic]):::data --> R2[extract numeric<br/>gold answer]:::proc --> R3[(rl_prompts.jsonl<br/>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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
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 `<think>…</think><answer>N</answer>` 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 `<answer>4</answer>` 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 |
|
||||
<br>
|
||||
|
||||
➡️ Next: [Stage 1 — Pretraining](02_pretraining.md).
|
||||
@@ -0,0 +1,102 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
H[(pile_train.h5)]:::store --> IT[get_batch_iterator<br/>random windows]:::proc
|
||||
IT --> FWD{{forward<br/>bf16 autocast}}:::model
|
||||
FWD --> L[cross-entropy loss]:::loss
|
||||
L --> BWD[backward<br/>x grad_accum]:::model
|
||||
BWD --> CL[clip grad norm 1.0]:::proc
|
||||
CL --> ST[AdamW step<br/>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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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).
|
||||
@@ -0,0 +1,95 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
H[(sft_packed.h5<br/>tokens + loss_mask)]:::store --> B[batch:<br/>tokens, mask]:::proc
|
||||
B --> M{{Transformer<br/>logits}}:::model
|
||||
M --> SH[shift: predict t+1]:::proc
|
||||
SH --> CE[token cross-entropy]:::loss
|
||||
MASK([loss_mask = 1 on<br/>assistant tokens]):::data --> CE
|
||||
CE --> AVG[average over<br/>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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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
|
||||
`<answer>…</answer>` 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).
|
||||
@@ -0,0 +1,105 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
P([preference pair<br/>prompt + chosen / rejected]):::data --> BB{{SFT backbone<br/>forward_hidden}}:::model
|
||||
BB --> LT[take last real token]:::proc
|
||||
LT --> RH[reward head<br/>Linear→1]:::model
|
||||
RH --> RC([r_chosen]):::rl
|
||||
RH --> RR([r_rejected]):::rl
|
||||
RC --> BT[Bradley-Terry<br/>-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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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).
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
P([chosen / rejected]):::data --> POL{{policy<br/>trainable}}:::model
|
||||
P --> REF{{reference<br/>frozen SFT copy}}:::ckpt
|
||||
POL --> LPP[seq log-probs<br/>π_chosen, π_rejected]:::proc
|
||||
REF --> LPR[seq log-probs<br/>ref_chosen, ref_rejected]:::proc
|
||||
LPP --> D[DPO loss<br/>-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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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).
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
PR([GSM8K prompts]):::data --> RO[rollout<br/>generate_with_logprobs]:::proc
|
||||
RO --> SC{score: verifier<br/>or reward model}:::rl
|
||||
SC --> KL[+ per-token<br/>KL-to-ref penalty]:::rl
|
||||
KL --> GAE[compute_gae<br/>advantages + returns]:::proc
|
||||
GAE --> UP{{clipped update<br/>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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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.
|
||||
+102
@@ -0,0 +1,102 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
PR([prompt]):::data --> G[sample a GROUP<br/>of G answers]:::proc
|
||||
G --> V{verifier reward<br/>per answer}:::rl
|
||||
V --> A[group advantage<br/>r - mean / std]:::proc
|
||||
A --> L[clipped surrogate<br/>+ 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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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).
|
||||
@@ -0,0 +1,99 @@
|
||||
<!-- omit in toc -->
|
||||
# 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.
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
CK[(stage checkpoint)]:::ckpt --> GEN[batched_generate<br/>greedy, length-bucketed]:::proc
|
||||
Q([GSM8K question]):::data --> GEN
|
||||
GEN --> EX[extract_answer<br/>answer / #### / last number]:::proc
|
||||
EX --> CMP{== gold?}:::eval
|
||||
GOLD([gold number]):::data --> CMP
|
||||
CMP --> ACC[accuracy across<br/>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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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 `<answer>…</answer>` 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).
|
||||
@@ -0,0 +1,92 @@
|
||||
<!-- omit in toc -->
|
||||
# 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).
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
CK[(any checkpoint)]:::ckpt --> LD[load_model_from_ckpt<br/>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<br/>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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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).
|
||||
+108
@@ -0,0 +1,108 @@
|
||||
<!-- omit in toc -->
|
||||
# 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:
|
||||
|
||||
|
||||
|
||||
<details>
|
||||
<summary>Mermaid source (live, editable)</summary>
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
PILE([The Pile<br/>9.8B tokens]):::data --> PRE{{Pretrain<br/>~400M base}}:::model
|
||||
PRE --> BASE[(base_pretrained.pt)]:::ckpt
|
||||
BASE --> SFT{{SFT<br/>Alpaca · Dolly · GSM8K}}:::model
|
||||
SFT --> SFTCK[(sft.pt)]:::ckpt
|
||||
SFTCK --> RM{{Reward Model<br/>Bradley-Terry}}:::rl
|
||||
SFTCK --> DPO{{DPO / ORPO / KTO<br/>preference}}:::rl
|
||||
RM --> RMCK[(reward.pt)]:::ckpt
|
||||
RMCK -->|reward signal| PPO{{PPO<br/>GAE + clip + KL}}:::rl
|
||||
SFTCK --> PPO
|
||||
SFTCK --> GRPO{{GRPO / RLVR<br/>group-relative}}:::rl
|
||||
PPO --> EVAL([GSM8K eval<br/>+ 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;
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 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 `<think>/<answer>` 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.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 181 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 322 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 570 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 91 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 117 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 103 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 151 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 139 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 123 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 138 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 106 KiB |
@@ -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/<name>.mmd`, then from the repo root:
|
||||
|
||||
```bash
|
||||
bash scripts/render_diagrams.sh
|
||||
```
|
||||
|
||||
That re-renders every `src/*.mmd` to `docs/diagrams/<name>.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
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 310 KiB |
@@ -0,0 +1,20 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart TD
|
||||
PILE([The Pile<br/>9.8B tokens]):::data --> PRE{{Pretrain<br/>~400M base}}:::model
|
||||
PRE --> BASE[(base_pretrained.pt)]:::ckpt
|
||||
BASE --> SFT{{SFT<br/>Alpaca · Dolly · GSM8K}}:::model
|
||||
SFT --> SFTCK[(sft.pt)]:::ckpt
|
||||
SFTCK --> RM{{Reward Model<br/>Bradley-Terry}}:::rl
|
||||
SFTCK --> DPO{{DPO / ORPO / KTO<br/>preference}}:::rl
|
||||
RM --> RMCK[(reward.pt)]:::ckpt
|
||||
RMCK -->|reward signal| PPO{{PPO<br/>GAE + clip + KL}}:::rl
|
||||
SFTCK --> PPO
|
||||
SFTCK --> GRPO{{GRPO / RLVR<br/>group-relative}}:::rl
|
||||
PPO --> EVAL([GSM8K eval<br/>+ 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;
|
||||
@@ -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<br/>+ tiktoken r50k_base]:::proc --> P3[(pile_train.h5<br/>flat tokens)]:::store
|
||||
end
|
||||
subgraph SF[" 2 · SFT "]
|
||||
S1([Alpaca · Dolly · GSM8K]):::data --> S2[render chat template<br/>mask prompt tokens]:::proc --> S3[pack to 1024]:::proc --> S4[(sft_packed.h5<br/>tokens + loss_mask)]:::store
|
||||
end
|
||||
subgraph PF[" 3 · Preferences "]
|
||||
F1([HH-RLHF · UltraFeedback]):::data --> F2[split prompt /<br/>chosen / rejected]:::proc --> F3[(preferences.jsonl)]:::store
|
||||
end
|
||||
subgraph RL[" 4 · RL prompts "]
|
||||
R1([GSM8K · arithmetic]):::data --> R2[extract numeric<br/>gold answer]:::proc --> R3[(rl_prompts.jsonl<br/>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;
|
||||
@@ -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<br/>random windows]:::proc
|
||||
IT --> FWD{{forward<br/>bf16 autocast}}:::model
|
||||
FWD --> L[cross-entropy loss]:::loss
|
||||
L --> BWD[backward<br/>x grad_accum]:::model
|
||||
BWD --> CL[clip grad norm 1.0]:::proc
|
||||
CL --> ST[AdamW step<br/>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;
|
||||
@@ -0,0 +1,13 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart LR
|
||||
H[(sft_packed.h5<br/>tokens + loss_mask)]:::store --> B[batch:<br/>tokens, mask]:::proc
|
||||
B --> M{{Transformer<br/>logits}}:::model
|
||||
M --> SH[shift: predict t+1]:::proc
|
||||
SH --> CE[token cross-entropy]:::loss
|
||||
MASK([loss_mask = 1 on<br/>assistant tokens]):::data --> CE
|
||||
CE --> AVG[average over<br/>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;
|
||||
@@ -0,0 +1,15 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart LR
|
||||
P([preference pair<br/>prompt + chosen / rejected]):::data --> BB{{SFT backbone<br/>forward_hidden}}:::model
|
||||
BB --> LT[take last real token]:::proc
|
||||
LT --> RH[reward head<br/>Linear→1]:::model
|
||||
RH --> RC([r_chosen]):::rl
|
||||
RH --> RR([r_rejected]):::rl
|
||||
RC --> BT[Bradley-Terry<br/>-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;
|
||||
@@ -0,0 +1,14 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart LR
|
||||
P([chosen / rejected]):::data --> POL{{policy<br/>trainable}}:::model
|
||||
P --> REF{{reference<br/>frozen SFT copy}}:::ckpt
|
||||
POL --> LPP[seq log-probs<br/>π_chosen, π_rejected]:::proc
|
||||
REF --> LPR[seq log-probs<br/>ref_chosen, ref_rejected]:::proc
|
||||
LPP --> D[DPO loss<br/>-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;
|
||||
@@ -0,0 +1,17 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart LR
|
||||
PR([GSM8K prompts]):::data --> RO[rollout<br/>generate_with_logprobs]:::proc
|
||||
RO --> SC{score: verifier<br/>or reward model}:::rl
|
||||
SC --> KL[+ per-token<br/>KL-to-ref penalty]:::rl
|
||||
KL --> GAE[compute_gae<br/>advantages + returns]:::proc
|
||||
GAE --> UP{{clipped update<br/>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;
|
||||
@@ -0,0 +1,16 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart LR
|
||||
PR([prompt]):::data --> G[sample a GROUP<br/>of G answers]:::proc
|
||||
G --> V{verifier reward<br/>per answer}:::rl
|
||||
V --> A[group advantage<br/>r - mean / std]:::proc
|
||||
A --> L[clipped surrogate<br/>+ 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;
|
||||
@@ -0,0 +1,12 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart LR
|
||||
CK[(stage checkpoint)]:::ckpt --> GEN[batched_generate<br/>greedy, length-bucketed]:::proc
|
||||
Q([GSM8K question]):::data --> GEN
|
||||
GEN --> EX[extract_answer<br/>answer / #### / last number]:::proc
|
||||
EX --> CMP{== gold?}:::eval
|
||||
GOLD([gold number]):::data --> CMP
|
||||
CMP --> ACC[accuracy across<br/>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;
|
||||
@@ -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<br/>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<br/>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;
|
||||
@@ -0,0 +1,20 @@
|
||||
%%{init: {'look':'handDrawn','theme':'base','themeVariables':{'fontFamily':'Comic Sans MS, cursive'}}}%%
|
||||
flowchart TD
|
||||
PILE([The Pile<br/>9.8B tokens]):::data --> PRE{{Pretrain<br/>~400M base}}:::model
|
||||
PRE --> BASE[(base_pretrained.pt)]:::ckpt
|
||||
BASE --> SFT{{SFT<br/>Alpaca · Dolly · GSM8K}}:::model
|
||||
SFT --> SFTCK[(sft.pt)]:::ckpt
|
||||
SFTCK --> RM{{Reward Model<br/>Bradley-Terry}}:::rl
|
||||
SFTCK --> DPO{{DPO / ORPO / KTO}}:::rl
|
||||
RM --> RMCK[(reward.pt)]:::ckpt
|
||||
RMCK -->|reward signal| PPO{{PPO<br/>GAE + clip + KL}}:::rl
|
||||
SFTCK --> PPO
|
||||
SFTCK --> GRPO{{GRPO / RLVR<br/>group-relative}}:::rl
|
||||
PPO --> EVAL([GSM8K eval<br/>+ 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;
|
||||
@@ -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_{<t})
|
||||
\]
|
||||
|
||||
The Transformer does not learn "truth" directly. It learns a distribution over the next token. A
|
||||
large amount of useful behavior appears because the next-token task forces the model to compress
|
||||
syntax, facts, formats, style, and reasoning traces into its weights.
|
||||
|
||||
## Where each concept lives in this repo
|
||||
|
||||
| Concept | Why it matters | Main code |
|
||||
|---|---|---|
|
||||
| Tokenization | Text must become integer ids before the model can train. | `scripts/prepare_pretrain_data.py`, `src/post_training/chat_template.py` |
|
||||
| Fixed context windows | Training examples are contiguous windows of tokens. | `data_loader/data_loader.py` |
|
||||
| Token and position embeddings | Integer ids become vectors and positions. | `src/models/transformer.py` |
|
||||
| Causal self-attention | Each token mixes information from previous tokens only. | `src/models/attention.py` |
|
||||
| Transformer block | Attention plus MLP, with pre-norm residual structure. | `src/models/transformer_block.py` |
|
||||
| Logits | Hidden states become unnormalized vocabulary scores. | `src/models/transformer.py` |
|
||||
| Cross-entropy | The base objective rewards probability on the true next token. | `src/models/transformer.py`, `src/post_training/sft.py` |
|
||||
| AdamW and LR schedule | Optimization details that decide whether training is stable. | `src/post_training/optim.py` |
|
||||
| Gradient accumulation | Simulates a larger batch under memory limits. | `scripts/pretrain_base.py`, `scripts/train_sft.py` |
|
||||
| Generation | The model feeds sampled tokens back into itself. | `src/models/transformer.py`, `src/post_training/inference.py` |
|
||||
|
||||
## Learning path
|
||||
|
||||
Read in this order:
|
||||
|
||||
1. [Tokenization & Data Shapes](tokenization.md) - how text becomes batches.
|
||||
2. [Decoder-Only Transformer](transformer.md) - the model skeleton.
|
||||
3. [Attention, Masks & Heads](attention.md) - the core operation.
|
||||
4. [Objectives, Losses & Perplexity](objectives.md) - what the model is optimized to do.
|
||||
5. [Optimization & Training Systems](optimization.md) - how the loop stays stable.
|
||||
6. [Generation & Sampling](generation.md) - how logits become text.
|
||||
|
||||
Then continue to the pipeline pages:
|
||||
|
||||
- [Data handling](../01_data_pipeline.md)
|
||||
- [Pretraining](../02_pretraining.md)
|
||||
- [SFT](../03_sft.md)
|
||||
- [Reward Model](../04_reward_model.md)
|
||||
- [DPO / ORPO / KTO](../05_dpo.md)
|
||||
- [PPO](../06_ppo.md)
|
||||
- [GRPO / RLVR](../07_grpo.md)
|
||||
|
||||
## Mental model
|
||||
|
||||
The base training loop is compact:
|
||||
|
||||
\[
|
||||
\text{text} \to \text{token ids} \to \text{embeddings} \to \text{Transformer blocks}
|
||||
\to \text{logits} \to \text{cross-entropy} \to \nabla_\theta
|
||||
\]
|
||||
|
||||
Post-training mostly changes the data and the loss:
|
||||
|
||||
- SFT keeps next-token prediction, but masks the loss to assistant tokens.
|
||||
- Reward modeling changes the output from vocabulary logits to one scalar score.
|
||||
- DPO compares sequence log-probabilities for chosen and rejected responses.
|
||||
- PPO and GRPO sample completions, score them, and update the policy with a constrained RL objective.
|
||||
|
||||
The same backbone is reused throughout. That is the most important design idea in this repo.
|
||||
|
||||
## Primary references
|
||||
|
||||
- [Attention Is All You Need](https://arxiv.org/abs/1706.03762) introduced the Transformer and scaled dot-product attention.
|
||||
- [The Pile](https://arxiv.org/abs/2101.00027) describes the 825 GiB text corpus used by the pretraining path.
|
||||
- [Neural Machine Translation of Rare Words with Subword Units](https://arxiv.org/abs/1508.07909) introduced the BPE subword idea that modern tokenizers build on.
|
||||
- [Decoupled Weight Decay Regularization](https://arxiv.org/abs/1711.05101) motivates AdamW.
|
||||
- [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155) is the classic SFT -> reward model -> PPO RLHF recipe.
|
||||
- [Direct Preference Optimization](https://arxiv.org/abs/2305.18290) motivates the DPO page.
|
||||
- [DeepSeekMath](https://arxiv.org/abs/2402.03300) introduces GRPO in the mathematical reasoning setting.
|
||||
@@ -0,0 +1,168 @@
|
||||
# Attention, Masks & Heads
|
||||
|
||||
Self-attention is the operation that lets each token decide which previous tokens matter.
|
||||
|
||||
In a decoder-only language model, position \(t\) can use information from positions \(0..t\), but not
|
||||
from positions \(t+1..T-1\). That restriction is what makes next-token training honest.
|
||||
|
||||
## The attention equation
|
||||
|
||||
For an input tensor \(X \in \mathbb{R}^{B \times T \times C}\), one attention head learns three
|
||||
linear projections:
|
||||
|
||||
\[
|
||||
Q = X W_Q,\quad K = X W_K,\quad V = X W_V
|
||||
\]
|
||||
|
||||
where \(Q,K,V \in \mathbb{R}^{B \times T \times D}\).
|
||||
|
||||
Scaled dot-product attention is:
|
||||
|
||||
\[
|
||||
\text{Attention}(Q,K,V)
|
||||
= \text{softmax}\left(\frac{QK^T}{\sqrt{D}} + M\right)V
|
||||
\]
|
||||
|
||||
The mask \(M\) is `0` for allowed positions and \(-\infty\) for future positions.
|
||||
|
||||
## Query, key, value intuition
|
||||
|
||||
For each token:
|
||||
|
||||
- query: "what am I looking for?"
|
||||
- key: "what information do I contain?"
|
||||
- value: "what content should I pass forward if selected?"
|
||||
|
||||
The dot product \(q_t \cdot k_s\) measures how much token \(t\) wants information from token \(s\).
|
||||
|
||||
## Why divide by \(\sqrt{D}\)?
|
||||
|
||||
If the components of queries and keys have roughly unit variance, their dot product has variance
|
||||
proportional to \(D\). Larger head sizes would create large logits, making softmax overly sharp and
|
||||
gradients weak. The scale factor keeps attention logits in a more stable range:
|
||||
|
||||
\[
|
||||
\frac{q_t \cdot k_s}{\sqrt{D}}
|
||||
\]
|
||||
|
||||
## Causal masking
|
||||
|
||||
For \(T=5\), the allowed attention pattern is:
|
||||
|
||||
\[
|
||||
\begin{bmatrix}
|
||||
1 & 0 & 0 & 0 & 0 \\
|
||||
1 & 1 & 0 & 0 & 0 \\
|
||||
1 & 1 & 1 & 0 & 0 \\
|
||||
1 & 1 & 1 & 1 & 0 \\
|
||||
1 & 1 & 1 & 1 & 1
|
||||
\end{bmatrix}
|
||||
\]
|
||||
|
||||
The repo stores this as a lower-triangular buffer:
|
||||
|
||||
```python
|
||||
self.register_buffer("tril", torch.tril(torch.ones(context_length, context_length)))
|
||||
```
|
||||
|
||||
Then masks future locations before softmax:
|
||||
|
||||
```python
|
||||
attn_weights = q @ k.transpose(-2, -1) * scale_factor
|
||||
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
|
||||
attn_weights = F.softmax(attn_weights, dim=-1)
|
||||
out = attn_weights @ v
|
||||
```
|
||||
|
||||
Because the future logits become \(-\infty\), their softmax probability becomes zero.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
X["x (B,T,C)"] --> Q["query projection"]
|
||||
X --> K["key projection"]
|
||||
X --> V["value projection"]
|
||||
Q --> S["QK^T / sqrt(D)"]
|
||||
K --> S
|
||||
S --> M["causal mask"]
|
||||
M --> P["softmax weights"]
|
||||
V --> O["weights @ V"]
|
||||
P --> O
|
||||
```
|
||||
|
||||
## Multi-head attention
|
||||
|
||||
One head has one attention pattern. Multiple heads let the model learn several patterns in parallel:
|
||||
|
||||
- syntax dependencies;
|
||||
- repeated names or entities;
|
||||
- local phrase structure;
|
||||
- delimiter and format tracking;
|
||||
- arithmetic or code-like dependencies.
|
||||
|
||||
The repo creates `n_head` independent `Head` modules:
|
||||
|
||||
```python
|
||||
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)
|
||||
```
|
||||
|
||||
Then concatenates head outputs:
|
||||
|
||||
```python
|
||||
x = torch.cat([h(x) for h in self.heads], dim=-1)
|
||||
x = self.proj(x)
|
||||
```
|
||||
|
||||
If \(H\) heads each emit width \(D=C/H\), concatenation returns width \(C\):
|
||||
|
||||
\[
|
||||
\text{Concat}(\text{head}_1,\ldots,\text{head}_H) \in \mathbb{R}^{B \times T \times C}
|
||||
\]
|
||||
|
||||
The final projection mixes information across heads.
|
||||
|
||||
## Attention cost
|
||||
|
||||
The attention score matrix has shape:
|
||||
|
||||
\[
|
||||
B \times T \times T
|
||||
\]
|
||||
|
||||
per head. With \(H\) heads, the core score storage is roughly:
|
||||
|
||||
\[
|
||||
O(BHT^2)
|
||||
\]
|
||||
|
||||
This is why context length is expensive. Doubling \(T\) roughly quadruples the attention matrix size.
|
||||
The repo's educational attention implementation is intentionally readable and materializes these
|
||||
matrices directly.
|
||||
|
||||
## What attention can and cannot do
|
||||
|
||||
Attention mixes information between positions. It does not by itself:
|
||||
|
||||
- create a probability distribution over the vocabulary;
|
||||
- know token order without positional information;
|
||||
- perform nonlinear transformation beyond weighted averaging.
|
||||
|
||||
Those jobs come from position embeddings, the MLP, layer norms, residual paths, and the final LM head.
|
||||
|
||||
## Debugging attention mentally
|
||||
|
||||
If training behaves oddly, ask these questions:
|
||||
|
||||
1. Is the mask causal, or can tokens see the answer?
|
||||
2. Are `q`, `k`, and `v` projected from the same normalized input?
|
||||
3. Is `head_size = n_embed // n_head` an integer?
|
||||
4. Does concatenating all heads return exactly `n_embed` channels?
|
||||
5. Is sequence length `T` less than or equal to `context_length`?
|
||||
|
||||
## Next
|
||||
|
||||
Attention creates hidden states. The loss decides how hidden states become learning signal. Continue
|
||||
to [Objectives, Losses & Perplexity](objectives.md).
|
||||
@@ -0,0 +1,140 @@
|
||||
# Generation & Sampling
|
||||
|
||||
Training predicts the next token for known text. Generation uses the model's own sampled token as the
|
||||
next input.
|
||||
|
||||
That feedback loop is why small probability differences can produce very different completions.
|
||||
|
||||
## Autoregressive loop
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
P["prompt ids"] --> CROP["keep last context_length ids"]
|
||||
CROP --> M["Transformer"]
|
||||
M --> LOGITS["last-position logits"]
|
||||
LOGITS --> PROBS["softmax / sampling"]
|
||||
PROBS --> NEXT["next token id"]
|
||||
NEXT --> APPEND["append to sequence"]
|
||||
APPEND --> CROP
|
||||
```
|
||||
|
||||
The simple implementation in `src/models/transformer.py`:
|
||||
|
||||
```python
|
||||
for _ in range(max_new_tokens):
|
||||
idx_cond = idx[:, -self.context_length:]
|
||||
logits, _ = self(idx_cond)
|
||||
logits = logits[:, -1, :]
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
idx_next = torch.multinomial(probs, num_samples=1)
|
||||
idx = torch.cat((idx, idx_next), dim=1)
|
||||
```
|
||||
|
||||
Only the final position is used because that position has attended to the whole current context.
|
||||
|
||||
## Greedy decoding vs sampling
|
||||
|
||||
Greedy decoding chooses:
|
||||
|
||||
\[
|
||||
\arg\max_i p_i
|
||||
\]
|
||||
|
||||
Sampling draws:
|
||||
|
||||
\[
|
||||
x \sim \text{Categorical}(p)
|
||||
\]
|
||||
|
||||
Greedy decoding is deterministic and often repetitive. Sampling is stochastic and can produce more
|
||||
varied text.
|
||||
|
||||
This repo's simple base `generate` method samples directly from the full softmax distribution.
|
||||
The post-training inference utilities add more controls.
|
||||
|
||||
## Temperature
|
||||
|
||||
Temperature rescales logits before softmax:
|
||||
|
||||
\[
|
||||
p_i =
|
||||
\frac{\exp(z_i / \tau)}
|
||||
{\sum_j \exp(z_j / \tau)}
|
||||
\]
|
||||
|
||||
Effects:
|
||||
|
||||
- \(\tau < 1\): sharper distribution, safer but less diverse;
|
||||
- \(\tau = 1\): unchanged;
|
||||
- \(\tau > 1\): flatter distribution, more diverse but more error-prone.
|
||||
|
||||
Temperature operates on logits, not probabilities.
|
||||
|
||||
## Top-k and top-p
|
||||
|
||||
Many generation systems restrict the candidate set before sampling.
|
||||
|
||||
Top-k keeps only the \(k\) highest-probability tokens.
|
||||
|
||||
Top-p, also called nucleus sampling, keeps the smallest set of tokens whose cumulative probability is
|
||||
at least \(p\).
|
||||
|
||||
These controls are not part of the base model architecture. They are decoding policies layered on top
|
||||
of the model's logits.
|
||||
|
||||
## Context cropping
|
||||
|
||||
The model has a fixed maximum context length:
|
||||
|
||||
```python
|
||||
idx_cond = idx[:, -self.context_length:]
|
||||
```
|
||||
|
||||
If the conversation grows beyond that, the oldest tokens are dropped. The model cannot attend to text
|
||||
outside the retained context window.
|
||||
|
||||
This is why context length is a product constraint, not just a training hyperparameter.
|
||||
|
||||
## Stop tokens
|
||||
|
||||
The tokenizer's EOT token is:
|
||||
|
||||
```python
|
||||
EOT_ID = 50256
|
||||
```
|
||||
|
||||
During training, EOT appears between documents and after assistant messages. During inference, a chat
|
||||
loop can stop when EOT appears or when a formatted answer is complete.
|
||||
|
||||
If the model was never trained to emit a clear stop token or answer delimiter, decoding has to guess
|
||||
when to stop.
|
||||
|
||||
## Why generated text can drift
|
||||
|
||||
During teacher-forced training, every input prefix comes from the dataset. During generation, prefixes
|
||||
come from the model itself. If the model samples a poor token early, future predictions condition on
|
||||
that poor token.
|
||||
|
||||
This distribution shift is one reason post-training matters:
|
||||
|
||||
- SFT teaches the answer format;
|
||||
- reward/preference methods push the model toward preferred completions;
|
||||
- RLVR/GRPO can reward final answers that satisfy an external verifier.
|
||||
|
||||
## Useful generation diagnostics
|
||||
|
||||
| Symptom | Likely cause | Check |
|
||||
|---|---|---|
|
||||
| Repeats forever | distribution too sharp or no learned stop behavior | lower max tokens, check EOT handling, adjust sampling |
|
||||
| Ignores instruction | base model not SFT-trained enough | test an SFT checkpoint |
|
||||
| Wrong answer format | SFT data format mismatch | inspect chat template and masks |
|
||||
| Random-looking text | undertrained model or high temperature | compare train/dev loss and lower temperature |
|
||||
| Crashes on long prompts | prompt exceeds context or device memory | crop context and inspect `context_length` |
|
||||
|
||||
## Next
|
||||
|
||||
You now have the foundation for the full pipeline. Continue to:
|
||||
|
||||
- [Data handling](../01_data_pipeline.md)
|
||||
- [Pretraining](../02_pretraining.md)
|
||||
- [SFT](../03_sft.md)
|
||||
@@ -0,0 +1,246 @@
|
||||
# Objectives, Losses & Perplexity
|
||||
|
||||
The architecture defines what the model can compute. The objective defines what behavior training
|
||||
rewards.
|
||||
|
||||
For a decoder-only language model, the base objective is next-token prediction.
|
||||
|
||||
## From logits to probability
|
||||
|
||||
At position \(t\), the model emits logits:
|
||||
|
||||
\[
|
||||
z_t \in \mathbb{R}^{V}
|
||||
\]
|
||||
|
||||
Softmax turns logits into a distribution:
|
||||
|
||||
\[
|
||||
p_\theta(x_{t+1}=i \mid x_{\leq t})
|
||||
= \frac{\exp(z_{t,i})}{\sum_{j=1}^{V}\exp(z_{t,j})}
|
||||
\]
|
||||
|
||||
The target is one integer token id \(y_t\). Cross-entropy for that position is:
|
||||
|
||||
\[
|
||||
\ell_t = -\log p_\theta(y_t \mid x_{\leq t})
|
||||
\]
|
||||
|
||||
The batch loss is the mean over positions:
|
||||
|
||||
\[
|
||||
\mathcal{L}_{\text{LM}}
|
||||
= -\frac{1}{BT}\sum_{b=1}^{B}\sum_{t=1}^{T}
|
||||
\log p_\theta(y_{b,t} \mid x_{b,\leq t})
|
||||
\]
|
||||
|
||||
## The shift
|
||||
|
||||
The model receives:
|
||||
|
||||
\[
|
||||
x = [t_0,t_1,\ldots,t_{T-1}]
|
||||
\]
|
||||
|
||||
and predicts:
|
||||
|
||||
\[
|
||||
y = [t_1,t_2,\ldots,t_T]
|
||||
\]
|
||||
|
||||
In `src/models/transformer.py`, the simple `forward` path computes cross-entropy over all positions:
|
||||
|
||||
```python
|
||||
logits, loss = model(idx, targets)
|
||||
flat_logits = logits.view(B * T, C)
|
||||
targets = targets.view(B * T).long()
|
||||
loss = F.cross_entropy(flat_logits, targets)
|
||||
```
|
||||
|
||||
The post-training SFT path makes the shift explicit because it needs a mask:
|
||||
|
||||
```python
|
||||
logits = logits[:, :-1, :]
|
||||
targets = tokens[:, 1:]
|
||||
mask = loss_mask[:, 1:].to(logits.dtype)
|
||||
```
|
||||
|
||||
## Perplexity
|
||||
|
||||
Perplexity is exponentiated cross-entropy:
|
||||
|
||||
\[
|
||||
\text{PPL} = \exp(\mathcal{L})
|
||||
\]
|
||||
|
||||
If the model assigns high probability to the true next tokens, loss falls and perplexity falls.
|
||||
|
||||
Interpretation:
|
||||
|
||||
- loss near \(\log(V)\) means the model is close to uniform random guessing;
|
||||
- lower loss means the model is concentrating probability on plausible next tokens;
|
||||
- validation loss matters more than training loss for generalization.
|
||||
|
||||
With `V = 50304`:
|
||||
|
||||
\[
|
||||
\log(V) \approx 10.83
|
||||
\]
|
||||
|
||||
So an untrained model often starts near that value.
|
||||
|
||||
## SFT masked loss
|
||||
|
||||
SFT still uses next-token cross-entropy, but only assistant completion tokens count:
|
||||
|
||||
\[
|
||||
\mathcal{L}_{\text{SFT}} =
|
||||
\frac{\sum_{b,t} m_{b,t}\,\ell_{b,t}}
|
||||
{\sum_{b,t} m_{b,t}}
|
||||
\]
|
||||
|
||||
where \(m_{b,t}=1\) for assistant tokens and \(0\) for prompt tokens.
|
||||
|
||||
Implementation in `src/post_training/sft.py`:
|
||||
|
||||
```python
|
||||
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)
|
||||
```
|
||||
|
||||
This distinction is crucial. The model should learn how to answer the prompt, not how to predict the
|
||||
prompt itself.
|
||||
|
||||
## Sequence log-probabilities
|
||||
|
||||
Preference optimization and RL need the log-probability of an entire response, not only one token.
|
||||
|
||||
For response tokens \(a_1,\ldots,a_L\) after prompt \(p\):
|
||||
|
||||
\[
|
||||
\log \pi_\theta(a \mid p)
|
||||
= \sum_{t=1}^{L}
|
||||
\log \pi_\theta(a_t \mid p, a_{<t})
|
||||
\]
|
||||
|
||||
The repo computes this in `src/post_training/rollout.py` with `sequence_logprobs`. It applies a
|
||||
response mask and sums token log-probs over answer positions.
|
||||
|
||||
That one primitive is reused by:
|
||||
|
||||
- DPO, ORPO, and KTO;
|
||||
- PPO policy ratios;
|
||||
- GRPO policy ratios;
|
||||
- KL measurements against a frozen reference model.
|
||||
|
||||
## DPO objective
|
||||
|
||||
DPO uses preference pairs: chosen response \(y_w\) and rejected response \(y_l\). It compares the policy
|
||||
against a frozen reference model:
|
||||
|
||||
\[
|
||||
\Delta_\pi =
|
||||
\log \pi_\theta(y_w \mid x) - \log \pi_\theta(y_l \mid x)
|
||||
\]
|
||||
|
||||
\[
|
||||
\Delta_{\text{ref}} =
|
||||
\log \pi_{\text{ref}}(y_w \mid x) - \log \pi_{\text{ref}}(y_l \mid x)
|
||||
\]
|
||||
|
||||
\[
|
||||
\mathcal{L}_{\text{DPO}} =
|
||||
-\log \sigma\left(\beta(\Delta_\pi - \Delta_{\text{ref}})\right)
|
||||
\]
|
||||
|
||||
In `src/post_training/dpo.py`:
|
||||
|
||||
```python
|
||||
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()
|
||||
```
|
||||
|
||||
Intuition: make the chosen response more likely than the rejected response, but measure the change
|
||||
relative to the reference model so the policy does not drift without constraint.
|
||||
|
||||
## PPO objective in one picture
|
||||
|
||||
PPO samples responses, scores them, and updates the policy using the ratio between new and old action
|
||||
probabilities:
|
||||
|
||||
\[
|
||||
r_t(\theta) =
|
||||
\frac{\pi_\theta(a_t \mid s_t)}
|
||||
{\pi_{\text{old}}(a_t \mid s_t)}
|
||||
= \exp(\log \pi_\theta - \log \pi_{\text{old}})
|
||||
\]
|
||||
|
||||
The clipped policy objective is:
|
||||
|
||||
\[
|
||||
\mathcal{L}_{\text{PPO}} =
|
||||
-\mathbb{E}_t
|
||||
\left[
|
||||
\min
|
||||
\left(
|
||||
r_t(\theta) A_t,
|
||||
\text{clip}(r_t(\theta),1-\epsilon,1+\epsilon) A_t
|
||||
\right)
|
||||
\right]
|
||||
\]
|
||||
|
||||
The repo implements that in `src/post_training/ppo.py`:
|
||||
|
||||
```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)
|
||||
```
|
||||
|
||||
The clip prevents one update from moving too far from the sampled policy.
|
||||
|
||||
## GRPO objective in one picture
|
||||
|
||||
GRPO avoids a learned value function. For each prompt, it samples a group of responses and normalizes
|
||||
their rewards within the group:
|
||||
|
||||
\[
|
||||
A_i = \frac{r_i - \text{mean}(r_1,\ldots,r_G)}
|
||||
{\text{std}(r_1,\ldots,r_G)+\epsilon}
|
||||
\]
|
||||
|
||||
That advantage says: "Was this answer better or worse than its siblings for the same prompt?"
|
||||
|
||||
In `src/post_training/grpo.py`:
|
||||
|
||||
```python
|
||||
r = rewards.view(-1, group_size)
|
||||
adv = (r - r.mean(1, keepdim=True)) / (r.std(1, keepdim=True) + eps)
|
||||
```
|
||||
|
||||
This is useful for verifiable-reward reasoning tasks because it removes PPO's value head and critic
|
||||
training loop.
|
||||
|
||||
## Objective comparison
|
||||
|
||||
| Stage | Data | Main signal | Learns |
|
||||
|---|---|---|---|
|
||||
| Pretraining | raw token stream | next-token CE | language modeling |
|
||||
| SFT | prompt/answer examples | masked next-token CE | instruction following format |
|
||||
| Reward model | chosen/rejected pairs | Bradley-Terry preference loss | scalar preference scoring |
|
||||
| DPO | chosen/rejected pairs | sequence log-prob preference loss | preference alignment without RL rollout |
|
||||
| PPO | sampled responses + reward | clipped policy gradient | reward-seeking behavior under KL control |
|
||||
| GRPO | grouped sampled responses + verifier | group-relative clipped policy gradient | verifier-driven reasoning without critic |
|
||||
|
||||
## Next
|
||||
|
||||
Loss gives direction. Optimization decides whether the model can follow it reliably. Continue to
|
||||
[Optimization & Training Systems](optimization.md).
|
||||
@@ -0,0 +1,220 @@
|
||||
# Optimization & Training Systems
|
||||
|
||||
Once the loss is defined, training is an engineering problem: move parameters in the right direction
|
||||
without numerical instability, memory blowups, or throughput collapse.
|
||||
|
||||
The main ingredients in this repo are:
|
||||
|
||||
- AdamW;
|
||||
- linear warmup plus cosine learning-rate decay;
|
||||
- gradient accumulation;
|
||||
- gradient clipping;
|
||||
- bf16 autocast;
|
||||
- DistributedDataParallel for multi-GPU training.
|
||||
|
||||
## The training step
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
B["batch"] --> F["forward under autocast"]
|
||||
F --> L["loss / grad_accum"]
|
||||
L --> BW["backward"]
|
||||
BW --> GA{"more microsteps?"}
|
||||
GA -- yes --> B
|
||||
GA -- no --> CLIP["clip grad norm"]
|
||||
CLIP --> LR["set scheduled LR"]
|
||||
LR --> STEP["AdamW step"]
|
||||
STEP --> ZERO["zero gradients"]
|
||||
```
|
||||
|
||||
The pretraining loop in `scripts/pretrain_base.py` uses this pattern:
|
||||
|
||||
```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
|
||||
loss.backward()
|
||||
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
|
||||
optimizer.step()
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
```
|
||||
|
||||
Dividing the loss by `grad_accum` keeps the gradient scale the same as if the full effective batch had
|
||||
fit in memory.
|
||||
|
||||
## AdamW
|
||||
|
||||
Adam keeps exponential moving averages of gradients and squared gradients:
|
||||
|
||||
\[
|
||||
m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t
|
||||
\]
|
||||
|
||||
\[
|
||||
v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2
|
||||
\]
|
||||
|
||||
After bias correction, parameters are updated approximately by:
|
||||
|
||||
\[
|
||||
\theta_{t+1}
|
||||
= \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}
|
||||
\]
|
||||
|
||||
AdamW decouples weight decay from the gradient update:
|
||||
|
||||
\[
|
||||
\theta_{t+1}
|
||||
= \theta_t - \eta \left(
|
||||
\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}
|
||||
+ \lambda \theta_t
|
||||
\right)
|
||||
\]
|
||||
|
||||
The repo applies weight decay only to matrix-like parameters:
|
||||
|
||||
```python
|
||||
if p.dim() >= 2:
|
||||
decay.append(p)
|
||||
else:
|
||||
no_decay.append(p)
|
||||
```
|
||||
|
||||
This is the standard GPT recipe: decay large weight matrices, but not biases, LayerNorm scales, or
|
||||
one-dimensional parameters.
|
||||
|
||||
## Learning-rate warmup and cosine decay
|
||||
|
||||
The learning rate is small at the start, ramps up, then decays:
|
||||
|
||||
\[
|
||||
\eta(s) =
|
||||
\eta_{\max}\frac{s+1}{S_{\text{warmup}}}
|
||||
\quad \text{if } s < S_{\text{warmup}}
|
||||
\]
|
||||
|
||||
After warmup:
|
||||
|
||||
\[
|
||||
\eta(s) =
|
||||
\eta_{\min}
|
||||
+ \frac{1}{2}(1+\cos(\pi p))(\eta_{\max}-\eta_{\min})
|
||||
\]
|
||||
|
||||
where:
|
||||
|
||||
\[
|
||||
p = \frac{s-S_{\text{warmup}}}{S_{\max}-S_{\text{warmup}}}
|
||||
\]
|
||||
|
||||
Implementation in `src/post_training/optim.py`:
|
||||
|
||||
```python
|
||||
if step < warmup_steps:
|
||||
return lr * (step + 1) / max(1, warmup_steps)
|
||||
progress = (step - warmup_steps) / max(1, max_steps - warmup_steps)
|
||||
coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
return min_lr + coeff * (lr - min_lr)
|
||||
```
|
||||
|
||||
Warmup prevents early unstable updates while weights are still poorly calibrated. Cosine decay reduces
|
||||
step size as training approaches the end of the budget.
|
||||
|
||||
## Gradient accumulation
|
||||
|
||||
If one batch is too large for GPU memory, split it into microbatches:
|
||||
|
||||
\[
|
||||
B_{\text{effective}} = B_{\text{micro}} \times N_{\text{accum}} \times N_{\text{gpus}}
|
||||
\]
|
||||
|
||||
Example:
|
||||
|
||||
- microbatch size: 8;
|
||||
- accumulation steps: 12;
|
||||
- GPUs: 2.
|
||||
|
||||
\[
|
||||
B_{\text{effective}} = 8 \times 12 \times 2 = 192
|
||||
\]
|
||||
|
||||
The optimizer steps once after all microbatches have contributed gradients.
|
||||
|
||||
## Gradient clipping
|
||||
|
||||
Gradient clipping limits the global norm:
|
||||
|
||||
\[
|
||||
g \leftarrow g \cdot \min\left(1, \frac{c}{\|g\|_2}\right)
|
||||
\]
|
||||
|
||||
If the gradient norm is below the threshold \(c\), nothing changes. If it is too large, the whole
|
||||
gradient vector is scaled down. This is a stability guard, especially useful in RL and long-sequence
|
||||
training.
|
||||
|
||||
## bf16 autocast
|
||||
|
||||
`bf16` uses fewer bits than fp32, but keeps an 8-bit exponent like fp32. That makes it much more
|
||||
forgiving than fp16 for deep learning training.
|
||||
|
||||
The repo uses autocast for forward computation:
|
||||
|
||||
```python
|
||||
with amp_autocast(cfg.amp_dtype, ctx.device):
|
||||
logits, _ = model(tokens)
|
||||
loss = sft_loss(logits, tokens, mask)
|
||||
```
|
||||
|
||||
Model parameters usually remain fp32. Many matrix multiplications run in bf16, improving memory and
|
||||
throughput on supported GPUs.
|
||||
|
||||
## DistributedDataParallel
|
||||
|
||||
DDP creates one process per GPU. Each process:
|
||||
|
||||
1. owns a full copy of the model;
|
||||
2. receives a different shard or random stream of data;
|
||||
3. computes gradients locally;
|
||||
4. synchronizes gradients across processes before the optimizer step.
|
||||
|
||||
With gradient accumulation, synchronization is needed only on the last microstep. The repo uses
|
||||
`model.no_sync()` for earlier microsteps to avoid unnecessary communication.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
R0["rank 0 GPU"] --> G["gradient all-reduce"]
|
||||
R1["rank 1 GPU"] --> G
|
||||
G --> U0["rank 0 optimizer step"]
|
||||
G --> U1["rank 1 optimizer step"]
|
||||
```
|
||||
|
||||
## What to watch during training
|
||||
|
||||
| Metric | Healthy behavior | Problem signal |
|
||||
|---|---|---|
|
||||
| train loss | falls steadily | flat near random baseline |
|
||||
| dev loss | falls, then stabilizes | rises while train loss falls |
|
||||
| grad norm | finite, bounded after clipping | NaN or repeated huge spikes |
|
||||
| tokens/sec | stable for same config | sudden drop or dataloader stall |
|
||||
| KL in RL stages | bounded | runaway drift from reference |
|
||||
| reward in RL stages | rises with variance | zero signal for many iterations |
|
||||
|
||||
## Memory levers
|
||||
|
||||
If a config does not fit, reduce in this order:
|
||||
|
||||
1. `batch_size`;
|
||||
2. `context_length`;
|
||||
3. `n_blocks`;
|
||||
4. `n_embed`;
|
||||
5. `n_head` only if it still divides `n_embed`.
|
||||
|
||||
Context length is especially expensive because attention uses a \(T \times T\) score matrix.
|
||||
|
||||
## Next
|
||||
|
||||
After training, the model still only emits logits. Generation turns those logits into text. Continue
|
||||
to [Generation & Sampling](generation.md).
|
||||
@@ -0,0 +1,159 @@
|
||||
# Tokenization & Data Shapes
|
||||
|
||||
The Transformer never sees characters or words. It sees integer token ids. A tokenizer is the boundary
|
||||
between language and tensors.
|
||||
|
||||
This repo uses OpenAI's `r50k_base` tokenizer through `tiktoken`, with:
|
||||
|
||||
- vocabulary size `50304`;
|
||||
- end-of-text token `<|endoftext|>` with id `50256`;
|
||||
- plain text role markers for chat because the tokenizer has no custom chat tokens.
|
||||
|
||||
## Why subword tokenization exists
|
||||
|
||||
A word-level vocabulary cannot handle rare names, typos, code identifiers, URLs, and new terms without
|
||||
exploding in size. A character-level vocabulary handles everything but makes sequences long. Subword
|
||||
tokenization is the compromise: frequent words can be one token, rare words can be decomposed.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
T["The tokenizer sees text"] --> A["common words -> short token sequences"]
|
||||
T --> B["rare words -> subword pieces"]
|
||||
T --> C["punctuation/code -> reusable fragments"]
|
||||
A --> IDS["integer ids"]
|
||||
B --> IDS
|
||||
C --> IDS
|
||||
```
|
||||
|
||||
The original BPE idea is simple: start with small units, repeatedly merge frequent adjacent pairs, and
|
||||
end with a fixed vocabulary of reusable pieces. The repo does not train its own tokenizer; it reuses
|
||||
`r50k_base`.
|
||||
|
||||
## Pretraining shape: one long token stream
|
||||
|
||||
For pretraining, documents are converted into one flat array:
|
||||
|
||||
\[
|
||||
[d_1, \text{EOT}, d_2, \text{EOT}, \ldots, d_N, \text{EOT}]
|
||||
\]
|
||||
|
||||
`scripts/prepare_pretrain_data.py` streams Pile shards, tokenizes documents, appends EOT, and writes
|
||||
the result to HDF5:
|
||||
|
||||
```python
|
||||
for ids in enc.encode_ordinary_batch(docs):
|
||||
buf.extend(ids)
|
||||
buf.append(EOT_ID)
|
||||
if len(buf) >= WRITE_CHUNK:
|
||||
flush()
|
||||
```
|
||||
|
||||
The training loader slices random windows of length `context_length + 1`. The first `context_length`
|
||||
tokens are inputs. The next `context_length` tokens are targets shifted by one position:
|
||||
|
||||
\[
|
||||
x = [t_0, t_1, \ldots, t_{T-1}]
|
||||
\]
|
||||
|
||||
\[
|
||||
y = [t_1, t_2, \ldots, t_T]
|
||||
\]
|
||||
|
||||
That shift is the entire next-token prediction task.
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
H[(pile_train.h5 flat ids)] --> W["sample T+1-token window"]
|
||||
W --> X["input: tokens 0..T-1"]
|
||||
W --> Y["target: tokens 1..T"]
|
||||
X --> M["Transformer"]
|
||||
M --> L["cross-entropy against target"]
|
||||
Y --> L
|
||||
```
|
||||
|
||||
## SFT shape: tokens plus a loss mask
|
||||
|
||||
SFT examples are conversations. We want the model to learn the assistant answer, not memorize the user
|
||||
prompt. So the data contains two aligned arrays:
|
||||
|
||||
- `tokens`: token ids;
|
||||
- `loss_mask`: `1` for assistant completion tokens, `0` for prompt tokens.
|
||||
|
||||
The chat template is plain text:
|
||||
|
||||
```text
|
||||
<|user|>
|
||||
{question}<|endoftext|><|assistant|>
|
||||
{answer}<|endoftext|>
|
||||
```
|
||||
|
||||
The key implementation is in `src/post_training/chat_template.py`:
|
||||
|
||||
```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))
|
||||
ids.append(EOT_ID)
|
||||
mask.append(1 if is_completion else 0)
|
||||
```
|
||||
|
||||
The mask aligns with the token ids:
|
||||
|
||||
| Span | Example | Mask |
|
||||
|---|---|---|
|
||||
| user marker | `<|user|>` | 0 |
|
||||
| user question | `What is 2+2?` | 0 |
|
||||
| assistant marker | `<|assistant|>` | 0 |
|
||||
| assistant answer | `<answer>4</answer>` | 1 |
|
||||
| assistant EOT | `<|endoftext|>` | 1 |
|
||||
|
||||
## Preference shape: prompt, chosen, rejected
|
||||
|
||||
Preference learning uses pairs:
|
||||
|
||||
```json
|
||||
{"prompt": "...", "chosen": "...", "rejected": "..."}
|
||||
```
|
||||
|
||||
Both responses share the same prompt. That matters because DPO and reward modeling should compare
|
||||
answer quality, not prompt difficulty.
|
||||
|
||||
For a batch, the loader creates two tokenized sequences:
|
||||
|
||||
\[
|
||||
\text{chosen ids} = \text{chat}(prompt, chosen)
|
||||
\]
|
||||
|
||||
\[
|
||||
\text{rejected ids} = \text{chat}(prompt, rejected)
|
||||
\]
|
||||
|
||||
The chosen and rejected sides are padded to the same length for batching. The repo tracks true sequence
|
||||
lengths so the reward model can read the last real token instead of a padding token.
|
||||
|
||||
## RL prompt shape: prompt plus verifiable gold answer
|
||||
|
||||
PPO and GRPO need prompts that can be scored after generation:
|
||||
|
||||
```json
|
||||
{"prompt": "Jan has 3 apples...", "gold": "12"}
|
||||
```
|
||||
|
||||
The verifier extracts the model's final answer and compares it to `gold`. This is called verifiable
|
||||
reward because the training signal does not require a human labeler or a learned reward model.
|
||||
|
||||
## Common shape bugs
|
||||
|
||||
| Bug | Symptom | Prevention |
|
||||
|---|---|---|
|
||||
| Target not shifted | Model learns to copy the current token. | Always predict `tokens[:, 1:]` from `tokens[:, :-1]`. |
|
||||
| Prompt tokens included in SFT loss | Model wastes capacity predicting user input. | Use `loss_mask` and average over masked positions only. |
|
||||
| Missing EOT | Model does not learn when to stop. | Include EOT after documents and assistant messages. |
|
||||
| Preference prompt mismatch | Reward/DPO compares different tasks. | Normalize to shared prompt plus chosen/rejected responses. |
|
||||
| Padding used as reward position | Reward model trains on meaningless pad hidden states. | Track `seq_lengths` and gather the last real token. |
|
||||
|
||||
## Next
|
||||
|
||||
Once text is tokenized, the model needs to turn ids into vectors. Continue to
|
||||
[Decoder-Only Transformer](transformer.md).
|
||||
@@ -0,0 +1,234 @@
|
||||
# Decoder-Only Transformer
|
||||
|
||||
This repo implements a GPT-style decoder-only Transformer. "Decoder-only" means:
|
||||
|
||||
- the model reads a single token sequence;
|
||||
- each position can attend only to previous positions;
|
||||
- the output at every position is a distribution over the next token.
|
||||
|
||||
It is the right architecture for autoregressive language modeling:
|
||||
|
||||
\[
|
||||
p_\theta(x_t \mid x_{<t})
|
||||
\]
|
||||
|
||||
## The forward pass
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
IDS["token ids (B, T)"] --> TOK["token embedding (B, T, C)"]
|
||||
IDS --> POS["position ids 0..T-1"]
|
||||
POS --> PE["position embedding (T, C)"]
|
||||
TOK --> SUM["token + position"]
|
||||
PE --> SUM
|
||||
SUM --> B1["Block 1"]
|
||||
B1 --> B2["Block 2"]
|
||||
B2 --> BN["... N blocks"]
|
||||
BN --> LN["final LayerNorm"]
|
||||
LN --> HEAD["lm_head Linear(C -> vocab)"]
|
||||
HEAD --> LOGITS["logits (B, T, V)"]
|
||||
```
|
||||
|
||||
The implementation is in `src/models/transformer.py`:
|
||||
|
||||
```python
|
||||
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)
|
||||
```
|
||||
|
||||
Shape symbols used throughout the docs:
|
||||
|
||||
| Symbol | Meaning |
|
||||
|---|---|
|
||||
| `B` | batch size |
|
||||
| `T` | sequence length / context length |
|
||||
| `C` | embedding width, `n_embed` |
|
||||
| `H` | number of attention heads |
|
||||
| `D` | head width, usually `C / H` |
|
||||
| `V` | vocabulary size |
|
||||
|
||||
## Embeddings
|
||||
|
||||
Token ids are categorical. The embedding table is a learned lookup:
|
||||
|
||||
\[
|
||||
E_{\text{tok}} \in \mathbb{R}^{V \times C}
|
||||
\]
|
||||
|
||||
For token id \(x_t\), the token vector is:
|
||||
|
||||
\[
|
||||
e_t = E_{\text{tok}}[x_t]
|
||||
\]
|
||||
|
||||
The model also learns absolute position embeddings:
|
||||
|
||||
\[
|
||||
E_{\text{pos}} \in \mathbb{R}^{T_{\max} \times C}
|
||||
\]
|
||||
|
||||
The input to the first block is:
|
||||
|
||||
\[
|
||||
h_t^{(0)} = E_{\text{tok}}[x_t] + E_{\text{pos}}[t]
|
||||
\]
|
||||
|
||||
In code:
|
||||
|
||||
```python
|
||||
tok_embedding = self.token_embed(idx)
|
||||
pos_embedding = self.position_embed(self.pos_idxs[:T])
|
||||
return tok_embedding + pos_embedding
|
||||
```
|
||||
|
||||
The position embedding is necessary because attention alone is permutation-equivariant: without
|
||||
position information, the model would not know whether a token occurred first, last, or in the middle.
|
||||
|
||||
## Transformer block
|
||||
|
||||
Each block in `src/models/transformer_block.py` uses pre-norm residual structure:
|
||||
|
||||
```python
|
||||
x = x + self.attn(self.ln1(x))
|
||||
x = x + self.mlp(self.ln2(x))
|
||||
```
|
||||
|
||||
Mathematically:
|
||||
|
||||
\[
|
||||
u = x + \text{MHA}(\text{LN}(x))
|
||||
\]
|
||||
|
||||
\[
|
||||
y = u + \text{MLP}(\text{LN}(u))
|
||||
\]
|
||||
|
||||
This gives each block two jobs:
|
||||
|
||||
- attention moves information across token positions;
|
||||
- the MLP transforms each position independently.
|
||||
|
||||
## Why residual connections matter
|
||||
|
||||
A residual block learns an update, not a full replacement:
|
||||
|
||||
\[
|
||||
y = x + f(x)
|
||||
\]
|
||||
|
||||
If a layer is not useful yet, it can learn a small update and let information pass through. This makes
|
||||
deep stacks trainable because gradients have a direct path backward through the addition.
|
||||
|
||||
## Why LayerNorm appears before sublayers
|
||||
|
||||
LayerNorm normalizes each token vector across its feature dimension:
|
||||
|
||||
\[
|
||||
\text{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta
|
||||
\]
|
||||
|
||||
For a token vector \(x \in \mathbb{R}^{C}\):
|
||||
|
||||
\[
|
||||
\mu = \frac{1}{C} \sum_{i=1}^{C} x_i
|
||||
\]
|
||||
|
||||
\[
|
||||
\sigma^2 = \frac{1}{C} \sum_{i=1}^{C} (x_i - \mu)^2
|
||||
\]
|
||||
|
||||
This repo uses pre-norm (`LN -> sublayer -> residual`) instead of post-norm (`sublayer -> residual ->
|
||||
LN`). Pre-norm is common in GPT-like models because it tends to make deeper stacks easier to optimize.
|
||||
|
||||
## MLP / feed-forward network
|
||||
|
||||
The block MLP in `src/models/mlp.py` is:
|
||||
|
||||
```python
|
||||
self.hidden = nn.Linear(n_embed, 4 * n_embed)
|
||||
self.relu = nn.ReLU()
|
||||
self.proj = nn.Linear(4 * n_embed, n_embed)
|
||||
```
|
||||
|
||||
For each position independently:
|
||||
|
||||
\[
|
||||
\text{MLP}(x) = W_2 \, \text{ReLU}(W_1 x + b_1) + b_2
|
||||
\]
|
||||
|
||||
where:
|
||||
|
||||
- \(W_1\) expands from \(C\) to \(4C\);
|
||||
- \(W_2\) projects from \(4C\) back to \(C\).
|
||||
|
||||
Attention lets tokens communicate. The MLP gives each token vector nonlinear compute after that
|
||||
communication.
|
||||
|
||||
## Logits and the language-model head
|
||||
|
||||
After the final block and final norm:
|
||||
|
||||
\[
|
||||
z_t = W_{\text{lm}} h_t + b_{\text{lm}}
|
||||
\]
|
||||
|
||||
The result \(z_t \in \mathbb{R}^{V}\) is a vector of logits: one unnormalized score per token in the
|
||||
vocabulary.
|
||||
|
||||
The probability distribution comes from softmax:
|
||||
|
||||
\[
|
||||
p_\theta(x_{t+1}=i \mid x_{\leq t}) =
|
||||
\frac{\exp(z_{t,i})}{\sum_{j=1}^{V}\exp(z_{t,j})}
|
||||
\]
|
||||
|
||||
## Parameter count intuition
|
||||
|
||||
Ignoring biases and norms, a rough per-block parameter estimate is:
|
||||
|
||||
\[
|
||||
\text{attention} \approx 4C^2
|
||||
\]
|
||||
|
||||
because Q, K, V, and the output projection are each \(C \times C\).
|
||||
|
||||
\[
|
||||
\text{MLP} \approx 8C^2
|
||||
\]
|
||||
|
||||
because \(C \to 4C\) and \(4C \to C\).
|
||||
|
||||
So each block is roughly:
|
||||
|
||||
\[
|
||||
12C^2
|
||||
\]
|
||||
|
||||
The embedding and LM head add:
|
||||
|
||||
\[
|
||||
V C + C V
|
||||
\]
|
||||
|
||||
This repo does not tie token embeddings and output embeddings, so the input embedding and `lm_head`
|
||||
are separate parameter matrices.
|
||||
|
||||
## Repo-specific architecture notes
|
||||
|
||||
| Choice | Repo implementation | Consequence |
|
||||
|---|---|---|
|
||||
| Absolute learned positions | `nn.Embedding(context_length, n_embed)` | Simple and readable; fixed max context. |
|
||||
| Causal attention mask | lower-triangular buffer in each head | Prevents future-token leakage. |
|
||||
| MLP activation | ReLU | Educationally simple; many production GPT models use GELU/SwiGLU variants. |
|
||||
| Dropout | not present in base modules | Less code noise; regularization comes mostly from data and optimizer choices. |
|
||||
| Weight tying | not used | Easier to read; more parameters than tied embeddings. |
|
||||
| Post-training heads | use `forward_hidden` | Reward/value heads reuse the same backbone. |
|
||||
|
||||
## Next
|
||||
|
||||
The most important sublayer is attention. Continue to [Attention, Masks & Heads](attention.md).
|
||||
@@ -0,0 +1,6 @@
|
||||
# Command cheatsheet
|
||||
|
||||
The condensed command reference for the whole pipeline (included verbatim from
|
||||
[`POST_TRAINING.md`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/POST_TRAINING.md)).
|
||||
|
||||
--8<-- "POST_TRAINING.md"
|
||||
@@ -0,0 +1,65 @@
|
||||
# Configure with JSON
|
||||
|
||||
Every training stage is configured by a small, human-editable JSON file under
|
||||
[`configs/`](https://github.com/FareedKhan-dev/train-llm-from-scratch/tree/main/configs). You edit one
|
||||
file and see only that stage's knobs — the model architecture and runtime live in the shared
|
||||
`configs/base.json`.
|
||||
|
||||
```
|
||||
configs/
|
||||
base.json # shared model + runtime (vocab, n_embed, device, amp_dtype, ...)
|
||||
pretrain.json # one file per stage — only that stage's hyperparameters
|
||||
sft.json reward.json dpo.json ppo.json grpo.json
|
||||
smoke/ # tiny CPU variant (small model + few steps) for quick tests
|
||||
base.json pretrain.json sft.json ...
|
||||
```
|
||||
|
||||
## How a config is resolved
|
||||
|
||||
[`config/loader.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/config/loader.py)
|
||||
merges four layers, **lowest precedence first**:
|
||||
|
||||
```
|
||||
dataclass defaults < configs/base.json < configs/<stage>.json < CLI --field overrides
|
||||
```
|
||||
|
||||
The dataclasses in `config/post_training_config.py` are the typed schema; the JSON is the editable
|
||||
source; the CLI wins for quick one-offs. JSON `null` maps to Python `None` (the clean way to set a
|
||||
`str | None` field such as `amp_dtype`).
|
||||
|
||||
!!! tip "Smoke configs auto-shrink the model"
|
||||
A stage JSON in a sub-folder uses that folder's `base.json`. So `configs/smoke/sft.json` picks up
|
||||
`configs/smoke/base.json` (tiny model, CPU) automatically — handy for a fast end-to-end test.
|
||||
|
||||
## Editing & inspecting
|
||||
|
||||
=== "Edit the JSON"
|
||||
|
||||
```jsonc
|
||||
// configs/sft.json
|
||||
{
|
||||
"pretrained_ckpt": "/ephemeral/ckpts/base_pretrained.pt",
|
||||
"data_path": "/ephemeral/data/sft_packed.h5",
|
||||
"lr": 1e-5,
|
||||
"epochs": 3,
|
||||
"batch_size": 16
|
||||
}
|
||||
```
|
||||
|
||||
=== "Override on the CLI"
|
||||
|
||||
```bash
|
||||
# --field beats the JSON
|
||||
python scripts/train_sft.py --config configs/sft.json --lr 2e-5 --batch_size 8
|
||||
```
|
||||
|
||||
=== "Print the resolved config"
|
||||
|
||||
```bash
|
||||
python scripts/train_sft.py --config configs/sft.json --print-config
|
||||
# dumps the fully merged config as JSON, then exits
|
||||
```
|
||||
|
||||
Every trainer accepts `--config <path>` (defaulting to its stage file), all `--field` overrides, and
|
||||
`--print-config`. The legacy pretraining path (`scripts/train_transformer.py` + `config/config.py`) is
|
||||
untouched and still works as the README teaches.
|
||||
@@ -0,0 +1,67 @@
|
||||
# Train (UI & CLI)
|
||||
|
||||
You can run every stage two ways: from the **command line** (full control, best for long jobs) or from
|
||||
the **Streamlit control panel** (forms, one-click launch, live logs, charts — see [The UI](ui.md)).
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
pip install -e ".[train]" # editable install — no more PYTHONPATH=.
|
||||
export HF_HOME=/ephemeral/hf_cache
|
||||
```
|
||||
|
||||
## The pipeline, end to end (CLI)
|
||||
|
||||
Each command reads its stage JSON from `configs/` (override anything with `--field`, or point at another
|
||||
file with `--config`). Use `python` for one GPU and `torchrun` for many.
|
||||
|
||||
=== "1. Data"
|
||||
|
||||
```bash
|
||||
python scripts/prepare_pretrain_data.py --split val --out /ephemeral/data/pile_dev.h5
|
||||
python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out /ephemeral/data/pile_train.h5
|
||||
python scripts/prepare_sft_data.py
|
||||
python scripts/prepare_preference_data.py --source both
|
||||
python scripts/prepare_rl_prompts.py
|
||||
```
|
||||
|
||||
=== "2. Pretrain"
|
||||
|
||||
```bash
|
||||
# single GPU
|
||||
python scripts/pretrain_base.py --config configs/pretrain.json
|
||||
# both GPUs (effective batch = batch_size * grad_accum * num_gpus)
|
||||
torchrun --standalone --nproc_per_node=2 scripts/pretrain_base.py --config configs/pretrain.json
|
||||
```
|
||||
|
||||
=== "3. Align"
|
||||
|
||||
```bash
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_reward.py
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_ppo.py --reward_source verifier
|
||||
torchrun --standalone --nproc_per_node=2 scripts/train_grpo.py
|
||||
```
|
||||
|
||||
The whole alignment chain in one shot:
|
||||
|
||||
```bash
|
||||
bash scripts/run_posttraining.sh # SFT → RM → DPO → PPO → GRPO → eval table
|
||||
```
|
||||
|
||||
## Multi-GPU notes
|
||||
|
||||
- `torchrun --standalone --nproc_per_node=N` launches N data-parallel ranks (DDP + bf16). Only rank 0
|
||||
logs and checkpoints.
|
||||
- On the dev box (2× H100, no NVLink) the educational attention materializes a `(B, n_head, T, T)` tensor
|
||||
per block, so memory scales with sequence length — at context 1024 use `--batch_size 8 --grad_accum 12`
|
||||
and recover the effective batch via accumulation. Set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
|
||||
|
||||
## Where outputs go
|
||||
|
||||
- Checkpoints → `/ephemeral/ckpts/<stage>.pt` (each carries its own resolved `cfg`).
|
||||
- Metrics → `/ephemeral/logs/<stage>_<timestamp>.jsonl` (one JSON per logged step). The UI plots these
|
||||
live; you can also `--use_wandb true` to mirror to Weights & Biases.
|
||||
|
||||
Then [evaluate](../08_evaluation.md) on GSM8K and [chat](../09_inference.md) with any checkpoint.
|
||||
@@ -0,0 +1,36 @@
|
||||
# The control-panel UI
|
||||
|
||||
A polished **Streamlit** app that wraps the whole pipeline — theory, one-click training, live logs,
|
||||
metric charts, evaluation, and chat — so you can drive everything without memorizing commands.
|
||||
|
||||
```bash
|
||||
pip install -e ".[ui]"
|
||||
streamlit run ui/app.py
|
||||
```
|
||||
|
||||
## What's inside
|
||||
|
||||
- **Home** — the pipeline diagram, live job status, and GPU status at a glance.
|
||||
- **Data** — launch the `prepare_*` scripts and watch the files appear under `/ephemeral/data`.
|
||||
- **A page per stage** (Pretrain, SFT, Reward, DPO, PPO, GRPO). Each page gives you:
|
||||
1. the **theory + hand-drawn diagram** for that stage (pulled straight from these docs),
|
||||
2. a **config form** that writes the stage's `configs/*.json`,
|
||||
3. a **Launch** button that runs the real training script as a background job, and
|
||||
4. a **live log tail** + **metric charts** (loss / reward / KL / accuracy) read from the JSONL.
|
||||
- **Evaluate** — run GSM8K accuracy on any checkpoint and inspect sample generations.
|
||||
- **Chat** — talk to any checkpoint in-process, with temperature / top-p / top-k controls.
|
||||
|
||||
## How launching works
|
||||
|
||||
The app builds the same command you'd type (`torchrun … scripts/train_sft.py --config …`), runs it as a
|
||||
detached background process, and streams its log file into the page — so jobs keep running even if you
|
||||
navigate away or refresh. A **Stop** button cleanly terminates the job (and all `torchrun` workers).
|
||||
|
||||
!!! warning "One GPU job at a time"
|
||||
Two multi-GPU `torchrun` jobs would fight over the same cards and OOM. The UI's **GPU-busy guard**
|
||||
disables *Launch* for a GPU stage while another GPU job is running (e.g. while pretraining occupies
|
||||
both H100s). CPU/smoke runs and the in-process **Chat** page always stay available.
|
||||
|
||||
!!! tip "Smoke mode"
|
||||
Each stage page has a *smoke* toggle that uses the tiny `configs/smoke/*.json` (small model, few
|
||||
steps, CPU) so you can validate the whole flow in seconds before committing to a real run.
|
||||
@@ -0,0 +1,12 @@
|
||||
window.MathJax = {
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex"
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,21 @@
|
||||
flowchart LR
|
||||
PILE([The Pile<br/>raw text]):::data --> PRE{{Pretrain<br/>base ~400M}}:::model
|
||||
PRE --> SFT{{SFT<br/>instruct}}:::model
|
||||
SFT --> RM{{Reward Model<br/>Bradley-Terry}}:::rl
|
||||
SFT --> DPO{{DPO / ORPO / KTO}}:::rl
|
||||
SFT --> GRPO{{GRPO / RLVR}}:::rl
|
||||
RM -->|reward signal| PPO{{PPO}}:::rl
|
||||
PPO --> EVAL([GSM8K eval<br/>+ chat]):::eval
|
||||
DPO --> EVAL
|
||||
GRPO --> EVAL
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 192 KiB |
@@ -0,0 +1,16 @@
|
||||
flowchart LR
|
||||
A1([Pretrain: The Pile]):::data --> A2[tiktoken r50k_base<br/>stream + encode]:::proc --> A3[(pile_train.h5<br/>flat int32 tokens)]:::store
|
||||
B1([SFT: Alpaca · Dolly · GSM8K]):::data --> B2[chat template<br/>+ mask the prompt]:::proc --> B3[(sft_packed.h5<br/>tokens + loss_mask)]:::store
|
||||
C1([Prefs: HH-RLHF · UltraFeedback]):::data --> C2[split prompt /<br/>chosen / rejected]:::proc --> C3[(preferences.jsonl)]:::store
|
||||
D1([RL: GSM8K · arithmetic]):::data --> D2[keep the gold answer]:::proc --> D3[(rl_prompts.jsonl)]:::store
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 745 KiB |
@@ -0,0 +1,17 @@
|
||||
flowchart LR
|
||||
T1["raw text<br/>'Once upon a time'"]:::data --> T2["tiktoken<br/>r50k_base"]:::proc
|
||||
T2 --> T3["token ids<br/>[7454, 2402, 257, 640]"]:::store
|
||||
T3 --> T4["append <|endoftext|><br/>(id 50256)"]:::proc
|
||||
T4 --> T5["fixed windows of<br/>context_length + 1"]:::store
|
||||
T5 --> T6["x = window[:-1]<br/>y = window[1:]"]:::out
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 189 KiB |
@@ -0,0 +1,16 @@
|
||||
flowchart LR
|
||||
M1["x (B, T, n_embed)"]:::data --> M2["Linear<br/>n_embed -> 4 · n_embed"]:::proc
|
||||
M2 --> M3["ReLU"]:::proc
|
||||
M3 --> M4["Linear<br/>4 · n_embed -> n_embed"]:::proc
|
||||
M4 --> M5["out (B, T, n_embed)"]:::out
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 162 KiB |
@@ -0,0 +1,21 @@
|
||||
flowchart LR
|
||||
X["x : token vectors<br/>(B, T, C)"]:::data --> Q["Q = x Wq"]:::proc
|
||||
X --> K["K = x Wk"]:::proc
|
||||
X --> V["V = x Wv"]:::proc
|
||||
Q --> S["scores = Q Kᵀ / √d"]:::rl
|
||||
K --> S
|
||||
S --> M["causal mask<br/>+ softmax"]:::rl
|
||||
M --> O["weights · V"]:::rl
|
||||
V --> O
|
||||
O --> H["head output<br/>(B, T, head_size)"]:::out
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 219 KiB |
@@ -0,0 +1,20 @@
|
||||
flowchart LR
|
||||
X["x (B, T, C)"]:::data --> H1["head 1"]:::proc
|
||||
X --> H2["head 2"]:::proc
|
||||
X --> H3["... (n_head heads)"]:::proc
|
||||
H1 --> CC["concat heads<br/>(B, T, C)"]:::rl
|
||||
H2 --> CC
|
||||
H3 --> CC
|
||||
CC --> PJ["output projection<br/>Linear C -> C"]:::proc
|
||||
PJ --> HO["multi-head output"]:::out
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 229 KiB |
@@ -0,0 +1,17 @@
|
||||
flowchart LR
|
||||
X["x"]:::data --> LN1["LayerNorm"]:::proc --> AT["Multi-Head Attention"]:::rl --> A1(("+")):::op
|
||||
X --> A1
|
||||
A1 --> LN2["LayerNorm"]:::proc --> ML["MLP"]:::proc --> A2(("+")):::op
|
||||
A1 --> A2
|
||||
A2 --> BO["block output"]:::out
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 126 KiB |
@@ -0,0 +1,21 @@
|
||||
flowchart LR
|
||||
I["token ids (B, T)"]:::data --> TE["token embedding"]:::proc
|
||||
I --> PE["position embedding"]:::proc
|
||||
TE --> SUM(("+")):::op
|
||||
PE --> SUM
|
||||
SUM --> BLK["N_BLOCKS ×<br/>Transformer Block"]:::model
|
||||
BLK --> LNF["final LayerNorm"]:::proc
|
||||
LNF --> LMH["lm_head<br/>Linear -> vocab_size"]:::proc
|
||||
LMH --> LOG["logits (B, T, vocab)"]:::out
|
||||
LOG -. targets present .-> CE["cross-entropy loss"]:::loss
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 153 KiB |
@@ -0,0 +1,14 @@
|
||||
flowchart LR
|
||||
D[(pile_train.h5)]:::store --> IT["get_batch_iterator<br/>random windows"]:::proc --> FW["forward<br/>(bf16 autocast)"]:::model --> CE["cross-entropy"]:::loss --> BW["backward<br/>× grad_accum"]:::proc --> CL["clip grad norm 1.0"]:::proc --> ST["AdamW step<br/>cosine LR + warmup"]:::model --> CK[(checkpoint)]:::ckpt
|
||||
ST -. next step .-> IT
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 167 KiB |
@@ -0,0 +1,15 @@
|
||||
flowchart LR
|
||||
S0[(sft_packed.h5<br/>tokens + loss_mask)]:::store --> S1["Transformer forward"]:::model --> S2["shift: predict t+1"]:::proc --> S3["per-token cross-entropy"]:::loss
|
||||
S0 --> S4["loss_mask = 1 only on<br/>assistant tokens"]:::proc --> S5["mean over masked tokens"]:::loss
|
||||
S3 --> S5 --> S6["AdamW step"]:::model
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 196 KiB |
@@ -0,0 +1,14 @@
|
||||
flowchart LR
|
||||
R0["prompt + chosen / rejected"]:::data --> R1["SFT backbone<br/>forward_hidden"]:::model --> R2["take last real token"]:::proc --> R3["reward head<br/>Linear -> 1"]:::proc
|
||||
R3 --> R4["r_chosen , r_rejected"]:::rl --> R5["Bradley-Terry loss<br/>-log σ(r_chosen - r_rejected)"]:::loss --> R6["AdamW step"]:::model
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 168 KiB |
@@ -0,0 +1,16 @@
|
||||
flowchart LR
|
||||
P0["chosen / rejected pair"]:::data --> P1["policy (trainable)<br/>sequence log-probs"]:::model
|
||||
P0 --> P2["reference (frozen SFT)<br/>sequence log-probs"]:::ckpt
|
||||
P1 --> P3["DPO loss<br/>-log σ(β · Δ)"]:::loss
|
||||
P2 --> P3 --> P4["AdamW step"]:::model
|
||||
|
||||
classDef data fill:#cdeccd,stroke:#2e7d32,stroke-width:2px,color:#143d1a;
|
||||
classDef store fill:#cfeee9,stroke:#1f9e8f,stroke-width:2px,color:#0c3a34;
|
||||
classDef proc fill:#cfe3fb,stroke:#1565c0,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:#ffd4d4,stroke:#c0392b,stroke-width:2px,color:#5c1212;
|
||||
classDef ckpt fill:#ececec,stroke:#666666,stroke-width:2px,color:#222222;
|
||||
classDef eval fill:#e7d6fb,stroke:#8e44ad,stroke-width:2px,color:#3d1a5a;
|
||||
classDef out fill:#e7d6fb,stroke:#7b3fbf,stroke-width:2px,color:#3d1a5a;
|
||||
classDef op fill:#ffffff,stroke:#888888,stroke-width:2px,color:#333333;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 239 KiB |
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user