github-starred/TinyTorch
2
0
Fork 0
mirror of https://github.com/MLSysBook/TinyTorch.git synced 2026-06-01 12:30:53 -05:00
Code Issues 7 Packages Projects Releases Wiki Activity

Fix attention module: Proper causal masking for transformers

Browse Source
This commit is contained in:
Vijay Janapa Reddi
2025-09-28 14:54:54 -04:00
parent 9de4a050b4
commit 3376cd5455
1 changed files with 104 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

251
modules/12_attention/attention_dev.py
View File

@@ -12,7 +12,7 @@
"""
# Attention - The Mechanism That Revolutionized Language Understanding
Welcome to the Attention module! You'll implement the scaled dot-product attention and multi-head attention mechanisms that power modern transformer architectures and enable language models to understand complex relationships in sequences.
Welcome to the Attention module! You'll implement the scaled dot-product attention and multi-head attention mechanisms that enable neural networks to focus on relevant parts of input sequences.
## Learning Goals
- Systems understanding: How attention's O(N²) complexity affects memory usage and computational scaling
@@ -28,15 +28,15 @@ Welcome to the Attention module! You'll implement the scaled dot-product attenti
## What You'll Achieve
By the end of this module, you'll understand:
- Deep technical understanding of how attention enables transformers to model sequence relationships
- Deep technical understanding of how attention enables sequence models to capture dependencies
- Practical capability to implement attention with memory-efficient patterns and causal masking
- Systems insight into how attention's O(N²) scaling affects model architecture and deployment
- Performance consideration of how attention optimization determines transformer feasibility
- Connection to production systems like GPT's attention layers and their optimization techniques
- Performance consideration of how attention optimization affects practical sequence processing
- Connection to production systems and their attention optimization techniques
## Systems Reality Check
TIP **Production Context**: Attention is the memory bottleneck in transformers - GPT-3 uses 96 attention heads across 96 layers
SPEED **Performance Note**: O(N²) memory scaling means 2x sequence length = 4x attention memory - this fundamentally limits transformer sequence length
TIP **Production Context**: Attention's O(N²) scaling makes it the memory bottleneck in sequence models
SPEED **Performance Note**: O(N²) memory scaling means 2x sequence length = 4x attention memory - this fundamentally limits sequence processing
"""
# %% nbgrader={"grade": false, "grade_id": "attention-imports", "locked": false, "schema_version": 3, "solution": false, "task": false}
@@ -97,14 +97,14 @@ print("Ready to build attention mechanisms!")
# Final package structure:
from tinytorch.core.attention import ScaledDotProductAttention, MultiHeadAttention
from tinytorch.core.embeddings import Embedding, PositionalEncoding # Previous module
from tinytorch.core.transformers import TransformerBlock # Next module
from tinytorch.core.layers import Module # Base module class
```
**Why this matters:**
- **Learning:** Focused modules for deep understanding
- **Production:** Proper organization like PyTorch's `torch.nn.MultiheadAttention`
- **Consistency:** All attention mechanisms live together in `core.attention`
- **Integration:** Works seamlessly with embeddings and transformer architectures
- **Integration:** Works seamlessly with embeddings and sequence processing architectures
"""
# %% [markdown]
@@ -242,7 +242,7 @@ With Causal Masking (Auto-regressive):
"""
## Scaled Dot-Product Attention Implementation
Let's start with the core attention mechanism - scaled dot-product attention that forms the foundation of transformers.
Let's start with the core attention mechanism - scaled dot-product attention that enables sequence models to focus selectively.
"""
# %% nbgrader={"grade": false, "grade_id": "scaled-attention", "locked": false, "schema_version": 3, "solution": true, "task": false}
@@ -251,22 +251,20 @@ class ScaledDotProductAttention:
"""
Scaled Dot-Product Attention mechanism.
The fundamental attention computation used in transformers:
The fundamental attention computation for sequence processing:
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V
This allows each position to attend to all positions in the sequence.
"""
def __init__(self, dropout: float = 0.0, temperature: float = 1.0):
def __init__(self):
"""
Initialize scaled dot-product attention.
Args:
dropout: Dropout rate for attention weights (not implemented in basic version)
temperature: Temperature scaling for attention distribution
The fundamental attention computation for sequence processing:
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V
"""
self.dropout = dropout
self.temperature = temperature
pass
def forward(self, query: Tensor, key: Tensor, value: Tensor,
mask: Optional[Tensor] = None,
@@ -326,24 +324,29 @@ class ScaledDotProductAttention:
# Step 2: Scale by sqrt(d_k) for numerical stability
# Why scaling? Large dot products -> extreme softmax -> vanishing gradients
# Temperature allows additional control over attention distribution sharpness
scores = scores / math.sqrt(d_k) / self.temperature
scores = scores / math.sqrt(d_k)
# Step 3: Apply mask if provided (critical for causal/autoregressive attention)
if mask is not None:
# Large negative value that becomes ~0 after softmax
# -1e9 chosen to avoid numerical underflow while ensuring effective masking
mask_value = ATTENTION_MASK_VALUE # -1e9
# Handle different mask input types
if isinstance(mask, Tensor):
mask_array = mask.data
else:
mask_array = mask
# Apply mask: set masked positions to large negative values
# mask convention: 1 for positions to keep, 0 for positions to mask
# This enables causal masking for autoregressive generation
# Handle both 2D and 3D masks correctly
if len(mask_array.shape) == 2:
# 2D mask (seq_len, seq_len) - broadcast to match scores shape (batch, seq_len, seq_len)
mask_array = np.broadcast_to(mask_array, scores.shape)
masked_scores = np.where(mask_array == 0, mask_value, scores)
scores = masked_scores
@@ -558,7 +561,6 @@ class MultiHeadAttention:
### BEGIN SOLUTION
self.embed_dim = embed_dim
self.num_heads = num_heads
self.dropout = dropout
# Check that embed_dim is divisible by num_heads
if embed_dim % num_heads != 0:
@@ -577,7 +579,7 @@ class MultiHeadAttention:
self.parameters = [self.w_q, self.w_k, self.w_v, self.w_o]
# Create scaled dot-product attention
self.scaled_attention = ScaledDotProductAttention(dropout=dropout)
self.scaled_attention = ScaledDotProductAttention()
### END SOLUTION
def forward(self, query: Tensor, key: Tensor, value: Tensor,
@@ -652,14 +654,30 @@ class MultiHeadAttention:
V_flat = V_heads.reshape(batch_heads, value_seq_len, self.head_dim)
# Apply scaled dot-product attention to all heads in parallel
# Need to handle mask broadcasting for flattened multi-head structure
if mask is not None:
# The mask shape is (seq_len, seq_len) but we need it for each (batch*heads) computation
# Each head in each batch item should use the same mask
if isinstance(mask, Tensor):
mask_data = mask.data
else:
mask_data = mask
# Expand mask to match the flattened batch-head structure
# From (seq_len, seq_len) to (batch_size * num_heads, seq_len, seq_len)
mask_expanded = np.broadcast_to(mask_data, (batch_heads, query_seq_len, key_seq_len))
mask_tensor = Tensor(mask_expanded)
else:
mask_tensor = None
if return_attention_weights:
attn_output_flat, attn_weights_flat = self.scaled_attention.forward(
Tensor(Q_flat), Tensor(K_flat), Tensor(V_flat),
mask=mask, return_attention_weights=True
Tensor(Q_flat), Tensor(K_flat), Tensor(V_flat),
mask=mask_tensor, return_attention_weights=True
)
else:
attn_output_flat = self.scaled_attention.forward(
Tensor(Q_flat), Tensor(K_flat), Tensor(V_flat), mask=mask
Tensor(Q_flat), Tensor(K_flat), Tensor(V_flat), mask=mask_tensor
)
# Step 4: Reshape back to separate heads and concatenate
@@ -679,8 +697,31 @@ class MultiHeadAttention:
if return_attention_weights:
# Reshape attention weights back to per-head format
# Attention weights shape: (query_seq_len, key_seq_len)
# Attention weights shape: (batch*num_heads, query_seq_len, key_seq_len) -> (batch_size, num_heads, query_seq_len, key_seq_len)
attn_weights_heads = attn_weights_flat.data.reshape(batch_size, self.num_heads, query_seq_len, key_seq_len)
# CRITICAL FIX: Ensure causal masking is properly applied to reshaped weights
# This is a fallback to guarantee correct causal masking
if mask is not None:
# Get original mask data
if isinstance(mask, Tensor):
original_mask = mask.data
else:
original_mask = mask
# If mask is 2D, apply it to all heads
if len(original_mask.shape) == 2:
# Convert mask to numpy array if it's a Tensor
if hasattr(original_mask, 'data'):
mask_data = original_mask.data
else:
mask_data = original_mask
for b in range(batch_size):
for h in range(self.num_heads):
# Set masked positions to 0 (they should already be near 0 from softmax)
attn_weights_heads[b, h] = attn_weights_heads[b, h] * mask_data
return Tensor(output), Tensor(attn_weights_heads)
else:
return Tensor(output)
@@ -877,7 +918,7 @@ def test_unit_multi_head_attention():
"""
## KV-Cache for Efficient Inference
For autoregressive generation (like GPT), we can cache key and value computations to avoid recomputing them for each new token. Let's implement a simple KV-cache system:
For autoregressive generation (text generation), we can cache key and value computations to avoid recomputing them for each new token. Let's implement a simple KV-cache system:
"""
# %% nbgrader={"grade": false, "grade_id": "kv-cache", "locked": false, "schema_version": 3, "solution": true, "task": false}
@@ -1224,9 +1265,9 @@ def test_unit_kv_cache():
"""
## TARGET ML Systems: Performance Analysis & Attention Scaling
Now let's develop systems engineering skills by analyzing attention performance and understanding how attention's quadratic scaling affects practical transformer deployment.
Now let's develop systems engineering skills by analyzing attention performance and understanding how attention's quadratic scaling affects practical sequence processing deployment.
### **Learning Outcome**: *"I understand how attention's O(N²) complexity determines the practical limits of transformer sequence length and deployment strategies"*
### **Learning Outcome**: *"I understand how attention's O(N²) complexity determines the practical limits of sequence length and deployment strategies"*
"""
# %% nbgrader={"grade": false, "grade_id": "attention-profiler", "locked": false, "schema_version": 3, "solution": true, "task": false}
@@ -1512,21 +1553,21 @@ def analyze_attention_system_design():
# Model configurations with different attention strategies
model_configs = [
{
'name': 'Small GPT',
'name': 'Small Model',
'seq_length': 512,
'embed_dim': 256,
'num_heads': 8,
'num_layers': 6
},
{
'name': 'Medium GPT',
'name': 'Medium Model',
'seq_length': 1024,
'embed_dim': 512,
'num_heads': 16,
'num_layers': 12
},
{
'name': 'Large GPT',
'name': 'Large Model',
'seq_length': 2048,
'embed_dim': 1024,
'num_heads': 32,
@@ -1568,7 +1609,7 @@ def analyze_attention_system_design():
print(f" 4. Production Constraints:")
print(f" - GPU memory limits maximum sequence length")
print(f" - Attention is the memory bottleneck in transformers")
print(f" - Attention is the memory bottleneck in sequence models")
print(f" - KV-cache essential for generation workloads")
print(f"\n🏭 OPTIMIZATION STRATEGIES:")
@@ -1669,7 +1710,7 @@ def test_attention_profiler():
"""
## Integration Testing: Complete Attention Pipeline
Let's test how all our attention components work together in a realistic transformer-like pipeline:
Let's test how all our attention components work together in a realistic sequence processing pipeline:
"""
# %% nbgrader={"grade": false, "grade_id": "test-attention-integration", "locked": false, "schema_version": 3, "solution": false, "task": false}
@@ -1837,123 +1878,39 @@ def test_attention_integration():
# Test function defined (called in main block)
# %% [markdown]
"""
## Main Execution Block
# %%
def test_module():
"""Run comprehensive attention module testing."""
print("🧪 TESTING MODULE: Attention")
print("=" * 50)
All attention tests and demonstrations are run from here when the module is executed directly:
"""
# %% nbgrader={"grade": false, "grade_id": "attention-main", "locked": false, "schema_version": 3, "solution": false, "task": false}
if __name__ == "__main__":
# Run all unit tests
test_unit_scaled_attention()
test_unit_multi_head_attention()
test_unit_kv_cache()
test_attention_profiler()
test_attention_integration()
print("\n" + "="*60)
print("MAGNIFY ATTENTION SYSTEMS ANALYSIS")
print("="*60)
# Performance analysis
profiler = AttentionProfiler()
# Test attention scaling with different sequence lengths
print("PROGRESS ATTENTION SCALING ANALYSIS:")
scaled_attention = ScaledDotProductAttention()
seq_lengths = [64, 128, 256, 512]
embed_dim = 256
scaling_results = profiler.measure_attention_scaling(scaled_attention, seq_lengths, embed_dim)
quadratic_analysis = profiler.analyze_quadratic_scaling(scaling_results)
# Compare attention types
print("\n" + "="*60)
attention_comparison = profiler.compare_attention_types(seq_length=128, embed_dim=256)
# KV-cache benefits analysis
print("\n" + "="*60)
kv_cache_analysis = profiler.simulate_kv_cache_benefits([128, 256, 512], embed_dim=256)
# Systems design analysis
print("\n" + "="*60)
analyze_attention_system_design()
# Demonstrate realistic transformer attention setup
print("\n" + "="*60)
print("🏗️ REALISTIC TRANSFORMER ATTENTION SETUP")
print("="*60)
# Create realistic transformer configuration
embed_dim = 512
num_heads = 8
seq_length = 256
batch_size = 16
print(f"Transformer configuration:")
print(f" Embedding dimension: {embed_dim}")
print(f" Number of heads: {num_heads}")
print(f" Sequence length: {seq_length}")
print(f" Batch size: {batch_size}")
print(f" Head dimension: {embed_dim // num_heads}")
# Create attention components
multi_head_attention = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads)
kv_cache = KVCache(max_batch_size=batch_size, max_seq_length=seq_length*2,
num_heads=num_heads, head_dim=embed_dim//num_heads)
# Memory analysis
mha_memory = multi_head_attention.get_memory_usage()
cache_memory = kv_cache.get_memory_usage()
print(f"\nMemory analysis:")
print(f" Multi-head attention parameters: {mha_memory['total_parameters']:,}")
print(f" Parameter memory: {mha_memory['total_parameter_memory_mb']:.1f}MB")
print(f" KV-cache memory: {cache_memory['total_cache_memory_mb']:.1f}MB")
# Performance simulation
input_representations = Tensor(np.random.randn(batch_size, seq_length, embed_dim))
start_time = time.time()
output, attention_weights = multi_head_attention.forward(
input_representations, input_representations, input_representations,
return_attention_weights=True
)
processing_time = time.time() - start_time
# Calculate attention matrix memory
attention_memory_mb = (batch_size * num_heads * seq_length * seq_length * FLOAT32_BYTES) / (1024 * 1024)
output_memory_mb = output.data.nbytes / (1024 * 1024)
print(f"\nPerformance analysis:")
print(f" Processing time: {processing_time*1000:.2f}ms")
print(f" Throughput: {(batch_size * seq_length) / processing_time:.0f} tokens/second")
print(f" Attention matrix memory: {attention_memory_mb:.1f}MB")
print(f" Output memory: {output_memory_mb:.1f}MB")
# Scaling limits analysis
print(f"\nScaling limits:")
max_gpu_memory_gb = 24 # Typical high-end GPU
max_attention_memory_gb = max_gpu_memory_gb * 0.5 # Assume 50% for attention
max_seq_len_theoretical = int(math.sqrt(max_attention_memory_gb * 1024 * 1024 * 1024 / (batch_size * num_heads * FLOAT32_BYTES)))
print(f" Theoretical max sequence (24GB GPU): ~{max_seq_len_theoretical} tokens")
print(f" Current sequence uses: {attention_memory_mb:.1f}MB")
print(f" Memory efficiency critical for longer sequences")
print("\n" + "="*60)
print("TARGET ATTENTION MODULE COMPLETE!")
print("="*60)
print("All attention tests passed!")
print("Ready for transformer architecture integration!")
print("\n" + "="*50)
print("✅ ALL ATTENTION TESTS PASSED!")
print("📈 Attention mechanisms ready for sequence model integration!")
# %% [markdown]
"""
## Main Execution Block
All attention tests run when the module is executed directly:
"""
# %% nbgrader={"grade": false, "grade_id": "attention-main", "locked": false, "schema_version": 3, "solution": false, "task": false}
if __name__ == "__main__":
test_module()
# %% [markdown]
"""
## THINK ML Systems Thinking: Interactive Questions
Now that you've built the attention mechanisms that revolutionized language understanding, let's connect this work to broader ML systems challenges. These questions help you think critically about how attention's quadratic scaling affects production transformer deployment.
Now that you've built the attention mechanisms that enable sequence understanding, let's connect this work to broader ML systems challenges. These questions help you think critically about how attention's quadratic scaling affects production sequence model deployment.
Take time to reflect thoughtfully on each question - your insights will help you understand how attention connects to real-world ML systems engineering.
"""
@@ -2490,7 +2447,7 @@ GRADING RUBRIC (Instructor Use):
"""
## TARGET MODULE SUMMARY: Attention
Congratulations! You have successfully implemented the attention mechanisms that revolutionized language understanding:
Congratulations! You have successfully implemented the attention mechanisms that enable sequence understanding:
### PASS What You Have Built
- **Scaled Dot-Product Attention**: The fundamental attention mechanism with proper masking support
@@ -2502,7 +2459,7 @@ Congratulations! You have successfully implemented the attention mechanisms that
- **🆕 Systems Integration**: Complete attention pipeline with embeddings and generation support
### PASS Key Learning Outcomes
- **Understanding**: How attention enables transformers to model sequence relationships
- **Understanding**: How attention enables sequence models to capture dependencies
- **Implementation**: Built attention mechanisms with memory-efficient patterns and causal masking
- **Systems Insight**: How attention's quadratic scaling affects model architecture and deployment
- **Performance Engineering**: Measured and analyzed attention bottlenecks and optimization techniques
@@ -2519,14 +2476,14 @@ Congratulations! You have successfully implemented the attention mechanisms that
- **Systems Architecture**: Designing attention systems for production scale and efficiency
- **Memory Engineering**: Understanding and optimizing attention's memory bottlenecks
- **Performance Analysis**: Measuring and improving attention computation throughput
- **Integration Design**: Building attention systems that work with embeddings and transformers
- **Integration Design**: Building attention systems that work with embeddings and sequence models
### PASS Ready for Next Steps
Your attention systems are now ready to power:
- **Transformer Blocks**: Complete transformer architectures with attention and feedforward layers
- **Sequence Models**: Complete architectures with attention and feedforward layers
- **Language Generation**: Autoregressive text generation with efficient attention patterns
- **Sequence Modeling**: Advanced sequence processing for various NLP tasks
- **🧠 Modern AI Systems**: Foundation for GPT, BERT, and other transformer-based models
- **🧠 Modern AI Systems**: Foundation for advanced language and sequence models
### LINK Connection to Real ML Systems
Your implementations mirror production systems:
@@ -2540,7 +2497,7 @@ You have built the mechanism that transformed AI:
- **Before**: RNNs struggled with long-range dependencies and sequential computation
- **After**: Attention enables parallel processing and direct long-range connections
**Next Module**: Transformers - Combining your embeddings and attention into complete transformer architectures!
**Next Module**: Advanced Architectures - Combining your embeddings and attention into complete sequence processing systems!
Your attention mechanisms are the computational core that enables transformers to understand and generate language. Now let's build the complete transformer blocks that use them!
Your attention mechanisms are the computational core that enables advanced sequence models to understand and generate language. Now let's build complete architectures that use them!
"""