distil-trainer 0.1.10__py3-none-any.whl

distil_trainer/pruning/depth_pruning.py

@@ -0,0 +1,261 @@
+"""Depth pruning (layer reduction) for transformer models."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from typing import Literal
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+
+from sentence_transformers import SentenceTransformer
+
+from distil_trainer.pruning.importance import ImportanceEstimator
+
+logger = logging.getLogger(__name__)
+
+
+class DepthPruner:
+    """
+    Handles layer removal from transformer models.
+
+    Depth pruning removes entire transformer layers to reduce model size
+    while preserving the overall architecture.
+
+    Example:
+        >>> pruner = DepthPruner(model)
+        >>> importance = pruner.compute_layer_importance(calibration_data)
+        >>> layers_to_keep = pruner.select_layers_to_keep(num_layers=8)
+        >>> pruned_model = pruner.prune(layers_to_keep=layers_to_keep)
+    """
+
+    def __init__(
+        self,
+        model: SentenceTransformer | PreTrainedModel,
+        calibration_data: list[str] | None = None,
+    ):
+        """
+        Initialize the DepthPruner.
+
+        Args:
+            model: The model to prune.
+            calibration_data: Optional calibration data for importance estimation.
+        """
+        self.model = model
+        self.calibration_data = calibration_data
+        self._transformer = self._get_transformer_module()
+        self._original_num_layers = self._get_num_layers()
+
+    def _get_transformer_module(self) -> nn.Module:
+        """Get the underlying transformer module."""
+        if isinstance(self.model, SentenceTransformer):
+            # SentenceTransformer wraps a transformer model
+            for module in self.model.modules():
+                if hasattr(module, "encoder") and hasattr(module.encoder, "layer"):
+                    return module
+                if hasattr(module, "layers"):
+                    return module
+            # Try the first module which is typically the transformer
+            return self.model._first_module()
+        else:
+            return self.model
+
+    def _get_encoder_layers(self) -> nn.ModuleList:
+        """Get the encoder layers from the transformer."""
+        transformer = self._transformer
+
+        # Try different attribute names used by different models
+        if hasattr(transformer, "encoder") and hasattr(transformer.encoder, "layer"):
+            return transformer.encoder.layer
+        if hasattr(transformer, "encoder") and hasattr(transformer.encoder, "layers"):
+            return transformer.encoder.layers
+        if hasattr(transformer, "layers"):
+            return transformer.layers
+        if hasattr(transformer, "auto_model"):
+            auto_model = transformer.auto_model
+            if hasattr(auto_model, "encoder") and hasattr(auto_model.encoder, "layer"):
+                return auto_model.encoder.layer
+            if hasattr(auto_model, "layers"):
+                return auto_model.layers
+
+        raise ValueError("Could not find encoder layers in model")
+
+    def _get_num_layers(self) -> int:
+        """Get the number of layers in the model."""
+        return len(self._get_encoder_layers())
+
+    def compute_layer_importance(
+        self,
+        calibration_data: list[str] | None = None,
+        method: Literal["activation", "gradient", "cosine_similarity", "lm_loss"] = "activation",
+        num_samples: int = 1024,
+    ) -> dict[int, float]:
+        """
+        Compute importance scores for each layer.
+
+        Args:
+            calibration_data: Sentences to use for calibration.
+            method: Importance estimation method.
+            num_samples: Number of samples to use.
+
+        Returns:
+            Dictionary mapping layer index to importance score.
+        """
+        data = calibration_data or self.calibration_data
+        if data is None:
+            raise ValueError("Calibration data required for importance estimation")
+
+        estimator = ImportanceEstimator(self.model)
+
+        if method == "activation":
+            importance = estimator.activation_based_layer_importance(data[:num_samples])
+        elif method == "cosine_similarity":
+            importance = estimator.drop_layer_importance(data[:num_samples])
+        else:
+            importance = estimator.activation_based_layer_importance(data[:num_samples])
+
+        return importance
+
+    def select_layers_to_keep(
+        self,
+        num_layers: int | None = None,
+        ratio: float | None = None,
+        strategy: Literal["first", "last", "even", "importance"] = "importance",
+        importance_scores: dict[int, float] | None = None,
+    ) -> list[int]:
+        """
+        Select which layers to keep based on the strategy.
+
+        Args:
+            num_layers: Number of layers to keep.
+            ratio: Ratio of layers to keep (alternative to num_layers).
+            strategy: Layer selection strategy.
+            importance_scores: Precomputed importance scores (for 'importance' strategy).
+
+        Returns:
+            List of layer indices to keep.
+        """
+        total_layers = self._original_num_layers
+
+        if num_layers is None and ratio is not None:
+            num_layers = int(total_layers * ratio)
+        if num_layers is None:
+            raise ValueError("Either num_layers or ratio must be specified")
+
+        num_layers = min(num_layers, total_layers)
+
+        if strategy == "first":
+            return list(range(num_layers))
+
+        elif strategy == "last":
+            return list(range(total_layers - num_layers, total_layers))
+
+        elif strategy == "even":
+            # Evenly distribute layers
+            if num_layers == 1:
+                return [0]
+            step = (total_layers - 1) / (num_layers - 1)
+            return [int(round(i * step)) for i in range(num_layers)]
+
+        elif strategy == "importance":
+            if importance_scores is None:
+                if self.calibration_data is None:
+                    # Fall back to even distribution
+                    logger.warning("No importance scores or calibration data, using even distribution")
+                    return self.select_layers_to_keep(num_layers=num_layers, strategy="even")
+                importance_scores = self.compute_layer_importance()
+
+            # Sort layers by importance and keep the most important ones
+            sorted_layers = sorted(importance_scores.items(), key=lambda x: x[1], reverse=True)
+            layers_to_keep = sorted([layer_idx for layer_idx, _ in sorted_layers[:num_layers]])
+            return layers_to_keep
+
+        else:
+            raise ValueError(f"Unknown strategy: {strategy}")
+
+    def prune(
+        self,
+        layers_to_keep: list[int] | None = None,
+        num_layers_to_keep: int | None = None,
+        layers_to_drop: list[int] | None = None,
+        layer_selection: Literal["first", "last", "even", "importance", "custom"] = "custom",
+    ) -> SentenceTransformer | PreTrainedModel:
+        """
+        Create a new model with only the specified layers.
+
+        Args:
+            layers_to_keep: Explicit list of layer indices to keep.
+            num_layers_to_keep: Number of layers to keep (alternative).
+            layers_to_drop: Layers to drop (alternative).
+            layer_selection: Strategy for selecting layers when using num_layers_to_keep.
+
+        Returns:
+            New model with reduced layers.
+        """
+        # Determine which layers to keep
+        if layers_to_keep is not None:
+            final_layers_to_keep = layers_to_keep
+        elif layers_to_drop is not None:
+            all_layers = set(range(self._original_num_layers))
+            final_layers_to_keep = sorted(all_layers - set(layers_to_drop))
+        elif num_layers_to_keep is not None:
+            if layer_selection == "custom":
+                layer_selection = "importance"
+            final_layers_to_keep = self.select_layers_to_keep(
+                num_layers=num_layers_to_keep,
+                strategy=layer_selection,
+            )
+        else:
+            raise ValueError("Must specify layers_to_keep, num_layers_to_keep, or layers_to_drop")
+
+        logger.info(f"Keeping layers: {final_layers_to_keep}")
+        logger.info(f"Reducing from {self._original_num_layers} to {len(final_layers_to_keep)} layers")
+
+        # Create a deep copy of the model
+        pruned_model = copy.deepcopy(self.model)
+
+        # Get the encoder layers
+        if isinstance(pruned_model, SentenceTransformer):
+            transformer = pruned_model._first_module()
+            if hasattr(transformer, "auto_model"):
+                auto_model = transformer.auto_model
+            else:
+                auto_model = transformer
+        else:
+            auto_model = pruned_model
+
+        # Find and replace the layers
+        encoder_layers = None
+        if hasattr(auto_model, "encoder") and hasattr(auto_model.encoder, "layer"):
+            encoder_layers = auto_model.encoder.layer
+            parent = auto_model.encoder
+            attr_name = "layer"
+        elif hasattr(auto_model, "encoder") and hasattr(auto_model.encoder, "layers"):
+            encoder_layers = auto_model.encoder.layers
+            parent = auto_model.encoder
+            attr_name = "layers"
+        elif hasattr(auto_model, "layers"):
+            encoder_layers = auto_model.layers
+            parent = auto_model
+            attr_name = "layers"
+
+        if encoder_layers is None:
+            raise ValueError("Could not find encoder layers in model")
+
+        # Create new layer list with only kept layers
+        new_layers = nn.ModuleList([encoder_layers[i] for i in final_layers_to_keep])
+        setattr(parent, attr_name, new_layers)
+
+        # Update config if available
+        if hasattr(auto_model, "config"):
+            auto_model.config.num_hidden_layers = len(final_layers_to_keep)
+
+        # Log statistics
+        original_params = sum(p.numel() for p in self.model.parameters())
+        pruned_params = sum(p.numel() for p in pruned_model.parameters())
+        reduction = 1 - (pruned_params / original_params)
+        logger.info(f"Parameter reduction: {original_params:,} -> {pruned_params:,} ({reduction:.1%})")
+
+        return pruned_model

distil_trainer/pruning/importance.py

@@ -0,0 +1,365 @@
+"""Importance estimation for pruning."""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+import torch
+from torch import nn
+from tqdm import tqdm
+from transformers import PreTrainedModel
+
+from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger(__name__)
+
+
+class ImportanceEstimator:
+    """
+    Methods for estimating component importance.
+
+    Supports importance estimation for:
+    - Layers (for depth pruning)
+    - Hidden dimensions (for width pruning)
+    - Intermediate dimensions (for MLP pruning)
+    - Attention heads (for head pruning)
+
+    Example:
+        >>> estimator = ImportanceEstimator(model)
+        >>> layer_importance = estimator.activation_based_layer_importance(sentences)
+        >>> hidden_importance = estimator.hidden_dimension_importance(sentences)
+    """
+
+    def __init__(self, model: SentenceTransformer | PreTrainedModel):
+        """
+        Initialize the ImportanceEstimator.
+
+        Args:
+            model: The model to estimate importance for.
+        """
+        self.model = model
+        self.device = next(model.parameters()).device
+
+    def activation_based_layer_importance(
+        self,
+        sentences: list[str],
+        batch_size: int = 32,
+    ) -> dict[int, float]:
+        """
+        Estimate layer importance based on activation magnitudes.
+
+        Higher activation magnitude suggests higher importance.
+
+        Args:
+            sentences: Sentences to use for estimation.
+            batch_size: Batch size for inference.
+
+        Returns:
+            Dictionary mapping layer index to importance score.
+        """
+        self.model.eval()
+
+        # Get the transformer model
+        if isinstance(self.model, SentenceTransformer):
+            transformer = self.model._first_module()
+            if hasattr(transformer, "auto_model"):
+                auto_model = transformer.auto_model
+            else:
+                auto_model = transformer
+            tokenizer = self.model.tokenizer
+        else:
+            auto_model = self.model
+            tokenizer = None
+
+        # Get encoder layers
+        encoder_layers = self._get_encoder_layers(auto_model)
+        num_layers = len(encoder_layers)
+
+        # Track activations per layer
+        layer_activations = {i: 0.0 for i in range(num_layers)}
+        num_samples = 0
+
+        # Register hooks to capture activations
+        handles = []
+        activation_sums = {}
+
+        def make_hook(layer_idx):
+            def hook(module, input, output):
+                # Handle different output formats
+                if isinstance(output, tuple):
+                    hidden_states = output[0]
+                else:
+                    hidden_states = output
+                activation_sums[layer_idx] = activation_sums.get(layer_idx, 0) + hidden_states.abs().mean().item()
+            return hook
+
+        for i, layer in enumerate(encoder_layers):
+            handle = layer.register_forward_hook(make_hook(i))
+            handles.append(handle)
+
+        try:
+            with torch.no_grad():
+                for i in range(0, len(sentences), batch_size):
+                    batch_sentences = sentences[i:i + batch_size]
+
+                    if isinstance(self.model, SentenceTransformer):
+                        self.model.encode(batch_sentences, convert_to_tensor=True)
+                    else:
+                        if tokenizer:
+                            inputs = tokenizer(
+                                batch_sentences,
+                                padding=True,
+                                truncation=True,
+                                return_tensors="pt",
+                            ).to(self.device)
+                            self.model(**inputs)
+
+                    num_samples += 1
+
+            # Average the activation sums
+            for layer_idx in range(num_layers):
+                layer_activations[layer_idx] = activation_sums.get(layer_idx, 0) / max(num_samples, 1)
+
+        finally:
+            for handle in handles:
+                handle.remove()
+
+        return layer_activations
+
+    def drop_layer_importance(
+        self,
+        sentences: list[str],
+        batch_size: int = 32,
+    ) -> dict[int, float]:
+        """
+        Estimate layer importance by measuring embedding similarity when layer is dropped.
+
+        Lower similarity when dropped = higher importance.
+
+        Args:
+            sentences: Sentences to use for estimation.
+            batch_size: Batch size for inference.
+
+        Returns:
+            Dictionary mapping layer index to importance score.
+        """
+        if not isinstance(self.model, SentenceTransformer):
+            raise ValueError("drop_layer_importance only works with SentenceTransformer models")
+
+        # Get reference embeddings
+        with torch.no_grad():
+            reference_embeddings = self.model.encode(
+                sentences[:min(len(sentences), 500)],
+                convert_to_tensor=True,
+                show_progress_bar=False,
+            )
+
+        # Get transformer model
+        transformer = self.model._first_module()
+        if hasattr(transformer, "auto_model"):
+            auto_model = transformer.auto_model
+        else:
+            auto_model = transformer
+
+        encoder_layers = self._get_encoder_layers(auto_model)
+        num_layers = len(encoder_layers)
+
+        importance = {}
+
+        for layer_idx in range(num_layers):
+            # Temporarily remove the layer
+            original_layer = encoder_layers[layer_idx]
+
+            # Replace with identity
+            encoder_layers[layer_idx] = nn.Identity()
+
+            try:
+                with torch.no_grad():
+                    dropped_embeddings = self.model.encode(
+                        sentences[:min(len(sentences), 500)],
+                        convert_to_tensor=True,
+                        show_progress_bar=False,
+                    )
+
+                # Compute cosine similarity
+                similarity = torch.nn.functional.cosine_similarity(
+                    reference_embeddings, dropped_embeddings, dim=1
+                ).mean().item()
+
+                # Lower similarity = higher importance
+                importance[layer_idx] = 1 - similarity
+
+            finally:
+                # Restore the layer
+                encoder_layers[layer_idx] = original_layer
+
+        return importance
+
+    def hidden_dimension_importance(
+        self,
+        sentences: list[str],
+        method: Literal["activation", "magnitude", "taylor"] = "activation",
+        batch_size: int = 32,
+    ) -> torch.Tensor:
+        """
+        Estimate importance of each hidden dimension.
+
+        Args:
+            sentences: Sentences to use for estimation.
+            method: Importance estimation method.
+            batch_size: Batch size for inference.
+
+        Returns:
+            Tensor of importance scores for each hidden dimension.
+        """
+        self.model.eval()
+
+        if isinstance(self.model, SentenceTransformer):
+            transformer = self.model._first_module()
+            if hasattr(transformer, "auto_model"):
+                auto_model = transformer.auto_model
+            else:
+                auto_model = transformer
+        else:
+            auto_model = self.model
+
+        hidden_size = auto_model.config.hidden_size
+        importance = torch.zeros(hidden_size, device=self.device)
+
+        if method == "activation":
+            # Track output activations
+            num_samples = 0
+
+            with torch.no_grad():
+                for i in range(0, len(sentences), batch_size):
+                    batch_sentences = sentences[i:i + batch_size]
+
+                    if isinstance(self.model, SentenceTransformer):
+                        embeddings = self.model.encode(
+                            batch_sentences,
+                            convert_to_tensor=True,
+                            show_progress_bar=False,
+                        )
+                        importance += embeddings.abs().mean(dim=0)
+                    else:
+                        # For regular models, would need tokenizer
+                        pass
+
+                    num_samples += 1
+
+            importance = importance / max(num_samples, 1)
+
+        elif method == "magnitude":
+            # Use weight magnitudes from final layer
+            encoder_layers = self._get_encoder_layers(auto_model)
+            last_layer = encoder_layers[-1]
+
+            # Sum up weight magnitudes
+            for name, param in last_layer.named_parameters():
+                if param.dim() >= 2:
+                    # Aggregate across the hidden dimension
+                    importance += param.abs().sum(dim=0)[:hidden_size] if param.size(-1) >= hidden_size else torch.zeros(hidden_size, device=self.device)
+
+        return importance
+
+    def intermediate_dimension_importance(
+        self,
+        sentences: list[str],
+        method: Literal["activation", "magnitude"] = "activation",
+        batch_size: int = 32,
+    ) -> torch.Tensor:
+        """
+        Estimate importance of each intermediate (MLP) dimension.
+
+        Args:
+            sentences: Sentences to use for estimation.
+            method: Importance estimation method.
+            batch_size: Batch size for inference.
+
+        Returns:
+            Tensor of importance scores for each intermediate dimension.
+        """
+        self.model.eval()
+
+        if isinstance(self.model, SentenceTransformer):
+            transformer = self.model._first_module()
+            if hasattr(transformer, "auto_model"):
+                auto_model = transformer.auto_model
+            else:
+                auto_model = transformer
+        else:
+            auto_model = self.model
+
+        intermediate_size = auto_model.config.intermediate_size
+        importance = torch.zeros(intermediate_size, device=self.device)
+
+        if method == "magnitude":
+            encoder_layers = self._get_encoder_layers(auto_model)
+
+            for layer in encoder_layers:
+                if hasattr(layer, "intermediate") and hasattr(layer.intermediate, "dense"):
+                    weight = layer.intermediate.dense.weight
+                    importance += weight.abs().mean(dim=1)
+
+            importance = importance / len(encoder_layers)
+
+        return importance
+
+    def attention_head_importance(
+        self,
+        sentences: list[str],
+        method: Literal["activation", "attention_entropy"] = "activation",
+        batch_size: int = 32,
+    ) -> torch.Tensor:
+        """
+        Estimate importance of each attention head.
+
+        Args:
+            sentences: Sentences to use for estimation.
+            method: Importance estimation method.
+            batch_size: Batch size for inference.
+
+        Returns:
+            Tensor of importance scores for each attention head.
+        """
+        self.model.eval()
+
+        if isinstance(self.model, SentenceTransformer):
+            transformer = self.model._first_module()
+            if hasattr(transformer, "auto_model"):
+                auto_model = transformer.auto_model
+            else:
+                auto_model = transformer
+        else:
+            auto_model = self.model
+
+        num_heads = auto_model.config.num_attention_heads
+        importance = torch.zeros(num_heads, device=self.device)
+
+        if method == "activation":
+            encoder_layers = self._get_encoder_layers(auto_model)
+            head_dim = auto_model.config.hidden_size // num_heads
+
+            for layer in encoder_layers:
+                if hasattr(layer, "attention") and hasattr(layer.attention, "self"):
+                    self_attn = layer.attention.self
+                    if hasattr(self_attn, "query"):
+                        weight = self_attn.query.weight
+                        # Reshape to [num_heads, head_dim, hidden_size]
+                        weight_per_head = weight.view(num_heads, head_dim, -1)
+                        importance += weight_per_head.abs().mean(dim=(1, 2))
+
+            importance = importance / len(encoder_layers)
+
+        return importance
+
+    def _get_encoder_layers(self, model) -> nn.ModuleList:
+        """Get encoder layers from the model."""
+        if hasattr(model, "encoder") and hasattr(model.encoder, "layer"):
+            return model.encoder.layer
+        if hasattr(model, "encoder") and hasattr(model.encoder, "layers"):
+            return model.encoder.layers
+        if hasattr(model, "layers"):
+            return model.layers
+        raise ValueError("Could not find encoder layers")