mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-05-29 16:27:22 -05:00
ee9f559b8c
PROBLEM: - nbdev requires #| export directive on EACH cell to export when using # %% markers - Cell markers inside class definitions split classes across multiple cells - Only partial classes were being exported to tinytorch package - Missing matmul, arithmetic operations, and activation classes in exports SOLUTION: 1. Removed # %% cell markers INSIDE class definitions (kept classes as single units) 2. Added #| export to imports cell at top of each module 3. Added #| export before each exportable class definition in all 20 modules 4. Added __call__ method to Sigmoid for functional usage 5. Fixed numpy import (moved to module level from __init__) MODULES FIXED: - 01_tensor: Tensor class with all operations (matmul, arithmetic, shape ops) - 02_activations: Sigmoid, ReLU, Tanh, GELU, Softmax classes - 03_layers: Linear, Dropout classes - 04_losses: MSELoss, CrossEntropyLoss, BinaryCrossEntropyLoss classes - 05_autograd: Function, AddBackward, MulBackward, MatmulBackward, SumBackward - 06_optimizers: Optimizer, SGD, Adam, AdamW classes - 07_training: CosineSchedule, Trainer classes - 08_dataloader: Dataset, TensorDataset, DataLoader classes - 09_spatial: Conv2d, MaxPool2d, AvgPool2d, SimpleCNN classes - 10-20: All exportable classes in remaining modules TESTING: - Test functions use 'if __name__ == "__main__"' guards - Tests run in notebooks but NOT on import - Rosenblatt Perceptron milestone working perfectly RESULT: ✅ All 20 modules export correctly ✅ Perceptron (1957) milestone functional ✅ Clean separation: development (modules/source) vs package (tinytorch)
2158 lines
89 KiB
Plaintext
2158 lines
89 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0fa7ad93",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"# Module 13: Transformers - Complete Transformer Architecture\n",
|
||
"\n",
|
||
"Welcome to Module 13! You're about to build the complete transformer architecture that powers modern language models like GPT.\n",
|
||
"\n",
|
||
"## 🔗 Prerequisites & Progress\n",
|
||
"**You've Built**: Tensors, activations, layers, attention mechanisms, embeddings, and all foundational components\n",
|
||
"**You'll Build**: TransformerBlock, complete GPT architecture, and autoregressive generation\n",
|
||
"**You'll Enable**: Full language model training and text generation capabilities\n",
|
||
"\n",
|
||
"**Connection Map**:\n",
|
||
"```\n",
|
||
"Attention + Layers + Embeddings → Transformers → GPT Architecture\n",
|
||
"(sequence processing) (building blocks) (complete model) (language generation)\n",
|
||
"```\n",
|
||
"\n",
|
||
"## Learning Objectives\n",
|
||
"By the end of this module, you will:\n",
|
||
"1. Implement complete TransformerBlock with attention, MLP, and layer normalization\n",
|
||
"2. Build full GPT architecture with multiple transformer blocks\n",
|
||
"3. Add autoregressive text generation capability\n",
|
||
"4. Understand parameter scaling in large language models\n",
|
||
"5. Test transformer components and generation pipeline\n",
|
||
"\n",
|
||
"Let's get started!\n",
|
||
"\n",
|
||
"## 📦 Where This Code Lives in the Final Package\n",
|
||
"\n",
|
||
"**Learning Side:** You work in `modules/13_transformers/transformers_dev.py` \n",
|
||
"**Building Side:** Code exports to `tinytorch.models.transformer`\n",
|
||
"\n",
|
||
"```python\n",
|
||
"# How to use this module:\n",
|
||
"from tinytorch.models.transformer import TransformerBlock, GPT\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Why this matters:**\n",
|
||
"- **Learning:** Complete transformer system showcasing how all components work together\n",
|
||
"- **Production:** Matches PyTorch's transformer implementation with proper model organization\n",
|
||
"- **Consistency:** All transformer components and generation logic in models.transformer\n",
|
||
"- **Integration:** Demonstrates the power of modular design by combining all previous modules"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c5e5dae4",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "imports",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"#| default_exp models.transformer\n",
|
||
"#| export\n",
|
||
"\n",
|
||
"import numpy as np\n",
|
||
"import math\n",
|
||
"from typing import Optional, List\n",
|
||
"\n",
|
||
"# Minimal implementations for development - in practice these import from previous modules\n",
|
||
"class Tensor:\n",
|
||
" \"\"\"Minimal Tensor class for transformer development - imports from Module 01 in practice.\"\"\"\n",
|
||
" def __init__(self, data, requires_grad=False):\n",
|
||
" self.data = np.array(data)\n",
|
||
" self.shape = self.data.shape\n",
|
||
" self.size = self.data.size\n",
|
||
" self.requires_grad = requires_grad\n",
|
||
" self.grad = None\n",
|
||
"\n",
|
||
" def __add__(self, other):\n",
|
||
" if isinstance(other, Tensor):\n",
|
||
" return Tensor(self.data + other.data)\n",
|
||
" return Tensor(self.data + other)\n",
|
||
"\n",
|
||
" def __mul__(self, other):\n",
|
||
" if isinstance(other, Tensor):\n",
|
||
" return Tensor(self.data * other.data)\n",
|
||
" return Tensor(self.data * other)\n",
|
||
"\n",
|
||
" def matmul(self, other):\n",
|
||
" return Tensor(np.dot(self.data, other.data))\n",
|
||
"\n",
|
||
" def sum(self, axis=None, keepdims=False):\n",
|
||
" return Tensor(self.data.sum(axis=axis, keepdims=keepdims))\n",
|
||
"\n",
|
||
" def mean(self, axis=None, keepdims=False):\n",
|
||
" return Tensor(self.data.mean(axis=axis, keepdims=keepdims))\n",
|
||
"\n",
|
||
" def reshape(self, *shape):\n",
|
||
" return Tensor(self.data.reshape(shape))\n",
|
||
"\n",
|
||
" def __repr__(self):\n",
|
||
" return f\"Tensor(data={self.data}, shape={self.shape})\"\n",
|
||
"\n",
|
||
"class Linear:\n",
|
||
" \"\"\"Minimal Linear layer - imports from Module 03 in practice.\"\"\"\n",
|
||
" def __init__(self, in_features, out_features, bias=True):\n",
|
||
" # Xavier/Glorot initialization\n",
|
||
" std = math.sqrt(2.0 / (in_features + out_features))\n",
|
||
" self.weight = Tensor(np.random.normal(0, std, (in_features, out_features)))\n",
|
||
" self.bias = Tensor(np.zeros(out_features)) if bias else None\n",
|
||
"\n",
|
||
" def forward(self, x):\n",
|
||
" output = x.matmul(self.weight)\n",
|
||
" if self.bias is not None:\n",
|
||
" output = output + self.bias\n",
|
||
" return output\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" params = [self.weight]\n",
|
||
" if self.bias is not None:\n",
|
||
" params.append(self.bias)\n",
|
||
" return params\n",
|
||
"\n",
|
||
"class MultiHeadAttention:\n",
|
||
" \"\"\"Minimal MultiHeadAttention - imports from Module 12 in practice.\"\"\"\n",
|
||
" def __init__(self, embed_dim, num_heads):\n",
|
||
" assert embed_dim % num_heads == 0\n",
|
||
" self.embed_dim = embed_dim\n",
|
||
" self.num_heads = num_heads\n",
|
||
" self.head_dim = embed_dim // num_heads\n",
|
||
"\n",
|
||
" self.q_proj = Linear(embed_dim, embed_dim)\n",
|
||
" self.k_proj = Linear(embed_dim, embed_dim)\n",
|
||
" self.v_proj = Linear(embed_dim, embed_dim)\n",
|
||
" self.out_proj = Linear(embed_dim, embed_dim)\n",
|
||
"\n",
|
||
" def forward(self, x, mask=None):\n",
|
||
" batch_size, seq_len, embed_dim = x.shape\n",
|
||
"\n",
|
||
" # Linear projections\n",
|
||
" Q = self.q_proj.forward(x)\n",
|
||
" K = self.k_proj.forward(x)\n",
|
||
" V = self.v_proj.forward(x)\n",
|
||
"\n",
|
||
" # Reshape for multi-head attention\n",
|
||
" Q = Q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)\n",
|
||
" K = K.reshape(batch_size, seq_len, self.num_heads, self.head_dim)\n",
|
||
" V = V.reshape(batch_size, seq_len, self.num_heads, self.head_dim)\n",
|
||
"\n",
|
||
" # Transpose to (batch_size, num_heads, seq_len, head_dim)\n",
|
||
" Q = Tensor(np.transpose(Q.data, (0, 2, 1, 3)))\n",
|
||
" K = Tensor(np.transpose(K.data, (0, 2, 1, 3)))\n",
|
||
" V = Tensor(np.transpose(V.data, (0, 2, 1, 3)))\n",
|
||
"\n",
|
||
" # Scaled dot-product attention\n",
|
||
" scores = Tensor(np.matmul(Q.data, np.transpose(K.data, (0, 1, 3, 2))))\n",
|
||
" scores = scores * (1.0 / math.sqrt(self.head_dim))\n",
|
||
"\n",
|
||
" # Apply causal mask for autoregressive generation\n",
|
||
" if mask is not None:\n",
|
||
" scores = Tensor(scores.data + mask.data)\n",
|
||
"\n",
|
||
" # Softmax\n",
|
||
" attention_weights = self._softmax(scores)\n",
|
||
"\n",
|
||
" # Apply attention to values\n",
|
||
" out = Tensor(np.matmul(attention_weights.data, V.data))\n",
|
||
"\n",
|
||
" # Transpose back and reshape\n",
|
||
" out = Tensor(np.transpose(out.data, (0, 2, 1, 3)))\n",
|
||
" out = out.reshape(batch_size, seq_len, embed_dim)\n",
|
||
"\n",
|
||
" # Final linear projection\n",
|
||
" return self.out_proj.forward(out)\n",
|
||
"\n",
|
||
" def _softmax(self, x):\n",
|
||
" \"\"\"Numerically stable softmax.\"\"\"\n",
|
||
" exp_x = Tensor(np.exp(x.data - np.max(x.data, axis=-1, keepdims=True)))\n",
|
||
" return Tensor(exp_x.data / np.sum(exp_x.data, axis=-1, keepdims=True))\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" params = []\n",
|
||
" params.extend(self.q_proj.parameters())\n",
|
||
" params.extend(self.k_proj.parameters())\n",
|
||
" params.extend(self.v_proj.parameters())\n",
|
||
" params.extend(self.out_proj.parameters())\n",
|
||
" return params\n",
|
||
"\n",
|
||
"class Embedding:\n",
|
||
" \"\"\"Minimal Embedding layer - imports from Module 11 in practice.\"\"\"\n",
|
||
" def __init__(self, vocab_size, embed_dim):\n",
|
||
" self.vocab_size = vocab_size\n",
|
||
" self.embed_dim = embed_dim\n",
|
||
" # Initialize with small random values\n",
|
||
" self.weight = Tensor(np.random.normal(0, 0.02, (vocab_size, embed_dim)))\n",
|
||
"\n",
|
||
" def forward(self, indices):\n",
|
||
" # Simple embedding lookup\n",
|
||
" return Tensor(self.weight.data[indices.data])\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" return [self.weight]\n",
|
||
"\n",
|
||
"def gelu(x):\n",
|
||
" \"\"\"GELU activation function.\"\"\"\n",
|
||
" return Tensor(0.5 * x.data * (1 + np.tanh(np.sqrt(2 / np.pi) * (x.data + 0.044715 * x.data**3))))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "946c33e2",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 1. Introduction: What are Transformers?\n",
|
||
"\n",
|
||
"Transformers are the revolutionary architecture that powers modern AI language models like GPT, ChatGPT, and Claude. The key breakthrough is **self-attention**, which allows every token in a sequence to directly interact with every other token, creating rich contextual understanding.\n",
|
||
"\n",
|
||
"### The Transformer Revolution\n",
|
||
"\n",
|
||
"Before transformers, language models used RNNs or CNNs that processed text sequentially or locally. Transformers changed everything by processing all positions in parallel while maintaining global context.\n",
|
||
"\n",
|
||
"### Complete GPT Architecture Overview\n",
|
||
"\n",
|
||
"```\n",
|
||
"Input: \"Hello world\" → [Token IDs: 15496, 1917]\n",
|
||
" ↓\n",
|
||
" ┌─────────────────────────────────────────────┐\n",
|
||
" │ EMBEDDING LAYER │\n",
|
||
" │ ┌─────────────┐ ┌──────────────────────┐ │\n",
|
||
" │ │Token Embed │ + │ Positional Embed │ │\n",
|
||
" │ │[15496→vec] │ │[pos_0→vec, pos_1→vec]│ │\n",
|
||
" │ └─────────────┘ └──────────────────────┘ │\n",
|
||
" └─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
" ┌─────────────────────────────────────────────┐\n",
|
||
" │ TRANSFORMER BLOCK 1 │\n",
|
||
" │ │\n",
|
||
" │ Input → LayerNorm → MultiHeadAttention │\n",
|
||
" │ ↓ ↓ │\n",
|
||
" │ └────── Residual Add ←────┘ │\n",
|
||
" │ ↓ │\n",
|
||
" │ Result → LayerNorm → MLP (Feed Forward) │\n",
|
||
" │ ↓ ↓ │\n",
|
||
" │ └──── Residual Add ←──┘ │\n",
|
||
" └─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
" ┌─────────────────────────────────────────────┐\n",
|
||
" │ TRANSFORMER BLOCK 2 │\n",
|
||
" │ ... (same structure) │\n",
|
||
" └─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
" ... (more blocks)\n",
|
||
" ↓\n",
|
||
" ┌─────────────────────────────────────────────┐\n",
|
||
" │ OUTPUT HEAD │\n",
|
||
" │ Final LayerNorm → Linear → Vocabulary Logits│\n",
|
||
" └─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
"Output: [Prob(\"Hello\"), Prob(\"world\"), Prob(\"!\"), ...]\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Why Transformers Dominate\n",
|
||
"\n",
|
||
"**Parallel Processing**: Unlike RNNs that process tokens one by one, transformers process all positions simultaneously. This makes training much faster.\n",
|
||
"\n",
|
||
"**Global Context**: Every token can directly attend to every other token in the sequence, capturing long-range dependencies that RNNs struggle with.\n",
|
||
"\n",
|
||
"**Scalability**: Performance predictably improves with more parameters and data. This enabled the scaling laws that led to GPT-3, GPT-4, and beyond.\n",
|
||
"\n",
|
||
"**Residual Connections**: Allow training very deep networks (100+ layers) by providing gradient highways.\n",
|
||
"\n",
|
||
"### The Building Blocks We'll Implement\n",
|
||
"\n",
|
||
"1. **LayerNorm**: Stabilizes training by normalizing activations\n",
|
||
"2. **Multi-Layer Perceptron (MLP)**: Provides non-linear transformation\n",
|
||
"3. **TransformerBlock**: Combines attention + MLP with residuals\n",
|
||
"4. **GPT**: Complete model with embeddings and generation capability"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f8388844",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 2. Foundations: Essential Transformer Mathematics\n",
|
||
"\n",
|
||
"### Layer Normalization: The Stability Engine\n",
|
||
"\n",
|
||
"Layer Normalization is crucial for training deep transformer networks. Unlike batch normalization (which normalizes across the batch), layer norm normalizes across the feature dimension for each individual sample.\n",
|
||
"\n",
|
||
"```\n",
|
||
"Mathematical Formula:\n",
|
||
"output = (x - μ) / σ * γ + β\n",
|
||
"\n",
|
||
"where:\n",
|
||
" μ = mean(x, axis=features) # Mean across feature dimension\n",
|
||
" σ = sqrt(var(x) + ε) # Standard deviation + small epsilon\n",
|
||
" γ = learnable scale parameter # Initialized to 1.0\n",
|
||
" β = learnable shift parameter # Initialized to 0.0\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Why Layer Norm Works:**\n",
|
||
"- **Independence**: Each sample normalized independently (good for variable batch sizes)\n",
|
||
"- **Stability**: Prevents internal covariate shift that breaks training\n",
|
||
"- **Gradient Flow**: Helps gradients flow better through deep networks\n",
|
||
"\n",
|
||
"### Residual Connections: The Gradient Highway\n",
|
||
"\n",
|
||
"Residual connections are the secret to training deep networks. They create \"gradient highways\" that allow information to flow directly through the network.\n",
|
||
"\n",
|
||
"```\n",
|
||
"Residual Pattern in Transformers:\n",
|
||
"┌─────────────────────────────────────────────┐\n",
|
||
"│ Pre-Norm Architecture (Modern Standard): │\n",
|
||
"│ │\n",
|
||
"│ x → LayerNorm → MultiHeadAttention → + x │\n",
|
||
"│ │ ↑ │\n",
|
||
"│ │ residual connection │ │\n",
|
||
"│ └──────────────────────────────────────┘ │\n",
|
||
"│ │ │\n",
|
||
"│ x → LayerNorm → MLP → + x │\n",
|
||
"│ │ ↑ ↑ │\n",
|
||
"│ │ residual connection │ │\n",
|
||
"│ └───────────────────────┘ │\n",
|
||
"└─────────────────────────────────────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Gradient Flow Visualization:**\n",
|
||
"```\n",
|
||
"Backward Pass Without Residuals: With Residuals:\n",
|
||
"Loss Loss\n",
|
||
" │ gradients get smaller │ gradients stay strong\n",
|
||
" ↓ at each layer ↓ via residual paths\n",
|
||
"Layer N ← tiny gradients Layer N ← strong gradients\n",
|
||
" │ │ ↗ (direct path)\n",
|
||
" ↓ ↓ ↗\n",
|
||
"Layer 2 ← vanishing Layer 2 ← strong gradients\n",
|
||
" │ │ ↗\n",
|
||
" ↓ ↓ ↗\n",
|
||
"Layer 1 ← gone! Layer 1 ← strong gradients\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Feed-Forward Network (MLP): The Thinking Layer\n",
|
||
"\n",
|
||
"The MLP provides the actual \"thinking\" in each transformer block. It's a simple two-layer network with a specific expansion pattern.\n",
|
||
"\n",
|
||
"```\n",
|
||
"MLP Architecture:\n",
|
||
"Input (embed_dim) → Linear → GELU → Linear → Output (embed_dim)\n",
|
||
" 512 2048 2048 512\n",
|
||
" (4x expansion)\n",
|
||
"\n",
|
||
"Mathematical Formula:\n",
|
||
"FFN(x) = Linear₂(GELU(Linear₁(x)))\n",
|
||
" = W₂ · GELU(W₁ · x + b₁) + b₂\n",
|
||
"\n",
|
||
"where:\n",
|
||
" W₁: (embed_dim, 4*embed_dim) # Expansion matrix\n",
|
||
" W₂: (4*embed_dim, embed_dim) # Contraction matrix\n",
|
||
" GELU: smooth activation function (better than ReLU for language)\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Why 4x Expansion?**\n",
|
||
"- **Capacity**: More parameters = more representation power\n",
|
||
"- **Non-linearity**: GELU activation creates complex transformations\n",
|
||
"- **Information Bottleneck**: Forces the model to compress useful information\n",
|
||
"\n",
|
||
"### The Complete Transformer Block Data Flow\n",
|
||
"\n",
|
||
"```\n",
|
||
"Input Tensor (batch, seq_len, embed_dim)\n",
|
||
" ↓\n",
|
||
" ┌─────────────────────────────────────┐\n",
|
||
" │ ATTENTION SUB-LAYER │\n",
|
||
" │ │\n",
|
||
" │ x₁ = LayerNorm(x₀) │\n",
|
||
" │ attention_out = MultiHeadAttn(x₁) │\n",
|
||
" │ x₂ = x₀ + attention_out (residual) │\n",
|
||
" └─────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
" ┌─────────────────────────────────────┐\n",
|
||
" │ MLP SUB-LAYER │\n",
|
||
" │ │\n",
|
||
" │ x₃ = LayerNorm(x₂) │\n",
|
||
" │ mlp_out = MLP(x₃) │\n",
|
||
" │ x₄ = x₂ + mlp_out (residual) │\n",
|
||
" └─────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
"Output Tensor (batch, seq_len, embed_dim)\n",
|
||
"```\n",
|
||
"\n",
|
||
"**Key Insight**: Each sub-layer (attention and MLP) gets a \"clean\" normalized input but adds its contribution to the residual stream. This creates a stable training dynamic."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "aa924c73",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 3. Implementation: Building Transformer Components\n",
|
||
"\n",
|
||
"Now we'll implement each transformer component with a clear understanding of their role in the overall architecture. We'll follow the pattern: **Explanation → Implementation → Test** for each component.\n",
|
||
"\n",
|
||
"Each component serves a specific purpose:\n",
|
||
"- **LayerNorm**: Stabilizes training and normalizes activations\n",
|
||
"- **MLP**: Provides non-linear transformation and \"thinking\" capacity\n",
|
||
"- **TransformerBlock**: Combines attention with MLP using residual connections\n",
|
||
"- **GPT**: Complete autoregressive language model for text generation"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "3dc23c53",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### Understanding Layer Normalization\n",
|
||
"\n",
|
||
"Layer Normalization is the foundation of stable transformer training. Unlike batch normalization, it normalizes each sample independently across its feature dimensions.\n",
|
||
"\n",
|
||
"#### Why Layer Norm is Essential\n",
|
||
"\n",
|
||
"Without normalization, deep networks suffer from \"internal covariate shift\" - the distribution of inputs to each layer changes during training, making learning unstable.\n",
|
||
"\n",
|
||
"#### Layer Norm Visualization\n",
|
||
"\n",
|
||
"```\n",
|
||
"Input Tensor: (batch=2, seq=3, features=4)\n",
|
||
"┌──────────────────────────────────────────┐\n",
|
||
"│ Sample 1: [[1.0, 2.0, 3.0, 4.0], │\n",
|
||
"│ [5.0, 6.0, 7.0, 8.0], │\n",
|
||
"│ [9.0, 10., 11., 12.]] │\n",
|
||
"│ │\n",
|
||
"│ Sample 2: [[13., 14., 15., 16.], │\n",
|
||
"│ [17., 18., 19., 20.], │\n",
|
||
"│ [21., 22., 23., 24.]] │\n",
|
||
"└──────────────────────────────────────────┘\n",
|
||
" ↓ Layer Norm (across features for each position)\n",
|
||
"┌──────────────────────────────────────────┐\n",
|
||
"│ Each position normalized to mean=0, std=1│\n",
|
||
"│ Sample 1: [[-1.34, -0.45, 0.45, 1.34], │\n",
|
||
"│ [-1.34, -0.45, 0.45, 1.34], │\n",
|
||
"│ [-1.34, -0.45, 0.45, 1.34]] │\n",
|
||
"│ │\n",
|
||
"│ Sample 2: [[-1.34, -0.45, 0.45, 1.34], │\n",
|
||
"│ [-1.34, -0.45, 0.45, 1.34], │\n",
|
||
"│ [-1.34, -0.45, 0.45, 1.34]] │\n",
|
||
"└──────────────────────────────────────────┘\n",
|
||
" ↓ Apply learnable scale (γ) and shift (β)\n",
|
||
"┌──────────────────────────────────────────┐\n",
|
||
"│ Final Output: γ * normalized + β │\n",
|
||
"└──────────────────────────────────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Key Properties\n",
|
||
"- **Per-sample normalization**: Each sequence position normalized independently\n",
|
||
"- **Learnable parameters**: γ (scale) and β (shift) allow the model to recover any desired distribution\n",
|
||
"- **Gradient friendly**: Helps gradients flow smoothly through deep networks"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "4c26bf73",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "layer-norm",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"class LayerNorm:\n",
|
||
" \"\"\"\n",
|
||
" Layer Normalization for transformer blocks.\n",
|
||
"\n",
|
||
" Normalizes across the feature dimension (last axis) for each sample independently,\n",
|
||
" unlike batch normalization which normalizes across the batch dimension.\n",
|
||
" \"\"\"\n",
|
||
"\n",
|
||
" def __init__(self, normalized_shape, eps=1e-5):\n",
|
||
" \"\"\"\n",
|
||
" Initialize LayerNorm with learnable parameters.\n",
|
||
"\n",
|
||
" TODO: Set up normalization parameters\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Store the shape to normalize over (usually embed_dim)\n",
|
||
" 2. Initialize learnable scale (gamma) and shift (beta) parameters\n",
|
||
" 3. Set small epsilon for numerical stability\n",
|
||
"\n",
|
||
" EXAMPLE:\n",
|
||
" >>> ln = LayerNorm(512) # For 512-dimensional embeddings\n",
|
||
" >>> x = Tensor(np.random.randn(2, 10, 512)) # (batch, seq, features)\n",
|
||
" >>> normalized = ln.forward(x)\n",
|
||
" >>> # Each (2, 10) sample normalized independently across 512 features\n",
|
||
"\n",
|
||
" HINTS:\n",
|
||
" - gamma should start at 1.0 (identity scaling)\n",
|
||
" - beta should start at 0.0 (no shift)\n",
|
||
" - eps prevents division by zero in variance calculation\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" self.normalized_shape = normalized_shape\n",
|
||
" self.eps = eps\n",
|
||
"\n",
|
||
" # Learnable parameters: scale and shift\n",
|
||
" self.gamma = Tensor(np.ones(normalized_shape)) # Scale parameter\n",
|
||
" self.beta = Tensor(np.zeros(normalized_shape)) # Shift parameter\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def forward(self, x):\n",
|
||
" \"\"\"\n",
|
||
" Apply layer normalization.\n",
|
||
"\n",
|
||
" TODO: Implement layer normalization formula\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Compute mean and variance across the last dimension\n",
|
||
" 2. Normalize: (x - mean) / sqrt(variance + eps)\n",
|
||
" 3. Apply learnable scale and shift: gamma * normalized + beta\n",
|
||
"\n",
|
||
" MATHEMATICAL FORMULA:\n",
|
||
" y = (x - μ) / σ * γ + β\n",
|
||
" where μ = mean(x), σ = sqrt(var(x) + ε)\n",
|
||
"\n",
|
||
" HINT: Use keepdims=True to maintain tensor dimensions for broadcasting\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" # Compute statistics across last dimension (features)\n",
|
||
" mean = x.mean(axis=-1, keepdims=True)\n",
|
||
"\n",
|
||
" # Compute variance: E[(x - μ)²]\n",
|
||
" diff = Tensor(x.data - mean.data)\n",
|
||
" variance = Tensor((diff.data ** 2).mean(axis=-1, keepdims=True))\n",
|
||
"\n",
|
||
" # Normalize\n",
|
||
" std = Tensor(np.sqrt(variance.data + self.eps))\n",
|
||
" normalized = Tensor((x.data - mean.data) / std.data)\n",
|
||
"\n",
|
||
" # Apply learnable transformation\n",
|
||
" output = normalized * self.gamma + self.beta\n",
|
||
" return output\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" \"\"\"Return learnable parameters.\"\"\"\n",
|
||
" return [self.gamma, self.beta]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "33272f95",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### 🔬 Unit Test: Layer Normalization\n",
|
||
"This test validates our LayerNorm implementation works correctly.\n",
|
||
"**What we're testing**: Normalization statistics and parameter learning\n",
|
||
"**Why it matters**: Essential for transformer stability and training\n",
|
||
"**Expected**: Mean ≈ 0, std ≈ 1 after normalization, learnable parameters work"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "42c87208",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": true,
|
||
"grade_id": "test-layer-norm",
|
||
"locked": true,
|
||
"points": 10
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def test_unit_layer_norm():\n",
|
||
" \"\"\"🔬 Test LayerNorm implementation.\"\"\"\n",
|
||
" print(\"🔬 Unit Test: Layer Normalization...\")\n",
|
||
"\n",
|
||
" # Test basic normalization\n",
|
||
" ln = LayerNorm(4)\n",
|
||
" x = Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]) # (2, 4)\n",
|
||
"\n",
|
||
" normalized = ln.forward(x)\n",
|
||
"\n",
|
||
" # Check output shape\n",
|
||
" assert normalized.shape == (2, 4)\n",
|
||
"\n",
|
||
" # Check normalization properties (approximately)\n",
|
||
" # For each sample, mean should be close to 0, std close to 1\n",
|
||
" for i in range(2):\n",
|
||
" sample_mean = np.mean(normalized.data[i])\n",
|
||
" sample_std = np.std(normalized.data[i])\n",
|
||
" assert abs(sample_mean) < 1e-5, f\"Mean should be ~0, got {sample_mean}\"\n",
|
||
" assert abs(sample_std - 1.0) < 1e-4, f\"Std should be ~1, got {sample_std}\"\n",
|
||
"\n",
|
||
" # Test parameter shapes\n",
|
||
" params = ln.parameters()\n",
|
||
" assert len(params) == 2\n",
|
||
" assert params[0].shape == (4,) # gamma\n",
|
||
" assert params[1].shape == (4,) # beta\n",
|
||
"\n",
|
||
" print(\"✅ LayerNorm works correctly!\")\n",
|
||
"\n",
|
||
"test_unit_layer_norm()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4eb1e55a",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### Understanding the Multi-Layer Perceptron (MLP)\n",
|
||
"\n",
|
||
"The MLP is where the \"thinking\" happens in each transformer block. It's a simple feed-forward network that provides non-linear transformation capacity.\n",
|
||
"\n",
|
||
"#### The Role of MLP in Transformers\n",
|
||
"\n",
|
||
"While attention handles relationships between tokens, the MLP processes each position independently, adding computational depth and non-linearity.\n",
|
||
"\n",
|
||
"#### MLP Architecture and Information Flow\n",
|
||
"\n",
|
||
"```\n",
|
||
"Information Flow Through MLP:\n",
|
||
"\n",
|
||
"Input: (batch, seq_len, embed_dim=512)\n",
|
||
" ↓\n",
|
||
"┌─────────────────────────────────────────────┐\n",
|
||
"│ Linear Layer 1: Expansion │\n",
|
||
"│ Weight: (512, 2048) Bias: (2048,) │\n",
|
||
"│ Output: (batch, seq_len, 2048) │\n",
|
||
"└─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
"┌─────────────────────────────────────────────┐\n",
|
||
"│ GELU Activation │\n",
|
||
"│ Smooth, differentiable activation │\n",
|
||
"│ Better than ReLU for language modeling │\n",
|
||
"└─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
"┌─────────────────────────────────────────────┐\n",
|
||
"│ Linear Layer 2: Contraction │\n",
|
||
"│ Weight: (2048, 512) Bias: (512,) │\n",
|
||
"│ Output: (batch, seq_len, 512) │\n",
|
||
"└─────────────────────────────────────────────┘\n",
|
||
" ↓\n",
|
||
"Output: (batch, seq_len, embed_dim=512)\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Why 4x Expansion?\n",
|
||
"\n",
|
||
"```\n",
|
||
"Parameter Count Analysis:\n",
|
||
"\n",
|
||
"Embed Dim: 512\n",
|
||
"MLP Hidden: 2048 (4x expansion)\n",
|
||
"\n",
|
||
"Parameters:\n",
|
||
"- Linear1: 512 × 2048 + 2048 = 1,050,624\n",
|
||
"- Linear2: 2048 × 512 + 512 = 1,049,088\n",
|
||
"- Total MLP: ~2.1M parameters\n",
|
||
"\n",
|
||
"For comparison:\n",
|
||
"- Attention (same embed_dim): ~1.5M parameters\n",
|
||
"- MLP has MORE parameters → more computational capacity\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### GELU vs ReLU\n",
|
||
"\n",
|
||
"```\n",
|
||
"Activation Function Comparison:\n",
|
||
"\n",
|
||
"ReLU(x) = max(0, x) # Hard cutoff at 0\n",
|
||
" ┌────\n",
|
||
" │\n",
|
||
" ─────┘\n",
|
||
" 0\n",
|
||
"\n",
|
||
"GELU(x) ≈ x * Φ(x) # Smooth, probabilistic\n",
|
||
" ╭────\n",
|
||
" ╱\n",
|
||
" ───╱\n",
|
||
" ╱\n",
|
||
" 0\n",
|
||
"\n",
|
||
"GELU is smoother and provides better gradients for language modeling.\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c5acb8f3",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "mlp",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"class MLP:\n",
|
||
" \"\"\"\n",
|
||
" Multi-Layer Perceptron (Feed-Forward Network) for transformer blocks.\n",
|
||
"\n",
|
||
" Standard pattern: Linear -> GELU -> Linear with expansion ratio of 4:1.\n",
|
||
" This provides the non-linear transformation in each transformer block.\n",
|
||
" \"\"\"\n",
|
||
"\n",
|
||
" def __init__(self, embed_dim, hidden_dim=None, dropout_prob=0.1):\n",
|
||
" \"\"\"\n",
|
||
" Initialize MLP with two linear layers.\n",
|
||
"\n",
|
||
" TODO: Set up the feed-forward network layers\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. First layer expands from embed_dim to hidden_dim (usually 4x larger)\n",
|
||
" 2. Second layer projects back to embed_dim\n",
|
||
" 3. Use GELU activation (smoother than ReLU, preferred in transformers)\n",
|
||
"\n",
|
||
" EXAMPLE:\n",
|
||
" >>> mlp = MLP(512) # Will create 512 -> 2048 -> 512 network\n",
|
||
" >>> x = Tensor(np.random.randn(2, 10, 512))\n",
|
||
" >>> output = mlp.forward(x)\n",
|
||
" >>> assert output.shape == (2, 10, 512)\n",
|
||
"\n",
|
||
" HINT: Standard transformer MLP uses 4x expansion (hidden_dim = 4 * embed_dim)\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" if hidden_dim is None:\n",
|
||
" hidden_dim = 4 * embed_dim # Standard 4x expansion\n",
|
||
"\n",
|
||
" self.embed_dim = embed_dim\n",
|
||
" self.hidden_dim = hidden_dim\n",
|
||
"\n",
|
||
" # Two-layer feed-forward network\n",
|
||
" self.linear1 = Linear(embed_dim, hidden_dim)\n",
|
||
" self.linear2 = Linear(hidden_dim, embed_dim)\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def forward(self, x):\n",
|
||
" \"\"\"\n",
|
||
" Forward pass through MLP.\n",
|
||
"\n",
|
||
" TODO: Implement the feed-forward computation\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. First linear transformation: embed_dim -> hidden_dim\n",
|
||
" 2. Apply GELU activation (smooth, differentiable)\n",
|
||
" 3. Second linear transformation: hidden_dim -> embed_dim\n",
|
||
"\n",
|
||
" COMPUTATION FLOW:\n",
|
||
" x -> Linear -> GELU -> Linear -> output\n",
|
||
"\n",
|
||
" HINT: GELU activation is implemented above as a function\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" # First linear layer with expansion\n",
|
||
" hidden = self.linear1.forward(x)\n",
|
||
"\n",
|
||
" # GELU activation\n",
|
||
" hidden = gelu(hidden)\n",
|
||
"\n",
|
||
" # Second linear layer back to original size\n",
|
||
" output = self.linear2.forward(hidden)\n",
|
||
"\n",
|
||
" return output\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" \"\"\"Return all learnable parameters.\"\"\"\n",
|
||
" params = []\n",
|
||
" params.extend(self.linear1.parameters())\n",
|
||
" params.extend(self.linear2.parameters())\n",
|
||
" return params"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "054236fd",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### 🔬 Unit Test: MLP (Feed-Forward Network)\n",
|
||
"This test validates our MLP implementation works correctly.\n",
|
||
"**What we're testing**: Shape preservation and parameter counting\n",
|
||
"**Why it matters**: MLP provides the non-linear transformation in transformers\n",
|
||
"**Expected**: Input/output shapes match, correct parameter count"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8849696d",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": true,
|
||
"grade_id": "test-mlp",
|
||
"locked": true,
|
||
"points": 10
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def test_unit_mlp():\n",
|
||
" \"\"\"🔬 Test MLP implementation.\"\"\"\n",
|
||
" print(\"🔬 Unit Test: MLP (Feed-Forward Network)...\")\n",
|
||
"\n",
|
||
" # Test MLP with standard 4x expansion\n",
|
||
" embed_dim = 64\n",
|
||
" mlp = MLP(embed_dim)\n",
|
||
"\n",
|
||
" # Test forward pass\n",
|
||
" batch_size, seq_len = 2, 10\n",
|
||
" x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))\n",
|
||
" output = mlp.forward(x)\n",
|
||
"\n",
|
||
" # Check shape preservation\n",
|
||
" assert output.shape == (batch_size, seq_len, embed_dim)\n",
|
||
"\n",
|
||
" # Check hidden dimension is 4x\n",
|
||
" assert mlp.hidden_dim == 4 * embed_dim\n",
|
||
"\n",
|
||
" # Test parameter counting\n",
|
||
" params = mlp.parameters()\n",
|
||
" expected_params = 4 # 2 weights + 2 biases\n",
|
||
" assert len(params) == expected_params\n",
|
||
"\n",
|
||
" # Test custom hidden dimension\n",
|
||
" custom_mlp = MLP(embed_dim, hidden_dim=128)\n",
|
||
" assert custom_mlp.hidden_dim == 128\n",
|
||
"\n",
|
||
" print(\"✅ MLP works correctly!\")\n",
|
||
"\n",
|
||
"test_unit_mlp()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "dac755a4",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### Understanding the Complete Transformer Block\n",
|
||
"\n",
|
||
"The TransformerBlock is the core building unit of GPT and other transformer models. It combines self-attention with feed-forward processing using a carefully designed residual architecture.\n",
|
||
"\n",
|
||
"#### Pre-Norm vs Post-Norm Architecture\n",
|
||
"\n",
|
||
"Modern transformers use \"pre-norm\" architecture where LayerNorm comes BEFORE the sub-layers, not after. This provides better training stability.\n",
|
||
"\n",
|
||
"```\n",
|
||
"Pre-Norm Architecture (What We Implement):\n",
|
||
"┌─────────────────────────────────────────────────────────┐\n",
|
||
"│ INPUT (x) │\n",
|
||
"│ │ │\n",
|
||
"│ ┌───────────────┴───────────────┐ │\n",
|
||
"│ │ │ │\n",
|
||
"│ ▼ │ │\n",
|
||
"│ LayerNorm │ │\n",
|
||
"│ │ │ │\n",
|
||
"│ ▼ │ │\n",
|
||
"│ MultiHeadAttention │ │\n",
|
||
"│ │ │ │\n",
|
||
"│ └───────────────┬───────────────┘ │\n",
|
||
"│ │ (residual connection) │\n",
|
||
"│ ▼ │\n",
|
||
"│ x + attention │\n",
|
||
"│ │ │\n",
|
||
"│ ┌───────────────┴───────────────┐ │\n",
|
||
"│ │ │ │\n",
|
||
"│ ▼ │ │\n",
|
||
"│ LayerNorm │ │\n",
|
||
"│ │ │ │\n",
|
||
"│ ▼ │ │\n",
|
||
"│ MLP │ │\n",
|
||
"│ │ │ │\n",
|
||
"│ └───────────────┬───────────────┘ │\n",
|
||
"│ │ (residual connection) │\n",
|
||
"│ ▼ │\n",
|
||
"│ x + mlp │\n",
|
||
"│ │ │\n",
|
||
"│ ▼ │\n",
|
||
"│ OUTPUT │\n",
|
||
"└─────────────────────────────────────────────────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Why Pre-Norm Works Better\n",
|
||
"\n",
|
||
"**Training Stability**: LayerNorm before operations provides clean, normalized inputs to attention and MLP layers.\n",
|
||
"\n",
|
||
"**Gradient Flow**: Residual connections carry gradients directly from output to input, bypassing the normalized operations.\n",
|
||
"\n",
|
||
"**Deeper Networks**: Pre-norm enables training much deeper networks (100+ layers) compared to post-norm.\n",
|
||
"\n",
|
||
"#### Information Processing in Transformer Block\n",
|
||
"\n",
|
||
"```\n",
|
||
"Step-by-Step Data Transformation:\n",
|
||
"\n",
|
||
"1. Input Processing:\n",
|
||
" x₀: (batch, seq_len, embed_dim) # Original input\n",
|
||
"\n",
|
||
"2. Attention Sub-layer:\n",
|
||
" x₁ = LayerNorm(x₀) # Normalize input\n",
|
||
" attn_out = MultiHeadAttn(x₁) # Self-attention\n",
|
||
" x₂ = x₀ + attn_out # Residual connection\n",
|
||
"\n",
|
||
"3. MLP Sub-layer:\n",
|
||
" x₃ = LayerNorm(x₂) # Normalize again\n",
|
||
" mlp_out = MLP(x₃) # Feed-forward\n",
|
||
" x₄ = x₂ + mlp_out # Final residual\n",
|
||
"\n",
|
||
"4. Output:\n",
|
||
" return x₄ # Ready for next block\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Residual Stream Concept\n",
|
||
"\n",
|
||
"Think of the residual connections as a \"stream\" that carries information through the network:\n",
|
||
"\n",
|
||
"```\n",
|
||
"Residual Stream Flow:\n",
|
||
"\n",
|
||
"Layer 1: [original embeddings] ─┐\n",
|
||
" ├─→ + attention info ─┐\n",
|
||
"Attention adds information ──────┘ │\n",
|
||
" ├─→ + MLP info ─┐\n",
|
||
"MLP adds information ───────────────────────────────────┘ │\n",
|
||
" │\n",
|
||
"Layer 2: carries accumulated information ──────────────────────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"Each layer adds information to this stream rather than replacing it, creating a rich representation."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "6ad0f601",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "transformer-block",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"class TransformerBlock:\n",
|
||
" \"\"\"\n",
|
||
" Complete Transformer Block with self-attention, MLP, and residual connections.\n",
|
||
"\n",
|
||
" This is the core building block of GPT and other transformer models.\n",
|
||
" Each block processes the input sequence and passes it to the next block.\n",
|
||
" \"\"\"\n",
|
||
"\n",
|
||
" def __init__(self, embed_dim, num_heads, mlp_ratio=4, dropout_prob=0.1):\n",
|
||
" \"\"\"\n",
|
||
" Initialize a complete transformer block.\n",
|
||
"\n",
|
||
" TODO: Set up all components of the transformer block\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Multi-head self-attention for sequence modeling\n",
|
||
" 2. First layer normalization (pre-norm architecture)\n",
|
||
" 3. MLP with specified expansion ratio\n",
|
||
" 4. Second layer normalization\n",
|
||
"\n",
|
||
" TRANSFORMER BLOCK ARCHITECTURE:\n",
|
||
" x → LayerNorm → MultiHeadAttention → + (residual) →\n",
|
||
" LayerNorm → MLP → + (residual) → output\n",
|
||
"\n",
|
||
" EXAMPLE:\n",
|
||
" >>> block = TransformerBlock(embed_dim=512, num_heads=8)\n",
|
||
" >>> x = Tensor(np.random.randn(2, 10, 512)) # (batch, seq, embed)\n",
|
||
" >>> output = block.forward(x)\n",
|
||
" >>> assert output.shape == (2, 10, 512)\n",
|
||
"\n",
|
||
" HINT: We use pre-norm architecture (LayerNorm before attention/MLP)\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" self.embed_dim = embed_dim\n",
|
||
" self.num_heads = num_heads\n",
|
||
"\n",
|
||
" # Multi-head self-attention\n",
|
||
" self.attention = MultiHeadAttention(embed_dim, num_heads)\n",
|
||
"\n",
|
||
" # Layer normalizations (pre-norm architecture)\n",
|
||
" self.ln1 = LayerNorm(embed_dim) # Before attention\n",
|
||
" self.ln2 = LayerNorm(embed_dim) # Before MLP\n",
|
||
"\n",
|
||
" # Feed-forward network\n",
|
||
" hidden_dim = int(embed_dim * mlp_ratio)\n",
|
||
" self.mlp = MLP(embed_dim, hidden_dim)\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def forward(self, x, mask=None):\n",
|
||
" \"\"\"\n",
|
||
" Forward pass through transformer block.\n",
|
||
"\n",
|
||
" TODO: Implement the complete transformer block computation\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Apply layer norm, then self-attention, then add residual\n",
|
||
" 2. Apply layer norm, then MLP, then add residual\n",
|
||
" 3. Return the transformed sequence\n",
|
||
"\n",
|
||
" COMPUTATION FLOW:\n",
|
||
" x → ln1 → attention → + x → ln2 → mlp → + → output\n",
|
||
"\n",
|
||
" RESIDUAL CONNECTIONS:\n",
|
||
" These are crucial for training deep networks - they allow gradients\n",
|
||
" to flow directly through the network during backpropagation.\n",
|
||
"\n",
|
||
" HINT: Store intermediate results to add residual connections properly\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" # First sub-layer: Multi-head self-attention with residual connection\n",
|
||
" # Pre-norm: LayerNorm before attention\n",
|
||
" normed1 = self.ln1.forward(x)\n",
|
||
" attention_out = self.attention.forward(normed1, mask)\n",
|
||
"\n",
|
||
" # Residual connection\n",
|
||
" x = x + attention_out\n",
|
||
"\n",
|
||
" # Second sub-layer: MLP with residual connection\n",
|
||
" # Pre-norm: LayerNorm before MLP\n",
|
||
" normed2 = self.ln2.forward(x)\n",
|
||
" mlp_out = self.mlp.forward(normed2)\n",
|
||
"\n",
|
||
" # Residual connection\n",
|
||
" output = x + mlp_out\n",
|
||
"\n",
|
||
" return output\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" \"\"\"Return all learnable parameters.\"\"\"\n",
|
||
" params = []\n",
|
||
" params.extend(self.attention.parameters())\n",
|
||
" params.extend(self.ln1.parameters())\n",
|
||
" params.extend(self.ln2.parameters())\n",
|
||
" params.extend(self.mlp.parameters())\n",
|
||
" return params"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "736d101d",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### 🔬 Unit Test: Transformer Block\n",
|
||
"This test validates our complete TransformerBlock implementation.\n",
|
||
"**What we're testing**: Shape preservation, residual connections, parameter counting\n",
|
||
"**Why it matters**: This is the core component that will be stacked to create GPT\n",
|
||
"**Expected**: Input/output shapes match, all components work together"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "65540a0f",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": true,
|
||
"grade_id": "test-transformer-block",
|
||
"locked": true,
|
||
"points": 15
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def test_unit_transformer_block():\n",
|
||
" \"\"\"🔬 Test TransformerBlock implementation.\"\"\"\n",
|
||
" print(\"🔬 Unit Test: Transformer Block...\")\n",
|
||
"\n",
|
||
" # Test transformer block\n",
|
||
" embed_dim = 64\n",
|
||
" num_heads = 4\n",
|
||
" block = TransformerBlock(embed_dim, num_heads)\n",
|
||
"\n",
|
||
" # Test forward pass\n",
|
||
" batch_size, seq_len = 2, 8\n",
|
||
" x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))\n",
|
||
" output = block.forward(x)\n",
|
||
"\n",
|
||
" # Check shape preservation\n",
|
||
" assert output.shape == (batch_size, seq_len, embed_dim)\n",
|
||
"\n",
|
||
" # Test with causal mask (for autoregressive generation)\n",
|
||
" mask = Tensor(np.triu(np.ones((seq_len, seq_len)) * -np.inf, k=1))\n",
|
||
" masked_output = block.forward(x, mask)\n",
|
||
" assert masked_output.shape == (batch_size, seq_len, embed_dim)\n",
|
||
"\n",
|
||
" # Test parameter counting\n",
|
||
" params = block.parameters()\n",
|
||
" expected_components = 4 # attention, ln1, ln2, mlp parameters\n",
|
||
" assert len(params) > expected_components # Should have parameters from all components\n",
|
||
"\n",
|
||
" # Test different configurations\n",
|
||
" large_block = TransformerBlock(embed_dim=128, num_heads=8, mlp_ratio=2)\n",
|
||
" assert large_block.mlp.hidden_dim == 256 # 128 * 2\n",
|
||
"\n",
|
||
" print(\"✅ TransformerBlock works correctly!\")\n",
|
||
"\n",
|
||
"test_unit_transformer_block()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "17ad8926",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### Understanding the Complete GPT Architecture\n",
|
||
"\n",
|
||
"GPT (Generative Pre-trained Transformer) is the complete language model that combines all our components into a text generation system. It's designed for **autoregressive** generation - predicting the next token based on all previous tokens.\n",
|
||
"\n",
|
||
"#### GPT's Autoregressive Nature\n",
|
||
"\n",
|
||
"GPT generates text one token at a time, using all previously generated tokens as context:\n",
|
||
"\n",
|
||
"```\n",
|
||
"Autoregressive Generation Process:\n",
|
||
"\n",
|
||
"Step 1: \"The cat\" → model predicts → \"sat\"\n",
|
||
"Step 2: \"The cat sat\" → model predicts → \"on\"\n",
|
||
"Step 3: \"The cat sat on\" → model predicts → \"the\"\n",
|
||
"Step 4: \"The cat sat on the\" → model predicts → \"mat\"\n",
|
||
"\n",
|
||
"Result: \"The cat sat on the mat\"\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Complete GPT Architecture\n",
|
||
"\n",
|
||
"```\n",
|
||
"┌─────────────────────────────────────────────────────────────┐\n",
|
||
"│ GPT ARCHITECTURE │\n",
|
||
"│ │\n",
|
||
"│ Input: Token IDs [15496, 1917, ...] │\n",
|
||
"│ │ │\n",
|
||
"│ ┌──────────────────┴──────────────────┐ │\n",
|
||
"│ │ EMBEDDING LAYER │ │\n",
|
||
"│ │ ┌─────────────┐ ┌─────────────────┐│ │\n",
|
||
"│ │ │Token Embed │+│Position Embed ││ │\n",
|
||
"│ │ │vocab→vector ││ │sequence→vector ││ │\n",
|
||
"│ │ └─────────────┘ └─────────────────┘│ │\n",
|
||
"│ └──────────────────┬──────────────────┘ │\n",
|
||
"│ │ │\n",
|
||
"│ ┌──────────────────┴──────────────────┐ │\n",
|
||
"│ │ TRANSFORMER BLOCK 1 │ │\n",
|
||
"│ │ ┌─────────┐ ┌─────────┐ ┌───────┐ │ │\n",
|
||
"│ │ │LayerNorm│→│Attention│→│ +x │ │ │\n",
|
||
"│ │ └─────────┘ └─────────┘ └───┬───┘ │ │\n",
|
||
"│ │ │ │ │\n",
|
||
"│ │ ┌─────────┐ ┌─────────┐ ┌───▼───┐ │ │\n",
|
||
"│ │ │LayerNorm│→│ MLP │→│ +x │ │ │\n",
|
||
"│ │ └─────────┘ └─────────┘ └───────┘ │ │\n",
|
||
"│ └──────────────────┬──────────────────┘ │\n",
|
||
"│ │ │\n",
|
||
"│ ... (more transformer blocks) ... │\n",
|
||
"│ │ │\n",
|
||
"│ ┌──────────────────┴──────────────────┐ │\n",
|
||
"│ │ OUTPUT HEAD │ │\n",
|
||
"│ │ ┌─────────┐ ┌─────────────────────┐ │ │\n",
|
||
"│ │ │LayerNorm│→│Linear(embed→vocab) │ │ │\n",
|
||
"│ │ └─────────┘ └─────────────────────┘ │ │\n",
|
||
"│ └──────────────────┬──────────────────┘ │\n",
|
||
"│ │ │\n",
|
||
"│ Output: Vocabulary Logits [0.1, 0.05, 0.8, ...] │\n",
|
||
"└─────────────────────────────────────────────────────────────┘\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Causal Masking for Autoregressive Training\n",
|
||
"\n",
|
||
"During training, GPT sees the entire sequence but must not \"cheat\" by looking at future tokens:\n",
|
||
"\n",
|
||
"```\n",
|
||
"Causal Attention Mask:\n",
|
||
"\n",
|
||
"Sequence: [\"The\", \"cat\", \"sat\", \"on\"]\n",
|
||
"Positions: 0 1 2 3\n",
|
||
"\n",
|
||
"Attention Matrix (what each position can see):\n",
|
||
" 0 1 2 3\n",
|
||
" 0 [ ✓ ✗ ✗ ✗ ] # \"The\" only sees itself\n",
|
||
" 1 [ ✓ ✓ ✗ ✗ ] # \"cat\" sees \"The\" and itself\n",
|
||
" 2 [ ✓ ✓ ✓ ✗ ] # \"sat\" sees \"The\", \"cat\", itself\n",
|
||
" 3 [ ✓ ✓ ✓ ✓ ] # \"on\" sees all previous tokens\n",
|
||
"\n",
|
||
"Implementation: Upper triangular matrix with -∞\n",
|
||
"[[ 0, -∞, -∞, -∞],\n",
|
||
" [ 0, 0, -∞, -∞],\n",
|
||
" [ 0, 0, 0, -∞],\n",
|
||
" [ 0, 0, 0, 0]]\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Generation Temperature Control\n",
|
||
"\n",
|
||
"Temperature controls the randomness of generation:\n",
|
||
"\n",
|
||
"```\n",
|
||
"Temperature Effects:\n",
|
||
"\n",
|
||
"Original logits: [1.0, 2.0, 3.0]\n",
|
||
"\n",
|
||
"Temperature = 0.1 (Conservative):\n",
|
||
"Scaled: [10.0, 20.0, 30.0] → Sharp distribution\n",
|
||
"Probs: [0.00, 0.00, 1.00] → Always picks highest\n",
|
||
"\n",
|
||
"Temperature = 1.0 (Balanced):\n",
|
||
"Scaled: [1.0, 2.0, 3.0] → Moderate distribution\n",
|
||
"Probs: [0.09, 0.24, 0.67] → Weighted sampling\n",
|
||
"\n",
|
||
"Temperature = 2.0 (Creative):\n",
|
||
"Scaled: [0.5, 1.0, 1.5] → Flatter distribution\n",
|
||
"Probs: [0.18, 0.33, 0.49] → More random\n",
|
||
"```\n",
|
||
"\n",
|
||
"#### Model Scaling and Parameters\n",
|
||
"\n",
|
||
"```\n",
|
||
"GPT Model Size Scaling:\n",
|
||
"\n",
|
||
"Tiny GPT (our implementation):\n",
|
||
"- embed_dim: 64, layers: 2, heads: 4\n",
|
||
"- Parameters: ~50K\n",
|
||
"- Use case: Learning and experimentation\n",
|
||
"\n",
|
||
"GPT-2 Small:\n",
|
||
"- embed_dim: 768, layers: 12, heads: 12\n",
|
||
"- Parameters: 117M\n",
|
||
"- Use case: Basic text generation\n",
|
||
"\n",
|
||
"GPT-3:\n",
|
||
"- embed_dim: 12,288, layers: 96, heads: 96\n",
|
||
"- Parameters: 175B\n",
|
||
"- Use case: Advanced language understanding\n",
|
||
"\n",
|
||
"GPT-4 (estimated):\n",
|
||
"- embed_dim: ~16,384, layers: ~120, heads: ~128\n",
|
||
"- Parameters: ~1.7T\n",
|
||
"- Use case: Reasoning and multimodal tasks\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "586f2e46",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "gpt",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"class GPT:\n",
|
||
" \"\"\"\n",
|
||
" Complete GPT (Generative Pre-trained Transformer) model.\n",
|
||
"\n",
|
||
" This combines embeddings, positional encoding, multiple transformer blocks,\n",
|
||
" and a language modeling head for text generation.\n",
|
||
" \"\"\"\n",
|
||
"\n",
|
||
" def __init__(self, vocab_size, embed_dim, num_layers, num_heads, max_seq_len=1024):\n",
|
||
" \"\"\"\n",
|
||
" Initialize complete GPT model.\n",
|
||
"\n",
|
||
" TODO: Set up all components of the GPT architecture\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Token embedding layer to convert tokens to vectors\n",
|
||
" 2. Positional embedding to add position information\n",
|
||
" 3. Stack of transformer blocks (the main computation)\n",
|
||
" 4. Final layer norm and language modeling head\n",
|
||
"\n",
|
||
" GPT ARCHITECTURE:\n",
|
||
" tokens → embedding → + pos_embedding →\n",
|
||
" transformer_blocks → layer_norm → lm_head → logits\n",
|
||
"\n",
|
||
" EXAMPLE:\n",
|
||
" >>> model = GPT(vocab_size=1000, embed_dim=256, num_layers=6, num_heads=8)\n",
|
||
" >>> tokens = Tensor(np.random.randint(0, 1000, (2, 10))) # (batch, seq)\n",
|
||
" >>> logits = model.forward(tokens)\n",
|
||
" >>> assert logits.shape == (2, 10, 1000) # (batch, seq, vocab)\n",
|
||
"\n",
|
||
" HINTS:\n",
|
||
" - Positional embeddings are learned, not fixed sinusoidal\n",
|
||
" - Final layer norm stabilizes training\n",
|
||
" - Language modeling head shares weights with token embedding (tie_weights)\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" self.vocab_size = vocab_size\n",
|
||
" self.embed_dim = embed_dim\n",
|
||
" self.num_layers = num_layers\n",
|
||
" self.num_heads = num_heads\n",
|
||
" self.max_seq_len = max_seq_len\n",
|
||
"\n",
|
||
" # Token and positional embeddings\n",
|
||
" self.token_embedding = Embedding(vocab_size, embed_dim)\n",
|
||
" self.position_embedding = Embedding(max_seq_len, embed_dim)\n",
|
||
"\n",
|
||
" # Stack of transformer blocks\n",
|
||
" self.blocks = []\n",
|
||
" for _ in range(num_layers):\n",
|
||
" block = TransformerBlock(embed_dim, num_heads)\n",
|
||
" self.blocks.append(block)\n",
|
||
"\n",
|
||
" # Final layer normalization\n",
|
||
" self.ln_f = LayerNorm(embed_dim)\n",
|
||
"\n",
|
||
" # Language modeling head (projects to vocabulary)\n",
|
||
" self.lm_head = Linear(embed_dim, vocab_size, bias=False)\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def forward(self, tokens):\n",
|
||
" \"\"\"\n",
|
||
" Forward pass through GPT model.\n",
|
||
"\n",
|
||
" TODO: Implement the complete GPT forward pass\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Get token embeddings and positional embeddings\n",
|
||
" 2. Add them together (broadcasting handles different shapes)\n",
|
||
" 3. Pass through all transformer blocks sequentially\n",
|
||
" 4. Apply final layer norm and language modeling head\n",
|
||
"\n",
|
||
" COMPUTATION FLOW:\n",
|
||
" tokens → embed + pos_embed → blocks → ln_f → lm_head → logits\n",
|
||
"\n",
|
||
" CAUSAL MASKING:\n",
|
||
" For autoregressive generation, we need to prevent tokens from\n",
|
||
" seeing future tokens. This is handled by the attention mask.\n",
|
||
"\n",
|
||
" HINT: Create position indices as range(seq_len) for positional embedding\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" batch_size, seq_len = tokens.shape\n",
|
||
"\n",
|
||
" # Token embeddings\n",
|
||
" token_emb = self.token_embedding.forward(tokens)\n",
|
||
"\n",
|
||
" # Positional embeddings\n",
|
||
" positions = Tensor(np.arange(seq_len).reshape(1, seq_len))\n",
|
||
" pos_emb = self.position_embedding.forward(positions)\n",
|
||
"\n",
|
||
" # Combine embeddings\n",
|
||
" x = token_emb + pos_emb\n",
|
||
"\n",
|
||
" # Create causal mask for autoregressive generation\n",
|
||
" mask = self._create_causal_mask(seq_len)\n",
|
||
"\n",
|
||
" # Pass through transformer blocks\n",
|
||
" for block in self.blocks:\n",
|
||
" x = block.forward(x, mask)\n",
|
||
"\n",
|
||
" # Final layer normalization\n",
|
||
" x = self.ln_f.forward(x)\n",
|
||
"\n",
|
||
" # Language modeling head\n",
|
||
" logits = self.lm_head.forward(x)\n",
|
||
"\n",
|
||
" return logits\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def _create_causal_mask(self, seq_len):\n",
|
||
" \"\"\"Create causal mask to prevent attending to future positions.\"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" # Upper triangular matrix filled with -inf\n",
|
||
" mask = np.triu(np.ones((seq_len, seq_len)) * -np.inf, k=1)\n",
|
||
" return Tensor(mask)\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def generate(self, prompt_tokens, max_new_tokens=50, temperature=1.0):\n",
|
||
" \"\"\"\n",
|
||
" Generate text autoregressively.\n",
|
||
"\n",
|
||
" TODO: Implement autoregressive text generation\n",
|
||
"\n",
|
||
" APPROACH:\n",
|
||
" 1. Start with prompt tokens\n",
|
||
" 2. For each new position:\n",
|
||
" - Run forward pass to get logits\n",
|
||
" - Sample next token from logits\n",
|
||
" - Append to sequence\n",
|
||
" 3. Return generated sequence\n",
|
||
"\n",
|
||
" AUTOREGRESSIVE GENERATION:\n",
|
||
" At each step, the model predicts the next token based on all\n",
|
||
" previous tokens. This is how GPT generates coherent text.\n",
|
||
"\n",
|
||
" EXAMPLE:\n",
|
||
" >>> model = GPT(vocab_size=100, embed_dim=64, num_layers=2, num_heads=4)\n",
|
||
" >>> prompt = Tensor([[1, 2, 3]]) # Some token sequence\n",
|
||
" >>> generated = model.generate(prompt, max_new_tokens=5)\n",
|
||
" >>> assert generated.shape[1] == 3 + 5 # original + new tokens\n",
|
||
"\n",
|
||
" HINT: Use np.random.choice with temperature for sampling\n",
|
||
" \"\"\"\n",
|
||
" ### BEGIN SOLUTION\n",
|
||
" current_tokens = Tensor(prompt_tokens.data.copy())\n",
|
||
"\n",
|
||
" for _ in range(max_new_tokens):\n",
|
||
" # Get logits for current sequence\n",
|
||
" logits = self.forward(current_tokens)\n",
|
||
"\n",
|
||
" # Get logits for last position (next token prediction)\n",
|
||
" last_logits = logits.data[:, -1, :] # (batch_size, vocab_size)\n",
|
||
"\n",
|
||
" # Apply temperature scaling\n",
|
||
" scaled_logits = last_logits / temperature\n",
|
||
"\n",
|
||
" # Convert to probabilities (softmax)\n",
|
||
" exp_logits = np.exp(scaled_logits - np.max(scaled_logits, axis=-1, keepdims=True))\n",
|
||
" probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)\n",
|
||
"\n",
|
||
" # Sample next token\n",
|
||
" next_token = np.array([[np.random.choice(self.vocab_size, p=probs[0])]])\n",
|
||
"\n",
|
||
" # Append to sequence\n",
|
||
" current_tokens = Tensor(np.concatenate([current_tokens.data, next_token], axis=1))\n",
|
||
"\n",
|
||
" return current_tokens\n",
|
||
" ### END SOLUTION\n",
|
||
"\n",
|
||
" def parameters(self):\n",
|
||
" \"\"\"Return all learnable parameters.\"\"\"\n",
|
||
" params = []\n",
|
||
" params.extend(self.token_embedding.parameters())\n",
|
||
" params.extend(self.position_embedding.parameters())\n",
|
||
"\n",
|
||
" for block in self.blocks:\n",
|
||
" params.extend(block.parameters())\n",
|
||
"\n",
|
||
" params.extend(self.ln_f.parameters())\n",
|
||
" params.extend(self.lm_head.parameters())\n",
|
||
"\n",
|
||
" return params"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "c9a7758f",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"### 🔬 Unit Test: GPT Model\n",
|
||
"This test validates our complete GPT implementation.\n",
|
||
"**What we're testing**: Model forward pass, shape consistency, generation capability\n",
|
||
"**Why it matters**: This is the complete language model that ties everything together\n",
|
||
"**Expected**: Correct output shapes, generation works, parameter counting"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e4ba240a",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": true,
|
||
"grade_id": "test-gpt",
|
||
"locked": true,
|
||
"points": 20
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def test_unit_gpt():\n",
|
||
" \"\"\"🔬 Test GPT model implementation.\"\"\"\n",
|
||
" print(\"🔬 Unit Test: GPT Model...\")\n",
|
||
"\n",
|
||
" # Test small GPT model\n",
|
||
" vocab_size = 100\n",
|
||
" embed_dim = 64\n",
|
||
" num_layers = 2\n",
|
||
" num_heads = 4\n",
|
||
"\n",
|
||
" model = GPT(vocab_size, embed_dim, num_layers, num_heads)\n",
|
||
"\n",
|
||
" # Test forward pass\n",
|
||
" batch_size, seq_len = 2, 8\n",
|
||
" tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))\n",
|
||
" logits = model.forward(tokens)\n",
|
||
"\n",
|
||
" # Check output shape\n",
|
||
" expected_shape = (batch_size, seq_len, vocab_size)\n",
|
||
" assert logits.shape == expected_shape\n",
|
||
"\n",
|
||
" # Test generation\n",
|
||
" prompt = Tensor(np.random.randint(0, vocab_size, (1, 5)))\n",
|
||
" generated = model.generate(prompt, max_new_tokens=3)\n",
|
||
"\n",
|
||
" # Check generation shape\n",
|
||
" assert generated.shape == (1, 8) # 5 prompt + 3 new tokens\n",
|
||
"\n",
|
||
" # Test parameter counting\n",
|
||
" params = model.parameters()\n",
|
||
" assert len(params) > 10 # Should have many parameters from all components\n",
|
||
"\n",
|
||
" # Test different model sizes\n",
|
||
" larger_model = GPT(vocab_size=200, embed_dim=128, num_layers=4, num_heads=8)\n",
|
||
" test_tokens = Tensor(np.random.randint(0, 200, (1, 10)))\n",
|
||
" larger_logits = larger_model.forward(test_tokens)\n",
|
||
" assert larger_logits.shape == (1, 10, 200)\n",
|
||
"\n",
|
||
" print(\"✅ GPT model works correctly!\")\n",
|
||
"\n",
|
||
"test_unit_gpt()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "1ecc1961",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 4. Integration: Complete Transformer Workflow\n",
|
||
"\n",
|
||
"Now that we've built all the components, let's see how they work together in a complete language modeling pipeline. This demonstrates the full power of the transformer architecture.\n",
|
||
"\n",
|
||
"### The Language Modeling Pipeline\n",
|
||
"\n",
|
||
"```\n",
|
||
"Complete Workflow Visualization:\n",
|
||
"\n",
|
||
"1. Text Input:\n",
|
||
" \"hello world\" → Tokenization → [15496, 1917]\n",
|
||
"\n",
|
||
"2. Model Processing:\n",
|
||
" [15496, 1917]\n",
|
||
" ↓ Token Embedding\n",
|
||
" [[0.1, 0.5, ...], [0.3, -0.2, ...]] # Vector representations\n",
|
||
" ↓ + Position Embedding\n",
|
||
" [[0.2, 0.7, ...], [0.1, -0.4, ...]] # With position info\n",
|
||
" ↓ Transformer Block 1\n",
|
||
" [[0.3, 0.2, ...], [0.5, -0.1, ...]] # After attention + MLP\n",
|
||
" ↓ Transformer Block 2\n",
|
||
" [[0.1, 0.9, ...], [0.7, 0.3, ...]] # Further processed\n",
|
||
" ↓ Final LayerNorm + LM Head\n",
|
||
" [[0.1, 0.05, 0.8, ...], [...]] # Probability over vocab\n",
|
||
"\n",
|
||
"3. Generation:\n",
|
||
" Model predicts next token: \"!\" (token 33)\n",
|
||
" New sequence: \"hello world!\"\n",
|
||
"```\n",
|
||
"\n",
|
||
"This integration demo will show:\n",
|
||
"- **Character-level tokenization** for simplicity\n",
|
||
"- **Forward pass** through all components\n",
|
||
"- **Autoregressive generation** in action\n",
|
||
"- **Temperature effects** on creativity"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "04f8fd5c",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "integration-demo",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def demonstrate_transformer_integration():\n",
|
||
" \"\"\"\n",
|
||
" Demonstrate complete transformer pipeline.\n",
|
||
"\n",
|
||
" This simulates training a small language model on a simple vocabulary.\n",
|
||
" \"\"\"\n",
|
||
" print(\"🔗 Integration Demo: Complete Language Model Pipeline\")\n",
|
||
" print(\"Building a mini-GPT for character-level text generation\")\n",
|
||
"\n",
|
||
" # Create a small vocabulary (character-level)\n",
|
||
" vocab = list(\"abcdefghijklmnopqrstuvwxyz .\")\n",
|
||
" vocab_size = len(vocab)\n",
|
||
" char_to_idx = {char: i for i, char in enumerate(vocab)}\n",
|
||
" idx_to_char = {i: char for i, char in enumerate(vocab)}\n",
|
||
"\n",
|
||
" print(f\"Vocabulary size: {vocab_size}\")\n",
|
||
" print(f\"Characters: {''.join(vocab)}\")\n",
|
||
"\n",
|
||
" # Create model\n",
|
||
" model = GPT(\n",
|
||
" vocab_size=vocab_size,\n",
|
||
" embed_dim=64,\n",
|
||
" num_layers=2,\n",
|
||
" num_heads=4,\n",
|
||
" max_seq_len=32\n",
|
||
" )\n",
|
||
"\n",
|
||
" # Sample text encoding\n",
|
||
" text = \"hello world.\"\n",
|
||
" tokens = [char_to_idx[char] for char in text]\n",
|
||
" input_tokens = Tensor(np.array([tokens]))\n",
|
||
"\n",
|
||
" print(f\"\\nOriginal text: '{text}'\")\n",
|
||
" print(f\"Tokenized: {tokens}\")\n",
|
||
" print(f\"Input shape: {input_tokens.shape}\")\n",
|
||
"\n",
|
||
" # Forward pass\n",
|
||
" logits = model.forward(input_tokens)\n",
|
||
" print(f\"Output logits shape: {logits.shape}\")\n",
|
||
" print(f\"Each position predicts next token from {vocab_size} possibilities\")\n",
|
||
"\n",
|
||
" # Generation demo\n",
|
||
" prompt_text = \"hello\"\n",
|
||
" prompt_tokens = [char_to_idx[char] for char in prompt_text]\n",
|
||
" prompt = Tensor(np.array([prompt_tokens]))\n",
|
||
"\n",
|
||
" print(f\"\\nGeneration demo:\")\n",
|
||
" print(f\"Prompt: '{prompt_text}'\")\n",
|
||
"\n",
|
||
" generated = model.generate(prompt, max_new_tokens=8, temperature=1.0)\n",
|
||
" generated_text = ''.join([idx_to_char[idx] for idx in generated.data[0]])\n",
|
||
"\n",
|
||
" print(f\"Generated: '{generated_text}'\")\n",
|
||
" print(\"(Note: Untrained model produces random text)\")\n",
|
||
"\n",
|
||
" return model\n",
|
||
"\n",
|
||
"demonstrate_transformer_integration()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "0c53c926",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 5. Systems Analysis: Parameter Scaling and Memory\n",
|
||
"\n",
|
||
"Transformer models scale dramatically with size, leading to both opportunities and challenges. Let's analyze the computational and memory requirements to understand why training large language models requires massive infrastructure.\n",
|
||
"\n",
|
||
"### The Scaling Laws Revolution\n",
|
||
"\n",
|
||
"One of the key discoveries in modern AI is that transformer performance follows predictable scaling laws:\n",
|
||
"\n",
|
||
"```\n",
|
||
"Scaling Laws Pattern:\n",
|
||
"Performance ∝ Parameters^α × Data^β × Compute^γ\n",
|
||
"\n",
|
||
"where α ≈ 0.7, β ≈ 0.8, γ ≈ 0.5\n",
|
||
"\n",
|
||
"This means:\n",
|
||
"- 10× more parameters → ~5× better performance\n",
|
||
"- 10× more data → ~6× better performance\n",
|
||
"- 10× more compute → ~3× better performance\n",
|
||
"```\n",
|
||
"\n",
|
||
"### Memory Scaling Analysis\n",
|
||
"\n",
|
||
"Memory requirements grow in different ways for different components:\n",
|
||
"\n",
|
||
"```\n",
|
||
"Memory Scaling by Component:\n",
|
||
"\n",
|
||
"1. Parameter Memory (Linear with model size):\n",
|
||
" - Embeddings: vocab_size × embed_dim\n",
|
||
" - Transformer blocks: ~4 × embed_dim²\n",
|
||
" - Total: O(embed_dim²)\n",
|
||
"\n",
|
||
"2. Attention Memory (Quadratic with sequence length):\n",
|
||
" - Attention matrices: batch × heads × seq_len²\n",
|
||
" - This is why long context is expensive!\n",
|
||
" - Total: O(seq_len²)\n",
|
||
"\n",
|
||
"3. Activation Memory (Linear with batch size):\n",
|
||
" - Forward pass activations for backprop\n",
|
||
" - Scales with: batch × seq_len × embed_dim\n",
|
||
" - Total: O(batch_size)\n",
|
||
"```\n",
|
||
"\n",
|
||
"### The Attention Memory Wall\n",
|
||
"\n",
|
||
"```\n",
|
||
"Attention Memory Wall Visualization:\n",
|
||
"\n",
|
||
"Sequence Length vs Memory Usage:\n",
|
||
"\n",
|
||
"1K tokens: [▓] 16 MB # Manageable\n",
|
||
"2K tokens: [▓▓▓▓] 64 MB # 4× memory (quadratic!)\n",
|
||
"4K tokens: [▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓] 256 MB # 16× memory\n",
|
||
"8K tokens: [████████████████████████████████] 1 GB # 64× memory\n",
|
||
"16K tokens: ████████████████████████████████████████████████████████████████ 4 GB\n",
|
||
"32K tokens: ████████████████████████████████████████████████████████████████████████████████████████████████████████████████ 16 GB\n",
|
||
"\n",
|
||
"This is why:\n",
|
||
"- GPT-3 context: 2K tokens\n",
|
||
"- GPT-4 context: 8K tokens (32K in turbo)\n",
|
||
"- Claude-3: 200K tokens (requires special techniques!)\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "039199a8",
|
||
"metadata": {
|
||
"lines_to_next_cell": 1,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "analyze-scaling",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def analyze_parameter_scaling():\n",
|
||
" \"\"\"📊 Analyze how parameter count scales with model dimensions.\"\"\"\n",
|
||
" print(\"📊 Analyzing Parameter Scaling in Transformers...\")\n",
|
||
" print(\"Understanding why model size affects performance and cost\\n\")\n",
|
||
"\n",
|
||
" # Test different model sizes\n",
|
||
" configs = [\n",
|
||
" {\"name\": \"Tiny\", \"embed_dim\": 64, \"num_layers\": 2, \"num_heads\": 4},\n",
|
||
" {\"name\": \"Small\", \"embed_dim\": 128, \"num_layers\": 4, \"num_heads\": 8},\n",
|
||
" {\"name\": \"Medium\", \"embed_dim\": 256, \"num_layers\": 8, \"num_heads\": 16},\n",
|
||
" {\"name\": \"Large\", \"embed_dim\": 512, \"num_layers\": 12, \"num_heads\": 16},\n",
|
||
" ]\n",
|
||
"\n",
|
||
" vocab_size = 50000 # Typical vocabulary size\n",
|
||
"\n",
|
||
" for config in configs:\n",
|
||
" model = GPT(\n",
|
||
" vocab_size=vocab_size,\n",
|
||
" embed_dim=config[\"embed_dim\"],\n",
|
||
" num_layers=config[\"num_layers\"],\n",
|
||
" num_heads=config[\"num_heads\"]\n",
|
||
" )\n",
|
||
"\n",
|
||
" # Count parameters\n",
|
||
" total_params = 0\n",
|
||
" for param in model.parameters():\n",
|
||
" total_params += param.size\n",
|
||
"\n",
|
||
" # Calculate memory requirements (4 bytes per float32 parameter)\n",
|
||
" memory_mb = (total_params * 4) / (1024 * 1024)\n",
|
||
"\n",
|
||
" print(f\"{config['name']} Model:\")\n",
|
||
" print(f\" Parameters: {total_params:,}\")\n",
|
||
" print(f\" Memory: {memory_mb:.1f} MB\")\n",
|
||
" print(f\" Embed dim: {config['embed_dim']}, Layers: {config['num_layers']}\")\n",
|
||
" print()\n",
|
||
"\n",
|
||
" print(\"💡 Parameter scaling is roughly quadratic with embedding dimension\")\n",
|
||
" print(\"🚀 Real GPT-3 has 175B parameters, requiring ~350GB memory!\")\n",
|
||
"\n",
|
||
"analyze_parameter_scaling()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "a249a5a0",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "analyze-attention-memory",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def analyze_attention_memory():\n",
|
||
" \"\"\"📊 Analyze attention memory complexity with sequence length.\"\"\"\n",
|
||
" print(\"📊 Analyzing Attention Memory Complexity...\")\n",
|
||
" print(\"Why long context is expensive and how it scales\\n\")\n",
|
||
"\n",
|
||
" embed_dim = 512\n",
|
||
" num_heads = 8\n",
|
||
" batch_size = 4\n",
|
||
"\n",
|
||
" # Test different sequence lengths\n",
|
||
" sequence_lengths = [128, 256, 512, 1024, 2048]\n",
|
||
"\n",
|
||
" print(\"Attention Matrix Memory Usage:\")\n",
|
||
" print(\"Seq Len | Attention Matrix Size | Memory (MB)\")\n",
|
||
" print(\"-\" * 45)\n",
|
||
"\n",
|
||
" for seq_len in sequence_lengths:\n",
|
||
" # Attention matrix is (batch_size, num_heads, seq_len, seq_len)\n",
|
||
" attention_elements = batch_size * num_heads * seq_len * seq_len\n",
|
||
"\n",
|
||
" # 4 bytes per float32\n",
|
||
" memory_bytes = attention_elements * 4\n",
|
||
" memory_mb = memory_bytes / (1024 * 1024)\n",
|
||
"\n",
|
||
" print(f\"{seq_len:6d} | {seq_len}×{seq_len} × {batch_size}×{num_heads} | {memory_mb:8.1f}\")\n",
|
||
"\n",
|
||
" print()\n",
|
||
" print(\"💡 Attention memory grows quadratically with sequence length\")\n",
|
||
" print(\"🚀 This is why techniques like FlashAttention are crucial for long sequences\")\n",
|
||
"\n",
|
||
"analyze_attention_memory()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "253b8e90",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\"",
|
||
"lines_to_next_cell": 1
|
||
},
|
||
"source": [
|
||
"## 🧪 Module Integration Test\n",
|
||
"\n",
|
||
"Final validation that everything works together correctly."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "d9431f80",
|
||
"metadata": {
|
||
"nbgrader": {
|
||
"grade": true,
|
||
"grade_id": "module-integration",
|
||
"locked": true,
|
||
"points": 25
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": [
|
||
"def test_module():\n",
|
||
" \"\"\"\n",
|
||
" Comprehensive test of entire module functionality.\n",
|
||
"\n",
|
||
" This final test runs before module summary to ensure:\n",
|
||
" - All unit tests pass\n",
|
||
" - Functions work together correctly\n",
|
||
" - Module is ready for integration with TinyTorch\n",
|
||
" \"\"\"\n",
|
||
" print(\"🧪 RUNNING MODULE INTEGRATION TEST\")\n",
|
||
" print(\"=\" * 50)\n",
|
||
"\n",
|
||
" # Run all unit tests\n",
|
||
" print(\"Running unit tests...\")\n",
|
||
" test_unit_layer_norm()\n",
|
||
" test_unit_mlp()\n",
|
||
" test_unit_transformer_block()\n",
|
||
" test_unit_gpt()\n",
|
||
"\n",
|
||
" print(\"\\nRunning integration scenarios...\")\n",
|
||
"\n",
|
||
" # Test complete transformer training scenario\n",
|
||
" print(\"🔬 Integration Test: Full Training Pipeline...\")\n",
|
||
"\n",
|
||
" # Create model and data\n",
|
||
" vocab_size = 50\n",
|
||
" embed_dim = 64\n",
|
||
" num_layers = 2\n",
|
||
" num_heads = 4\n",
|
||
"\n",
|
||
" model = GPT(vocab_size, embed_dim, num_layers, num_heads)\n",
|
||
"\n",
|
||
" # Test batch processing\n",
|
||
" batch_size = 3\n",
|
||
" seq_len = 16\n",
|
||
" tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))\n",
|
||
"\n",
|
||
" # Forward pass\n",
|
||
" logits = model.forward(tokens)\n",
|
||
" assert logits.shape == (batch_size, seq_len, vocab_size)\n",
|
||
"\n",
|
||
" # Test generation with different temperatures\n",
|
||
" prompt = Tensor(np.random.randint(0, vocab_size, (1, 8)))\n",
|
||
"\n",
|
||
" # Conservative generation\n",
|
||
" conservative = model.generate(prompt, max_new_tokens=5, temperature=0.1)\n",
|
||
" assert conservative.shape == (1, 13)\n",
|
||
"\n",
|
||
" # Creative generation\n",
|
||
" creative = model.generate(prompt, max_new_tokens=5, temperature=2.0)\n",
|
||
" assert creative.shape == (1, 13)\n",
|
||
"\n",
|
||
" # Test parameter counting consistency\n",
|
||
" total_params = sum(param.size for param in model.parameters())\n",
|
||
" assert total_params > 1000 # Should have substantial parameters\n",
|
||
"\n",
|
||
" print(\"✅ Full transformer pipeline works!\")\n",
|
||
"\n",
|
||
" print(\"\\n\" + \"=\" * 50)\n",
|
||
" print(\"🎉 ALL TESTS PASSED! Module ready for export.\")\n",
|
||
" print(\"Run: tito module complete 13_transformers\")\n",
|
||
"\n",
|
||
"test_module()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "28f5c8ca",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"if __name__ == \"__main__\":\n",
|
||
" print(\"🚀 Running Transformers module...\")\n",
|
||
" test_module()\n",
|
||
" print(\"✅ Module validation complete!\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "440ab431",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 🤔 ML Systems Thinking: Transformer Architecture\n",
|
||
"\n",
|
||
"Now that you've built a complete transformer model, let's reflect on the systems implications and design decisions."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "2f6986b9",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "systems-q1",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a465d45c",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"### Question 1: Attention Complexity Analysis\n",
|
||
"You implemented multi-head attention that computes attention matrices of size (batch, heads, seq_len, seq_len).\n",
|
||
"\n",
|
||
"**a) Memory Scaling**: For GPT-4 scale (context length 8192, batch size 16, 96 attention heads):\n",
|
||
"- Attention matrix elements: _____ (calculate: 16 × 96 × 8192 × 8192)\n",
|
||
"- Memory in GB (4 bytes/float): _____ GB per layer\n",
|
||
"- For 96 layers: _____ GB total just for attention matrices\n",
|
||
"\n",
|
||
"**b) Why Quadratic Matters**: If processing costs $0.01 per GB, what's the cost difference between:\n",
|
||
"- 1K context: $_____\n",
|
||
"- 8K context: $_____\n",
|
||
"- 32K context: $_____\n",
|
||
"\n",
|
||
"*Think about: Why long-context models are expensive, and why FlashAttention matters*"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "fb3a7788",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "systems-q2",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "a2f18da0",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"### Question 2: Parameter Distribution Analysis\n",
|
||
"Your GPT model has parameters in embeddings, transformer blocks, and the language head.\n",
|
||
"\n",
|
||
"**a) Parameter Breakdown**: For a model with vocab_size=50K, embed_dim=1024, num_layers=24:\n",
|
||
"- Token embedding: _____ parameters (vocab_size × embed_dim)\n",
|
||
"- Each transformer block: approximately _____ parameters\n",
|
||
"- Language head: _____ parameters\n",
|
||
"- Total model: approximately _____ parameters\n",
|
||
"\n",
|
||
"**b) Memory During Training**: Training requires storing:\n",
|
||
"- Parameters (model weights)\n",
|
||
"- Gradients (same size as parameters)\n",
|
||
"- Optimizer states (2-3× parameters for Adam)\n",
|
||
"- Activations (depends on batch size and sequence length)\n",
|
||
"\n",
|
||
"For your calculated model size, estimate total training memory: _____ GB\n",
|
||
"\n",
|
||
"*Consider: Why training large models requires hundreds of GPUs*"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "259119cb",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "systems-q3",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "680a951e",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"### Question 3: Autoregressive Generation Bottlenecks\n",
|
||
"Your generate() method runs the full model forward pass for each new token.\n",
|
||
"\n",
|
||
"**a) Generation Inefficiency**: To generate 100 tokens with a 24-layer model:\n",
|
||
"- Token 1: _____ layer computations (24 layers × 1 position)\n",
|
||
"- Token 2: _____ layer computations (24 layers × 2 positions)\n",
|
||
"- Token 100: _____ layer computations (24 layers × 100 positions)\n",
|
||
"- Total: _____ layer computations\n",
|
||
"\n",
|
||
"**b) KV-Cache Optimization**: With KV-caching, each new token only needs:\n",
|
||
"- _____ layer computations (just the new position)\n",
|
||
"- This reduces computation by approximately _____× for 100 tokens\n",
|
||
"\n",
|
||
"*Think about: Why inference optimization matters for production deployment*"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e99d5ae3",
|
||
"metadata": {
|
||
"lines_to_next_cell": 0,
|
||
"nbgrader": {
|
||
"grade": false,
|
||
"grade_id": "systems-q4",
|
||
"solution": true
|
||
}
|
||
},
|
||
"outputs": [],
|
||
"source": []
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "678beea4",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"### Question 4: Pre-norm vs Post-norm Architecture\n",
|
||
"You implemented pre-norm (LayerNorm before attention/MLP) rather than post-norm (LayerNorm after).\n",
|
||
"\n",
|
||
"**a) Training Stability**: Pre-norm helps with gradient flow because:\n",
|
||
"- Residual connections pass _____ gradients directly through the network\n",
|
||
"- LayerNorm before operations provides _____ input distributions\n",
|
||
"- This enables training _____ networks compared to post-norm\n",
|
||
"\n",
|
||
"**b) Performance Trade-offs**:\n",
|
||
"- Pre-norm: Better training stability, but slightly _____ final performance\n",
|
||
"- Post-norm: Better performance when it trains, but requires _____ learning rates\n",
|
||
"- Most modern large models use _____ because scale requires stability\n",
|
||
"\n",
|
||
"*Consider: Why architectural choices become more important at scale*"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "4e8cc6dc",
|
||
"metadata": {
|
||
"cell_marker": "\"\"\""
|
||
},
|
||
"source": [
|
||
"## 🎯 MODULE SUMMARY: Transformers\n",
|
||
"\n",
|
||
"Congratulations! You've built the complete transformer architecture that powers modern language models!\n",
|
||
"\n",
|
||
"### Key Accomplishments\n",
|
||
"- Built LayerNorm for stable training across deep networks\n",
|
||
"- Implemented MLP (feed-forward) networks with GELU activation\n",
|
||
"- Created complete TransformerBlock with self-attention, residual connections, and pre-norm architecture\n",
|
||
"- Built full GPT model with embeddings, positional encoding, and autoregressive generation\n",
|
||
"- Analyzed parameter scaling and attention memory complexity\n",
|
||
"- All tests pass ✅ (validated by `test_module()`)\n",
|
||
"\n",
|
||
"### Ready for Next Steps\n",
|
||
"Your transformer implementation is the foundation for modern language models! This architecture enables:\n",
|
||
"- **Training**: Learn patterns from massive text datasets\n",
|
||
"- **Generation**: Produce coherent, contextual text\n",
|
||
"- **Transfer Learning**: Fine-tune for specific tasks\n",
|
||
"- **Scaling**: Grow to billions of parameters for emergent capabilities\n",
|
||
"\n",
|
||
"Export with: `tito module complete 13_transformers`\n",
|
||
"\n",
|
||
"**Next**: Module 14 will add KV-caching for efficient generation, optimizing the autoregressive inference you just implemented!"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3 (ipykernel)",
|
||
"language": "python",
|
||
"name": "python3"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 5
|
||
}
|