Fine-Tuning Small LLMs on a Desktop - Part 1: Setup and Environment

# test_setup.py
import sys
import torch
import transformers
import unsloth
import pandas as pd
import numpy as np

print("🚀 Installation Verification")
print("=" * 40)
print(f"Python version: {sys.version}")
print(f"PyTorch version: {torch.__version__}")
print(f"Transformers version: {transformers.__version__}")

# Check GPU support
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"MPS available: {torch.backends.mps.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"GPU count: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
        memory_gb = torch.cuda.get_device_properties(i).total_memory / 1e9
        print(f"    Memory: {memory_gb:.1f} GB")
elif torch.backends.mps.is_available():
    print("✅ Apple Silicon GPU (MPS) available")
    print("   Training will use Metal Performance Shaders")
else:
    print("⚠️ No GPU acceleration available - will use CPU")

# Test Unsloth import
try:
    from unsloth import FastLanguageModel
    print("✅ Unsloth imported successfully")
except ImportError as e:
    print(f"❌ Unsloth import failed: {e}")

# Test device selection
if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

print(f"🎯 Recommended device for training: {device}")

print("\n🎉 Setup verification complete!")

# test_complete_setup.py
import os
import sys
import requests
import subprocess
from pathlib import Path
import torch
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

def test_python_environment():
    """Test Python environment and packages"""
    print("🐍 Testing Python Environment")
    print("-" * 30)

    required_packages = [
        'torch', 'transformers', 'datasets', 'accelerate',
        'unsloth', 'bitsandbytes', 'peft', 'wandb'
    ]

    missing_packages = []
    for package in required_packages:
        try:
            __import__(package)
            print(f"✅ {package}")
        except ImportError:
            print(f"❌ {package}")
            missing_packages.append(package)

    if missing_packages:
        print(f"\n⚠️  Missing packages: {', '.join(missing_packages)}")
        return False
    return True

def test_gpu_setup():
    """Test GPU availability and configuration"""
    print("\n🖥️  Testing GPU Setup")
    print("-" * 30)

    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        print(f"✅ CUDA available with {gpu_count} GPU(s)")

        for i in range(gpu_count):
            name = torch.cuda.get_device_name(i)
            memory = torch.cuda.get_device_properties(i).total_memory / 1e9
            print(f"   GPU {i}: {name} ({memory:.1f} GB)")
        return True
    else:
        print("⚠️  CUDA not available - will use CPU")
        return True  # Not a failure, just different configuration

def test_docker_setup():
    """Test Docker and Docker Model Runner"""
    print("\n🐳 Testing Docker Setup")
    print("-" * 30)

    try:
        # Test basic Docker
        result = subprocess.run(['docker', '--version'],
                              capture_output=True, text=True)
        if result.returncode == 0:
            print(f"✅ Docker: {result.stdout.strip()}")
        else:
            print("❌ Docker not available")
            return False

        # Test Docker Model Runner
        result = subprocess.run(['docker', 'model', '--help'],
                              capture_output=True, text=True)
        if result.returncode == 0:
            print("✅ Docker Model Runner available")
        else:
            print("❌ Docker Model Runner not enabled")
            return False

        return True
    except FileNotFoundError:
        print("❌ Docker not found in PATH")
        return False

def test_ollama_setup():
    """Test Ollama installation and service"""
    print("\n🦙 Testing Ollama Setup")
    print("-" * 30)

    try:
        # Test Ollama binary
        result = subprocess.run(['ollama', '--version'],
                              capture_output=True, text=True)
        if result.returncode == 0:
            print(f"✅ Ollama: {result.stdout.strip()}")
        else:
            print("❌ Ollama not available")
            return False

        # Test Ollama service
        try:
            response = requests.get("http://localhost:11434/api/tags", timeout=5)
            if response.status_code == 200:
                print("✅ Ollama service running")
                models = response.json().get("models", [])
                print(f"   {len(models)} models available")
            else:
                print("⚠️  Ollama service not responding (run 'ollama serve')")
        except requests.exceptions.ConnectionError:
            print("⚠️  Ollama service not running (run 'ollama serve')")

        return True
    except FileNotFoundError:
        print("❌ Ollama not found in PATH")
        return False

def test_project_structure():
    """Test project directory structure"""
    print("\n📁 Testing Project Structure")
    print("-" * 30)

    required_dirs = [
        'data', 'models', 'notebooks', 'scripts',
        'configs', 'logs', 'tests', 'docker'
    ]

    missing_dirs = []
    for dir_name in required_dirs:
        if Path(dir_name).exists():
            print(f"✅ {dir_name}/")
        else:
            print(f"❌ {dir_name}/")
            missing_dirs.append(dir_name)

    if missing_dirs:
        print(f"\n⚠️  Missing directories: {', '.join(missing_dirs)}")
        return False
    return True

def test_model_loading():
    """Test loading a small model"""
    print("\n🤖 Testing Model Loading")
    print("-" * 30)

    try:
        from unsloth import FastLanguageModel

        print("Loading small test model...")
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name="unsloth/tinyllama-bnb-4bit",
            max_seq_length=512,
            dtype=None,
            load_in_4bit=True,
        )
        print("✅ Model loaded successfully")

        # Test tokenization
        test_text = "Hello, world!"
        tokens = tokenizer(test_text, return_tensors="pt")
        print(f"✅ Tokenization test passed ({len(tokens['input_ids'][0])} tokens)")

        return True
    except Exception as e:
        print(f"❌ Model loading failed: {e}")
        return False

def main():
    """Run all tests"""
    print("🧪 Complete Setup Testing")
    print("=" * 50)

    tests = [
        test_python_environment,
        test_gpu_setup,
        test_docker_setup,
        test_ollama_setup,
        test_project_structure,
        test_model_loading
    ]

    results = []
    for test in tests:
        try:
            result = test()
            results.append(result)
        except Exception as e:
            print(f"❌ Test failed with error: {e}")
            results.append(False)

    # Summary
    print("\n📊 Test Summary")
    print("=" * 50)
    passed = sum(results)
    total = len(results)

    print(f"Tests passed: {passed}/{total}")

    if passed == total:
        print("🎉 All tests passed! Your environment is ready for fine-tuning.")
    else:
        print("⚠️  Some tests failed. Please address the issues above.")
        print("\n💡 Common fixes:")
        print("- Install missing Python packages: pip install -r requirements.txt")
        print("- Enable Docker Model Runner in Docker Desktop settings")
        print("- Start Ollama service: ollama serve")
        print("- Create missing directories: mkdir -p data models notebooks scripts configs logs tests docker")

    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)

# Solution: Use smaller batch sizes and 4-bit quantization
load_in_4bit = True
per_device_train_batch_size = 1
gradient_accumulation_steps = 8

# Solution: Add MPS fallback for unsupported operations
import os
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'

# Alternative: Force CPU for problematic operations
device = "cpu"  # Instead of "mps" if issues persist