gpt2-124m-fineweb

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

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

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

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.