zen-family/tools/conversion/mlx-conversion/optimize.py

#!/usr/bin/env python3
"""
MLX Model Optimization Suite for Apple Silicon
Memory-efficient inference and performance tuning
"""

import argparse
import json
import os
from pathlib import Path
from typing import Optional, Dict, Any
import logging
import subprocess

import mlx.core as mx
import mlx.nn as nn
from mlx_lm import load

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class AppleSiliconOptimizer:
    """Optimize MLX models for Apple Silicon unified memory"""

    def __init__(self):
        self.device_info = self._get_device_info()
        logger.info(f"Device: {self.device_info}")

    def _get_device_info(self) -> Dict[str, Any]:
        """Get Apple Silicon device information"""
        info = {
            "device": str(mx.default_device()),
            "metal_available": mx.metal.is_available() if hasattr(mx, 'metal') else False
        }

        # Get system memory
        try:
            result = subprocess.run(
                ["sysctl", "hw.memsize"],
                capture_output=True,
                text=True
            )
            if result.returncode == 0:
                mem_bytes = int(result.stdout.split(":")[1].strip())
                info["unified_memory_gb"] = mem_bytes / (1024**3)
        except Exception:
            pass

        # Get chip info
        try:
            result = subprocess.run(
                ["sysctl", "-n", "machdep.cpu.brand_string"],
                capture_output=True,
                text=True
            )
            if result.returncode == 0:
                info["chip"] = result.stdout.strip()
        except Exception:
            pass

        return info

    def optimize_model(self, model_path: Path) -> Dict[str, Any]:
        """Apply optimizations to model for Apple Silicon"""
        logger.info(f"Optimizing model: {model_path}")

        optimizations = {}

        # 1. Memory mapping optimization
        self._optimize_memory_mapping(model_path)
        optimizations["memory_mapped"] = True

        # 2. Graph optimization
        self._optimize_compute_graph(model_path)
        optimizations["graph_optimized"] = True

        # 3. Batch size recommendations
        batch_config = self._recommend_batch_size(model_path)
        optimizations["batch_config"] = batch_config

        # 4. Cache optimization
        cache_config = self._optimize_cache(model_path)
        optimizations["cache_config"] = cache_config

        # Save optimization config
        config_path = model_path / "optimization_config.json"
        with open(config_path, "w") as f:
            json.dump(optimizations, f, indent=2)

        logger.info(f"Optimizations saved to {config_path}")
        return optimizations

    def _optimize_memory_mapping(self, model_path: Path):
        """Enable efficient memory mapping for weights"""
        weights_path = model_path / "weights.npz"

        if weights_path.exists():
            # Ensure weights are stored in optimal format
            import numpy as np

            weights = mx.load(str(weights_path))

            # Save with memory-efficient settings
            mx.save(
                str(weights_path),
                weights,
                compressed=True
            )

            logger.info("Memory mapping optimized for weights")

    def _optimize_compute_graph(self, model_path: Path):
        """Optimize compute graph for Metal Performance Shaders"""
        config_path = model_path / "config.json"

        if config_path.exists():
            with open(config_path) as f:
                config = json.load(f)

            # Add Metal optimization flags
            config["metal_kernel_cache"] = True
            config["fused_ops"] = True
            config["graph_optimization_level"] = 2

            with open(config_path, "w") as f:
                json.dump(config, f, indent=2)

            logger.info("Compute graph optimized for Metal")

    def _recommend_batch_size(self, model_path: Path) -> Dict[str, int]:
        """Recommend optimal batch sizes based on model and memory"""
        metadata_path = model_path / "metadata.json"

        if metadata_path.exists():
            with open(metadata_path) as f:
                metadata = json.load(f)

            model_size = metadata.get("original_size", "7B")
            quantized = metadata.get("quantized", False)
            q_bits = metadata.get("quantization_bits", 16)

            # Estimate model memory
            size_gb = float(model_size.rstrip("B"))
            if quantized:
                model_memory_gb = size_gb * (q_bits / 16)
            else:
                model_memory_gb = size_gb * 2

            available_memory = self.device_info.get("unified_memory_gb", 16)
            free_memory = available_memory - model_memory_gb - 2  # Reserve 2GB for system

            # Calculate batch sizes
            batch_config = {
                "inference_batch_size": max(1, int(free_memory / 0.5)),
                "training_batch_size": max(1, int(free_memory / 2)),
                "max_sequence_length": 4096 if free_memory > 4 else 2048,
                "kv_cache_size": min(int(free_memory * 0.3 * 1024), 4096)  # MB
            }

            logger.info(f"Recommended batch config: {batch_config}")
            return batch_config

        return {
            "inference_batch_size": 1,
            "training_batch_size": 1,
            "max_sequence_length": 2048,
            "kv_cache_size": 1024
        }

    def _optimize_cache(self, model_path: Path) -> Dict[str, Any]:
        """Configure optimal caching strategy"""
        cache_config = {
            "kv_cache_dtype": "float16",
            "attention_cache": True,
            "gradient_checkpointing": False,  # For inference
            "cache_implementation": "ring_buffer",
            "max_cache_tokens": 8192
        }

        return cache_config

    def profile_model(self, model_path: Path, prompt: str = "Hello") -> Dict[str, Any]:
        """Profile model performance"""
        from mlx_lm import load, generate
        import time
        import psutil
        import os

        logger.info("Starting performance profiling...")

        # Load model
        model, tokenizer = load(str(model_path))

        # Memory before
        process = psutil.Process(os.getpid())
        mem_before = process.memory_info().rss / (1024**3)  # GB

        # Warmup
        _ = generate(model, tokenizer, prompt, max_tokens=10)

        # Profile generation
        start = time.perf_counter()
        response = generate(model, tokenizer, prompt, max_tokens=100)
        elapsed = time.perf_counter() - start

        # Memory after
        mem_after = process.memory_info().rss / (1024**3)  # GB

        profile_data = {
            "generation_time": elapsed,
            "tokens_per_second": 100 / elapsed,
            "memory_used_gb": mem_after - mem_before,
            "peak_memory_gb": mem_after,
            "model_path": str(model_path)
        }

        logger.info(f"Profile results: {profile_data}")
        return profile_data


