c

Author: codelion

examples/mlx_kernel_optimization/README.md

@@ -1,193 +1,200 @@
-# MLX-LM Performance Optimization with OpenEvolve
+# MLX Training Performance Optimization with OpenEvolve
 
-This example demonstrates using OpenEvolve to optimize real MLX-LM inference and training performance on Apple Silicon, directly measuring speedups on the `Qwen2.5-0.5B-Instruct-bf16` model.
+This example demonstrates using OpenEvolve to optimize MLX training performance on Apple Silicon, focusing exclusively on accelerating neural network training workloads.
 
-## The New Approach: Real-World MLX-LM Optimization
+## The Training-Focused Approach: Real-World MLX Training Optimization
 
-Instead of synthetic matrix benchmarks, we now optimize **actual MLX-LM performance**:
+We now focus exclusively on **MLX training performance** optimization:
 
-✅ **Real model**: Qwen2.5-0.5B-Instruct-bf16 for fast but realistic testing  
-✅ **Real workloads**: Text generation (inference) and training simulation  
-✅ **Real metrics**: End-to-end speedup measurement vs original MLX  
-✅ **Practical focus**: Optimize for transformer attention and MLP patterns  
+✅ **Training Workloads**: Forward + backward passes with gradient computation  
+✅ **Realistic Models**: Transformer architectures with substantial matrix operations  
+✅ **Training Patterns**: Batch processing, MLP layers, attention computation  
+✅ **Clear Signal**: Consistent evaluation without inference noise  
+✅ **Practical Value**: Accelerate model development and research workflows  
 
-## Background
+## Why Training-Only Optimization?
 
-MLX is the fastest inference engine on Apple Silicon:
+### 1. **Cleaner Evaluation Signal**
 
-```
-Performance Comparison:
-pytorch_mps    : 1.190s avg, 42.0 tokens/s
-mlx            : 0.044s avg, 1135.8 tokens/s ⭐ 25x FASTER  
-llama_cpp      : 0.316s avg, 158.0 tokens/s
-```
-
-However, MLX's matrix multiplication can be further optimized through intelligent tiling strategies that better utilize Apple Silicon's architecture.
-
-## The Optimization Challenge
+Training provides much more consistent evaluation than inference:
 
-MLX-LM performance depends on efficient matrix multiplication for:
+```python
+# Training: Deterministic, substantial computation
+def training_step():
+    inputs = mx.random.randint(0, vocab_size, (batch_size, seq_len))  # Fixed size
+    logits = model(inputs)  # Deterministic forward pass
+    loss, grads = mx.value_and_grad(loss_fn)(model, inputs, targets)  # Gradient computation
+    optimizer.update(model, grads)  # Parameter updates
+```
 
-🧠 **Transformer Workloads**:
-- **Attention layers**: (batch×seq_len) × hidden_dim × hidden_dim
-- **MLP expansion**: (batch×seq_len) × hidden_dim × (4×hidden_dim)  
-- **MLP projection**: (batch×seq_len) × (4×hidden_dim) × hidden_dim
-- **Output projection**: (batch×seq_len) × hidden_dim × vocab_size
+**Benefits:**
+- No model loading overhead (1-2 second penalty eliminated)
+- No text generation variability 
+- Deterministic computation graphs
+- Consistent matrix dimensions across runs
+- More matrix operations per evaluation
 
-🏗️ **Apple Silicon Architecture**:
-- **M1/M2**: 16-element vector units, 12-20MB L2 cache
-- **M3/M4**: 32-element AMX units, 24-48MB shared cache
-- **All**: Unified memory with 200-400GB/s bandwidth
-- **Challenge**: Choose optimal tile sizes for each chip and workload
+### 2. **Training-Specific Matrix Patterns**
 
-## How OpenEvolve Optimizes MLX-LM
+Training has unique characteristics that benefit from specialized optimization:
 
-OpenEvolve evolves the `choose_tile_size()` function to:
+🧠 **Training Workload Patterns**:
+- **Larger Batch Sizes**: 16-32 vs 1-4 for inference
+- **Forward + Backward**: Double the matrix operations
+- **Gradient Computation**: Requires transpose operations
+- **Memory Pressure**: Activations + gradients + parameters
+- **Repeated Patterns**: Same operations across many training steps
 
-1. **Detect workload patterns** (attention vs MLP) mathematically
-2. **Adapt to Apple Silicon variant** (M1/M2/M3/M4 specific optimizations)
-3. **Balance memory hierarchy** (L1/L2 cache vs unified memory bandwidth)
-4. **Optimize for real transformer patterns** (not synthetic benchmarks)
+🎯 **Optimization Opportunities**:
+- **Batch-Aware Tiling**: Different strategies for larger batch dimensions
+- **Gradient-Friendly Patterns**: Consider transpose operations in backward pass
+- **Memory Hierarchy**: Balance cache usage with gradient storage
+- **Training Consistency**: Optimize for repeated execution patterns
 
-## Quick Start
+### 3. **Substantial Practical Value**
 
