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

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

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