mini-transformer-tinystories

Mini Transformer (From Scratch)

A decoder-only transformer built entirely from first principles in PyTorch, no pre-built attention layers, no HuggingFace Transformers library internals. Every component (multi-head self-attention, positional encoding, feedforward blocks, residual connections) is implemented from scratch and trained on the TinyStories dataset.

This is part of a larger project studying transformer internals, followed by implementing DeepSeek's MLA (Multi-head Latent Attention) and mHC (Manifold-Constrained Hyper-Connections) as architectural upgrades to this same baseline.

Architecture

• Decoder-only transformer, GPT-style
• Embedding dim: 256
• Attention heads: 8
• Layers: 6
• Feedforward hidden dim: 1024
• Context length: 128 tokens
• Tokenizer: GPT-2 (tiktoken), 50,257 vocab
• Total parameters: 336,488

Training

• Dataset: TinyStories (roneneldan/TinyStories), ~474M training tokens
• Steps: 10,000
• Batch size: 32
• Optimizer: AdamW, lr=3e-4
• Hardware: single RTX 3060 (6GB VRAM)
• Final train loss: None
• Final val loss: None

Usage

Load model.py for the architecture classes, then:

import torch
from model import MiniGPT

model = MiniGPT(vocab_size=50257, embed_dim=256, num_heads=8,
                 num_layers=6, hidden_dim=1024, max_seq_len=128)
ckpt = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(ckpt["model"])

Sample Output

Once upon a time, there was a little girl named Lily who loved to play outside. She had a lot of fun playing in the mud and splashing around...