ffn_pruner.py

"""
FFN Neuron Pruning

Physical removal of FFN neurons based on importance scores.
"""

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


class FFNPruner:
    """
    Prunes FFN neurons from transformer models.
    
    Removes low-importance neurons from intermediate FFN layers.
    """
    
    def __init__(self, model: nn.Module, model_type: str = 'auto'):
        self.model = model
        self.model_type = self._detect_model_type() if model_type == 'auto' else model_type
        self.ffn_modules = self._detect_ffn_modules()
    
    def _detect_model_type(self) -> str:
        """Detect model architecture."""
        model_class = type(self.model).__name__.lower()
        if 'bert' in model_class or 'roberta' in model_class:
            return 'bert'
        elif 'gpt' in model_class:
            return 'gpt2'
        else:
            return 'unknown'
    
    def _detect_ffn_modules(self) -> Dict[str, Dict]:
        """Detect FFN modules in the model."""
        ffn_modules = {}
        
        if self.model_type == 'bert':
            if hasattr(self.model, 'bert'):
                encoder = self.model.bert.encoder
            elif hasattr(self.model, 'encoder'):
                encoder = self.model.encoder
            else:
                return ffn_modules
            
            for i, layer in enumerate(encoder.layer):
                if hasattr(layer, 'intermediate') and hasattr(layer, 'output'):
                    ffn_modules[f'layer.{i}.ffn'] = {
                        'layer_idx': i,
                        'layer_obj': layer,
                        'intermediate': layer.intermediate.dense,
                        'output': layer.output.dense,
                        'intermediate_size': layer.intermediate.dense.out_features,
                        'hidden_size': layer.intermediate.dense.in_features,
                    }
        
        elif self.model_type == 'gpt2':
            if hasattr(self.model, 'transformer'):
                blocks = self.model.transformer.h
            elif hasattr(self.model, 'h'):
                blocks = self.model.h
            else:
                return ffn_modules
            
            for i, block in enumerate(blocks):
                if hasattr(block, 'mlp'):
                    mlp = block.mlp
                    if hasattr(mlp, 'c_fc') and hasattr(mlp, 'c_proj'):
                        c_fc = mlp.c_fc
                        # Conv1D stores dimensions as nf (out) and nx (in)
                        # Conv1D weight shape: [out_features, in_features]
                        if hasattr(c_fc, 'nf'):
                            intermediate_size = c_fc.nf
                            hidden_size = c_fc.nx
                        else:
                            intermediate_size = c_fc.out_features
                            hidden_size = c_fc.in_features
                        
                        ffn_modules[f'layer.{i}.ffn'] = {
                            'layer_idx': i,
                            'layer_obj': block,
                            'mlp': mlp,
                            'intermediate': c_fc,
                            'output': mlp.c_proj,
                            'intermediate_size': intermediate_size,
                            'hidden_size': hidden_size,
                        }
        
        return ffn_modules
    
    def prune_ffn_neurons(
        self,
        neurons_to_prune: Dict[str, List[int]],
        inplace: bool = False
    ) -> nn.Module:
        """
        Prune specified FFN neurons.
        
        Args:
            neurons_to_prune: Dict mapping layer names to list of neuron indices to prune
            inplace: Whether to modify the model in-place
        
        Returns:
            Pruned model
        """
        if not inplace:
            model = copy.deepcopy(self.model)
            pruner = FFNPruner(model, self.model_type)
        else:
            model = self.model
            pruner = self
        
        if self.model_type == 'bert':
            pruner._prune_bert_ffn(neurons_to_prune)
        elif self.model_type == 'gpt2':
            pruner._prune_gpt2_ffn(neurons_to_prune)
        else:
            raise ValueError(f"Unsupported model type: {self.model_type}")
        
        return model
    
    def _prune_bert_ffn(self, neurons_to_prune: Dict[str, List[int]]):
        """Prune BERT FFN neurons."""
        for layer_name, neuron_indices in neurons_to_prune.items():
            if layer_name not in self.ffn_modules:
                continue
            
            info = self.ffn_modules[layer_name]
            layer = info['layer_obj']
            intermediate_layer = layer.intermediate.dense
            output_layer = layer.output.dense
            
            intermediate_size = info['intermediate_size']
            hidden_size = info['hidden_size']
            
            # Create mask for neurons to keep
            keep_mask = torch.ones(intermediate_size, dtype=torch.bool)
            keep_mask[neuron_indices] = False
            keep_indices = torch.where(keep_mask)[0]
            
            new_intermediate_size = len(keep_indices)
            
            if new_intermediate_size == 0:
                print(f"  Warning: All neurons pruned in {layer_name}, skipping")
                continue
            
            # Prune intermediate layer (output dimension)
            old_weight = intermediate_layer.weight.data  # [intermediate_size, hidden_size]
            old_bias = intermediate_layer.bias.data if intermediate_layer.bias is not None else None
            
            new_weight = old_weight[keep_indices, :]
            new_bias = old_bias[keep_indices] if old_bias is not None else None
            
            # Create new intermediate layer
            new_intermediate = nn.Linear(hidden_size, new_intermediate_size, bias=old_bias is not None)
            new_intermediate.weight.data = new_weight
            if new_bias is not None:
                new_intermediate.bias.data = new_bias
            
            layer.intermediate.dense = new_intermediate
            
            # Prune output layer (input dimension)
            old_weight = output_layer.weight.data  # [hidden_size, intermediate_size]
            old_bias = output_layer.bias.data if output_layer.bias is not None else None
            
            new_weight = old_weight[:, keep_indices]
            
            # Create new output layer
            new_output = nn.Linear(new_intermediate_size, hidden_size, bias=old_bias is not None)
            new_output.weight.data = new_weight
            if old_bias is not None:
                new_output.bias.data = old_bias
            
            layer.output.dense = new_output
    
    def _prune_gpt2_ffn(self, neurons_to_prune: Dict[str, List[int]]):
        """Prune GPT-2 FFN neurons."""
        for layer_name, neuron_indices in neurons_to_prune.items():
            if layer_name not in self.ffn_modules:
                continue
            
            info = self.ffn_modules[layer_name]
            mlp = info['mlp']
            c_fc = info['intermediate']
            c_proj = info['output']
            
            intermediate_size = info['intermediate_size']
            hidden_size = info['hidden_size']
            
            # Create mask for neurons to keep
            keep_mask = torch.ones(intermediate_size, dtype=torch.bool)
            keep_mask[neuron_indices] = False
            keep_indices = torch.where(keep_mask)[0]
            
            new_intermediate_size = len(keep_indices)
            
            if new_intermediate_size == 0:
                print(f"  Warning: All neurons pruned in {layer_name}, skipping")
                continue
            
            # GPT-2 uses Conv1D layers
            from transformers.pytorch_utils import Conv1D
            
            # Prune c_fc (intermediate layer)
            if isinstance(c_fc, Conv1D):
                # Conv1D: weight shape is [nx, nf] = [in_features, out_features] 
                # This is transposed compared to Linear!
                # For c_fc: [hidden_size, intermediate_size] = [768, 3072]
                old_weight = c_fc.weight.data  # [hidden_size, intermediate_size]
                old_bias = c_fc.bias.data if c_fc.bias is not None else None
                
                # Prune along the output dimension (dim 1)
                new_weight = old_weight[:, keep_indices]
                new_bias = old_bias[keep_indices] if old_bias is not None else None
                
                # Create new Conv1D
                new_c_fc = Conv1D(new_intermediate_size, hidden_size)
                new_c_fc.weight.data = new_weight
                if new_bias is not None:
                    new_c_fc.bias.data = new_bias
                
                mlp.c_fc = new_c_fc
            else:
                # Regular Linear layer
                old_weight = c_fc.weight.data  # [intermediate_size, hidden_size]
                old_bias = c_fc.bias.data if c_fc.bias is not None else None
                
                new_weight = old_weight[keep_indices, :]
                new_bias = old_bias[keep_indices] if old_bias is not None else None
                
                new_c_fc = nn.Linear(hidden_size, new_intermediate_size, bias=old_bias is not None)
                new_c_fc.weight.data = new_weight
                if new_bias is not None:
                    new_c_fc.bias.data = new_bias
                
                mlp.c_fc = new_c_fc
            
            # Prune c_proj (output layer)
            if isinstance(c_proj, Conv1D):
                # Conv1D: weight shape is [nx, nf] = [in_features, out_features]
                # For c_proj: [intermediate_size, hidden_size] = [3072, 768]
                old_weight = c_proj.weight.data  # [intermediate_size, hidden_size]
                old_bias = c_proj.bias.data if c_proj.bias is not None else None
                
                # Prune along the input dimension (dim 0)
                new_weight = old_weight[keep_indices, :]
                
                new_c_proj = Conv1D(hidden_size, new_intermediate_size)
                new_c_proj.weight.data = new_weight
                if old_bias is not None:
                    new_c_proj.bias.data = old_bias
                
                mlp.c_proj = new_c_proj
            else:
                # Regular Linear layer
                old_weight = c_proj.weight.data  # [hidden_size, intermediate_size]
                old_bias = c_proj.bias.data if c_proj.bias is not None else None
                
                new_weight = old_weight[:, keep_indices]
                
                new_c_proj = nn.Linear(new_intermediate_size, hidden_size, bias=old_bias is not None)
                new_c_proj.weight.data = new_weight
                if old_bias is not None:
                    new_c_proj.bias.data = old_bias
                
                mlp.c_proj = new_c_proj


