attention_head_pruning.py

"""
Attention Head Pruning Pass

Analyzes importance of each attention head and prunes low-importance heads
to reduce model size and improve inference speed.

Now uses enhanced importance calculation supporting HuggingFace models.
"""

import torch
import torch.nn as nn
from typing import Dict, List, Tuple, Any
import numpy as np

# Import enhanced importance calculator
try:
    from .enhanced_head_importance import compute_head_importance_enhanced, AttentionHeadDetector
    ENHANCED_AVAILABLE = True
except ImportError:
    ENHANCED_AVAILABLE = False


def compute_head_importance_gradient(
    model: nn.Module,
    calibration_data: torch.Tensor,
    loss_fn=None
) -> Dict[str, List[float]]:
    """
    Compute importance score for each attention head using gradient magnitude.
    
    Args:
        model: PyTorch model with MultiheadAttention layers
        calibration_data: Sample data for computing importance
        loss_fn: Optional loss function (uses random target if None)
    
    Returns:
        Dictionary mapping attention layer names to head importance scores
    """
    # Use enhanced method if available
    if ENHANCED_AVAILABLE:
        importance_dict = compute_head_importance_enhanced(
            model, calibration_data, method='gradient', num_samples=10
        )
        # Convert numpy arrays to lists
        return {k: v.tolist() for k, v in importance_dict.items()}
    
    # Fallback to original implementation
    model.eval()
    head_importance = {}
    
    for name, module in model.named_modules():
        if isinstance(module, nn.MultiheadAttention):
            num_heads = module.num_heads
            importance_scores = [1.0 / num_heads] * num_heads
            head_importance[name] = importance_scores
    
    return head_importance


def compute_head_importance_entropy(
    model: nn.Module,
    calibration_data: torch.Tensor
) -> Dict[str, List[float]]:
    """
    Compute importance score based on attention pattern entropy.
    Lower entropy = more focused attention = potentially more important.
    
    Args:
        model: PyTorch model with MultiheadAttention layers or HuggingFace models
        calibration_data: Sample data for computing attention patterns
    
    Returns:
        Dictionary mapping attention layer names to head importance scores
    """
    # Use enhanced method if available
    if ENHANCED_AVAILABLE:
        importance_dict = compute_head_importance_enhanced(
            model, calibration_data, method='entropy', num_samples=10
        )
        # Convert numpy arrays to lists
        return {k: v.tolist() for k, v in importance_dict.items()}
    
    # Fallback to original implementation
    model.eval()
    head_importance = {}
    
    # Hook to capture attention weights
    attention_weights = {}
    
    def hook_fn(name):
        def hook(module, input, output):
            # MultiheadAttention returns (attn_output, attn_weights)
            if len(output) == 2:
                attention_weights[name] = output[1].detach()
        return hook
    
    # Register hooks
    hooks = []
    for name, module in model.named_modules():
        if isinstance(module, nn.MultiheadAttention):
            hook = module.register_forward_hook(hook_fn(name))
            hooks.append(hook)
    
    # Forward pass to capture attention weights
    with torch.no_grad():
        model(calibration_data)
    
    # Remove hooks
    for hook in hooks:
        hook.remove()
    
    # Compute entropy for each head
    for name, attn_w in attention_weights.items():
        # attn_w can be either:
        # - (batch, num_heads, seq_len, seq_len) for batch_first=True with return_attention_weights
        # - (seq_len, batch, num_heads) averaged weights
        # - (batch, seq_len, seq_len) if heads already averaged
        
        if attn_w.dim() == 4:
            # Standard format: (batch, num_heads, seq_len, seq_len)
            num_heads = attn_w.shape[1]
        elif attn_w.dim() == 3:
            # Might be (batch, seq_len, seq_len) - already averaged, use uniform importance
            # For now, skip this layer or assign uniform importance
            # We'll set num_heads based on the model
            module = dict(model.named_modules()).get(name)
            if module and hasattr(module, 'num_heads'):
                num_heads = module.num_heads
                # Use uniform importance
                importance_scores = [1.0 / num_heads] * num_heads
                importance_scores = [s / sum(importance_scores) for s in importance_scores]
                head_importance[name] = importance_scores
            continue
        else:
            # Unknown format, skip
            continue
        
        importance_scores = []
        
        for head_idx in range(num_heads):
            head_attn = attn_w[:, head_idx, :, :]
            
            # Compute entropy: -sum(p * log(p))
            # Lower entropy = more focused = more important
            epsilon = 1e-10
            head_attn_safe = head_attn + epsilon
            entropy = -torch.sum(head_attn_safe * torch.log(head_attn_safe), dim=-1)
            avg_entropy = entropy.mean().item()
            
            # Convert entropy to importance (inverse relationship)
            importance = 1.0 / (avg_entropy + 1.0)
            importance_scores.append(importance)
        
        # Normalize importance scores
        total = sum(importance_scores)
        if total > 0:
            importance_scores = [s / total for s in importance_scores]
        
        head_importance[name] = importance_scores
    
    return head_importance