-### Install Dependencies
-```bash
-pip install -r requirements.txt
-```
-
-### Run Real MLX-LM Optimization
-```bash
-python ../../openevolve-run.py initial_program.py evaluator.py --config config.yaml --iterations 200
-```
+Training optimization provides real benefits:
+- **Faster Research Iteration**: Quicker model development cycles
+- **Cost Reduction**: Lower compute costs for training runs  
+- **Better Hardware Utilization**: More efficient use of Apple Silicon
+- **Scalability**: Benefits increase with larger models and datasets
 
-### Resume from Checkpoint
-```bash
-# If interrupted, resume with:
-python ../../openevolve-run.py initial_program.py evaluator.py --config config.yaml --checkpoint ./openevolve_output/mlx_lm_optimization_db/checkpoints/checkpoint_XX --iterations 100
-```
+## Technical Implementation
 
-## What Gets Optimized
+### Matrix Operation Focus
 
-The evolution targets two key functions:
+The evolution targets the key functions used in training:
 
-### 1. Tile Size Selection
 ```python
 def choose_tile_size(M, N, K, device_info):
     """
-    Choose optimal tile sizes for MLX matrix multiplication
-    
-    Args:
-        M, N, K: Matrix dimensions (C = A @ B where A is M×K, B is K×N)
-        device_info: Apple Silicon characteristics (chip, memory, etc.)
-        
-    Returns:
-        (tile_M, tile_N, tile_K): Optimal tile sizes for this workload
+    Optimize for training-specific patterns:
+    - Batch-heavy matrices (large M dimension)
+    - MLP expansion/projection (4x hidden dimension scaling)
+    - Attention computation (square-ish matrices)
+    - Gradient computation (consider transpose patterns)
     """
-    # This function gets evolved by OpenEvolve!
-    # From simple heuristics to sophisticated Apple Silicon optimization
-```
 
-### 2. Optimized Matrix Multiplication
-```python
 def optimized_matmul(A, B, tile_M, tile_N, tile_K):
     """
-    Perform tiled matrix multiplication with optimized memory access patterns
-    
-    Must be numerically correct while maximizing Apple Silicon performance
+    Implement tiled multiplication optimized for:
+    - Training memory access patterns
+    - Apple Silicon architecture
+    - Cache efficiency with gradient storage
     """
-    # This function implements the actual tiled computation
 ```
 
-## Expected Results
-
-OpenEvolve should discover optimizations that provide:
-
-📈 **Inference Speedup**: 5-15% faster text generation  
-📈 **Training Speedup**: 10-25% faster training steps  
-🎯 **Targeted Optimization**: Better performance on larger batches and longer sequences  
-🏗️ **Architecture Awareness**: M3/M4 perform better than M1/M2  
-
-## Real-World Integration
+### Enhanced Training Evaluation
 
-Once optimized, integrate with any MLX-LM workflow:
+The evaluator creates realistic training scenarios:
 
 ```python
-from mlx_lm import load, generate
-from mlx_lm_openevolve import enable_optimizations
-
-# Enable OpenEvolve optimizations
-enable_optimizations("./openevolve_output/best/best_program.py")
+class EnhancedTrainingModel(nn.Module):
+    """
+    Transformer-like model with substantial matrix operations:
+    - Multiple MLP layers (4x expansion/projection)
+    - Attention-like operations  
+    - Large output projections
+    - Forward + backward passes
+    """
 
-# Your existing code gets automatic speedups!
-model, tokenizer = load("mlx-community/Qwen2.5-0.5B-Instruct-bf16")
-text = generate(model, tokenizer, prompt="Hello world", verbose=True)
+# Training Configuration
+batch_size = 32      # Realistic training batch
+seq_len = 512        # Longer sequences
+hidden_dim = 1024    # Large hidden dimension
+vocab_size = 6000    # Substantial vocabulary
 ```
 
-## Advanced: Understanding the Evaluation
+## Quick Start
 
-The new evaluator directly measures MLX-LM performance:
+### Install Dependencies
+```bash
+pip install -r requirements.txt
+```
 
-### Inference Test
-1. Load Qwen2.5-0.5B-Instruct-bf16 model
-2. Generate text with original MLX
-3. Generate same text with optimized MLX
-4. Measure speedup ratio
+### Run Training-Focused Optimization
+```bash
+python ../../openevolve-run.py initial_program.py evaluator.py --config config.yaml --iterations 200
+```
 
-### Training Test  
-1. Create realistic training scenario with transformer layers
-2. Run training steps with original MLX
-3. Run same steps with optimized MLX
-4. Measure training speedup ratio
+### Resume from Checkpoint
+```bash
+# If interrupted, resume with:
+python ../../openevolve-run.py initial_program.py evaluator.py --config config.yaml --checkpoint ./openevolve_output/mlx_training_optimization_db/checkpoints/checkpoint_XX --iterations 100
+```
+
+## Expected Results
 
-### Combined Score
-- **70% weight**: Inference speedup (most common use case)
-- **30% weight**: Training speedup (development workflows)
-- **Bonus**: Consistent optimization across both workloads
+The training-focused approach should discover optimizations providing:
 
