File size: 4,133 Bytes
2fd159e 1be11eb 2fd159e 1be11eb 2fd159e 1be11eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | ---
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
|