github-starred/TinyTorch
2
0
Fork 0
mirror of https://github.com/MLSysBook/TinyTorch.git synced 2026-08-03 08:56:10 -05:00
Code Issues 7 Packages Projects Releases Wiki Activity

refactor: Update attention module to match tokenization style

Browse Source 
- Clean import structure following TinyTorch dependency chain
- Add proper export declarations for key functions and classes
- Standardize NBGrader cell structure and testing patterns
- Enhance ASCII diagrams with improved formatting
- Align documentation style with tokenization module standards
- Maintain all core functionality and educational value
This commit is contained in:
Vijay Janapa Reddi
2025-10-25 15:26:33 -04:00
parent d4b1d7c279
commit 7c8b94b59a
2 changed files with 87 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+24 -24
View File
@@ -7,7 +7,7 @@
[![Documentation](https://img.shields.io/badge/docs-jupyter_book-orange.svg)](https://mlsysbook.github.io/TinyTorch/)
![Status](https://img.shields.io/badge/status-active-success.svg)
> 🚧 **Work in Progress** - We're actively developing TinyTorch for Spring 2025! Core modules (01-09) are complete and tested. Transformer modules (10-14) in active development right now.
> 🚧 **Work in Progress** - Actively developing TinyTorch for Spring 2025! All 20 core modules (01-20) are implemented but still being debugged and tested. Core foundation modules (01-09) are stable. Transformer and optimization modules (10-20) are functional but undergoing refinement. Join us in building the future of ML systems education.
## 📖 Table of Contents
- [Why TinyTorch?](#why-tinytorch)
@@ -224,15 +224,15 @@ tito checkpoint status
tito checkpoint timeline
```
**Checkpoint Progression:**
- **01-02**: Foundation (Tensors, Activations)
**Module Progression:**
- **01-02**: Foundation (Tensor, Activations)
- **03-07**: Core Networks (Layers, Losses, Autograd, Optimizers, Training)
- **08-09**: Computer Vision (DataLoaders, Spatial ops - unlocks CIFAR-10 @ 75%+)
- **08-09**: Computer Vision (DataLoader, Spatial ops - unlocks CIFAR-10 @ 75%+)
- **10-14**: Language Models (Tokenization, Embeddings, Attention, Transformers, KV-Caching)
- **15-19**: System Optimization (Profiling, Acceleration, Quantization, Compression, Benchmarking)
- **20**: Capstone (Complete end-to-end ML systems)
Each checkpoint asks: **"Can I build this capability from scratch?"** with hands-on validation.
Each module asks: **"Can I build this capability from scratch?"** with hands-on validation.
### Module Completion Workflow
@@ -270,46 +270,46 @@ tito module complete 01_tensor
As you complete modules, unlock historical ML milestones demonstrating YOUR implementations:
### 🧠 01. Perceptron (1957) - After Module 04
### 🧠 01. Perceptron (1957) - After Module 03
```bash
cd milestones/01_perceptron_1957
cd milestones/01_1957_perceptron
python perceptron_trained.py
# Rosenblatt's first trainable neural network
# YOUR Linear layer + Sigmoid recreates history!
```
**Requirements**: Modules 01-04 (Tensor, Activations, Layers, Losses)
**Requirements**: Modules 01-03 (Tensor, Activations, Layers)
**Achievement**: Binary classification with gradient descent
---
### ⚡ 02. XOR Crisis (1969) - After Module 06
### ⚡ 02. XOR Crisis (1969) - After Module 05
```bash
cd milestones/02_xor_crisis_1969
cd milestones/02_1969_xor_crisis
python xor_solved.py
# Solve Minsky's XOR challenge with hidden layers
# YOUR autograd enables multi-layer learning!
```
**Requirements**: Modules 01-06 (+ Autograd, Optimizers)
**Requirements**: Modules 01-05 (+ Autograd)
**Achievement**: Non-linear problem solving
---
### 🔢 03. MLP Revival (1986) - After Module 07
```bash
cd milestones/03_mlp_revival_1986
cd milestones/03_1986_mlp_revival
python mlp_digits.py # 8x8 digit classification
python mlp_mnist.py # Full MNIST dataset
# Backpropagation revolution on real vision!
# YOUR training loops achieve 95%+ accuracy
```
**Requirements**: Modules 01-07 (+ Training)
**Requirements**: Modules 01-07 (+ Optimizers, Training)
**Achievement**: Real computer vision with MLPs
---
### 🖼️ 04. CNN Revolution (1998) - After Module 09
```bash
cd milestones/04_cnn_revolution_1998
cd milestones/04_1998_cnn_revolution
python cnn_digits.py # Spatial features on digits
python lecun_cifar10.py # Natural images (CIFAR-10)
# LeCun's CNNs achieve 75%+ on CIFAR-10!
@@ -322,7 +322,7 @@ python lecun_cifar10.py # Natural images (CIFAR-10)
### 🤖 05. Transformer Era (2017) - After Module 13
```bash
cd milestones/05_transformer_era_2017
cd milestones/05_2017_transformer_era
python vaswani_shakespeare.py
# Attention mechanisms for language modeling
# YOUR attention implementation generates text!
@@ -334,7 +334,7 @@ python vaswani_shakespeare.py
### ⚡ 06. Systems Age (2024) - After Module 19
```bash
cd milestones/06_systems_age_2024
cd milestones/06_2024_systems_age
python optimize_models.py
# Profile, optimize, and benchmark YOUR framework
# Compete on TinyMLPerf leaderboard!
@@ -372,14 +372,14 @@ pytest tests/
```
**Current Status**:
- ✅ **20 complete modules** (01 Tensor → 20 Capstone)
- ✅ **20 modules implemented** (01 Tensor → 20 Capstone) - all code exists
- ✅ **6 historical milestones** (1957 Perceptron → 2024 Systems Age)
- ✅ **Capability-based checkpoints** tracking learning progress
- ✅ **Complete optimization pipeline** from profiling to benchmarking
- ✅ **TinyMLPerf competition framework** for performance excellence
- ✅ **Foundation modules stable** (01-09): Tensor through Spatial operations
- 🚧 **Transformer modules functional** (10-14): Tokenization through KV-Caching - undergoing testing
- 🚧 **Optimization modules functional** (15-20): Profiling through Capstone - undergoing testing
- ✅ **KISS principle design** for clear, maintainable code
- ✅ **Essential-only features**: Focus on what's used in production ML systems
- 🚧 **Active development**: Transformer integration (modules 10-14) on `transformers-integration` branch
- 🎯 **Target: Spring 2025** - Active debugging and refinement in progress
## 📚 Documentation & Resources
@@ -455,11 +455,11 @@ Special thanks to students and contributors who helped refine this educational f
### What Makes TinyTorch Different?
- ✅ **Essential-only features** - Focus on what's actually used in production
- ✅ **Complete implementation** - Build every component from scratch
- ✅ **Real achievements** - Train CNNs on CIFAR-10 to 75%+ accuracy
- 🚧 **Complete implementation** - Build every component from scratch (20 modules in development)
- 🎯 **Real achievements** - Train CNNs on CIFAR-10 to 75%+ accuracy (target)
- ✅ **Systems thinking** - Understand memory, performance, and scaling
- ✅ **Production relevance** - Learn patterns from PyTorch and TensorFlow
- ✅ **Immediate validation** - 20 capability checkpoints track progress
- ✅ **Progressive learning** - 20 modules from tensors to transformers to optimization
### Your Learning Journey
1. **Week 1-2**: Foundation (Tensors, Activations, Layers)
modules/source/12_attention/attention_dev.py
+63 -91
View File
@@ -44,7 +44,7 @@ Let's get started!
## 📦 Where This Code Lives in the Final Package
**Learning Side:** You work in `modules/12_attention/attention_dev.py`
**Learning Side:** You work in `modules/12_attention/attention_dev.py`
**Building Side:** Code exports to `tinytorch.core.attention`
```python
@@ -59,76 +59,16 @@ from tinytorch.core.attention import scaled_dot_product_attention, MultiHeadAtte
- **Integration:** Works seamlessly with embeddings for complete sequence processing pipelines
"""
# %% nbgrader={"grade": false, "grade_id": "imports", "locked": false, "solution": true}
# %%
import numpy as np
import math
import time
from typing import Optional, Tuple, List
# Import dependencies from other modules
# Import dependencies from previous modules - following TinyTorch dependency chain
from tinytorch.core.tensor import Tensor
from tinytorch.core.layers import Linear
# Note: Keeping simplified implementations for reference during development
class _SimplifiedTensor:
"""Simplified tensor for attention operations development."""
def __init__(self, data, requires_grad=False):
self.data = np.array(data, dtype=np.float32)
self.shape = self.data.shape
self.requires_grad = requires_grad
self.grad = None
def __repr__(self):
return f"Tensor(shape={self.shape}, data=\n{self.data})"
def __add__(self, other):
if isinstance(other, Tensor):
return Tensor(self.data + other.data)
return Tensor(self.data + other)
def __mul__(self, other):
if isinstance(other, Tensor):
return Tensor(self.data * other.data)
return Tensor(self.data * other)
def sum(self, axis=None):
return Tensor(np.sum(self.data, axis=axis))
def mean(self, axis=None):
return Tensor(np.mean(self.data, axis=axis))
def matmul(self, other):
return Tensor(np.matmul(self.data, other.data))
def softmax(self, axis=-1):
"""Apply softmax along specified axis."""
# Subtract max for numerical stability
shifted = self.data - np.max(self.data, axis=axis, keepdims=True)
exp_values = np.exp(shifted)
return Tensor(exp_values / np.sum(exp_values, axis=axis, keepdims=True))
# Simplified Linear layer for development
class _SimplifiedLinear:
"""Simplified linear layer for attention projections."""
def __init__(self, in_features, out_features):
self.in_features = in_features
self.out_features = out_features
# Initialize weights and bias (simplified Xavier initialization)
self.weight = Tensor(np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features))
self.bias = Tensor(np.zeros(out_features))
def forward(self, x):
"""Forward pass: y = xW + b"""
output = x.matmul(self.weight)
# Add bias (broadcast across batch and sequence dimensions)
return Tensor(output.data + self.bias.data)
def parameters(self):
"""Return list of parameters for this layer."""
return [self.weight, self.bias]
# %% [markdown]
"""
## Part 1: Introduction - What is Attention?
@@ -307,7 +247,8 @@ Step-by-Step Attention Computation:
```
"""
# %% nbgrader={"grade": false, "grade_id": "attention-function", "locked": false, "solution": true}
# %% nbgrader={"grade": false, "grade_id": "attention-function", "solution": true}
#| export
def scaled_dot_product_attention(Q: Tensor, K: Tensor, V: Tensor, mask: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
"""
Compute scaled dot-product attention.
@@ -443,7 +384,9 @@ def test_unit_scaled_dot_product_attention():
print("✅ scaled_dot_product_attention works correctly!")
test_unit_scaled_dot_product_attention()
# Run test immediately when developing this module
if __name__ == "__main__":
test_unit_scaled_dot_product_attention()
# %% [markdown]
"""
@@ -470,22 +413,46 @@ Multi-head attention runs multiple attention "heads" in parallel, each learning
### Understanding Multi-Head Architecture
```
Single-Head vs Multi-Head Attention:
SINGLE HEAD (Limited):
Input → [Linear] → Q,K,V → [Attention] → Output
512×512 512×512 512
MULTI-HEAD (Rich):
Input → [Linear] → Q₁,K₁,V₁ → [Attention₁] → Head₁ (64 dims)
→ [Linear] → Q₂,K₂,V₂ → [Attention₂] → Head₂ (64 dims)
→ [Linear] → Q₃,K₃,V₃ → [Attention₃] → Head₃ (64 dims)
...
→ [Linear] → Q₈,K₈,V₈ → [Attention₈] → Head₈ (64 dims)
↓
[Concatenate]
↓
[Linear Mix] → Output (512)
┌─────────────────────────────────────────────────────────────────────────┐
│ SINGLE-HEAD vs MULTI-HEAD ATTENTION ARCHITECTURE │
├─────────────────────────────────────────────────────────────────────────┤
│ │
│ SINGLE HEAD ATTENTION (Limited Representation): │
│ ┌─────────────────────────────────────────────────────────────────────┐ │
│ │ Input (512) → [Linear] → Q,K,V (512) → [Attention] → Output (512) │ │
│ │ ↑ ↑ ↑ ↑ │ │
│ │ Single proj Full dimensions One head Limited focus │ │
│ └─────────────────────────────────────────────────────────────────────┘ │
│ │
│ MULTI-HEAD ATTENTION (Rich Parallel Processing): │
│ ┌─────────────────────────────────────────────────────────────────────┐ │
│ │ Input (512) │ │
│ │ ↓ │ │
│ │ [Q/K/V Projections] → 512 dimensions each │ │
│ │ ↓ │ │
│ │ [Split into 8 heads] → 8 × 64 dimensions per head │ │
│ │ ↓ │ │
│ │ Head₁: Q₁(64) ⊗ K₁(64) → Attention₁ → Output₁(64) │ Syntax focus │ │
│ │ Head₂: Q₂(64) ⊗ K₂(64) → Attention₂ → Output₂(64) │ Semantic │ │
│ │ Head₃: Q₃(64) ⊗ K₃(64) → Attention₃ → Output₃(64) │ Position │ │
│ │ Head₄: Q₄(64) ⊗ K₄(64) → Attention₄ → Output₄(64) │ Long-range │ │
│ │ Head₅: Q₅(64) ⊗ K₅(64) → Attention₅ → Output₅(64) │ Local deps │ │
│ │ Head₆: Q₆(64) ⊗ K₆(64) → Attention₆ → Output₆(64) │ Coreference │ │
│ │ Head₇: Q₇(64) ⊗ K₇(64) → Attention₇ → Output₇(64) │ Composition │ │
│ │ Head₈: Q₈(64) ⊗ K₈(64) → Attention₈ → Output₈(64) │ Global view │ │
│ │ ↓ │ │
│ │ [Concatenate] → 8 × 64 = 512 dimensions │ │
│ │ ↓ │ │
│ │ [Output Linear] → Final representation (512) │ │
│ └─────────────────────────────────────────────────────────────────────┘ │
│ │
│ Key Benefits of Multi-Head: │
│ • Parallel specialization across different relationship types │
│ • Same total parameters, distributed across multiple focused heads │
│ • Each head can learn distinct attention patterns │
│ • Enables rich, multifaceted understanding of sequences │
│ │
└─────────────────────────────────────────────────────────────────────────┘
```
### The Multi-Head Process Detailed
@@ -520,7 +487,7 @@ Each head can specialize in different patterns:
This parallelization allows the model to attend to different representation subspaces simultaneously.
"""
# %% nbgrader={"grade": false, "grade_id": "multihead-attention", "locked": false, "solution": true}
# %% nbgrader={"grade": false, "grade_id": "multihead-attention", "solution": true}
#| export
class MultiHeadAttention:
"""
@@ -718,7 +685,9 @@ def test_unit_multihead_attention():
print("✅ MultiHeadAttention works correctly!")
test_unit_multihead_attention()
# Run test immediately when developing this module
if __name__ == "__main__":
test_unit_multihead_attention()
# %% [markdown]
"""
@@ -770,7 +739,7 @@ Just for attention matrices!
```
"""
# %% nbgrader={"grade": false, "grade_id": "attention-complexity", "locked": false, "solution": true}
# %% nbgrader={"grade": false, "grade_id": "attention-complexity", "solution": true}
def analyze_attention_complexity():
"""📊 Analyze attention computational complexity and memory scaling."""
print("📊 Analyzing Attention Complexity...")
@@ -798,7 +767,7 @@ def analyze_attention_complexity():
print(f"\n💡 Attention memory scales as O(n²) with sequence length")
print(f"🚀 For seq_len=1024, attention matrix alone needs {(1024*1024*4)/1024/1024:.1f} MB")
# %% nbgrader={"grade": false, "grade_id": "attention-timing", "locked": false, "solution": true}
# %% nbgrader={"grade": false, "grade_id": "attention-timing", "solution": true}
def analyze_attention_timing():
"""📊 Measure attention computation time vs sequence length."""
print("\n📊 Analyzing Attention Timing...")
@@ -916,7 +885,7 @@ Example Attention Patterns in Language:
Let's see these patterns emerge in our implementation.
"""
# %% nbgrader={"grade": false, "grade_id": "attention-scenarios", "locked": false, "solution": true}
# %% nbgrader={"grade": false, "grade_id": "attention-scenarios", "solution": true}
def test_attention_scenarios():
"""Test attention mechanisms in realistic scenarios."""
print("🔬 Testing Attention Scenarios...")
@@ -998,7 +967,9 @@ def test_attention_scenarios():
print("\n✅ All attention scenarios work correctly!")
test_attention_scenarios()
# Run test immediately when developing this module
if __name__ == "__main__":
test_attention_scenarios()
# %% [markdown]
"""
@@ -1031,7 +1002,7 @@ The attention matrices you see here are the foundation of model interpretability
# %% [markdown]
"""
## 🧪 Module Integration Test
## 6. Module Integration Test
Final validation that everything works together correctly.
"""
@@ -1064,8 +1035,9 @@ def test_module():
print("🎉 ALL TESTS PASSED! Module ready for export.")
print("Run: tito module complete 12")
# Call before module summary
test_module()
# Run comprehensive module test when executed directly
if __name__ == "__main__":
test_module()
# %%
if __name__ == "__main__":
@@ -1133,7 +1105,7 @@ Congratulations! You've built the attention mechanism that revolutionized deep l
Your attention implementation is the core mechanism that enables modern language models!
Export with: `tito module complete 12`
**Next**: Module 13 will combine attention with feed-forward layers to build complete transformer blocks, leading to GPT-style language models!
**Next**: Module 13 will combine attention with feed-forward layers to build complete transformer blocks!
### What You Just Built Powers
- **GPT models**: Your attention mechanism is the exact pattern used in ChatGPT and GPT-4