def compute_head_importance_output_variance(
    model: nn.Module,
    calibration_data: torch.Tensor
) -> Dict[str, List[float]]:
    """
    Compute importance based on output contribution variance.
    Higher variance in output = more important head.
    
    Args:
        model: PyTorch model with MultiheadAttention layers
        calibration_data: Sample data for computing output variance
    
    Returns:
        Dictionary mapping attention layer names to head importance scores
    """
    model.eval()
    head_importance = {}
    
    # This is a placeholder implementation
    # Full implementation would capture per-head outputs and compute variance
    
    for name, module in model.named_modules():
        if isinstance(module, nn.MultiheadAttention):
            num_heads = module.num_heads
            
            # For MVP, use uniform importance
            importance_scores = [1.0 / num_heads] * num_heads
            
            head_importance[name] = importance_scores
    
    return head_importance


def prune_attention_heads(
    mir: Dict[str, Any],
    head_importance: Dict[str, List[float]],
    pruning_ratio: float = 0.3
) -> Dict[str, Any]:
    """
    Prune low-importance attention heads from MIR representation.
    
    Args:
        mir: MIR representation
        head_importance: Importance scores per head for each attention layer
        pruning_ratio: Fraction of heads to prune (0.0 to 1.0)
    
    Returns:
        Updated MIR with pruned heads
    """
    if not (0.0 <= pruning_ratio < 1.0):
        raise ValueError("Pruning ratio must be in [0.0, 1.0)")
    
    pruning_stats = {
        "total_heads_before": 0,
        "total_heads_after": 0,
        "layers_pruned": []
    }
    
    for node in mir["graph"]["nodes"]:
        if node["op_type"] == "attention_qkv_projection":
            origin_name = node["metadata"]["origin_name"]
            
            if origin_name in head_importance:
                scores = head_importance[origin_name]
                num_heads = node["params"]["num_heads"]
                
                pruning_stats["total_heads_before"] += num_heads
                
                # Determine heads to keep
                keep_count = max(1, int(num_heads * (1 - pruning_ratio)))
                heads_to_keep = sorted(
                    range(num_heads),
                    key=lambda i: scores[i],
                    reverse=True
                )[:keep_count]
                
                pruned_heads = [i for i in range(num_heads) if i not in heads_to_keep]
                
                # Update node parameters
                node["params"]["num_heads"] = keep_count
                node["params"]["pruned_heads"] = pruned_heads
                node["params"]["kept_heads"] = heads_to_keep
                
                # Update compression metadata
                if "compression" not in node:
                    node["compression"] = {}
                
                node["compression"]["head_pruning"] = {
                    "original_heads": num_heads,
                    "remaining_heads": keep_count,
                    "pruning_ratio": (num_heads - keep_count) / num_heads,
                    "pruned_head_indices": pruned_heads,
                    "kept_head_indices": heads_to_keep,
                    "importance_scores": scores
                }
                
                pruning_stats["total_heads_after"] += keep_count
                pruning_stats["layers_pruned"].append({
                    "layer": origin_name,
                    "original_heads": num_heads,
                    "remaining_heads": keep_count
                })
    
    # Add pruning statistics to MIR metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    
    mir["metadata"]["head_pruning_stats"] = pruning_stats
    
    return mir


def apply_head_pruning(
    mir: Dict[str, Any],
    model: nn.Module,
    calibration_data: torch.Tensor,
    pruning_ratio: float = 0.3,
    importance_method: str = "entropy"
) -> Dict[str, Any]:
    """
    Complete head pruning pipeline: compute importance and prune.
    
    Args:
        mir: MIR representation
        model: Original PyTorch model
        calibration_data: Data for computing head importance
        pruning_ratio: Fraction of heads to prune
        importance_method: Method for computing importance
            - "gradient": Gradient-based importance
            - "entropy": Attention pattern entropy
            - "variance": Output variance
    
    Returns:
        Updated MIR with pruned heads
    """
    # Compute head importance
    if importance_method == "gradient":
        head_importance = compute_head_importance_gradient(model, calibration_data)
    elif importance_method == "entropy":
        head_importance = compute_head_importance_entropy(model, calibration_data)
    elif importance_method == "variance":
        head_importance = compute_head_importance_output_variance(model, calibration_data)
    else:
        raise ValueError(f"Unknown importance method: {importance_method}")
    
    # Prune heads
    mir = prune_attention_heads(mir, head_importance, pruning_ratio)
    
    return mir


def get_pruning_statistics(mir: Dict[str, Any]) -> Dict[str, Any]:
    """
    Extract head pruning statistics from MIR.
    
    Args:
        mir: MIR representation with pruning metadata
    
    Returns:
        Dictionary with pruning statistics
    """
    stats = mir.get("metadata", {}).get("head_pruning_stats", {})
    
    if not stats:
        return {
            "pruning_applied": False,
            "message": "No head pruning metadata found"
        }
    
    total_before = stats.get("total_heads_before", 0)
    total_after = stats.get("total_heads_after", 0)
    
    return {
        "pruning_applied": True,
        "total_heads_before": total_before,
        "total_heads_after": total_after,
        "heads_pruned": total_before - total_after,
        "overall_pruning_ratio": (total_before - total_after) / total_before if total_before > 0 else 0,
        "layers_pruned": stats.get("layers_pruned", [])
    }