ffn_compression.py

"""
Feed-Forward Network (FFN) Compression Pass

Reduces the intermediate dimension of transformer FFN layers through
structured pruning or low-rank decomposition.
"""

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


def analyze_ffn_neuron_importance(
    model: nn.Module,
    calibration_data: torch.Tensor
) -> Dict[str, List[float]]:
    """
    Analyze importance of each neuron in FFN intermediate layers.
    
    Args:
        model: PyTorch model
        calibration_data: Sample data for analysis
    
    Returns:
        Dictionary mapping FFN layer names to neuron importance scores
    """
    model.eval()
    neuron_importance = {}
    
    # Capture activations
    activations = {}
    
    def hook_fn(name):
        def hook(module, input, output):
            activations[name] = output.detach()
        return hook
    
    # Register hooks for Linear layers in FFN
    hooks = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            # Check if this is an FFN layer (typically has 'ffn', 'mlp', 'fc' in name)
            if any(keyword in name.lower() for keyword in ['ffn', 'mlp', 'fc']):
                hook = module.register_forward_hook(hook_fn(name))
                hooks.append(hook)
    
    # Forward pass
    with torch.no_grad():
        model(calibration_data)
    
    # Remove hooks
    for hook in hooks:
        hook.remove()
    
    # Compute importance based on activation magnitude
    for name, act in activations.items():
        # act shape: (batch, seq_len, hidden_dim) or (batch, hidden_dim)
        
        # Compute L2 norm across batch and sequence dimensions
        if len(act.shape) == 3:
            importance = torch.norm(act, dim=(0, 1)).cpu().numpy()
        elif len(act.shape) == 2:
            importance = torch.norm(act, dim=0).cpu().numpy()
        else:
            importance = torch.norm(act.flatten(0, -2), dim=0).cpu().numpy()
        
        # Normalize
        if importance.sum() > 0:
            importance = importance / importance.sum()
        
        neuron_importance[name] = importance.tolist()
    
    return neuron_importance


def compress_ffn_width(
    mir: Dict[str, Any],
    compression_ratio: float = 0.5,
    neuron_importance: Dict[str, List[float]] = None
) -> Dict[str, Any]:
    """
    Compress FFN intermediate layers by reducing width.
    
    In transformer FFN: hidden_dim → ffn_dim → hidden_dim
    Typically ffn_dim = 4 * hidden_dim
    
    This reduces ffn_dim (e.g., 4x → 2x) based on neuron importance.
    
    Args:
        mir: MIR representation
        compression_ratio: Target compression ratio (0.0 to 1.0)
            0.5 means reduce to 50% of original size
        neuron_importance: Optional importance scores for neuron selection
    
    Returns:
        Updated MIR with compressed FFN layers
    """
    if not (0.0 < compression_ratio <= 1.0):
        raise ValueError("Compression ratio must be in (0.0, 1.0]")
    
    compression_stats = {
        "layers_compressed": [],
        "total_params_before": 0,
        "total_params_after": 0
    }
    
    for i, node in enumerate(mir["graph"]["nodes"]):
        if node["op_type"] == "linear":
            origin_name = node["metadata"]["origin_name"]
            
            # Check if this is an FFN intermediate layer
            # FFN layers typically have 'ffn', 'mlp', 'fc1', 'wi' in their name
            is_ffn_expansion = any(keyword in origin_name.lower() 
                                  for keyword in ['ffn.fc1', 'ffn.wi', 'mlp.fc1', 'mlp.c_fc'])
            
            if is_ffn_expansion:
                params = node.get("params", {})
                in_features = params.get("in", 0)
                out_features = params.get("out", 0)
                
                compression_stats["total_params_before"] += in_features * out_features
                
                # Compute new output dimension
                new_out = max(1, int(out_features * compression_ratio))
                
                # If importance scores provided, select top neurons
                if neuron_importance and origin_name in neuron_importance:
                    importance = neuron_importance[origin_name]
                    
                    # Select top neurons by importance
                    if len(importance) == out_features:
                        top_neurons = sorted(
                            range(out_features),
                            key=lambda i: importance[i],
                            reverse=True
                        )[:new_out]
                        
                        node["compression"]["kept_neurons"] = sorted(top_neurons)
                        node["compression"]["pruned_neurons"] = [
                            i for i in range(out_features) if i not in top_neurons
                        ]
                
                # Update parameters
                params["out"] = new_out
                node["params"] = params
                
                # Update compression metadata
                if "compression" not in node:
                    node["compression"] = {}
                
                node["compression"]["ffn_width_reduction"] = {
                    "original_width": out_features,
                    "compressed_width": new_out,
                    "compression_ratio": new_out / out_features,
                    "method": "width_reduction"
                }
                
                compression_stats["total_params_after"] += in_features * new_out
                compression_stats["layers_compressed"].append({
                    "layer": origin_name,
                    "original_width": out_features,
                    "compressed_width": new_out
                })
                
                # Also need to update the next linear layer (contraction) input dimension
                # Find the corresponding fc2/wo layer
                for j in range(i + 1, len(mir["graph"]["nodes"])):
                    next_node = mir["graph"]["nodes"][j]
                    if next_node["op_type"] == "linear":
                        next_origin = next_node["metadata"]["origin_name"]
                        
                        # Check if this is the corresponding contraction layer
                        is_ffn_contraction = any(
                            keyword in next_origin.lower() 
                            for keyword in ['ffn.fc2', 'ffn.wo', 'mlp.fc2', 'mlp.c_proj']
                        )
                        
                        if is_ffn_contraction:
                            # Update input dimension
                            next_params = next_node.get("params", {})
                            next_params["in"] = new_out
                            next_node["params"] = next_params
                            
                            # Add metadata
                            if "compression" not in next_node:
                                next_node["compression"] = {}
                            next_node["compression"]["ffn_input_adjusted"] = {
                                "original_in": out_features,
                                "adjusted_in": new_out
                            }
                            
                            break
                        
                        # Stop if we hit another FFN expansion (different layer)
                        if any(keyword in next_origin.lower() 
                              for keyword in ['ffn.fc1', 'ffn.wi', 'mlp.fc1']):
                            break
    
    # Add compression statistics to MIR metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    
    mir["metadata"]["ffn_compression_stats"] = compression_stats
    
    return mir