def get_neurons_to_prune_by_importance(
    importance_scores: Dict[str, np.ndarray],
    prune_ratio: float
) -> Dict[str, List[int]]:
    """
    Determine which neurons to prune based on importance scores.
    
    Args:
        importance_scores: Dict mapping layer names to importance scores
        prune_ratio: Fraction of neurons to prune (0.0 to 1.0)
    
    Returns:
        Dict mapping layer names to lists of neuron indices to prune
    """
    neurons_to_prune = {}
    
    for layer_name, scores in importance_scores.items():
        num_neurons = len(scores)
        num_to_prune = int(num_neurons * prune_ratio)
        
        if num_to_prune > 0:
            # Sort by importance (ascending)
            sorted_indices = np.argsort(scores)
            # Take the lowest importance neurons
            prune_indices = sorted_indices[:num_to_prune].tolist()
            neurons_to_prune[layer_name] = prune_indices
    
    return neurons_to_prune


def prune_ffn_by_importance(
    model: nn.Module,
    importance_scores: Dict[str, np.ndarray],
    prune_ratio: float,
    model_type: str = 'auto',
    inplace: bool = False
) -> nn.Module:
    """
    Prune FFN neurons based on importance scores.
    
    Args:
        model: Model to prune
        importance_scores: Neuron importance scores
        prune_ratio: Fraction of neurons to prune
        model_type: 'bert', 'gpt2', or 'auto'
        inplace: Whether to modify model in-place
    
    Returns:
        Pruned model
    """
    # Determine which neurons to prune
    neurons_to_prune = get_neurons_to_prune_by_importance(importance_scores, prune_ratio)
    
    # Create pruner and prune
    pruner = FFNPruner(model, model_type)
    pruned_model = pruner.prune_ffn_neurons(neurons_to_prune, inplace=inplace)
    
    return pruned_model


if __name__ == "__main__":
    print("FFN Neuron Pruning")
    print("=" * 60)
    print("\nPrunes low-importance FFN neurons from transformers.")
    print("Reduces ~50% of model parameters with minimal accuracy loss.")
    print("\nUsage:")
    print("  1. Compute importance scores using ffn_importance.py")
    print("  2. Call prune_ffn_by_importance() with desired ratio")
    print("  3. Evaluate pruned model accuracy")