-## Comparison to Synthetic Benchmarks
+📈 **Training Speedup**: 10-25% faster training steps  
+🎯 **Consistent Optimization**: Better signal-to-noise ratio for evolution  
+🔧 **Architecture-Aware**: M1/M2/M3/M4 specific optimizations  
+⚡ **Memory Efficient**: Optimized for training's memory pressure  
 
-| **Synthetic Matrix Benchmark** | **Real MLX-LM Optimization** |
-|--------------------------------|-------------------------------|
-| ❌ Artificial matrix sizes | ✅ Real transformer dimensions |
-| ❌ GFLOPS (doesn't reflect user experience) | ✅ End-to-end speedup (what users feel) |
-| ❌ Isolated operations | ✅ Full model inference/training |
-| ❌ May not transfer to real workloads | ✅ Directly optimizes actual use cases |
+## Evolution Discoveries
 
-## Expected Evolution Discoveries
+Based on training characteristics and Apple Silicon architecture, expect OpenEvolve to discover:
 
-Based on transformer architecture and Apple Silicon characteristics, expect OpenEvolve to discover:
+🧠 **Training Workload Classification**:
+```python
+is_batch_heavy = (M > 256)  # Large batch dimension
+is_mlp = (aspect_ratio_K > 1.5)  # MLP 4x expansion patterns
+is_gradient_computation = (transpose_pattern_detected)  # Backward pass
+```
 
-🧠 **Workload Classification**:
+🔧 **Apple Silicon Training Optimization**:
 ```python
-k_dominance = K / max(M, N)  # Detect MLP vs attention patterns
-aspect_ratio = max(M, N) / min(M, N)  # Handle rectangular matrices
+if "M4" in chip and is_batch_heavy:
+    base_tile = 128; vector_align = 32  # Large tiles for AMX units
+    memory_scale = 1.5  # Training can use more memory
+elif is_mlp and training_workload:
+    k_bias = 1.3  # Favor K dimension for MLP patterns
 ```
 
-🔧 **Chip-Specific Optimization**:
+⚡ **Training Memory Patterns**:
 ```python
-if "M4" in chip:
-    base_tile = 512; vector_align = 32  # Large tiles, AMX units
-elif "M1" in chip:
-    base_tile = 256; vector_align = 16  # Smaller tiles, older architecture
+# Optimize for training's repeated execution
+if total_elements > 1_000_000 and is_training:
+    scale = 1.1  # Larger tiles for substantial computation
+    cache_optimization = "training_friendly"  # Consider gradient storage
 ```
 
-⚡ **Memory Hierarchy Optimization**:
+## Integration with Training Workflows
+
+Once optimized, integrate with any MLX training code:
+
 ```python
-# Balance L2 cache utilization vs memory bandwidth
-cache_factor = device_info["l2_cache_mb"] / 16.0
-memory_factor = min(2.0, device_info["memory_gb"] / 16.0)
+import mlx.core as mx
+from optimized_kernels import enable_training_optimizations
+
+# Enable OpenEvolve training optimizations
+enable_training_optimizations("./openevolve_output/best/best_program.py")
+
+# Your existing training code gets automatic speedups!
+for epoch in range(num_epochs):
+    for batch in dataloader:
+        loss, grads = mx.value_and_grad(loss_fn)(model, batch)
+        optimizer.update(model, grads)  # Now faster!
 ```
 
-This represents a significant advance from generic matrix optimization to **transformer-aware, Apple Silicon-specific, real-world performance optimization**.
+## Comparison: Training vs Inference Optimization
+
+| **Inference Optimization** | **Training Optimization** |
+|------------------------------|---------------------------|
+| ❌ Noisy evaluation (model loading, text generation) | ✅ Clean evaluation (deterministic computation) |
+| ❌ Small matrices (batch=1-4) | ✅ Large matrices (batch=16-32) |
+| ❌ Variable workloads | ✅ Consistent patterns |
+| ❌ Complex pipeline overhead | ✅ Direct matrix operation focus |
+| ❌ Difficult signal extraction | ✅ Clear optimization signal |
 
 ## Research Impact
 
-This approach demonstrates:
+This training-focused approach demonstrates:
 
-1. **Practical AI Optimization**: Directly optimizing real AI workloads, not synthetic benchmarks
-2. **Hardware-Software Co-Design**: Evolving algorithms specifically for Apple Silicon architecture  
-3. **Measurable User Benefit**: End-to-end speedups that users actually experience
-4. **Automated Discovery**: Finding optimizations that would take experts months to develop manually
+1. **Practical AI Acceleration**: Directly optimizing the bottleneck of model development
+2. **Hardware-Software Co-Design**: Training-specific optimizations for Apple Silicon  
+3. **Clear Evaluation Methodology**: Robust metrics for evolutionary optimization
+4. **Real-World Application**: Immediate benefits for ML researchers and practitioners
 
-This moves beyond proof-of-concept to **production-ready AI performance optimization**.
+This moves from proof-of-concept to **production-ready training acceleration** that ML practitioners can immediately benefit from.