gpt2-124m-fineweb
GPT-2 124M pretrained from scratch on FineWeb-Edu, then supervised fine-tuned for instruction following. Trained on a single rented RTX 4090 for about $6.
Results
| Stage | Metric | Value |
|---|---|---|
| Pretraining | FineWeb-Edu val loss | 3.0771 |
| SFT | held-out val loss | 2.4581 |
For reference, OpenAI's GPT-2 124M scores roughly 3.29 on this validation set, trained on ~100B tokens. This run used 7.3B.
Benchmarks
| Task | n | Base | SFT | GPT-2 124M |
|---|---|---|---|---|
| HellaSwag (acc_norm) | 10,042 | 0.2966 | 0.3031 | 0.2955 |
| PIQA (acc_norm) | 1,838 | 0.6268 | 0.6143 | 0.628 |
| ARC-Easy (acc_norm) | 2,376 | 0.4541 | 0.4815 | 0.438 |
| ARC-Challenge (acc_norm) | 1,172 | 0.2526 | 0.2526 | 0.224 |
| OpenBookQA (acc_norm) | 500 | 0.2940 | 0.3100 | 0.274 |
| WinoGrande | 1,267 | 0.5257 | 0.5099 | 0.516 |
| LAMBADA | 5,153 | 0.2319 | 0.2567 | 0.326 |
Read these carefully. ARC-Challenge, OpenBookQA and WinoGrande sit at or near chance for GPT-2 itself, so differences there are noise, not results. What actually carries signal:
- Parity on HellaSwag and PIQA, and a genuine edge on ARC-Easy (+4.3 standard errors for the SFT model), on 7.3B tokens against GPT-2's ~100B. FineWeb-Edu's quality filtering is doing that work.
- LAMBADA is 9 points worse. At n=5,153 that is 16 standard errors, so it is not noise. LAMBADA is drawn from novels and requires long-range narrative coreference. FineWeb-Edu is filtered for educational content and contains almost no fiction. The filtering that bought parity everywhere else cost this.
Data quality is not a scalar you turn up. It is a choice about what your model will be bad at.
Reproduce with python hellaswag.py --ckpt sft_final.pt and
python evals.py --ckpt sft_final.pt.
Files
| File | Contents |
|---|---|
sft_final.pt |
Instruction-tuned model, weights only (498 MB) |
gpt2_124m_weights.pt |
Pretrained base, weights only (498 MB) |
pretrain_final.pt |
Pretrained base + optimizer state, resumable (1.4 GB) |
model.py |
Model definition, needed to load any checkpoint |
chat_format.py |
Chat tokenizer and template |
hellaswag.py, evals.py |
Benchmark scripts |
Architecture
Standard GPT-2 124M: 12 layers, 12 heads, 768 embedding dim, 1024 context, tied
input and output embeddings. vocab_size is 50304, padded up from GPT-2's 50257
for tensor-core alignment.
Training
Pretraining. 13,900 steps at 524,288 tokens per batch = 7.29B tokens, one
epoch of FineWeb-Edu sample-10BT. AdamW, peak LR 6e-4 with 715 warmup steps and
cosine decay to 6e-5, weight decay 0.1, grad clip 1.0. bf16 autocast, flash
attention via scaled_dot_product_attention, torch.compile. 15h02m on one RTX
4090 at 135,000 tokens/sec, roughly 70% model FLOPs utilization.
SFT. 2,837 steps over 3 epochs of 60,543 conversations from alpaca-cleaned
(51.7k single-turn) and no_robots (9.4k human-written, some multi-turn), giving
9.26M supervised tokens. LR 5e-5 cosine to zero, 100 warmup steps, no weight
decay. Loss is masked to assistant content and its closing <|end|> only. 24
minutes. Validation loss bottomed at 2.4477 mid-epoch-2 and finished at 2.4581.
Chat format
vocab_size is padded from 50257 to 50304, leaving 47 allocated but unused
embedding rows. The chat tokens live in that dead space, so SFT loaded the
pretrained checkpoint with no embedding resize and no new parameters:
<|system|> = 50257 <|user|> = 50258 <|assistant|> = 50259 <|end|> = 50260
Conversations render as:
<|system|>You are helpful.<|end|><|user|>What is 2+2?<|end|><|assistant|>4<|end|>
Generation should stop on <|end|> (50260) or <|endoftext|> (50256); the base
model learned the latter as a document boundary and still reaches for it.
Limitations
124M parameters, 7.3B tokens. It learned the shape of an answer well and the content unreliably.
Asked why the sky is blue, it correctly names Rayleigh scattering unprompted, then invents a "Rayleigh angle" and explains that scattering varies by time of night.
Asked how bread is made:
"Grilled bread is made from ingredients such as flour, sugar, eggs, and milk. It is baked in a flour-based mixture, called a bagel... until it reaches a temperature of between 350°F and 350°F (50°C and 65°C)."
Asked to list three uses for a paperclip, it produces a perfectly formatted numbered list of three items, all about face masks.
Arithmetic, code, recent facts and multi-step reasoning are all out of reach. Multi-turn context is weak because most SFT data is single-turn, so reset the conversation between unrelated questions. Narrative fiction is a particular weakness, as the LAMBADA result shows.
Not suitable for production. This is a reproduction built for learning.
Usage
Requires model.py and chat_format.py from this repo.
import torch
from model import GPT, GPTConfig
from chat_format import get_encoding, render_prompt, STOP_TOKENS
torch.serialization.add_safe_globals([GPTConfig])
ckpt = torch.load('sft_final.pt', map_location='cpu')
model = GPT(ckpt['config'])
model.load_state_dict(ckpt['model'])
model.eval()
enc = get_encoding()
ids = render_prompt([{'role': 'user', 'content': 'Explain why the sky is blue.'}], enc)
x = torch.tensor(ids).unsqueeze(0)
out = []
for _ in range(300):
logits, _ = model(x[:, -1024:])
probs = torch.softmax(logits[:, -1, :] / 0.8, dim=-1)
top_p, top_i = torch.topk(probs, 50, dim=-1)
nxt = torch.gather(top_i, -1, torch.multinomial(top_p, 1))
if nxt.item() in STOP_TOKENS:
break
out.append(nxt.item())
x = torch.cat((x, nxt), dim=1)
print(enc.decode(out))
Acknowledgements
Follows Andrej Karpathy's build-nanogpt.