AI for Everyone, Part 1 — Training a 16M-Parameter Language Model From Scratch in 15 Minutes

• Why build a small model yourself
• Setup
• The dataset — TinyStories
• Architecture — a decoder-only transformer
  • Without the causal mask, training is meaningless
  • Weight tying saves 40% of the parameters
  • Pre-norm placement
• The training loop
• Training log
• Generated samples
• What a perplexity of 8 means
• Summary
• 🧠 Comprehension quiz
• References

Why build a small model yourself

There is no shortage of writing about how to use large language models. What happens inside them, though, stays hidden behind an API. This series goes the other way: we train them ourselves.

The goal is not performance. The constraints are:

• Under 100 million parameters (in practice, far fewer)
• Under 20 minutes of training on a single GPU
• Only publicly available Hugging Face datasets
• Every log and artifact preserved

Small models make the whole structure legible. You can watch, in minutes rather than weeks, why attention is needed, what the loss function measures, and what failure looks like when it arrives.

This first part is a text generation model.

Setup

The actual hardware and versions used.

The dataset — TinyStories

roneneldan/TinyStories is a collection of short English fairy tales written using only vocabulary a three- or four-year-old would understand. That property makes it ideal for small-model experiments.

Ordinary web text has an enormous vocabulary, so a small model burns its capacity memorising words before it ever learns grammar. TinyStories deliberately restricts the vocabulary, which lets a small model spend its capacity on learning how sentences work instead.

Rather than training a tokenizer, we borrowed GPT-2's. Reusing one piece of a hub model is a perfectly good strategy.

from transformers import AutoTokenizer
from datasets import load_dataset

tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
ds = load_dataset("roneneldan/TinyStories", split="train[:40000]")

ids = []
for ex in ds:
    ids.extend(tok(ex["text"]).input_ids + [tok.eos_token_id])
    if len(ids) >= 4_000_000:
        break
data = torch.tensor(ids[:4_000_000], dtype=torch.long)

The tokens are concatenated into one long stream, and training samples are random slices of it. Some samples straddle document boundaries, but eos_token_id marks those boundaries, so the model learns where a story ends.

Architecture — a decoder-only transformer

This is the entire model: four layers, embedding dimension 256, eight heads.

CTX, DIM, LAYERS, HEADS = 128, 256, 4, 8

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.attn = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM)
        )

    def forward(self, x, mask):
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))


class TinyGPT(nn.Module):
    def __init__(self, vocab):
        super().__init__()
        self.emb = nn.Embedding(vocab, DIM)
        self.pos = nn.Embedding(CTX, DIM)
        self.blocks = nn.ModuleList(Block() for _ in range(LAYERS))
        self.ln = nn.LayerNorm(DIM)
        mask = torch.triu(torch.full((CTX, CTX), float("-inf")), 1)
        self.register_buffer("mask", mask)

    def forward(self, x):
        T = x.shape[1]
        h = self.emb(x) + self.pos(torch.arange(T, device=x.device))
        for b in self.blocks:
            h = b(h, self.mask[:T, :T])
        return self.ln(h) @ self.emb.weight.T  # weight tying

Three design decisions are worth pausing on.

Without the causal mask, training is meaningless

The upper-triangular matrix built by torch.triu(..., 1) fills everything above the diagonal with negative infinity. After softmax, those weights become zero — which is precisely what stops each position from seeing tokens that come after it.

Leave it out and the model is asked to predict the next token while already being shown that token as input. The loss collapses toward zero and generation produces nothing usable. A loss that falls far faster than expected is usually a sign that the answer is leaking.

Weight tying saves 40% of the parameters

The @ self.emb.weight.T on the final line is weight tying: instead of a separate output projection, we transpose and reuse the input embedding matrix.

GPT-2's vocabulary is 50,257 tokens. A separate output layer would cost 50257 × 256 ≈ 12.87 million extra parameters. The whole model is 16.06 million, so without tying it would be 28.93 million. Beyond the savings, mapping tokens to vectors and vectors back to tokens are inverse operations, so sharing one matrix is a natural fit.

Pre-norm placement

self.ln1(x) is applied before attention, and the residual adds to the un-normalised x. This pre-norm ordering is known to train more stably than post-norm in deeper models. At four layers the difference is small, but we followed the convention.

