ffn_importance.py

"""
FFN (Feed-Forward Network) Importance Analysis

Analyzes and compresses the feed-forward layers in transformers.
FFN layers typically account for 2/3 of transformer parameters.
"""

import torch
import torch.nn as nn
from typing import Dict, List, Tuple, Optional
import numpy as np
from collections import defaultdict


class FFNAnalyzer:
    """
    Analyzes importance of FFN (feed-forward network) neurons.
    
    FFN structure in transformers:
    - Linear1: [hidden_size, intermediate_size] (expansion, e.g., 768 -> 3072)
    - Activation: GELU/ReLU
    - Linear2: [intermediate_size, hidden_size] (projection back)
    
    Accounts for ~66% of parameters in BERT/GPT-2.
    """
    
    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.
        
        Returns:
            Dict mapping module names to FFN info
        """
        ffn_modules = {}
        
        if self.model_type == 'bert':
            # BERT: encoder.layer.X.intermediate and output
            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'):
                    intermediate = layer.intermediate.dense
                    output = layer.output.dense
                    
                    ffn_modules[f'layer.{i}.ffn'] = {
                        'layer_idx': i,
                        'intermediate': intermediate,
                        'output': output,
                        'intermediate_size': intermediate.out_features,
                        'hidden_size': intermediate.in_features,
                    }
        
        elif self.model_type == 'gpt2':
            # GPT-2: transformer.h.X.mlp
            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,
                            'intermediate': c_fc,
                            'output': mlp.c_proj,
                            'intermediate_size': intermediate_size,
                            'hidden_size': hidden_size,
                        }
        
        print(f"[FFN Analyzer] Detected {len(ffn_modules)} FFN modules")
        return ffn_modules
    
    def compute_activation_importance(
        self,
        calibration_data: torch.Tensor,
        num_samples: int = 10
    ) -> Dict[str, np.ndarray]:
        """
        Compute neuron importance based on activation magnitudes.
        
        Method: Neurons with consistently low activations are less important.
        
        Args:
            calibration_data: Input data for calibration
            num_samples: Number of samples to use
        
        Returns:
            Dict mapping layer names to importance scores [intermediate_size]
        """
        print(f"\n[FFN Activation Method] Computing neuron importance...")
        
        self.model.eval()
        activations = defaultdict(list)
        hooks = []
        
        def make_hook(layer_name):
            def hook(module, input, output):
                # Capture intermediate activations (after first linear layer)
                if isinstance(output, torch.Tensor):
                    activations[layer_name].append(output.detach().cpu().abs())
            return hook
        
        # Register hooks on intermediate layers
        for name, info in self.ffn_modules.items():
            intermediate = info['intermediate']
            hook = intermediate.register_forward_hook(make_hook(name))
            hooks.append(hook)
        
        # Forward passes
        batch_size = min(num_samples, calibration_data.shape[0])
        with torch.no_grad():
            for i in range(batch_size):
                sample = calibration_data[i:i+1]
                try:
                    _ = self.model(sample)
                except Exception as e:
                    print(f"  Warning: Forward pass {i} failed: {e}")
                    continue
        
        # Remove hooks
        for hook in hooks:
            hook.remove()
        
        # Compute importance from activations
        importance_scores = {}
        
        for name, info in self.ffn_modules.items():
            intermediate_size = info['intermediate_size']
            
            if name in activations and len(activations[name]) > 0:
                # Stack activations: [num_samples, batch, seq_len, intermediate_size]
                acts = torch.stack(activations[name])
                
                # Average magnitude across samples, batch, and sequence
                avg_activation = acts.mean(dim=(0, 1, 2))  # [intermediate_size]
                
                # Normalize to get importance scores
                importance = avg_activation.numpy()
                importance = importance / (importance.sum() + 1e-10)
                
                importance_scores[name] = importance
            else:
                # No activations captured, use uniform
                importance_scores[name] = np.ones(intermediate_size) / intermediate_size
        
        print(f"  Computed activation importance for {len(importance_scores)} FFN layers")
        return importance_scores
    
    def compute_weight_magnitude_importance(self) -> Dict[str, np.ndarray]:
        """
        Compute neuron importance based on weight magnitudes.
        
        Method: |W_in| * |W_out| for each neuron.
        Neurons with small weights contribute less to output.
        
        Returns:
            Dict mapping layer names to importance scores [intermediate_size]
        """
        print(f"\n[FFN Weight Magnitude Method] Computing neuron importance...")
        
        importance_scores = {}
        
        for name, info in self.ffn_modules.items():
            intermediate = info['intermediate']
            output = info['output']
            
            # Get weights
            W_in = intermediate.weight.data  # [intermediate_size, hidden_size]
            W_out = output.weight.data       # [hidden_size, intermediate_size]
            
            # Compute magnitude for each neuron
            in_magnitude = W_in.abs().sum(dim=1)   # [intermediate_size]
            out_magnitude = W_out.abs().sum(dim=0) # [intermediate_size]
            
            # Combined importance: product of input and output magnitudes
            importance = (in_magnitude * out_magnitude).cpu().numpy()
            
            # Normalize
            importance = importance / (importance.sum() + 1e-10)
            
            importance_scores[name] = importance
        
        print(f"  Computed weight magnitude importance for {len(importance_scores)} FFN layers")
        return importance_scores
    
    def compute_gradient_importance(
        self,
        calibration_data: torch.Tensor,
        num_samples: int = 10
    ) -> Dict[str, np.ndarray]:
        """
        Compute neuron importance based on gradients.
        
        Method: Measure gradient magnitude w.r.t. intermediate neurons.
        
        Args:
            calibration_data: Input data
            num_samples: Number of samples
        
        Returns:
            Dict mapping layer names to importance scores [intermediate_size]
        """
        print(f"\n[FFN Gradient Method] Computing neuron importance...")
        
        self.model.train()
        gradients = defaultdict(list)
        
        # Store intermediate outputs for gradient computation
        intermediate_outputs = {}
        hooks = []
        
        def make_hook(layer_name):
            def hook(module, input, output):
                if output.requires_grad:
                    intermediate_outputs[layer_name] = output
            return hook
        
        # Register hooks
        for name, info in self.ffn_modules.items():
            intermediate = info['intermediate']
            hook = intermediate.register_forward_hook(make_hook(name))
            hooks.append(hook)
        
        # Compute gradients
        batch_size = min(num_samples, calibration_data.shape[0])
        
        try:
            for i in range(batch_size):
                sample = calibration_data[i:i+1]
                
                outputs = self.model(sample)
                
                # Get logits
                if hasattr(outputs, 'logits'):
                    logits = outputs.logits
                elif hasattr(outputs, 'last_hidden_state'):
                    logits = outputs.last_hidden_state
                elif isinstance(outputs, tuple):
                    logits = outputs[0]
                else:
                    logits = outputs
                
                # Compute loss
                if isinstance(logits, torch.Tensor):
                    loss = logits.pow(2).mean()
                    loss.backward()
                    
                    # Capture gradients from intermediate outputs
                    for name, output in intermediate_outputs.items():
                        if output.grad is not None:
                            gradients[name].append(output.grad.detach().cpu().abs())
                    
                    self.model.zero_grad()
        
        except Exception as e:
            print(f"  Warning: Gradient computation failed: {e}")
        
        finally:
            for hook in hooks:
                hook.remove()
            self.model.zero_grad()
            self.model.eval()
        
        # Compute importance from gradients
        importance_scores = {}
        
        for name, info in self.ffn_modules.items():
            intermediate_size = info['intermediate_size']
            
            if name in gradients and len(gradients[name]) > 0:
                # Stack gradients and average
                grads = torch.stack(gradients[name])
                avg_grad = grads.mean(dim=(0, 1, 2))  # [intermediate_size]
                
                # Normalize
                importance = avg_grad.numpy()
                importance = importance / (importance.sum() + 1e-10)
                
                importance_scores[name] = importance
            else:
                importance_scores[name] = np.ones(intermediate_size) / intermediate_size
        
        print(f"  Computed gradient importance for {len(importance_scores)} FFN layers")
        return importance_scores
    
    def get_compression_statistics(self) -> Dict[str, any]:
        """Get statistics about FFN layers for compression analysis."""
        total_ffn_params = 0
        total_model_params = sum(p.numel() for p in self.model.parameters())
        
        for name, info in self.ffn_modules.items():
            intermediate = info['intermediate']
            output = info['output']
            
            ffn_params = sum(p.numel() for p in intermediate.parameters())
            ffn_params += sum(p.numel() for p in output.parameters())
            total_ffn_params += ffn_params
        
        return {
            'total_model_params': total_model_params,
            'total_ffn_params': total_ffn_params,
            'ffn_percentage': total_ffn_params / total_model_params * 100,
            'num_ffn_layers': len(self.ffn_modules),
        }


def compute_ffn_importance(
    model: nn.Module,
    calibration_data: torch.Tensor,
    method: str = 'activation',
    num_samples: int = 10,
    model_type: str = 'auto'
) -> Dict[str, np.ndarray]:
    """
    Compute FFN neuron importance scores.
    
    Args:
        model: Model to analyze
        calibration_data: Calibration input data
        method: 'activation', 'weight_magnitude', or 'gradient'
        num_samples: Number of calibration samples
        model_type: 'bert', 'gpt2', or 'auto'
    
    Returns:
        Dict mapping layer names to importance scores
    """
    analyzer = FFNAnalyzer(model, model_type)
    
    if method == 'activation':
        return analyzer.compute_activation_importance(calibration_data, num_samples)
    elif method == 'weight_magnitude':
        return analyzer.compute_weight_magnitude_importance()
    elif method == 'gradient':
        return analyzer.compute_gradient_importance(calibration_data, num_samples)
    else:
        raise ValueError(f"Unknown method: {method}")


if __name__ == "__main__":
    print("FFN Importance Analysis")
    print("=" * 60)
    print("\nAnalyzes feed-forward network layers in transformers.")
    print("FFN layers typically account for 66% of model parameters.")
    print("\nSupported methods:")
    print("  - activation: Based on activation magnitudes")
    print("  - weight_magnitude: Based on weight magnitudes")
    print("  - gradient: Based on gradient magnitudes")