utils_optimization.py

"""
Model optimization utilities: quantization, compilation, and performance enhancements.
"""
import torch
import torch.nn as nn
from typing import Optional, Dict, Any
import logging

logger = logging.getLogger(__name__)


def apply_quantization(
    model: nn.Module,
    quantization_bits: int = 8,
    quantization_type: str = "dynamic"
) -> nn.Module:
    """
    Apply quantization to the model.
    
    Args:
        model: PyTorch model to quantize
        quantization_bits: Number of bits (8 or 4)
        quantization_type: Type of quantization ('dynamic', 'static', 'qat')
    
    Returns:
        Quantized model
    """
    if quantization_bits == 8:
        if quantization_type == "dynamic":
            logger.info("Applying dynamic 8-bit quantization")
            model = torch.quantization.quantize_dynamic(
                model, {nn.Linear}, dtype=torch.qint8
            )
        elif quantization_type == "static":
            logger.info("Applying static 8-bit quantization")
            # Static quantization requires calibration data
            model.eval()
            model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
            torch.quantization.prepare(model, inplace=True)
            # Note: Calibration should be done with actual data
            torch.quantization.convert(model, inplace=True)
    elif quantization_bits == 4:
        logger.info("Applying 4-bit quantization (requires bitsandbytes)")
        try:
            import bitsandbytes as bnb
            # Apply 4-bit quantization to linear layers
            for name, module in model.named_modules():
                if isinstance(module, nn.Linear):
                    # Replace with 4-bit quantized linear
                    model._modules[name] = bnb.nn.Linear4bit(
                        module.in_features,
                        module.out_features,
                        bias=module.bias is not None
                    )
        except ImportError:
            logger.warning("bitsandbytes not available, skipping 4-bit quantization")
    
    return model


def apply_torch_compile(
    model: nn.Module,
    mode: str = "reduce-overhead",
    fullgraph: bool = False
) -> nn.Module:
    """
    Apply torch.compile for faster inference.
    
    Args:
        model: PyTorch model to compile
        mode: Compilation mode ('default', 'reduce-overhead', 'max-autotune')
        fullgraph: Whether to compile the entire graph
    
    Returns:
        Compiled model
    """
    if hasattr(torch, "compile"):
        logger.info(f"Compiling model with mode={mode}")
        try:
            model = torch.compile(model, mode=mode, fullgraph=fullgraph)
            logger.info("Model compilation successful")
        except Exception as e:
            logger.warning(f"Model compilation failed: {e}")
    else:
        logger.warning("torch.compile not available (requires PyTorch 2.0+)")
    
    return model


def apply_gradient_checkpointing(model: nn.Module, enable: bool = True):
    """
    Enable gradient checkpointing to save memory during training.
    
    Args:
        model: PyTorch model
        enable: Whether to enable gradient checkpointing
    """
    if enable:
        if hasattr(model, "gradient_checkpointing_enable"):
            model.gradient_checkpointing_enable()
            logger.info("Gradient checkpointing enabled")
        else:
            logger.warning("Model does not support gradient checkpointing")


def get_model_size_mb(model: nn.Module) -> float:
    """
    Calculate model size in megabytes.
    
    Args:
        model: PyTorch model
    
    Returns:
        Model size in MB
    """
    param_size = 0
    buffer_size = 0
    
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
    
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
    
    size_all_mb = (param_size + buffer_size) / 1024**2
    return size_all_mb


def count_parameters(model: nn.Module) -> Dict[str, int]:
    """
    Count trainable and total parameters.
    
    Args:
        model: PyTorch model
    
    Returns:
        Dictionary with parameter counts
    """
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    
    return {
        "trainable": trainable,
        "total": total,
        "trainable_millions": trainable / 1e6,
        "total_millions": total / 1e6,
    }


def optimize_model_for_inference(
    model: nn.Module,
    use_quantization: bool = False,
    quantization_bits: int = 8,
    use_torch_compile: bool = False,
    compile_mode: str = "reduce-overhead"
) -> nn.Module:
    """
    Apply multiple optimizations for inference.
    
    Args:
        model: PyTorch model
        use_quantization: Whether to apply quantization
        quantization_bits: Bits for quantization
        use_torch_compile: Whether to compile model
        compile_mode: Compilation mode
    
    Returns:
        Optimized model
    """
    model.eval()
    
    if use_quantization:
        model = apply_quantization(model, quantization_bits)
    
    if use_torch_compile:
        model = apply_torch_compile(model, mode=compile_mode)
    
    # Set to inference mode
    with torch.no_grad():
        model.eval()
    
    return model