---
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