| --- |
| language: en |
| license: mit |
| tags: |
| - pytorch |
| - language-model |
| - transformer |
| - decoder-only |
| - custom-architecture |
| - text-generation |
| - educational |
| - from-scratch |
| datasets: |
| - Skylion007/openwebtext |
| library_name: pytorch |
| --- |
| |
| # First5M — Decoder-Only Transformer Language Model |
|
|
| A ~5M parameter GPT-style language model built **entirely from scratch** using PyTorch. |
| Trained on [OpenWebText](https://huggingface.co/datasets/Skylion007/openwebtext) |
| following the architecture from ["Attention Is All You Need"](https://arxiv.org/abs/1706.03762) |
| (Vaswani et al., 2017). |
|
|
| This is an **educational project** — the goal is to understand every line of code |
| in a Transformer, not to build a production model. |
|
|
| ## Model Details |
|
|
| | Property | Value | |
| |---|---| |
| | Architecture | Decoder-only Transformer (Pre-LayerNorm) | |
| | Parameters | ~5M | |
| | Layers | 6 | |
| | Hidden Dimension (d_model) | 256 | |
| | Attention Heads | 4 | |
| | FFN Dimension (d_ff) | 1024 | |
| | Context Window | 256 tokens | |
| | Tokenizer | tiktoken GPT-2 BPE (50,257 vocab) | |
| | Training Data | OpenWebText (~328M tokens, 17999 steps) | |
| | Best Val Loss | 4.986894807815552 (PPL 146) | |
| | Positional Encoding | Sinusoidal (not learned) | |
| | Weight Tying | Yes (embedding = output head) | |
|
|
| ## Quick Start |
|
|
| **Requirements:** `pip install torch tiktoken huggingface_hub` |
|
|
| ```python |
| import torch |
| import tiktoken |
| from huggingface_hub import hf_hub_download |
| import importlib.util |
| |
| # Download model files |
| model_py_path = hf_hub_download("kaafivikrant/First5M", "model.py") |
| weights_path = hf_hub_download("kaafivikrant/First5M", "pytorch_model.pt") |
| |
| # Load the model class from model.py |
| spec = importlib.util.spec_from_file_location("model", model_py_path) |
| mod = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(mod) |
| |
| # Build model and load weights |
| config = mod.ModelConfig() |
| model = mod.TransformerLM(config) |
| state_dict = torch.load(weights_path, map_location="cpu", weights_only=True) |
| model.load_state_dict(state_dict, strict=False) # strict=False: lm_head is weight-tied |
| model.eval() |
| |
| # Generate text |
| enc = tiktoken.get_encoding("gpt2") |
| prompt = "The meaning of life is" |
| ids = torch.tensor([enc.encode(prompt)], dtype=torch.long) |
| out = model.generate(ids, max_new_tokens=100, temperature=0.8, top_k=50) |
| print(enc.decode(out[0].tolist())) |
| ``` |
|
|
| ## Architecture |
|
|
| ``` |
| Input token IDs [batch, seq_len] |
| | |
| Token Embedding [50257, 256] |
| + |
| Sinusoidal Positional Encoding |
| | |
| 6x Transformer Blocks: |
| |-- LayerNorm |
| |-- Multi-Head Self-Attention (4 heads x 64 dims, causal mask) |
| |-- Residual Add |
| |-- LayerNorm |
| |-- Feed-Forward (256 -> 1024 -> 256, GELU) |
| |-- Residual Add |
| | |
| Final LayerNorm |
| | |
| Output Head [256, 50257] (tied with embedding) |
| | |
| Logits [batch, seq_len, 50257] |
| ``` |
|
|
| ## Training Details |
|
|
| - **Dataset:** OpenWebText (Skylion007/openwebtext) — ~4.3B tokens total |
| - **Tokens Seen:** ~328M (~7.6% of dataset) |
| - **Optimizer:** AdamW (betas=0.9/0.95, weight_decay=0.1) |
| - **LR Schedule:** Cosine decay with linear warmup (500 steps) |
| - **Peak LR:** 3e-4, Min LR: 3e-5 |
| - **Batch Size:** 64 effective (16 x 4 gradient accumulation steps) |
| - **Hardware:** Apple M1, 16GB RAM, MPS backend |
| - **Training Time:** ~22 hours |
| |
| ## Generation Parameters |
| |
| The `generate()` method supports: |
| - `temperature`: Controls randomness (0.7-0.9 recommended) |
| - `top_k`: Limits sampling to top K tokens (40-50 recommended) |
| - `repetition_penalty`: Penalizes repeated tokens (1.2 default, 1.0 = off) |
|
|
| ## Limitations |
|
|
| This is a small educational model. It: |
| - Produces low-quality, often incoherent text (expected for 5M params) |
| - Has a tiny context window (256 tokens) |
| - Has NOT been instruction-tuned or aligned |
| - May produce repetitive, nonsensical, or inappropriate text |
| - Is NOT intended for any production use |
|
|
| ## Files |
|
|
| | File | Description | |
| |---|---| |
| | `model.py` | Standalone model class definitions (no dependencies beyond PyTorch) | |
| | `pytorch_model.pt` | Model weights (state dict) | |
| | `config.json` | Architecture hyperparameters | |
| | `tokenizer_config.json` | Tokenizer info (tiktoken GPT-2 BPE) | |
|
|
| ## License |
|
|
| MIT |
|
|
