taby-mm

A GPT-2 124M language model trained from scratch on 14.9B tokens of FineWeb-Edu with micrograd, a C++ and CUDA deep learning library written for this project. It uses OpenAI's GPT-2 architecture and tokenizer, and scores above OpenAI's GPT-2 124M on HellaSwag.

It is a base model. It continues text; it was not instruction-tuned or safety-tuned.

Results

HellaSwag, full validation set (10,042 examples), 0-shot, lm-evaluation-harness 0.4.13, hf backend, float32, batch size 1.

Model	`acc_norm`	`acc`
taby-mm	34.53% ± 0.47	30.85% ± 0.46
OpenAI GPT-2 124M	31.14% ± 0.46	28.92% ± 0.45

Both rows were measured with the same harness version, task and settings. taby-mm was scored on Apple mps, GPT-2 on CPU.

Validation loss on the held-out FineWeb-Edu shard fell from 10.968 at step 0 to 2.992 at the end of training.

Usage

from transformers import GPT2LMHeadModel, GPT2TokenizerFast

tokenizer = GPT2TokenizerFast.from_pretrained("ravit21/taby-mm")
model = GPT2LMHeadModel.from_pretrained("ravit21/taby-mm").eval()

ids = tokenizer("The water cycle is", return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=80, do_sample=True, top_k=50, temperature=0.8,
                     pad_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(out[0]))

Start a sentence rather than asking a question: the model continues the text it is given.

Model


Architecture	GPT-2 (decoder-only transformer), `GPT2LMHeadModel`
Parameters	124,439,808
Layers, heads, width	12, 12, 768
Context length	1,024 tokens
Vocabulary	50,257 (GPT-2 BPE, OpenAI's tokenizer unchanged)
Embeddings	Learned positions; output head tied to the token embedding
Activation	GELU, tanh approximation (`gelu_new`)
Weights	float32 safetensors

Training


Data	FineWeb-Edu, GPT-2 token shards from `karpathy/fineweb-edu-100B-gpt2-token-shards`: train shards 1 to 150, val shard 0, each checked against its published SHA-256
Tokens	14.94B, one pass, no repeated data
Steps	114,000 at 131,072 tokens per step (16 sequences x 1,024 tokens x 8 GPUs, no accumulation)
Optimizer	AdamW, beta1 0.9, beta2 0.99, weight decay 0.1, gradient clipping at 1.0
Learning rate	Linear warmup over 2,000 steps to 6e-4, cosine decay to 6e-5
Precision	bf16 compute
Initialisation	GPT-2 style: normal(0, 0.02), residual projections scaled by 1/sqrt(2 x layers)
Dropout	None
Seed	1337
Hardware	8x NVIDIA A100 40 GB SXM4, Lambda Cloud
Training time	3 h 38 min at about 1.18M tokens per second
Cost	About $70 of instance time (4 h 25 min including setup and checks)

The model was trained entirely in micrograd, then converted to the transformers layout with eval/export_hf.py: query, key and value are concatenated into c_attn, and the padded vocabulary used in training (50,304 rows) is cut back to 50,257. The converted model's log-probabilities were checked against micrograd's own scorer before evaluation.

Limitations

A 124M base model: fluent but often factually wrong, with no instruction following.
Trained on English educational web text; it reflects the biases and errors of that data.
No safety tuning. Do not use it where harmful or false output matters.
Context is limited to 1,024 tokens.

License

The model weights are released under the MIT license.

The training data, FineWeb-Edu by Hugging Face, is released under ODC-By 1.0. The tokenizer is OpenAI's GPT-2 tokenizer.

Loading…