Lemma

Lemma is a small LLM (small large language model lmao) pretrained on 10B tokens of fineweb-edu followed by post-training on the GSM8K dataset in an attempt to improve the mathematical reasoning capability of a small 124M parameter model.

This repo implements the entire end-to-end training pipeline in PyTorch covering data preprocessing, pretraining, SFT and DPO.

Why?

Wanted to try my hand at creating a full training pipeline from scratch.

Features

• Data preparation and pretraining scripts
• Supervised fine-tuning
• Direct preference optimization
• Distributed training
• Automatic checkpointing and resuming

Model Specs

• GPT-2 small-style decoder-only transformer
• 124M parameters
• 12 layers, 12 attention heads, 768 hidden size
• 1024 token context window
• Vocab size of 50,304 tokens. (Actual usable tokens are 50,257, the extra are padded for optimization)

Running it on your own

Clone the repo, cd into it and run

uv sync

Training

1. Pretraining

First download fineweb-edu

uv run scripts/prepare_fineweb.py

Then run

uv run pretrain.py # For a single gpu
torchrun --standalone --nproc_per_node=8 pretrain.py # For training across 8 gpus

2. Supervised Fine-Tuning

Prepare the data

uv run scripts/prepare_gsm8k.py

Then train

uv run sft.py # For a single gpu
torchrun --standalone --nproc_per_node=8 sft.py # For training across 8 gpus

3. Direct Preference Optimization

Prepare the pairs. You gotta change the split in the script to get the validation set. It's set to the train split by default. So run it twice with different splits.

uv run scripts/create_dpo_pairs.py # With a single gpu
torchrun --standalone --nproc_per_node=8 scripts/create_dpo_pairs.py # Across 8 gpus

Then train

uv run dpo.py # For a single gpu
torchrun --standalone --nproc_per_node=8 dpo.py # For training across 8 gpus

Evaluation

uv run eval.py # For a single gpu
torchrun --standalone --nproc_per_node=8 eval.py # For 8 gpus

Results

I trained the model on 8 RTX 5080s with 16GB VRAM each. The total cost of the whole training run was $18.50.

Pretraining

Metric	Value
Validation loss	3.40
Validation perplexity	29.96

Post-training

Stage	Accuracy	No-answer rate
SFT	1.52% (20/1319)	12.74% (168/1319)
DPO	2.43% (32/1319)	4.70% (62/1319)

The relatively small model size and limited compute budget (😞) naturally constrain performance on reasoning-heavy tasks such as GSM8K hence the low numbers on the accuracies. But the fact that we do see improvement, albeit small, I would consider this experiment a success.

Training Curves

Pretraining, SFT, and DPO loss curves across the full training pipeline.

Learning rate, gradient norm, tokens/sec, and MFU for each stage.

Train and validation preference accuracy during DPO.

Acknowledgements

Thank you to Andrej Karpathy for his nanoGPT lecture on which the whole pretraining script is based.