Building a 10M Parameter LLM from Scratch on Google Colab Free Tier
How I trained a decoder-only Transformer on the Tiny Shakespeare dataset using nothing but a free Google Colab T4 GPU.
Introduction
The best way to understand how large language models work is to build one yourself. Not by fine-tuning a pre-trained model, not by calling an API, but by writing every layer, every attention head, and every line of the training loop from scratch.
This is the story of Tiny_Shakespeare_LLM_from_scratch — a ~10 million parameter decoder-only Transformer built in PyTorch and trained end-to-end on the Tiny Shakespeare dataset using only the free tier of Google Colab.
No paid GPUs. No cloud credits. Just a free T4 and a few hours of training.
Why Build from Scratch?
Using libraries like Hugging Face transformers is convenient, but it hides the mechanics. Building from scratch forces you to confront the questions that matter:
- How does self-attention actually compute context?
- Why do we need positional embeddings?
- What does the residual stream look like at each layer?
- How does the causal mask prevent the model from cheating?
By the end of this project, the Transformer architecture stops being a black box. It becomes a sequence of tensor operations you wrote yourself.
The Dataset: Tiny Shakespeare
The model trains on the Tiny Shakespeare dataset — roughly 1.1 MB of text, about 1 million characters, drawn from the public-domain works of William Shakespeare. The vocabulary is character-level: 65 unique characters (letters, punctuation, newlines). This keeps the embedding table tiny and the whole training pipeline fast enough to run on a free Colab GPU.
The dataset is small enough to overfit in minutes, which makes it perfect for debugging. If your model cannot memorize Tiny Shakespeare, there is something wrong with your architecture or training loop.
Model Architecture
The model is a decoder-only Transformer — the same family as GPT, LLaMA, and Mistral. It generates text autoregressively: given a sequence of characters, it predicts the next one, appends it to the context, and repeats.
Here is the high-level architecture:
The decoder stack (right side of the original Transformer diagram) is the core of the model. We strip away the encoder entirely and keep only the masked self-attention and feed-forward blocks.
Hyperparameters
| Component | Value |
|---|---|
| Parameters | ~10.6M |
Layers (n_layer) |
6 |
Attention heads (n_head) |
6 |
Embedding dimension (n_embd) |
384 |
Context length (block_size) |
256 |
| Feed-forward dimension | 4 × 384 = 1536 |
| Dropout | 0.2 |
| Vocabulary size | 65 |
| Tokenization | Character-level |
These numbers are deliberately modest. A 6-layer, 6-head, 384-dimensional model is small enough to train on a free T4 in a couple of hours, yet large enough to produce coherent Shakespeare-like text after training.
The Forward Pass
Each forward pass does the following:
- Token Embedding — Maps each character index to a 384-dimensional vector.
- Positional Embedding — Adds a learned embedding for each position in the sequence, so the model knows the order of characters.
-
Transformer Blocks — Six identical blocks, each containing:
- LayerNorm → Multi-Head Causal Self-Attention → Residual Connection
- LayerNorm → Feed-Forward Network (MLP) → Residual Connection
- Final LayerNorm
- Language Modeling Head — A linear layer projecting from 384 dimensions to the 65-character vocabulary.
The causal mask in the self-attention layer ensures that position t can only attend to positions ≤ t. This is what makes the model autoregressive.
Training on Google Colab Free Tier
The entire training run fits comfortably within Google Colab's free T4 GPU (16 GB VRAM).
Training Configuration
| Setting | Value |
|---|---|
| Batch size | 64 |
| Learning rate | 3e-4 (with cosine decay) |
| Optimizer | AdamW |
| Weight decay | 0.1 |
| Warmup steps | 100 |
| Max steps | 5,000 |
| Gradient clipping | 1.0 |
| Mixed precision | Yes (torch.cuda.amp) |
A batch size of 64 with a context length of 256 means each forward pass processes 16,384 characters. With gradient accumulation, you can simulate larger batches if needed, but the free T4 handles this comfortably.
Memory Optimizations
To fit within Colab's memory limits:
-
Mixed precision training (
torch.cuda.amp.autocast) halves the memory footprint of activations. - Gradient clipping prevents exploding gradients.
- A small model (10M parameters) means the optimizer states are tiny compared to a 7B model.
Training for 5,000 steps takes roughly 30–45 minutes on a T4. The loss curve typically drops from ~4.2 (random) to ~1.5–1.8, producing text that is recognizably Shakespearean in structure, if not in meaning.
Sample Output
After training, the model generates text like this:
ROMEO:
What say'st thou, my lord? I will not be so:
The gentle heart is not a little word,
And yet the world is grown so bad, that men
Do call it virtue when they are most accursed.
It is not Shakespeare. But it is Shakespeare-shaped — and that is the point. The model learned grammar, punctuation, character names, and the rhythm of Elizabethan dialogue purely from next-character prediction.
Repository Structure
The repository is organized for clarity:
Tiny_Shakespeare_LLM_from_scratch/
├── data/
│ └── shakespeare.txt # Tiny Shakespeare dataset
├── model.py # Transformer architecture
├── train.py # Training loop
├── generate.py # Inference script
├── config.py # Hyperparameters
└── README.md
Every file is intentionally short and readable. The model definition is under 200 lines. The training loop is under 100.
What I Learned
- Attention is not magic. It is a weighted sum of value vectors, where the weights come from a scaled dot-product between queries and keys.
- The residual stream is the model's working memory. Every block reads from it and writes back to it.
- Character-level models are surprisingly capable. With only 65 tokens, the model learns syntax, names, and even some semantic structure.
- Free Colab is enough for real research. You do not need an A100 to understand how Transformers work.
Conclusion
This project is not about building the best language model. It is about building a language model — one you understand completely, from the embedding lookup to the final softmax.
If you want to run it yourself, clone the repository, open the notebook in Colab, and press play. The entire pipeline is designed to run on the free tier.
Repository: https://github.com/Devn913/Tiny_Shakespeare_LLM_from_scratch
Built with PyTorch on a free Google Colab T4.
Top comments (0)