head_pruner.py

"""
Physical Attention Head Pruning for HuggingFace Models

This module implements actual head removal (not just importance scoring).
Physically modifies model architecture to remove low-importance attention heads.
"""

import torch
import torch.nn as nn
from typing import Dict, List, Tuple, Optional, Set
import numpy as np
import copy


class AttentionHeadPruner:
    """
    Prunes attention heads from HuggingFace transformer models.
    
    Supports:
    - BERT-style models (separate Q/K/V projections)
    - GPT-2-style models (combined QKV projection)
    - Preserves model functionality after pruning
    """
    
    def __init__(self, model: nn.Module, model_type: str = 'auto'):
        """
        Initialize pruner.
        
        Args:
            model: HuggingFace model to prune
            model_type: 'bert', 'gpt2', or 'auto' (auto-detect)
        """
        self.model = model
        self.model_type = self._detect_model_type() if model_type == 'auto' else model_type
        self.pruned_heads = {}  # layer_idx -> set of pruned head indices
        
    def _detect_model_type(self) -> str:
        """Detect model architecture type."""
        model_class = type(self.model).__name__.lower()
        
        if 'bert' in model_class:
            return 'bert'
        elif 'gpt' in model_class or 'gpt2' in model_class:
            return 'gpt2'
        elif 'roberta' in model_class:
            return 'bert'  # RoBERTa uses BERT architecture
        else:
            return 'unknown'
    
    def prune_heads(
        self,
        heads_to_prune: Dict[int, List[int]],
        importance_scores: Optional[Dict[str, np.ndarray]] = None
    ) -> Dict[str, any]:
        """
        Prune specified attention heads from the model.
        
        Args:
            heads_to_prune: Dict mapping layer_idx -> list of head indices to prune
                           e.g., {0: [1, 5, 7], 1: [0, 3], ...}
            importance_scores: Optional importance scores for logging
        
        Returns:
            Dict with pruning statistics
        """
        print(f"\n[Head Pruner] Pruning attention heads...")
        print(f"  Model type: {self.model_type}")
        
        if self.model_type == 'bert':
            return self._prune_bert_heads(heads_to_prune, importance_scores)
        elif self.model_type == 'gpt2':
            return self._prune_gpt2_heads(heads_to_prune, importance_scores)
        else:
            raise ValueError(f"Unsupported model type: {self.model_type}")
    
    def _prune_bert_heads(
        self,
        heads_to_prune: Dict[int, List[int]],
        importance_scores: Optional[Dict[str, np.ndarray]] = None
    ) -> Dict[str, any]:
        """
        Prune heads from BERT-style models.
        
        BERT has separate Q, K, V projections:
        - self.query: Linear(hidden_size, all_head_size)
        - self.key: Linear(hidden_size, all_head_size)
        - self.value: Linear(hidden_size, all_head_size)
        - output.dense: Linear(all_head_size, hidden_size)
        
        To prune head i: remove columns from Q/K/V and rows from output.dense
        """
        stats = {
            'layers_pruned': 0,
            'total_heads_pruned': 0,
            'params_before': sum(p.numel() for p in self.model.parameters()),
            'params_after': 0,
        }
        
        # Get BERT encoder layers
        if hasattr(self.model, 'bert'):
            encoder = self.model.bert.encoder
        elif hasattr(self.model, 'encoder'):
            encoder = self.model.encoder
        else:
            raise ValueError("Cannot find encoder in model")
        
        # Prune each layer
        for layer_idx, head_indices in heads_to_prune.items():
            if layer_idx >= len(encoder.layer):
                print(f"  Warning: Layer {layer_idx} out of range, skipping")
                continue
            
            layer = encoder.layer[layer_idx]
            attention = layer.attention
            
            if not hasattr(attention, 'self'):
                print(f"  Warning: Layer {layer_idx} missing self-attention, skipping")
                continue
            
            self_attn = attention.self
            output_layer = attention.output.dense
            
            # Get dimensions
            num_attention_heads = self_attn.num_attention_heads
            attention_head_size = self_attn.attention_head_size
            all_head_size = num_attention_heads * attention_head_size
            
            # Convert head indices to keep (inverse of prune)
            heads_to_keep = [i for i in range(num_attention_heads) if i not in head_indices]
            
            if len(heads_to_keep) == 0:
                print(f"  Warning: Trying to prune all heads from layer {layer_idx}, skipping")
                continue
            
            print(f"  Layer {layer_idx}: Pruning {len(head_indices)} heads, keeping {len(heads_to_keep)}")
            
            # Prune Q, K, V projections (columns)
            self_attn.query = self._prune_linear_columns(
                self_attn.query, heads_to_keep, num_attention_heads
            )
            self_attn.key = self._prune_linear_columns(
                self_attn.key, heads_to_keep, num_attention_heads
            )
            self_attn.value = self._prune_linear_columns(
                self_attn.value, heads_to_keep, num_attention_heads
            )
            
            # Prune output projection (rows)
            attention.output.dense = self._prune_linear_rows(
                output_layer, heads_to_keep, num_attention_heads
            )
            
            # Update head count
            self_attn.num_attention_heads = len(heads_to_keep)
            
            # Track pruned heads
            if layer_idx not in self.pruned_heads:
                self.pruned_heads[layer_idx] = set()
            self.pruned_heads[layer_idx].update(head_indices)
            
            stats['layers_pruned'] += 1
            stats['total_heads_pruned'] += len(head_indices)
        
        stats['params_after'] = sum(p.numel() for p in self.model.parameters())
        stats['compression_ratio'] = stats['params_before'] / stats['params_after']
        
        print(f"\n  Pruning complete:")
        print(f"    Layers pruned: {stats['layers_pruned']}")
        print(f"    Total heads pruned: {stats['total_heads_pruned']}")
        print(f"    Parameters: {stats['params_before']:,} -> {stats['params_after']:,}")
        print(f"    Compression: {stats['compression_ratio']:.2f}x")
        
        return stats
    
    def _prune_gpt2_heads(
        self,
        heads_to_prune: Dict[int, List[int]],
        importance_scores: Optional[Dict[str, np.ndarray]] = None
    ) -> Dict[str, any]:
        """
        Prune heads from GPT-2-style models.
        
        GPT-2 has combined QKV projection:
        - c_attn: Linear(hidden_size, 3 * all_head_size)  # Combined Q, K, V
        - c_proj: Linear(all_head_size, hidden_size)
        
        To prune head i: remove corresponding sections from c_attn and rows from c_proj
        """
        stats = {
            'layers_pruned': 0,
            'total_heads_pruned': 0,
            'params_before': sum(p.numel() for p in self.model.parameters()),
            'params_after': 0,
        }
        
        # Get GPT-2 transformer blocks
        if hasattr(self.model, 'transformer'):
            blocks = self.model.transformer.h
        elif hasattr(self.model, 'h'):
            blocks = self.model.h
        else:
            raise ValueError("Cannot find transformer blocks in model")
        
        # Prune each layer
        for layer_idx, head_indices in heads_to_prune.items():
            if layer_idx >= len(blocks):
                print(f"  Warning: Layer {layer_idx} out of range, skipping")
                continue
            
            block = blocks[layer_idx]
            attn = block.attn
            
            if not hasattr(attn, 'c_attn'):
                print(f"  Warning: Layer {layer_idx} missing c_attn, skipping")
                continue
            
            # Get dimensions
            num_heads = attn.num_heads if hasattr(attn, 'num_heads') else 12  # Default for GPT-2
            embed_dim = attn.embed_dim if hasattr(attn, 'embed_dim') else 768
            head_dim = embed_dim // num_heads
            
            # Convert to heads to keep
            heads_to_keep = [i for i in range(num_heads) if i not in head_indices]
            
            if len(heads_to_keep) == 0:
                print(f"  Warning: Trying to prune all heads from layer {layer_idx}, skipping")
                continue
            
            print(f"  Layer {layer_idx}: Pruning {len(head_indices)} heads, keeping {len(heads_to_keep)}")
            
            # Prune c_attn (combined QKV projection)
            # c_attn has shape [embed_dim, 3 * embed_dim] (3x for Q, K, V)
            attn.c_attn = self._prune_gpt2_c_attn(
                attn.c_attn, heads_to_keep, num_heads, head_dim
            )
            
            # Prune c_proj (output projection)
            attn.c_proj = self._prune_linear_rows(
                attn.c_proj, heads_to_keep, num_heads
            )
            
            # Update head count
            if hasattr(attn, 'num_heads'):
                attn.num_heads = len(heads_to_keep)
            
            # Track pruned heads
            if layer_idx not in self.pruned_heads:
                self.pruned_heads[layer_idx] = set()
            self.pruned_heads[layer_idx].update(head_indices)
            
            stats['layers_pruned'] += 1
            stats['total_heads_pruned'] += len(head_indices)
        
        stats['params_after'] = sum(p.numel() for p in self.model.parameters())
        stats['compression_ratio'] = stats['params_before'] / stats['params_after']
        
        print(f"\n  Pruning complete:")
        print(f"    Layers pruned: {stats['layers_pruned']}")
        print(f"    Total heads pruned: {stats['total_heads_pruned']}")
        print(f"    Parameters: {stats['params_before']:,} -> {stats['params_after']:,}")
        print(f"    Compression: {stats['compression_ratio']:.2f}x")
        
        return stats
    
    def _prune_linear_columns(
        self,
        linear: nn.Linear,
        heads_to_keep: List[int],
        num_heads: int
    ) -> nn.Linear:
        """
        Prune columns from a linear layer (for Q/K/V projections).
        
        Args:
            linear: Linear layer to prune
            heads_to_keep: Indices of heads to keep
            num_heads: Total number of heads
        
        Returns:
            New linear layer with pruned columns
        """
        out_features, in_features = linear.weight.shape
        head_dim = out_features // num_heads
        
        # Select columns to keep (reshape by heads)
        indices = []
        for head_idx in heads_to_keep:
            start = head_idx * head_dim
            end = start + head_dim
            indices.extend(range(start, end))
        
        # Create new linear layer
        new_linear = nn.Linear(in_features, len(indices), bias=linear.bias is not None)
        new_linear.weight.data = linear.weight.data[indices, :]
        if linear.bias is not None:
            new_linear.bias.data = linear.bias.data[indices]
        
        return new_linear
    
    def _prune_linear_rows(
        self,
        linear: nn.Linear,
        heads_to_keep: List[int],
        num_heads: int
    ) -> nn.Linear:
        """
        Prune rows from a linear layer (for output projections).
        
        Args:
            linear: Linear layer to prune
            heads_to_keep: Indices of heads to keep
            num_heads: Total number of heads
        
        Returns:
            New linear layer with pruned rows
        """
        out_features, in_features = linear.weight.shape
        head_dim = in_features // num_heads
        
        # Select rows to keep
        indices = []
        for head_idx in heads_to_keep:
            start = head_idx * head_dim
            end = start + head_dim
            indices.extend(range(start, end))
        
        # Create new linear layer
        new_linear = nn.Linear(len(indices), out_features, bias=linear.bias is not None)
        new_linear.weight.data = linear.weight.data[:, indices]
        if linear.bias is not None:
            new_linear.bias.data = linear.bias.data
        
        return new_linear
    
    def _prune_gpt2_c_attn(
        self,
        c_attn: nn.Linear,
        heads_to_keep: List[int],
        num_heads: int,
        head_dim: int
    ) -> nn.Linear:
        """
        Prune GPT-2's combined QKV projection.
        
        c_attn projects to [Q, K, V] concatenated, so we need to prune
        from three separate sections.
        
        Note: GPT-2 uses Conv1D which has weight shape [out_features, in_features]
        but we treat it as Linear for consistency.
        """
        # Get weight shape - handle both Linear and Conv1D
        if hasattr(c_attn, 'weight'):
            weight = c_attn.weight
            bias = c_attn.bias if hasattr(c_attn, 'bias') else None
            
            # Check if this is Conv1D (out, in) or Linear (out, in)
            if weight.shape[0] == weight.shape[1] * 3:
                # Conv1D format: [3*embed_dim, embed_dim]
                out_features, in_features = weight.shape
            else:
                # Linear format: [embed_dim, 3*embed_dim]
                in_features, out_features = weight.shape
                weight = weight.t()  # Transpose to match Conv1D
                if bias is not None:
                    bias = bias.clone()
        else:
            raise ValueError("c_attn must have weight attribute")
        
        total_dim = out_features // 3  # Divide by 3 for Q, K, V
        
        # Prune each section (Q, K, V)
        indices = []
        for section in range(3):  # Q, K, V
            section_start = section * total_dim
            for head_idx in sorted(heads_to_keep):
                start = section_start + head_idx * head_dim
                end = start + head_dim
                if end <= section_start + total_dim:  # Safety check
                    indices.extend(range(start, end))
        
        # Create new linear layer with correct dimensions
        new_c_attn = nn.Linear(in_features, len(indices), bias=bias is not None)
        new_c_attn.weight.data = weight[indices, :].clone()
        if bias is not None:
            new_c_attn.bias.data = bias[indices].clone()
        
        return new_c_attn
    
    def get_heads_to_prune_by_importance(
        self,
        importance_scores: Dict[str, np.ndarray],
        prune_ratio: float = 0.25
    ) -> Dict[int, List[int]]:
        """
        Determine which heads to prune based on importance scores.
        
        Args:
            importance_scores: Dict mapping layer names to importance arrays
            prune_ratio: Fraction of heads to prune (0.25 = 25%)
        
        Returns:
            Dict mapping layer_idx -> list of head indices to prune
        """
        # Flatten all scores
        all_scores = []
        layer_mapping = {}  # (layer_idx, head_idx) -> score
        
        for layer_name, scores in importance_scores.items():
            # Extract layer index from name
            layer_idx = self._extract_layer_index(layer_name)
            if layer_idx is None:
                continue
            
            for head_idx, score in enumerate(scores):
                all_scores.append(score)
                layer_mapping[(layer_idx, head_idx)] = score
        
        # Determine threshold
        threshold = np.percentile(all_scores, prune_ratio * 100)
        
        # Group by layer
        heads_to_prune = {}
        for (layer_idx, head_idx), score in layer_mapping.items():
            if score < threshold:
                if layer_idx not in heads_to_prune:
                    heads_to_prune[layer_idx] = []
                heads_to_prune[layer_idx].append(head_idx)
        
        return heads_to_prune
    
    def _extract_layer_index(self, layer_name: str) -> Optional[int]:
        """Extract layer index from layer name."""
        import re
        
        # Try different patterns
        patterns = [
            r'layer\.(\d+)',  # encoder.layer.0
            r'\.h\.(\d+)\.',  # transformer.h.0.attn
            r'blocks\.(\d+)', # blocks.0
        ]
        
        for pattern in patterns:
            match = re.search(pattern, layer_name)
            if match:
                return int(match.group(1))
        
        return None


