Tiny-LLM-From-Scratch (~50M)
A ~50M parameter GPT-2-style decoder-only Transformer built entirely from scratch in PyTorch and pre-trained on FineWeb-Edu.
This model was created for educational purposes to demonstrate how a GPT-style LLM works
at the tensor/code level. Every component—tokenizer, data pipeline, attention, MLP,
training loop, and generation—is implemented from scratch without using the transformers library.
Architecture Details
| Parameter | Value |
|---|---|
| Architecture | Decoder-only Transformer (Pre-LayerNorm, GPT-2 style) |
| Total parameters | 50,677,760 (50.68 M) |
| Context length | 256 tokens |
| Vocabulary | 16,384 Byte-Level BPE tokens |
| Embedding dimension (C) | 640 |
| Attention heads (H) | 10 (head_dim = 64) |
| Transformer layers (N) | 6 |
| FFN hidden dim | 2,560 (4 x C) |
| Activation | GELU |
| Weight tying | No |
Training Details
| Setting | Value |
|---|---|
| Dataset | FineWeb-Edu sample-10BT |
| Training tokens | ~20 M |
| Tokenizer | Custom Byte-Level BPE (trained on 80k docs) |
| Optimiser | AdamW (lr=3e-4, weight_decay=0.1, betas=(0.9, 0.95)) |
| LR schedule | Cosine with 500-step linear warm-up |
| Effective batch size | 128 sequences (32 × 4 grad accum) |
| Steps | 5,000 |
| AMP | FP16 (PyTorch autocast + GradScaler) |
| Hardware | Google Colab T4 (16 GB) |
Sampling Capabilities
Supports autoregressive text generation with:
- Temperature scaling
- Top-K filtering
- Nucleus (Top-P) sampling
Usage (PyTorch)
from huggingface_hub import hf_hub_download
import torch, json
# Download files
config_path = hf_hub_download(repo_id="tiny-llm-50m", filename="config.json")
weights_path = hf_hub_download(repo_id="tiny-llm-50m", filename="pytorch_model.bin")
# Rebuild model (copy GPTConfig and GPT classes from the training notebook)
with open(config_path) as f:
cfg_dict = json.load(f)
config = GPTConfig(**cfg_dict)
model = GPT(config)
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()
Source Code
Full project source: Tiny-LLM-From-Scratch