The training loop

model = TinyGPT(len(tok)).to("cuda")
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

step = 0
while not run.over_budget():          # 15-minute budget
    ix = torch.randint(0, len(data) - CTX - 1, (64,))
    x = torch.stack([data[i:i + CTX] for i in ix]).to("cuda")
    y = torch.stack([data[i + 1:i + CTX + 1] for i in ix]).to("cuda")
    with torch.autocast("cuda", torch.bfloat16):
        loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
    opt.zero_grad(); loss.backward(); opt.step()
    step += 1

The whole trick is that x and y are offset by one position. If the input is "Once upon a", the target is "upon a time". Every position predicts simultaneously, so a single length-128 sample yields 128 training signals.

We use bfloat16 through torch.autocast. It saves memory and time on a 3090, and unlike float16, bf16 has the same exponent range as float32, so it stays stable without loss scaling.

Training log

The values as recorded.

[    0.6s] step      0 loss=168.4468 ppl=...
[   38.2s] step   1000 loss=5.2213  ppl=185.13
[  114.7s] step   3000 loss=3.6841  ppl=39.81
[  305.9s] step   6700 loss=3.0350  ppl=20.80
[  662.6s] step  16100 loss=2.2882  ppl=9.86
[  901.2s] step  22358 loss=2.0948  ppl=8.12

• 22,358 steps at batch 64 and context 128 — roughly 180 million tokens processed
• The starting loss of 168 is the value right after initialisation. A uniform distribution would give ln(50257) ≈ 10.8; an interaction between LayerNorm and weight tying produces a large first-step value that drops into the normal range immediately
• A final perplexity of 8.12 means the model has narrowed the next word to about 8 candidates on average

Generated samples

Sampled at temperature 0.8. Unedited.

PROMPT: Once upon a time
Once upon a time, there was a little boy named Timmy. Timmy liked to play
with his toy cars instead. One day, Timmy wanted to play with the red cars,
but he couldn't find it.

As Timmy was leaving a fight on the slide, he saw a big tree. Timmy thought
it looked like fun, so he started to tremble. He mixed the leaves and

PROMPT: Lily went to the
Lily went to the park with her mom. They saw a big tree, a hole and a hole.
They were curious and wanted to see what was inside.

They heard a hole in the hole. It was a zipper, but it was in the hole.

Let us read these honestly.

What works — the grammar is nearly flawless. Tense agreement, articles, and pronoun reference are correct. Character names stay consistent across a paragraph. Opened quotation marks get closed. Paragraph breaks fall naturally.

What does not — meaning falls apart. "a big tree, a hole and a hole" repeats itself; "leaving a fight on the slide" is grammatical but nonsensical; "he started to tremble" connects to nothing around it.

That contrast is the point. Grammar is a local pattern that a small model can learn statistically, while coherence is a long-range dependency that needs capacity and context. With a context of 128 tokens, the span the model can even refer back to is short. Scaling up improves exactly this first.

What a perplexity of 8 means

Perplexity is exp(cross_entropy). Intuitively it answers: how many candidates has the model narrowed the next token down to?

• Before training (uniform): 50,257
• After 1,000 steps: 185
• Final: 8.12

From 50,257 choices down to 8. But this number should not be compared against other models directly. Perplexity depends heavily on the tokenizer and the data distribution, and restricted-vocabulary data like TinyStories yields low values easily. It is only meaningful within the same data and the same tokenizer.

Summary

Sixteen million parameters, fifteen minutes, and pennies of electricity produced a model that writes grammatical English fairy tales.

Three takeaways. The causal mask exists to prevent answer leakage, and it is the first thing to suspect when loss drops abnormally fast. Weight tying saves nearly half the parameters in models with large vocabularies. And grammar and meaning are problems of different difficulty — a small model learns the former first.

The next part looks at how a single model can handle text and images at once. Treat pixels as tokens and the same transformer can both generate images and describe them.

🧠 Comprehension quiz

References

• TinyStories dataset
• TinyStories paper — How Small Can Language Models Be and Still Speak Coherent English?
• Attention Is All You Need
• Using the Output Embedding to Improve Language Models — weight tying