def prune_model_by_importance(
    model: nn.Module,
    importance_scores: Dict[str, np.ndarray],
    prune_ratio: float = 0.25,
    model_type: str = 'auto'
) -> Tuple[nn.Module, Dict[str, any]]:
    """
    Prune a model based on head importance scores.
    
    Args:
        model: Model to prune
        importance_scores: Head importance scores from compute_head_importance_enhanced
        prune_ratio: Fraction of heads to prune
        model_type: 'bert', 'gpt2', or 'auto'
    
    Returns:
        Tuple of (pruned_model, pruning_stats)
    """
    pruner = AttentionHeadPruner(model, model_type)
    
    # Determine which heads to prune
    heads_to_prune = pruner.get_heads_to_prune_by_importance(
        importance_scores, prune_ratio
    )
    
    # Prune the model
    stats = pruner.prune_heads(heads_to_prune, importance_scores)
    
    return model, stats


if __name__ == "__main__":
    print("Physical Attention Head Pruning")
    print("=" * 60)
    print("\nThis module implements actual head removal from HuggingFace models.")
    print("\nUsage:")
    print("""
    from passes.head_pruner import prune_model_by_importance
    from passes.enhanced_head_importance import compute_head_importance_enhanced
    
    # Compute importance
    importance = compute_head_importance_enhanced(
        model, calibration_data, method='entropy'
    )
    
    # Prune model
    pruned_model, stats = prune_model_by_importance(
        model, importance, prune_ratio=0.25
    )
    
    # stats contains: layers_pruned, total_heads_pruned, compression_ratio
    """)