加载中…
加载中…
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.
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.
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.
| 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 |
| 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.
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.