bmeyer2025 commited on
Commit
9d14757
Β·
verified Β·
1 Parent(s): ec47f0d

Upload DEVLOG.md with huggingface_hub

Browse files
Files changed (1) hide show
  1. DEVLOG.md +613 -0
DEVLOG.md ADDED
@@ -0,0 +1,613 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Tiny LLM Dev Log
2
+
3
+ Building a 10M-parameter language model from scratch in PyTorch to understand how modern LLMs work from the ground up.
4
+
5
+ **Goal**: Build a decoder-only transformer, train it on Shakespeare, modernize it with the same components used in LLaMA/Qwen/Mistral, and publish to HuggingFace.
6
+
7
+ **Why this matters**: Every LLM β€” GPT-4, Claude, Qwen β€” is a scaled-up version of what we're building here. Same architecture, same training loop, same attention mechanism. The only differences are scale (billions of params vs our 10M), data (internet vs Shakespeare), and engineering optimizations.
8
+
9
+ ---
10
+
11
+ ## The Core Idea
12
+
13
+ A language model predicts the next token. That's it. You give it "To be or not to" and it learns that "be" is the most likely next word. Why does this simple idea produce intelligence? Because to predict what comes next, the model has to learn grammar, context, character relationships, narrative structure β€” understanding emerges as a side effect of getting really good at prediction.
14
+
15
+ ---
16
+
17
+ ## Phase 1: Setup β€” 2026-03-28
18
+
19
+ **Hardware**: Apple Silicon Mac Mini M4 (16GB unified memory) running PyTorch 2.11 with MPS (Metal Performance Shaders) β€” Apple's GPU backend. We chose to train locally instead of on Colab because a ~10M param model on 1MB of data fits comfortably in 16GB.
20
+
21
+ **Dataset**: Tiny Shakespeare (~1.1MB, ~1M characters) β€” the standard toy dataset for learning transformer architectures. Small enough to train in minutes, complex enough to see real patterns emerge.
22
+
23
+ **GitHub repo**: [brianmeyer/tinyllm](https://github.com/brianmeyer/tinyllm)
24
+
25
+ ---
26
+
27
+ ## Phase 2: Building the Vanilla Transformer β€” 2026-03-28
28
+
29
+ ### Milestone 1: Tokenizer β€” turning text into numbers
30
+
31
+ Neural networks can't read text. They work with numbers. So step one is converting every character in Shakespeare into an integer.
32
+
33
+ We found 65 unique characters in the dataset (letters, punctuation, spaces, newlines). Each gets a number: `a=39, b=40, ...` etc. This is called **character-level tokenization** β€” the simplest possible approach. Real LLMs use **BPE** (Byte Pair Encoding) which groups common character sequences into single tokens (like "the" β†’ one token), but character-level is clearer for learning.
34
+
35
+ The key function is `get_batch()`. It grabs a random 256-character chunk from Shakespeare and creates a training pair:
36
+ - **Input (x)**: characters 0-255
37
+ - **Target (y)**: characters 1-256 (shifted by one)
38
+
39
+ At every position, the model's job: "given everything up to here, predict the next character."
40
+
41
+ **Numbers**: 65-char vocab, 1,003,854 training tokens, 111,540 validation tokens (90/10 split).
42
+
43
+ ### Milestone 2: Self-Attention β€” the core of the transformer
44
+
45
+ This is the mechanism that makes transformers work. Here's the intuition:
46
+
47
+ Imagine you're reading "The king picked up his crown." When you get to "his," how do you know it refers to "king" and not some other noun? You **attend** to the right word based on context. Self-attention is a differentiable version of this.
48
+
49
+ Every token produces three vectors:
50
+ - **Query (Q)**: "What am I looking for?"
51
+ - **Key (K)**: "What do I contain?"
52
+ - **Value (V)**: "What information should I share?"
53
+
54
+ The attention score between two tokens is `Q·K` (dot product). High score = "these tokens are relevant to each other." We scale by `√(head_size)` to prevent the scores from getting too large (which would make softmax too peaky).
55
+
56
+ **The causal mask** is crucial: we use a lower-triangular matrix to ensure each token can only attend to tokens that came *before* it. Token 5 can see tokens 0-5 but NOT tokens 6+. This is what makes it a language model that generates left-to-right, rather than a bidirectional encoder like BERT.
57
+
58
+ ```
59
+ Attention(Q, K, V) = softmax(Q @ K^T / √d_k) @ V
60
+ ```
61
+
62
+ ### Milestone 3: Transformer Block β€” stacking the pieces
63
+
64
+ A single attention head has limited capacity β€” it can only focus on one "type" of relationship at a time. **Multi-head attention** runs 6 heads in parallel, each learning to focus on different patterns (one might learn syntax, another might learn character names, etc.). Their outputs get concatenated and projected back down.
65
+
66
+ The **feed-forward network (FFN)** is two linear layers with a ReLU activation. Its job: after attention has gathered information from other positions, the FFN processes that information. Think of attention as "gather" and FFN as "think."
67
+
68
+ A **Block** combines them with two critical additions:
69
+ 1. **Residual connections**: `output = x + attention(x)`. The input gets added back to the output. This helps gradients flow during training β€” without it, deep networks are very hard to train.
70
+ 2. **Pre-norm**: We apply LayerNorm *before* attention and FFN (not after). This is the modern convention (LLaMA, GPT-3 era onwards) and trains more stably.
71
+
72
+ **One Block = 1.77M parameters.** We stack 6 of them.
73
+
74
+ ### Milestone 4: The Full GPT Model
75
+
76
+ The complete model:
77
+ 1. **Token embedding** (65 Γ— 384): converts each character ID to a 384-dimensional vector
78
+ 2. **Position embedding** (256 Γ— 384): adds a learned vector for each position (0-255)
79
+ 3. **6 transformer blocks**: the actual computation
80
+ 4. **Final LayerNorm**: stabilizes the output
81
+ 5. **Language model head** (384 β†’ 65): converts the final representation back to probabilities over the 65 characters
82
+
83
+ **Weight tying**: The language model head shares its weights with the token embedding. This means "the representation of character 'a' going into the model" is the same as "the model deciding to output character 'a'." This was in GPT-2 β€” it's a nice inductive bias and saves parameters.
84
+
85
+ **Total: 10.8M parameters.**
86
+
87
+ **Untrained loss**: 4.09. Theoretical random baseline for 65 characters: `ln(65) β‰ˆ 4.17`. The model starts knowing nothing β€” it's essentially randomly guessing among 65 characters.
88
+
89
+ ### Training the Vanilla Model
90
+
91
+ **Optimizer**: AdamW with lr=3e-4. AdamW is the standard for transformers β€” it's Adam (adaptive learning rates per parameter) with proper weight decay. The 3e-4 learning rate is Karpathy's recommendation for this scale.
92
+
93
+ **Batch size 64, block size 256**: Each training step processes 64 sequences of 256 characters = 16,384 characters per step.
94
+
95
+ **Loss curve:**
96
+
97
+ | Step | Train Loss | Val Loss | What's happening |
98
+ |------|-----------|---------|-----------------|
99
+ | 0 | 4.19 | 4.19 | Random guessing (1 in 65 chars) |
100
+ | 500 | 1.87 | 1.99 | Learned common characters, spaces, basic word shapes |
101
+ | 1000 | 1.44 | 1.65 | Learning word patterns, common Shakespeare phrases |
102
+ | 1500 | 1.29 | 1.54 | Diminishing returns, starting to fit real structure |
103
+ | 2000 | 1.21 | 1.49 | Val loss still improving |
104
+ | 2500 | 1.15 | 1.49 | Val loss starts plateauing |
105
+ | 3000 | 1.09 | 1.48 | **Best val loss** β€” sweet spot before overfitting |
106
+ | 3500 | 1.04 | 1.48 | Train still dropping, val flat |
107
+ | 4000 | 0.99 | 1.50 | Val ticks up β€” model starting to memorize |
108
+ | 4500 | 0.94 | 1.51 | Overfitting growing |
109
+ | 5000 | 0.88 | 1.54 | Final. Train-val gap = 0.66 |
110
+
111
+ **Total training time**: 88.1 minutes on M4 MPS (5000 steps, ~1 sec/step + eval overhead).
112
+
113
+ **Overfitting analysis**: Best val loss was **1.48 at step 3000**. After that, train loss kept dropping (the model was memorizing Shakespeare) but val loss crept back up. By step 5000 the train-val gap was 0.66 β€” significant. In a real project, you'd either stop early (at step 3000) or add more training data. For our learning purposes, 5000 steps is fine β€” we wanted to see the full curve including the overfitting regime.
114
+
115
+ **Speed**: ~1 second per training step on M4 MPS. First step took 88 seconds due to Metal shader compilation (one-time cost).
116
+
117
+ **Training crash** (first attempt): The first training run died silently after step 1500. Cause: running other GPU-intensive scripts in parallel overwhelmed MPS (16GB shared between CPU and GPU). Lesson: MPS doesn't crash cleanly like CUDA β€” it just kills the process with no error. Restarted from scratch with no other GPU work running.
118
+
119
+ ### Generated Shakespeare β€” Vanilla Model
120
+
121
+ After training, we generated text at different temperatures to see the quality/creativity tradeoff.
122
+
123
+ **Temperature = 0.5 (focused, conservative):**
124
+ ```
125
+ ROMEO:
126
+ I would be so straitly for thee for thy heart.
127
+
128
+ BENVOLIO:
129
+ By this and look on thee, who were thy son
130
+ As if thou couldst desire to thy love.
131
+ ```
132
+ Coherent sentences, proper character names, Shakespearean rhythm. Repetitive.
133
+
134
+ **Temperature = 0.8 (balanced):**
135
+ ```
136
+ ROMEO:
137
+ Thither the forest world they are.
138
+
139
+ MERCUTIO:
140
+ No better for the court.
141
+
142
+ MERCUTIO:
143
+ Let us always be for contented: have you not slander'd
144
+ therein like less behind than than offends it, and he
145
+ discharged in Verona his report.
146
+ ```
147
+ More varied, still mostly grammatical. Good default setting.
148
+
149
+ **Temperature = 1.0 (default):**
150
+ ```
151
+ ROMEO:
152
+ Long and Angelo: and yours, counterfeit that
153
+ Which then benefit your own solemnity.
154
+ ```
155
+ Gets creative β€” mixes characters from different plays (Angelo from Measure for Measure shows up in a Romeo scene). More errors but also more interesting.
156
+
157
+ **What temperature actually does**: it divides the logits (raw model scores) before softmax. Low temperature makes the probability distribution peakier (model picks its top choice), high temperature flattens it (model samples more randomly). temp=0 would be pure greedy decoding β€” always pick the most likely token.
158
+
159
+ **What this tells us**: The model has genuinely learned Shakespeare structure β€” character names, dialogue formatting, verse rhythm, vocabulary. At 10M parameters trained on 1MB for 88 minutes, that's remarkable. But it also makes mistakes: garbled phrases, crossed character contexts, meaningless filler. A bigger model with more data would fix these.
160
+
161
+ ---
162
+
163
+ ## Phase 3: Modernizing the Architecture β€” 2026-03-29
164
+
165
+ The vanilla transformer we built is architecturally similar to GPT-2 (2019). Modern LLMs (LLaMA, Qwen, Mistral β€” all 2023-2025) use four key improvements. We swap them in one at a time, training 2000 steps each time, to see what each change does.
166
+
167
+ ### Why one swap at a time?
168
+
169
+ If you change 4 things at once and the model gets better, you don't know which change helped. Controlled experiments β€” change one variable, measure the effect. This is how real ML research works.
170
+
171
+ ### Swap 1: LayerNorm β†’ RMSNorm
172
+
173
+ **What LayerNorm does**: normalizes activations by subtracting the mean and dividing by the standard deviation, then applies learned scale and bias.
174
+
175
+ **What RMSNorm does**: skips the mean subtraction, just divides by the root-mean-square. No bias parameter either.
176
+
177
+ **Why the change**: Turns out the mean subtraction doesn't help much. Removing it is simpler and ~7% faster with equivalent results. Zhang & Sennrich showed this in 2019.
178
+
179
+ **Used in**: LLaMA, LLaMA 2, LLaMA 3, Qwen, Qwen 2, Mistral, Gemma.
180
+
181
+ **Results** (2000 steps, crashed at step 1000 due to MPS β€” but we got enough data):
182
+
183
+ | Step | Vanilla val | RMSNorm val |
184
+ |------|-----------|------------|
185
+ | 500 | 1.99 | 1.99 |
186
+ | 1000 | 1.65 | 1.63 |
187
+
188
+ **Verdict**: Essentially identical. RMSNorm is a free upgrade β€” same quality, simpler code, fewer operations. This confirms why the entire industry adopted it: there's zero downside. In production LLMs serving millions of requests, the small efficiency gain compounds.
189
+
190
+ ### Swap 2: ReLU FFN β†’ SwiGLU
191
+
192
+ **What ReLU FFN does**: `Linear(384β†’1536) β†’ ReLU β†’ Linear(1536β†’384)`. ReLU sets all negative values to zero β€” a hard cutoff that destroys information.
193
+
194
+ **What SwiGLU does**: Uses three weight matrices instead of two:
195
+ - A **gate** matrix learns *what to let through*
196
+ - An **up** matrix provides *the values to let through*
197
+ - The gate uses SiLU (a smooth curve, no hard zeros), and its output multiplies the up values
198
+ - A **down** matrix projects back to the original dimension
199
+
200
+ **Why the change**: The gating mechanism gives the network much finer control over information flow. ReLU is a binary on/off switch; SwiGLU is a smooth dimmer. Shazeer's 2020 paper showed consistent 0.1-0.3 perplexity improvement.
201
+
202
+ **Param count trick**: Three matrices instead of two means more params. To keep it fair, we shrink the hidden dim: `int(2/3 Γ— 4 Γ— 384) = 1024` (vs 1536 for ReLU). Same total params, better architecture.
203
+
204
+ **Used in**: LLaMA (all versions), Qwen (all versions), Mistral, PaLM.
205
+
206
+ **Results** (full 2000 steps completed, 37.8 min):
207
+
208
+ | Step | Vanilla val | SwiGLU val | Difference |
209
+ |------|-----------|-----------|-----------|
210
+ | 500 | 1.99 | 1.88 | **-0.11** |
211
+ | 1000 | 1.65 | 1.58 | **-0.07** |
212
+ | 1500 | 1.54 | 1.51 | **-0.03** |
213
+ | 2000 | 1.49 | 1.50 | ~same |
214
+
215
+ **Verdict**: SwiGLU learned significantly faster at every early checkpoint. By step 2000 they converged β€” but that's because tiny Shakespeare is so small that both architectures eventually hit the same data-limited floor. On a real dataset with billions of tokens, that faster learning compounds into meaningfully better final quality. This is the swap that makes the biggest quality difference in practice.
216
+
217
+ **What surprised me**: The convergence at step 2000. I expected SwiGLU to stay ahead. But it makes sense β€” when your dataset is only 1MB, there's a ceiling on how good any architecture can get. The architectural advantage shows most clearly in how fast you get there, not how far you go. On larger datasets, SwiGLU would pull ahead and stay ahead.
218
+
219
+ ### Swap 3: Learned Positional Embeddings β†’ RoPE
220
+
221
+ **What learned pos embeddings do**: A lookup table of 256 vectors. Position 0 always adds vector[0], position 5 always adds vector[5]. The model has to learn what these vectors should be.
222
+
223
+ **What RoPE does**: Instead of *adding* position info to the token embedding, it *rotates* the Query and Key vectors in attention. Each position gets a different rotation angle. The angle decreases geometrically across dimensions β€” early dimensions rotate fast (capturing local patterns), later dimensions rotate slowly (capturing long-range structure).
224
+
225
+ **Why the change**: Three big advantages:
226
+ 1. **Relative positions**: Under RoPE, the attention score QΒ·K only depends on the *distance* between two tokens, not their absolute positions. "The word 3 positions back" is the same pattern regardless of whether you're at position 10 or position 200.
227
+ 2. **Length generalization**: The model can handle longer sequences than it was trained on (the rotation math works for any length). Learned pos embeddings can't β€” position 257 has no learned vector.
228
+ 3. **No extra parameters**: We *remove* the pos_emb table entirely (saves 256Γ—384 = 98,304 params). Position is encoded for free through rotations.
229
+
230
+ **The rotation math**: For each consecutive pair of dimensions (x₁, xβ‚‚):
231
+ ```
232
+ x₁' = x₁·cos(mΞΈ) - xβ‚‚Β·sin(mΞΈ)
233
+ xβ‚‚' = x₁·sin(mΞΈ) + xβ‚‚Β·cos(mΞΈ)
234
+ ```
235
+ where m = position and ΞΈ = frequency for that dimension pair. This is literally a 2D rotation matrix.
236
+
237
+ **Why only Q and K, not V?** RoPE's purpose is to make the attention *pattern* position-aware. Attention patterns come from QΒ·K. Values (V) carry the actual information content β€” they don't need position encoding.
238
+
239
+ **Used in**: LLaMA (all versions), Qwen (all versions), Mistral, GPT-NeoX, Gemma. This is THE standard for modern LLMs.
240
+
241
+ **This is the most important swap to understand** for working with production LLM architectures.
242
+
243
+ **Results** (full 2000 steps completed, 41.6 min):
244
+
245
+ | Step | Vanilla val | RoPE val | Difference |
246
+ |------|-----------|---------|-----------|
247
+ | 500 | 1.99 | 1.68 | **-0.31** |
248
+ | 1000 | 1.65 | 1.53 | **-0.12** |
249
+ | 1500 | 1.54 | 1.49 | **-0.05** |
250
+ | 2000 | 1.49 | 1.47 | **-0.02** |
251
+
252
+ **Verdict**: RoPE was the biggest single improvement β€” 0.31 better at step 500! It achieved this with 98K *fewer* parameters (no positional embedding table). The strong inductive bias of relative position encoding via rotations gives the model a massive head start. Like SwiGLU, the gap narrows as both models approach the data-limited floor, but RoPE gets there faster and ends up slightly ahead.
253
+
254
+ **Why RoPE outperformed even SwiGLU**: RoPE replaces a component the model had to *learn from scratch* (positional embeddings) with one that has strong mathematical structure *baked in* (rotations encode distance). SwiGLU is a better architecture but it still starts from random weights. RoPE starts with a correct inductive bias about how position should work.
255
+
256
+ ### Swap 4: Add KV Cache for Inference
257
+
258
+ **The problem**: Without caching, generating each new token requires reprocessing the *entire* sequence through the model. For a 200-token sequence, generating token 201 means re-running all 200 tokens through attention. Generating token 202 means re-running all 201. That's O(nΒ²) total compute.
259
+
260
+ **The solution**: After processing the prompt, *save* the Key and Value tensors for every layer. When generating a new token, only compute K and V for that one token, append to the cache, and run attention. This is O(n) total compute.
261
+
262
+ **For a 256-token context, that's up to 256Γ— less K/V computation per token generated.**
263
+
264
+ This doesn't change training at all (during training we process full sequences anyway). It's purely an inference optimization.
265
+
266
+ **Used in**: Every production LLM inference system. This is what makes real-time chat possible.
267
+
268
+ **Results** (benchmark using vanilla checkpoint, 200 generated tokens):
269
+
270
+ | | Speed | Throughput |
271
+ |--|-------|-----------|
272
+ | Without cache | 4.27s | 46.8 tok/s |
273
+ | With cache | 3.23s | 61.9 tok/s |
274
+ | **Speedup** | **1.3Γ—** | |
275
+
276
+ **Verdict**: 1.3Γ— faster on 256-token contexts. The modest speedup is because our sequences are short β€” KV cache shines on longer contexts. At 2048 tokens (typical for production LLMs), you'd see 10-50Γ— speedup. At 128K tokens (Claude/GPT-4 scale), it's the difference between "responds in 1 second" and "responds in 10 minutes."
277
+
278
+ ### Phase 3 Summary: All Swaps Compared
279
+
280
+ | Step | Vanilla | RMSNorm | SwiGLU | **RoPE** |
281
+ |------|---------|---------|--------|----------|
282
+ | 500 | 1.99 | 1.99 | 1.88 | **1.68** |
283
+ | 1000 | 1.65 | 1.63 | 1.58 | **1.53** |
284
+ | 1500 | 1.54 | β€” | 1.51 | **1.49** |
285
+ | 2000 | 1.49 | β€” | 1.50 | **1.47** |
286
+
287
+ **Rankings by impact**:
288
+ 1. **RoPE** β€” biggest quality improvement, fewer parameters, strong inductive bias
289
+ 2. **SwiGLU** β€” faster learning, better gradient flow
290
+ 3. **RMSNorm** β€” free efficiency upgrade, no quality change
291
+ 4. **KV Cache** β€” no quality change, 1.3Γ— faster inference (much more at scale)
292
+
293
+ All four are compatible β€” they modify different parts of the transformer block. They slot into different places:
294
+
295
+ ```
296
+ Transformer Block:
297
+ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
298
+ β”‚ RMSNorm ← Swap 1 (normalization) β”‚
299
+ β”‚ MultiHeadAttention β”‚
300
+ β”‚ └─ Q, K with RoPE ← Swap 3 β”‚
301
+ β”‚ └─ KV Cache ← Swap 4 β”‚
302
+ β”‚ + residual β”‚
303
+ β”‚ β”‚
304
+ β”‚ RMSNorm ← Swap 1 (normalization) β”‚
305
+ β”‚ SwiGLU ← Swap 2 (feed-forward) β”‚
306
+ β”‚ + residual β”‚
307
+ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
308
+ ```
309
+
310
+ Combined in `model_modern.py`, our model is architecturally identical to LLaMA/Qwen at tiny scale.
311
+
312
+ ### Full Modern Model Training β€” The Overfitting Disaster
313
+
314
+ Trained the combined modern model (all 4 swaps) for 5000 steps. **It started out amazing:**
315
+
316
+ | Step | Modern train | Modern val | Vanilla val (comparison) |
317
+ |------|-------------|-----------|------------------------|
318
+ | 0 | 4.25 | 4.25 | 4.19 |
319
+ | 500 | 1.37 | 1.59 | 1.99 |
320
+ | 1000 | 1.21 | 1.50 | 1.65 |
321
+ | **1500** | **1.12** | **1.47** | **1.54** |
322
+ | 2000 | 1.04 | 1.50 | 1.49 |
323
+ | 3000 | 0.88 | 1.55 | 1.48 |
324
+ | 5000 | 0.58 | 1.77 | 1.54 |
325
+
326
+ At step 500, modern was val 1.59 vs vanilla's 1.99 β€” **0.40 better**. At step 1500, it hit 1.47 (beating vanilla's all-time best of 1.48 that took 3000 steps to reach). But then it kept going and the val loss EXPLODED to 1.77. Train loss dropped to an absurdly low 0.58 β€” the model was essentially memorizing Shakespeare character by character.
327
+
328
+ **What the generated text looked like at step 5000:**
329
+ ```
330
+ ROMEO:
331
+ The theneveveveinourein treishathathatwhathon thishadishadishadin
332
+ madilllllllllllllllllllllllllllllllllllllllllllllllllll
333
+ ```
334
+
335
+ Pure garbage. Repetitive fragments. The model memorized training data so precisely that it couldn't generalize at all.
336
+
337
+ **Why this happened**: The modern architecture (RoPE + SwiGLU) is more powerful than vanilla. On a large dataset, that extra power means better generalization. On tiny Shakespeare (only 1MB), it means faster memorization. It's like giving a photographic-memory student a one-page cheat sheet β€” they'll memorize it word for word instead of understanding the concepts.
338
+
339
+ **The fix β€” early stopping**: Added best-checkpoint saving to `train_modern.py`. Now it saves the model at whichever eval step has the lowest val loss. The best model was at step 1500, not step 5000.
340
+
341
+ ### Plot twist: It wasn't overfitting β€” it was a RoPE position bug
342
+
343
+ The "garbage" output at step 5000 wasn't because the model was bad. When we ran the benchmark, the model produced `"ounounounounoun"` β€” but the val loss at step 1000 (1.51) should NOT produce garbage.
344
+
345
+ **Root cause**: A bug in the KV cache generation path. During generation, each new token is processed with `T=1`. Our code computed RoPE frequencies for `T=1`, which means **every generated token got position 0's rotation**. The model had no idea where any token was in the sequence.
346
+
347
+ ```python
348
+ # BUG: always computes rotation for position 0
349
+ cos, sin = precompute_rope_freqs(head_size, T, device) # T=1 during cache gen
350
+
351
+ # FIX: compute for actual position, slice to the right range
352
+ cos_full, sin_full = precompute_rope_freqs(head_size, cache_pos + T, device)
353
+ cos = cos_full[cache_pos : cache_pos + T] # correct position!
354
+ ```
355
+
356
+ During training (T=256, no cache), all positions got the right rotations β€” so the model trained perfectly. The bug only appeared during KV cache generation.
357
+
358
+ **After the fix**, the step 1000 checkpoint produces clean Shakespeare:
359
+ ```
360
+ ROMEO:
361
+ Here he hours, my lord, and not were my head,
362
+ Is is in not hither, and the my petty sights
363
+ And be pass'd the morning his throat, if you will
364
+ Throw your grave of those scarce comes the poor,
365
+ Where you
366
+ ```
367
+
368
+ **This is a great debugging lesson**: When a model has reasonable loss but produces garbage at inference, the bug is usually in the inference code, not the model. Check your positional encoding, attention masking, and caching logic before blaming the model.
369
+
370
+ We also discovered and fixed an **MPS memory leak** that was silently killing training processes after ~60-80 minutes. Fix: `torch.mps.empty_cache()` every 100 steps.
371
+
372
+ ### Modern model retrain β€” with fixes
373
+
374
+ With the RoPE bug fixed and MPS cache clearing added, retrained with early stopping. Best checkpoint saved automatically at lowest val loss.
375
+
376
+ | Step | Modern train | Modern val | Status |
377
+ |------|-------------|-----------|--------|
378
+ | 0 | 4.25 | 4.25 | |
379
+ | 500 | 1.39 | 1.60 | best, saved |
380
+ | 1000 | 1.23 | 1.51 | best, saved |
381
+
382
+ The model IS overfitting on tiny Shakespeare (modern architecture is too powerful for 1MB of data), but with early stopping we capture the sweet spot. The best checkpoint generates good Shakespeare text.
383
+
384
+ ---
385
+
386
+ ## Phase 4: Scaling β€” train_bpe.py
387
+
388
+ Three upgrades to train faster and smarter. Each one is something you'd use in a real production training run.
389
+
390
+ ### Upgrade 1: BPE Tokenization
391
+
392
+ Character-level tokenization was great for learning, but it's wasteful. The word "the" takes 3 tokens (t, h, e). The model wastes 3 attention steps processing one of the most common words in English.
393
+
394
+ **BPE (Byte Pair Encoding)** solves this. It's a compression algorithm that learns which character sequences appear often and merges them into single tokens. "the" β†’ one token. "tion" β†’ one token. Rare words still get broken into pieces, so it handles any input.
395
+
396
+ We use `tiktoken` with the GPT-2 tokenizer (50,257 tokens). The tradeoff: our embedding table grows from 65Γ—384 = 24,960 params to 50,257Γ—384 = 19.3M params. But each token carries ~4x more information, so sequences are ~4x shorter for the same text. Net win.
397
+
398
+ ### Upgrade 2: Mixed Precision Training
399
+
400
+ By default, PyTorch does all math in float32 (32 bits per number). Mixed precision runs most operations in float16 (16 bits) β€” literally half the data. This means:
401
+ - ~2Γ— faster matrix multiplications (the GPU processes twice as many numbers per cycle)
402
+ - ~Half the memory usage (bigger batches or bigger models fit)
403
+ - Minimal quality loss (sensitive operations like loss computation stay in float32)
404
+
405
+ On MPS (Apple Silicon), we use `torch.autocast(device_type='mps', dtype=torch.float16)`. CUDA uses a `GradScaler` for dynamic loss scaling, but MPS doesn't need it.
406
+
407
+ ### Upgrade 3: Gradient Accumulation
408
+
409
+ Bigger batches generally train better β€” the gradient estimate is less noisy because you're averaging over more examples. But bigger batches need more memory. On 16GB, we can't do batch_size=64 with a BPE model.
410
+
411
+ **Gradient accumulation** solves this: run 4 "micro-batches" of size 16, accumulate the gradients without updating the model, then do one optimizer step. The model sees 4Γ—16=64 examples per step, same as a batch_size=64, but only 16 examples are in memory at any time.
412
+
413
+ The key detail: we divide the loss by `ACCUM_STEPS` before calling `.backward()`. This ensures the accumulated gradient is properly averaged, not summed.
414
+
415
+ We also add **gradient clipping** (`max_norm=1.0`) β€” if any gradient gets too large, we scale the entire gradient vector down to keep its norm ≀ 1.0. This prevents training instability from gradient spikes.
416
+
417
+ *Results: [pending β€” will run after swap comparisons]*
418
+
419
+ ---
420
+
421
+ ### BPE Training β€” Divergence Disaster
422
+
423
+ First BPE training run (29.9M params, mixed precision float16, gradient accumulation):
424
+
425
+ | Step | Train | Val | Status |
426
+ |------|-------|-----|--------|
427
+ | 0 | 10.85 | 10.84 | Random over 50K vocab (ln(50257) β‰ˆ 10.83, perfect) |
428
+ | 500 | 7.68 | 8.84 | Learning! Best saved |
429
+ | 1000 | 9.46 | 11.79 | **DIVERGED β€” worse than random** |
430
+
431
+ **What happened**: Loss exploded from 8.84 to 11.79 between step 500 and 1000. The model didn't just stop learning β€” it got actively *worse* than random.
432
+
433
+ **Root cause**: Mixed precision float16 on MPS. Float16 has a very limited range (max ~65,504). With a 50,257-token vocabulary, the softmax over logits can produce values that overflow or underflow in float16. The gradients become NaN, and the model parameters get corrupted. On CUDA, `GradScaler` dynamically adjusts the loss scale to prevent this. On MPS, there's no GradScaler β€” we were running raw float16 with no safety net.
434
+
435
+ **The fix**:
436
+ 1. **Disable mixed precision** β€” use float32 (slower but numerically stable on MPS)
437
+ 2. **Lower learning rate** β€” 1e-4 instead of 3e-4 (larger model needs smaller steps)
438
+ 3. **Increase dropout to 0.3** β€” more regularization for the bigger 29.9M param model
439
+
440
+ **Lesson**: Mixed precision is a CUDA optimization. MPS doesn't have the same infrastructure (no GradScaler, different float16 behavior). Don't blindly copy CUDA training configs to MPS. Test on a few hundred steps first before committing to a long run.
441
+
442
+ ### The MPS Wall β€” Moving to Colab
443
+
444
+ After the BPE divergence, we tried retraining the modern model with fixes:
445
+ - Added `torch.mps.empty_cache()` every 100 steps
446
+ - Set `PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0`
447
+ - Increased dropout to 0.3
448
+
449
+ **None of it worked.** MPS kept killing the process after step 0-500. Every single time. We tried 4+ restarts with different configurations.
450
+
451
+ **Root cause**: MPS on macOS has a memory leak ([PyTorch issue #154329](https://github.com/pytorch/pytorch/issues/154329)). GPU memory slowly grows during training. On 16GB unified memory (shared between CPU, GPU, OS, and whatever else is running), the process gets OOM-killed by macOS after 60-80 minutes. There's no error message β€” the process just disappears.
452
+
453
+ **The irony**: The vanilla model trained fine on MPS (88 minutes, completed fully). But that was earlier in the session when less was in memory. The modern model runs kept dying because of accumulated system state.
454
+
455
+ **Decision: Move to Google Colab.** This is what the original guide recommended. CUDA on a T4 is rock-solid β€” no memory leaks, no silent kills. Created `train_colab.py` that runs all three training phases (vanilla, modern, BPE) in one shot on Colab.
456
+
457
+ **Lesson**: MPS is great for inference and short training runs. For anything over 30 minutes, use CUDA. Don't fight the hardware β€” use the right tool for the job. This cost us several hours of debugging that could have been spent learning.
458
+
459
+ ### Colab Training Results β€” All Three Phases
460
+
461
+ First run on free T4: completed but checkpoints lost to runtime disconnect. Second run on Colab Pro T4: completed with Google Drive saving. All three training phases ran back-to-back with zero crashes. Total wall time: ~3.1 hours. CUDA just works.
462
+
463
+ *Final numbers below are from the Colab Pro run.*
464
+
465
+ #### Part 1: Vanilla GPT (56.9 min, best val 1.4804)
466
+
467
+ | Step | Train | Val | Status |
468
+ |------|-------|-----|--------|
469
+ | 0 | 4.23 | 4.23 | best, saved |
470
+ | 500 | 1.79 | 1.93 | best, saved |
471
+ | 1000 | 1.41 | 1.63 | best, saved |
472
+ | 1500 | 1.28 | 1.55 | best, saved |
473
+ | 2000 | 1.19 | 1.50 | best, saved |
474
+ | 2500 | 1.13 | 1.48 | best, saved |
475
+ | **3000** | **1.08** | **1.48** | **best, saved** |
476
+ | 3500 | 1.02 | 1.49 | overfitting starts |
477
+ | 4000 | 0.97 | 1.51 | |
478
+ | 4500 | 0.92 | 1.54 | |
479
+ | 5000 | 0.86 | 1.56 | |
480
+
481
+ Best checkpoint at step 3000. After that, val loss climbs while train keeps dropping β€” classic overfitting on a small dataset.
482
+
483
+ #### Part 2: Modern GPT with dropout 0.3 (64.2 min, best val 1.4754)
484
+
485
+ | Step | Train | Val | Status |
486
+ |------|-------|-----|--------|
487
+ | 0 | 4.32 | 4.32 | best, saved |
488
+ | 500 | 1.47 | 1.67 | best, saved |
489
+ | 1000 | 1.29 | 1.53 | best, saved |
490
+ | 1500 | 1.21 | 1.50 | best, saved |
491
+ | 2000 | 1.14 | 1.48 | best, saved |
492
+ | **2500** | **1.09** | **1.48** | **best, saved** |
493
+ | 3000 | 1.05 | 1.48 | |
494
+ | 3500 | 1.00 | 1.48 | plateau |
495
+ | 4000 | 0.96 | 1.50 | overfitting starts |
496
+ | 4500 | 0.91 | 1.52 | |
497
+ | 5000 | 0.87 | 1.55 | |
498
+
499
+ **Modern beat vanilla**: best val 1.4754 vs 1.4804. Small margin, but modern got there 500 steps sooner (step 2500 vs step 3000). The dropout 0.3 was the key fix β€” it pushed the overfitting point from step 1500 (with dropout 0.2) to step 2500.
500
+
501
+ **Head-to-head at each step:**
502
+
503
+ | Step | Vanilla val | Modern val | Modern advantage |
504
+ |------|-----------|-----------|-----------------|
505
+ | 500 | 1.93 | 1.67 | **-0.26** |
506
+ | 1000 | 1.63 | 1.53 | **-0.10** |
507
+ | 2000 | 1.50 | 1.48 | **-0.02** |
508
+ | Best | 1.4804 (step 3000) | **1.4754** (step 2500) | **-0.005, 500 steps faster** |
509
+
510
+ #### Part 3: BPE + Modern + Gradient Accumulation (68.2 min, best val 4.6414)
511
+
512
+ | Step | Train | Val | Status |
513
+ |------|-------|-----|--------|
514
+ | 0 | 10.94 | 10.94 | Random over 50K vocab (ln(50257) β‰ˆ 10.83) |
515
+ | 500 | 4.32 | 4.85 | best, saved |
516
+ | **1000** | **3.64** | **4.64** | **best, saved** |
517
+ | 1500 | 3.15 | 4.77 | overfitting |
518
+ | 2000 | 2.71 | 4.93 | |
519
+ | 3000 | 2.04 | 5.43 | severe overfitting |
520
+
521
+ **BPE overfits even faster** β€” 29.9M params (mostly the 50K embedding table) on only 338K BPE tokens. The model memorizes the training data by step 1000. Best checkpoint saved at step 1000.
522
+
523
+ **Note on BPE loss numbers**: You can't directly compare BPE loss (4.64) to char-level loss (1.48) because they're predicting over different vocabularies. BPE perplexity = e^4.64 β‰ˆ 103 (choosing among ~103 tokens). Char-level perplexity = e^1.48 β‰ˆ 4.4 (choosing among ~4.4 characters). The BPE model is actually doing harder predictions β€” each token carries more information.
524
+
525
+ **Lesson**: BPE with a 50K vocab on 1MB of Shakespeare is a terrible ratio. The embedding table alone (50,257 Γ— 384 = 19.3M params) is larger than the rest of the model combined. You need millions of tokens to train those embeddings properly. This is why real LLMs train on trillions of tokens β€” BPE only pays off at scale.
526
+
527
+ ### Generated Samples β€” Colab T4
528
+
529
+ **Vanilla model, temp=0.8:**
530
+ ```
531
+ ROMEO:
532
+ Nay, be too be so head: but I am as betimes;
533
+ There is no man with her pleasure attentience,
534
+ She doth behold our queen arms.
535
+
536
+ PAULINA:
537
+ I'll not too woe to die for the law to the world,
538
+ I'll be old fas
539
+ ```
540
+
541
+ **Modern model, temp=0.8 (KV cached):**
542
+ ```
543
+ ROMEO:
544
+ A gallant-house! what says the woe?
545
+
546
+ MERCUTIO:
547
+ Good madam, my lord.
548
+
549
+ ROMEO:
550
+ Villain, for I do not say it is true,
551
+ Which hath a sin by him come to the crown,
552
+ That he is reports for me; for ever is he.
553
+ ```
554
+
555
+ Both produce recognizable Shakespeare with proper character names and dialogue formatting. The modern model's output is slightly more coherent β€” shorter sentences, cleaner dialogue turns.
556
+
557
+ ### Throughput Benchmarks β€” Colab T4
558
+
559
+ | Model | Throughput | Time for 300 tokens |
560
+ |-------|----------|-------------------|
561
+ | Vanilla (no cache) | **72.2 tok/s** | 4.16s |
562
+ | Modern (KV cache) | 40.7 tok/s | 7.37s |
563
+
564
+ **Surprise**: Vanilla is faster! The modern model's KV cache overhead (managing cache state, extra RoPE computation) outweighs the cache benefit at this tiny sequence length (256 tokens). KV cache pays off at longer contexts (2048+). At our scale, the extra complexity slows things down.
565
+
566
+ **Why is KV cache slower at our scale?** Three reasons:
567
+
568
+ 1. **The model is too small** β€” each forward pass takes microseconds. Cache management overhead (concatenating tensors, tracking positions, RoPE angle computation) is a larger fraction of total compute than the savings from not recomputing K/V.
569
+ 2. **Our sequences are short** β€” 256 tokens max. KV cache saves recomputing K/V for all previous tokens. At 256 tokens that's a small saving. At 128K tokens (Claude/GPT-4 scale) it would be ~500Γ— faster.
570
+ 3. **Python loop overhead** β€” our attention heads run in a Python for-loop. Production systems use fused CUDA kernels (Flash Attention) where KV cache is nearly free.
571
+
572
+ **The real lesson**: Architecture improvements that win at scale can lose at small scale. RoPE + SwiGLU + KV cache are designed for billion-parameter models processing thousands of tokens. At 10M params and 256 tokens, the simpler vanilla architecture has less overhead. But the modern architecture **trained better** β€” that's where the real value shows.
573
+
574
+ ---
575
+
576
+ ## Phase 5: Publish
577
+
578
+ - Code pushed to GitHub: [brianmeyer/tinyllm](https://github.com/brianmeyer/tinyllm)
579
+ - HuggingFace: `bmeyer2025/tiny-gpt-shakespeare` (pending checkpoint download from Colab)
580
+
581
+ ---
582
+
583
+ ## Errors and Lessons
584
+
585
+ | What happened | Why | What we learned |
586
+ |--------------|-----|----------------|
587
+ | Training died silently at step 1500 | MPS memory leak, GPU tests in parallel | MPS leaks memory over long runs. Fix: `torch.mps.empty_cache()` every 100 steps. Don't run multiple GPU jobs on 16GB. |
588
+ | Originally bundled all 4 swaps together | Rushing, skipped the guide | Controlled experiments: change one variable, measure the effect. This is how real ML research works. |
589
+ | Python output buffering hid progress | stdout buffered when piped to file | Always use `python -u` (unbuffered) for long-running scripts. |
590
+ | **Modern model generated garbage** | **RoPE position bug in KV cache** | During KV-cached generation, every new token got position 0's rotation instead of its actual position. Model trained fine but inference was broken. Fix: track `_cache_pos` and slice RoPE frequencies to the correct range. **Biggest lesson: when loss is good but output is bad, the bug is in inference, not training.** |
591
+ | Modern model overfitted catastrophically | 10M params too powerful for 1MB data | More powerful β‰  better. Modern architecture memorizes tiny datasets faster. Fix: early stopping + more data (BPE helps by making each token carry more information). |
592
+ | **BPE training diverged at step 1000** | **Mixed precision float16 on MPS** | Float16 softmax over 50K vocab overflows on MPS. CUDA has GradScaler to prevent this; MPS doesn't. Fix: use float32, lower lr. **Don't blindly copy CUDA configs to MPS.** |
593
+ | **MPS kills training after 60-80 min** | **MPS memory leak (PyTorch #154329)** | GPU memory grows slowly until macOS OOM-kills the process silently. `empty_cache()` and watermark ratio didn't fix it. **Move to CUDA (Colab) for training over 30 min.** |
594
+ | **Lost all Colab checkpoints** | **Colab runtime disconnected after training** | Free tier runtimes are ephemeral β€” everything in `/content/` is deleted when the runtime disconnects. 3+ hours of training, gone. **Always mount Google Drive and save checkpoints there.** |
595
+ | **Colab GPU quota exhausted** | **Used all free T4 hours on training** | Free tier gives ~4-6 GPU hours per 24h period. Our 3.2h training run used most of it. Had to wait for quota reset to retrain. **Plan your GPU time. Save to Drive first, not after.** |
596
+
597
+ ---
598
+
599
+ ## Key Concepts Glossary
600
+
601
+ **Attention**: Mechanism where every token computes how much to "care about" every other token. QΒ·K gives relevance scores, softmax normalizes them, then we take a weighted sum of V.
602
+
603
+ **Causal mask**: Lower-triangular matrix that prevents tokens from attending to future positions. This is what makes it a left-to-right language model.
604
+
605
+ **Pre-norm**: Applying normalization *before* attention/FFN (not after). More stable training. Used in all modern LLMs.
606
+
607
+ **Residual connections**: Adding the input back to the output of each sub-layer. Helps gradients flow in deep networks.
608
+
609
+ **Weight tying**: Sharing the token embedding matrix with the output projection. Saves params and acts as a regularizer.
610
+
611
+ **Loss (cross-entropy)**: How wrong the model's predictions are. Lower = better. For random guessing over 65 characters, loss starts at ln(65) β‰ˆ 4.17.
612
+
613
+ **Perplexity**: e^(loss). Intuition: "how many characters is the model effectively choosing between?" Perplexity of 65 = random. Perplexity of 5 = model has narrowed it down to ~5 plausible characters at each position.