class LoRAOptimizer:
    """Optimize LoRA adapters for efficient fine-tuning"""

    @staticmethod
    def create_efficient_lora(
        model_path: Path,
        rank: int = 8,
        target_modules: Optional[list] = None
    ) -> Dict[str, Any]:
        """Create memory-efficient LoRA configuration"""

        config = {
            "rank": rank,
            "alpha": rank * 2,
            "dropout": 0.05,
            "target_modules": target_modules or [
                "q_proj",
                "v_proj"  # Reduced from default to save memory
            ],
            "modules_to_save": [],
            "quantize_base": True,
            "gradient_checkpointing": True
        }

        adapter_path = model_path.parent / f"{model_path.name}-lora-r{rank}"
        adapter_path.mkdir(exist_ok=True)

        with open(adapter_path / "adapter_config.json", "w") as f:
            json.dump(config, f, indent=2)

        logger.info(f"Efficient LoRA config saved to {adapter_path}")
        return config

    @staticmethod
    def merge_lora(
        model_path: Path,
        adapter_path: Path,
        output_path: Path
    ):
        """Merge LoRA adapter with base model for faster inference"""
        from mlx_lm import load, save

        logger.info("Merging LoRA adapter with base model...")

        # Load model with adapter
        model, tokenizer = load(
            str(model_path),
            adapter_path=str(adapter_path)
        )

        # Save merged model
        save(model, tokenizer, str(output_path))

        logger.info(f"Merged model saved to {output_path}")


def optimize_for_deployment(model_path: Path) -> Path:
    """Full optimization pipeline for production deployment"""
    logger.info("Running full deployment optimization pipeline...")

    optimizer = AppleSiliconOptimizer()

    # 1. Apply optimizations
    optimizations = optimizer.optimize_model(model_path)

    # 2. Profile performance
    profile = optimizer.profile_model(model_path)

    # 3. Create deployment package
    deploy_path = model_path.parent / f"{model_path.name}-deployed"
    deploy_path.mkdir(exist_ok=True)

    # Copy optimized model
    import shutil
    shutil.copytree(model_path, deploy_path, dirs_exist_ok=True)

    # Save deployment metadata
    deployment_info = {
        "optimizations": optimizations,
        "performance_profile": profile,
        "device_info": optimizer.device_info,
        "deployment_ready": True
    }

    with open(deploy_path / "deployment.json", "w") as f:
        json.dump(deployment_info, f, indent=2)

    logger.info(f"Deployment-ready model at {deploy_path}")
    return deploy_path


def main():
    parser = argparse.ArgumentParser(
        description="Optimize MLX models for Apple Silicon"
    )
    parser.add_argument(
        "model",
        type=Path,
        help="Path to MLX model directory"
    )
    parser.add_argument(
        "--profile",
        action="store_true",
        help="Profile model performance"
    )
    parser.add_argument(
        "--create-lora",
        action="store_true",
        help="Create efficient LoRA adapter"
    )
    parser.add_argument(
        "--lora-rank",
        type=int,
        default=8,
        help="LoRA rank (default: 8)"
    )
    parser.add_argument(
        "--merge-lora",
        type=Path,
        help="Path to LoRA adapter to merge"
    )
    parser.add_argument(
        "--deploy",
        action="store_true",
        help="Run full deployment optimization"
    )

    args = parser.parse_args()

    if args.deploy:
        optimize_for_deployment(args.model)
    else:
        optimizer = AppleSiliconOptimizer()

        if args.profile:
            profile = optimizer.profile_model(args.model)
            print("\nPerformance Profile:")
            print(json.dumps(profile, indent=2))

        if args.create_lora:
            config = LoRAOptimizer.create_efficient_lora(
                args.model,
                rank=args.lora_rank
            )
            print("\nLoRA Configuration:")
            print(json.dumps(config, indent=2))

        if args.merge_lora:
            output = args.model.parent / f"{args.model.name}-merged"
            LoRAOptimizer.merge_lora(args.model, args.merge_lora, output)

        if not any([args.profile, args.create_lora, args.merge_lora]):
            # Run standard optimization
            optimizations = optimizer.optimize_model(args.model)
            print("\nOptimizations Applied:")
            print(json.dumps(optimizations, indent=2))


if __name__ == "__main__":
    main()