This is where the magic happens! In this part, we will deep dive into LoRA (Low-Rank Adaptation) fine-tuning and use MLX to train our model with incredible efficiency on Apple Silicon.
Understanding LoRA: The Game-Changing Technique
Imagine you are a master chef who wants to learn a new cuisine. Instead of forgetting everything you know and starting from scratch, you add new techniques and flavor profiles to your existing knowledge. That's exactly what LoRA (Low-Rank Adaptation) does for language models.
The Traditional Fine-Tuning Problem
Traditional fine-tuning updates all 1.7 billion parameters of our model. This means:
- ❌ Massive memory requirements
- ❌ Slow training
- ❌ Risk of "catastrophic forgetting" (losing general knowledge)
- ❌ Large model files
The LoRA Solution
LoRA adds small "adapter" layers that learn new behaviors while keeping the original model frozen:
- ✅ Minimal memory usage
- ✅ Fast training
- ✅ Preserves general knowledge
- ✅ Tiny adapter file size
- ✅ Can be combined or switched out easily
How LoRA Works Under the Hood
Think of the original model as a Swiss Army knife with all its tools welded in place. LoRA adds new attachments that can be snapped on or off.
MLX: Apple's Secret Weapon for AI
MLX is Apple's machine learning framework designed specifically for Apple Silicon. It's what makes our local fine-tuning possible and incredibly fast.
Why MLX is good for Local AI
- Unified Memory Architecture: M-series chips share memory between CPU and GPU, eliminating data transfer bottlenecks
- Optimized Computation: Hand-tuned for Apple Silicon's specific capabilities
- Memory Efficiency: Intelligent memory management for maximum model sizes
- Python Integration: Easy to use while being incredibly fast
Setting Up Our Fine-Tuning Pipeline
Let us build our fine-tuning system step by step, understanding each component.
Step 1: Configuration and Setup
First, let's create a comprehensive configuration system:
touch fine_tuning_config.py
# Create fine_tuning_config.py import os from pathlib import Path import mlx.core as mx class FineTuningConfig: """Centralized configuration for fine-tuning""" def __init__(self): # Model configuration self.base_model = "HuggingFaceTB/SmolLM2-1.7B-Instruct" self.adapter_path = "./adapters/email_sentiment" # Data paths self.train_data_path = "./data/mlx_format/train.jsonl" self.valid_data_path = "./data/mlx_format/valid.jsonl" # LoRA parameters self.lora_layers = 16 # Number of transformer layers to add LoRA to self.lora_rank = 16 # The 'r' in LoRA - higher = more capacity but slower self.lora_alpha = 32 # Scaling factor for LoRA adapters # Training parameters self.batch_size = 2 # Batch size (reduce if out of memory) self.learning_rate = 5e-5 # Learning rate self.max_iters = 1000 # Maximum training iterations self.steps_per_report = 10 # How often to print progress self.steps_per_eval = 200 # How often to run validation self.save_every = 400 # How often to save checkpoints # Hardware optimization self.use_gpu = mx.metal.is_available() self.max_sequence_length = 2048 # Create directories Path(self.adapter_path).mkdir(parents=True, exist_ok=True) def print_config(self): """Print current configuration""" print("🔧 Fine-tuning Configuration:") print(f" Base model: {self.base_model}") print(f" GPU available: {self.use_gpu}") print(f" LoRA rank: {self.lora_rank}") print(f" LoRA layers: {self.lora_layers}") print(f" Batch size: {self.batch_size}") print(f" Learning rate: {self.learning_rate}") print(f" Max iterations: {self.max_iters}") print(f" Adapter path: {self.adapter_path}") # Create and test config if __name__ == "__main__": config = FineTuningConfig() config.print_config() Step 2: Memory and Performance Monitoring
Before we start fine-tuning, let's create tools to monitor our system:
touch monitoring.py
# Create monitoring.py import time import mlx.core as mx from typing import Dict, List import psutil class PerformanceMonitor: """Monitor memory usage and training performance""" def __init__(self): self.start_time = time.time() self.metrics = [] def log_memory_usage(self, step: int, loss: float = None): """Log current memory and performance metrics""" # GPU memory (if available) gpu_memory = {} if mx.metal.is_available(): gpu_memory = { 'active_mb': mx.metal.get_active_memory() / 1e6, 'peak_mb': mx.metal.get_peak_memory() / 1e6 } # System memory system_memory = psutil.virtual_memory() # Training metrics elapsed = time.time() - self.start_time metrics = { 'step': step, 'elapsed_seconds': elapsed, 'loss': loss, 'gpu_active_mb': gpu_memory.get('active_mb', 0), 'gpu_peak_mb': gpu_memory.get('peak_mb', 0), 'system_memory_percent': system_memory.percent, 'system_memory_available_gb': system_memory.available / 1e9 } self.metrics.append(metrics) if step % 50 == 0: # Print every 50 steps self.print_status(metrics) return metrics def print_status(self, metrics: Dict): """Print current training status""" print(f"Step {metrics['step']:4d} | " f"Loss: {metrics['loss']:.4f} | " f"GPU: {metrics['gpu_active_mb']:.0f}MB | " f"Time: {metrics['elapsed_seconds']:.1f}s") def get_training_summary(self): """Get summary of training run""" if not self.metrics: return {} peak_gpu = max(m['gpu_peak_mb'] for m in self.metrics) total_time = self.metrics[-1]['elapsed_seconds'] final_loss = self.metrics[-1]['loss'] return { 'total_training_time': total_time, 'peak_gpu_memory_mb': peak_gpu, 'final_loss': final_loss, 'steps_completed': len(self.metrics) } Step 3: The Fine-Tuning Engine
Now let's create our main fine-tuning script using MLX-LM:
touch fine_tune_model.py
# Create fine_tune_model.py import subprocess import time import json import os from pathlib import Path from fine_tuning_config import FineTuningConfig from monitoring import PerformanceMonitor class MLXFineTuner: """Fine-tune models using MLX with LoRA""" def __init__(self, config: FineTuningConfig): self.config = config self.monitor = PerformanceMonitor() def validate_data(self): """Validate that training data exists and is properly formatted""" print("📊 Validating training data...") if not os.path.exists(self.config.train_data_path): raise FileNotFoundError(f"Training data not found: {self.config.train_data_path}") # Count training examples train_count = 0 with open(self.config.train_data_path, 'r') as f: for line in f: if line.strip(): train_count += 1 print(f"✅ Found {train_count} training examples") # Validate format with open(self.config.train_data_path, 'r') as f: first_line = f.readline() try: example = json.loads(first_line) if 'text' not in example: raise ValueError("Training data must have 'text' field") print("✅ Data format validated") except json.JSONDecodeError: raise ValueError("Training data must be valid JSONL format") return train_count def build_training_command(self): """Build the MLX-LM training command""" cmd = [ "python3", "-m", "mlx_lm", "lora", "--model", self.config.base_model, "--train", "--data", "./data/mlx_format", # Directory containing train.jsonl "--batch-size", str(self.config.batch_size), "--iters", str(self.config.max_iters), "--learning-rate", str(self.config.learning_rate), "--steps-per-report", str(self.config.steps_per_report), "--steps-per-eval", str(self.config.steps_per_eval), "--adapter-path", self.config.adapter_path, "--save-every", str(self.config.save_every) ] return cmd def run_fine_tuning(self): """Execute the fine-tuning process""" print("🚀 Starting LoRA fine-tuning with MLX...") print("=" * 60) # Validate everything is ready train_count = self.validate_data() self.config.print_config() # Build command cmd = self.build_training_command() print(f"\n📝 Command: {' '.join(cmd)}") # Start training start_time = time.time() print(f"\n🏃 Training started at {time.strftime('%H:%M:%S')}") print(f"📚 Training on {train_count} examples") print("💡 This typically takes 3-10 minutes on Apple Silicon M3") print("⏰ Progress will be reported every 10 steps\n") try: # Run the training command result = subprocess.run(cmd, capture_output=True, text=True, check=True) training_time = time.time() - start_time print("\n" + "="*60) print("🎉 Fine-tuning completed successfully!") print(f"⏱️ Total training time: {training_time:.1f} seconds") print(f"💾 Adapters saved to: {self.config.adapter_path}") # Save training metadata metadata = { 'model_name': self.config.base_model, 'training_time_seconds': training_time, 'training_examples': train_count, 'lora_rank': self.config.lora_rank, 'lora_layers': self.config.lora_layers, 'batch_size': self.config.batch_size, 'learning_rate': self.config.learning_rate, 'max_iters': self.config.max_iters, 'timestamp': time.time(), 'command_used': ' '.join(cmd) } metadata_path = f"{self.config.adapter_path}/training_metadata.json" with open(metadata_path, 'w') as f: json.dump(metadata, f, indent=2) print(f"📊 Training metadata saved to: {metadata_path}") # Parse and display training output self.parse_training_output(result.stdout) return True, metadata except subprocess.CalledProcessError as e: print("\n❌ Fine-tuning failed!") print(f"Error code: {e.returncode}") print(f"Error output: {e.stderr}") print(f"Standard output: {e.stdout}") return False, None def parse_training_output(self, output: str): """Parse and display key information from training output""" print("\n📈 Training Progress Summary:") print("-" * 40) lines = output.split('\n') # Look for key training metrics for line in lines: if 'Loss:' in line or 'Validation' in line: print(f" {line.strip()}") # Look for final metrics for line in reversed(lines): if 'Loss:' in line: print(f"\n🎯 Final training loss: {line.split('Loss:')[-1].strip()}") break def verify_training_output(self): """Verify that training produced the expected files""" print("\n🔍 Verifying training output...") adapter_path = Path(self.config.adapter_path) # Check for adapter files adapter_files = list(adapter_path.glob("*.safetensors")) + list(adapter_path.glob("*.npz")) if adapter_files: print(f"✅ Found adapter files: {[f.name for f in adapter_files]}") else: print("❌ No adapter files found") return False # Check for configuration config_file = adapter_path / "adapter_config.json" if config_file.exists(): print(f"✅ Found adapter config: {config_file}") # Display config contents with open(config_file, 'r') as f: config_data = json.load(f) print(f" LoRA rank: {config_data.get('r', 'unknown')}") print(f" LoRA alpha: {config_data.get('lora_alpha', 'unknown')}") else: print("⚠️ No adapter config found") # Calculate total size total_size = sum(f.stat().st_size for f in adapter_path.rglob('*') if f.is_file()) print(f"📁 Total adapter size: {total_size / 1e6:.1f} MB") return True def main(): """Main fine-tuning execution""" print("🤖 MLX LoRA Fine-Tuning Pipeline") print("=" * 50) # Create configuration config = FineTuningConfig() # Create fine-tuner fine_tuner = MLXFineTuner(config) # Run fine-tuning success, metadata = fine_tuner.run_fine_tuning() if success: # Verify output fine_tuner.verify_training_output() print("\n✨ Fine-tuning pipeline completed successfully!") print("\n🎯 Next steps:") print(" 1. Test your fine-tuned model") print(" 2. Run evaluation to measure performance") print(" 3. Build your application interface") return metadata else: print("\n💥 Fine-tuning failed. Please check the error messages above.") return None if __name__ == "__main__": metadata = main() 
Top comments (0)