llm-124m

A GPT-2 small (124M parameters) next-token model trained from scratch in plain PyTorch, plus a question-answer fine-tuned version. MIT licensed.

Code: https://github.com/alexander2323/llm-124m

Files

• out_124M/ckpt.pt pretrained weights, 475 MB
• out_124M/config.json architecture and training arguments
• out_124M/log.txt pretraining log
• out_chat/ckpt.pt question-answer fine-tuned weights, 475 MB
• out_chat/log.txt fine-tuning log
• code/ the training and inference scripts, same as the GitHub repo

Use

pip install torch tiktoken numpy
hf download alexander2323/llm-124m --local-dir .
cd code
python generate.py --ckpt ../out_124M/ckpt.pt --prompt "The French Revolution began in"
python chat.py --ckpt ../out_chat/ckpt.pt --question "What is the capital of France?"

Loading by hand:

import torch
from model import GPT, GPTConfig  # from code/

ckpt = torch.load("out_124M/ckpt.pt", map_location="cpu")
model = GPT(GPTConfig(**ckpt["config"]))
model.load_state_dict(ckpt["model"])
model.eval()

The tokenizer is the GPT-2 byte-pair tokenizer from tiktoken (gpt2). The fine-tuned model expects prompts in the form Question: ...\nAnswer:.

Architecture

Decoder-only transformer laid out exactly as GPT-2: 12 layers, 12 heads, 768 wide, 1024 token context, vocabulary padded to 50304, weights tied between the token embedding and the output head.

Training

Pretraining: 10 billion tokens, batch of 524,288 tokens, learning rate 6e-4 with 700 warmup steps and cosine decay to 6e-5, weight decay 0.1, AdamW. 4x H200 NVL, 1 h 36 min. Final validation loss 3.07.

Data mix, streamed and tokenized by code/prepare.py: FineWeb-Edu (sample-10BT), English Wikipedia (20231101), peS2o science papers (v2), and Pile of Law primary sources. The tokenized training set is not published; prepare.py rebuilds it from the public source datasets.

Fine-tuning: Dolly 15k plus Natural Questions Open (88k), 2 passes, learning rate 3e-5, 4 min. Validation loss 3.83 to 3.14.

Limitations

A 124M model trained on 10 billion tokens. It completes text plausibly and answers simple factual questions, and it also makes things up freely. No safety tuning.