def apply_ffn_compression(
    mir: Dict[str, Any],
    model: nn.Module,
    calibration_data: torch.Tensor,
    compression_ratio: float = 0.5,
    use_importance_analysis: bool = True
) -> Dict[str, Any]:
    """
    Complete FFN compression pipeline.
    
    Args:
        mir: MIR representation
        model: Original PyTorch model
        calibration_data: Data for importance analysis
        compression_ratio: Target compression ratio
        use_importance_analysis: Whether to use importance-based neuron selection
    
    Returns:
        Updated MIR with compressed FFN layers
    """
    neuron_importance = None
    
    if use_importance_analysis:
        neuron_importance = analyze_ffn_neuron_importance(model, calibration_data)
    
    mir = compress_ffn_width(mir, compression_ratio, neuron_importance)
    
    return mir


def get_ffn_compression_statistics(mir: Dict[str, Any]) -> Dict[str, Any]:
    """
    Extract FFN compression statistics from MIR.
    
    Args:
        mir: MIR representation with compression metadata
    
    Returns:
        Dictionary with compression statistics
    """
    stats = mir.get("metadata", {}).get("ffn_compression_stats", {})
    
    if not stats:
        return {
            "compression_applied": False,
            "message": "No FFN compression metadata found"
        }
    
    params_before = stats.get("total_params_before", 0)
    params_after = stats.get("total_params_after", 0)
    
    return {
        "compression_applied": True,
        "total_params_before": params_before,
        "total_params_after": params_after,
        "params_saved": params_before - params_after,
        "param_reduction_ratio": (params_before - params_after) / params_before if params_before > 0 else 0,
        "layers_compressed": stats.get("layers_compressed", [])
    }


def estimate_ffn_speedup(mir: Dict[str, Any]) -> Dict[str, Any]:
    """
    Estimate inference speedup from FFN compression.
    
    Args:
        mir: MIR representation with FFN compression
    
    Returns:
        Dictionary with speedup estimates
    """
    stats = mir.get("metadata", {}).get("ffn_compression_stats", {})
    
    if not stats:
        return {
            "estimated_speedup": 1.0,
            "note": "No compression applied"
        }
    
    params_before = stats.get("total_params_before", 1)
    params_after = stats.get("total_params_after", 1)
    
    # Simple linear speedup estimate based on parameter reduction
    # Real speedup depends on hardware, batch size, etc.
    param_ratio = params_after / params_before if params_before > 0 else 1.0
    estimated_speedup = 1.0 / param_ratio if param_ratio > 0 else 1.0
    
    return {
        "estimated_speedup": estimated_speedup,
        "parameter_reduction": 1.0 - param_ratio,
        "note": "Actual speedup may vary based on hardware and implementation"
    }