kaafivikrant/First5M
0205
First5M — Decoder-Only Transformer Language Model
A ~5M parameter GPT-style language model built entirely from scratch using PyTorch. Trained on OpenWebText following the architecture from "Attention Is All You Need" (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
Quick Start
Requirements: pip install torch tiktoken huggingface_hub
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
License
MIT