layer-scan 0.1.0__py3-none-any.whl

layer_scan/__init__.py

@@ -0,0 +1,3 @@
+"""layer-scan: Automated LLM layer duplication configuration scanner."""
+
+__version__ = "0.1.0"

layer_scan/backends/__init__.py

@@ -0,0 +1,9 @@
+"""Inference backends for layer-scan.
+
+Each backend wraps a specific inference framework and implements
+the forward_with_duplication interface.
+"""
+
+from layer_scan.backends.base import Backend
+
+__all__ = ["Backend"]

layer_scan/backends/base.py

@@ -0,0 +1,51 @@
+"""Abstract backend interface for layer-scan."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import torch
+
+from layer_scan.config import DuplicationConfig
+
+
+class Backend(ABC):
+    """Abstract base for inference backends.
+
+    A backend must provide:
+    1. Model loading and tokenization
+    2. Forward pass with optional layer duplication
+    3. Layer count information
+    """
+
+    @abstractmethod
+    def load(self, model_path: str, **kwargs) -> None:
+        """Load a model from the given path or HuggingFace ID."""
+
+    @abstractmethod
+    def get_total_layers(self) -> int:
+        """Return the total number of transformer decoder layers."""
+
+    @abstractmethod
+    def get_tokenizer(self):
+        """Return the model's tokenizer."""
+
+    @abstractmethod
+    def forward_with_duplication(
+        self,
+        text: str,
+        duplication_config: DuplicationConfig | None = None,
+    ) -> torch.Tensor:
+        """Run forward pass with optional layer duplication.
+
+        Args:
+            text: Input text to process.
+            duplication_config: If provided, duplicate layers [i..j-1].
+                If None, run the standard forward pass (baseline).
+
+        Returns:
+            Logits tensor of shape (vocab_size,) for the last token position.
+        """
+
+    def cleanup(self) -> None:
+        """Release resources (GPU memory, etc.). Optional."""

layer_scan/backends/exllamav2.py

@@ -0,0 +1,224 @@
+"""ExLlamaV2 backend — optimized for consumer GPUs with quantized models.
+
+This is the recommended backend for scanning large models (70B+) on
+consumer hardware (e.g., 2×RTX 4090). ExLlamaV2 supports GPTQ and EXL2
+quantization with excellent memory efficiency.
+
+Requires: pip install layer-scan[exllamav2]
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from layer_scan.backends.base import Backend
+from layer_scan.config import DuplicationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ExLlamaV2Backend(Backend):
+    """ExLlamaV2 backend with runtime layer duplication.
+
+    Uses ExLlamaV2's modular architecture to execute individual layers
+    in a custom order, enabling layer duplication without model copies.
+    """
+
+    def __init__(self) -> None:
+        self._model = None
+        self._tokenizer = None
+        self._cache = None
+        self._config = None
+        self._total_layers: int = 0
+
+    def load(self, model_path: str, **kwargs) -> None:
+        """Load an EXL2/GPTQ quantized model.
+
+        Kwargs:
+            gpu_split: List of GPU memory limits in MB (e.g., [22000, 22000]).
+            max_seq_len: Maximum sequence length (default: 4096).
+            rope_scale: RoPE scaling factor (default: 1.0).
+        """
+        try:
+            from exllamav2 import (
+                ExLlamaV2,
+                ExLlamaV2Cache,
+                ExLlamaV2Config,
+                ExLlamaV2Tokenizer,
+            )
+        except ImportError as e:
+            raise ImportError(
+                "ExLlamaV2 not installed. Install with: "
+                "pip install layer-scan[exllamav2]"
+            ) from e
+
+        logger.info("Loading ExLlamaV2 model from %s", model_path)
+
+        self._config = ExLlamaV2Config(model_path)
+        self._config.max_seq_len = kwargs.get("max_seq_len", 4096)
+
+        if "rope_scale" in kwargs:
+            self._config.scale_pos_emb = kwargs["rope_scale"]
+
+        self._model = ExLlamaV2(self._config)
+
+        gpu_split = kwargs.get("gpu_split")
+        if gpu_split:
+            self._model.load(gpu_split)
+        else:
+            self._model.load_autosplit()
+
+        self._tokenizer = ExLlamaV2Tokenizer(self._config)
+        self._cache = ExLlamaV2Cache(self._model, max_seq_len=self._config.max_seq_len)
+
+        # Count transformer layers (excluding embedding, head, norms)
+        self._total_layers = self._count_decoder_layers()
+        logger.info("ExLlamaV2 model loaded: %d decoder layers", self._total_layers)
+
+    def _count_decoder_layers(self) -> int:
+        """Count the number of decoder layers in the ExLlamaV2 model."""
+        count = 0
+        for module in self._model.modules:
+            module_name = type(module).__name__
+            if "Attention" in module_name or "MLP" in module_name:
+                # Each transformer block has attention + MLP
+                # We count by attention modules to get block count
+                if "Attention" in module_name:
+                    count += 1
+        return count
+
+    def get_total_layers(self) -> int:
+        if self._total_layers == 0:
+            raise RuntimeError("Model not loaded. Call load() first.")
+        return self._total_layers
+
+    def get_tokenizer(self):
+        """Return a tokenizer adapter compatible with HuggingFace interface."""
+        if self._tokenizer is None:
+            raise RuntimeError("Model not loaded. Call load() first.")
+        return _ExLlamaV2TokenizerAdapter(self._tokenizer)
+
+    @torch.no_grad()
+    def forward_with_duplication(
+        self,
+        text: str,
+        duplication_config: DuplicationConfig | None = None,
+    ) -> torch.Tensor:
+        """Run forward pass with optional layer duplication.
+
+        For ExLlamaV2, we use the module-level forward to control
+        execution order of individual transformer blocks.
+        """
+        input_ids = self._tokenizer.encode(text)
+        input_ids = input_ids.to(self._model.modules[0].device())
+
+        self._cache.current_seq_len = 0
+
+        if duplication_config is None:
+            # Standard forward pass
+            logits = self._model.forward(input_ids, self._cache)
+            return logits[0, -1, :]
+
+        # Duplicated forward: execute modules in custom order
+        return self._forward_duplicated(input_ids, duplication_config)
+
+    def _forward_duplicated(
+        self,
+        input_ids: torch.Tensor,
+        config: DuplicationConfig,
+    ) -> torch.Tensor:
+        """Execute ExLlamaV2 modules with layer duplication.
+
+        ExLlamaV2 exposes modules as a flat list. We need to map
+        decoder layer indices to module indices and build the
+        custom execution order.
+        """
+        # Build module execution order based on layer duplication config
+        # ExLlamaV2 modules: [embedding, layer0_attn, layer0_mlp, ..., norm, head]
+        exec_order = config.execution_order()
+
+        # Map layer indices to module ranges
+        layer_modules = self._get_layer_module_map()
+
+        # Execute embedding
+        hidden = self._model.modules[0].forward(input_ids, self._cache)
+
+        # Execute decoder layers in duplicated order
+        for layer_idx in exec_order:
+            for mod_idx in layer_modules[layer_idx]:
+                hidden = self._model.modules[mod_idx].forward(hidden, self._cache)
+
+        # Execute final norm and head
+        for mod_idx in self._get_post_layer_modules():
+            hidden = self._model.modules[mod_idx].forward(hidden, self._cache)
+
+        return hidden[0, -1, :]
+
+    def _get_layer_module_map(self) -> dict[int, list[int]]:
+        """Map decoder layer indices to ExLlamaV2 module indices."""
+        layer_map: dict[int, list[int]] = {}
+        current_layer = -1
+
+        for mod_idx, module in enumerate(self._model.modules):
+            name = type(module).__name__
+            if "Attention" in name:
+                current_layer += 1
+                layer_map[current_layer] = [mod_idx]
+            elif "MLP" in name and current_layer >= 0:
+                layer_map[current_layer].append(mod_idx)
+
+        return layer_map
+
+    def _get_post_layer_modules(self) -> list[int]:
+        """Get module indices for norm and lm_head (after all decoder layers)."""
+        post_modules = []
+        found_last_mlp = False
+
+        for mod_idx in range(len(self._model.modules) - 1, -1, -1):
+            name = type(self._model.modules[mod_idx]).__name__
+            if "MLP" in name or "Attention" in name:
+                if not found_last_mlp:
+                    found_last_mlp = True
+                    # Everything after this index is post-layer
+                    post_modules = list(range(mod_idx + 1, len(self._model.modules)))
+                break
+
+        return post_modules
+
+    def cleanup(self) -> None:
+        if self._model is not None:
+            del self._model
+            self._model = None
+        if self._cache is not None:
+            del self._cache
+            self._cache = None
+        self._tokenizer = None
+        self._total_layers = 0
+
+        import gc
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+class _ExLlamaV2TokenizerAdapter:
+    """Adapter to make ExLlamaV2 tokenizer compatible with HuggingFace interface.
+
+    layer-scan's scoring module expects tokenizer.encode(str) -> list[int].
+    """
+
+    def __init__(self, exl2_tokenizer) -> None:
+        self._tokenizer = exl2_tokenizer
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> list[int]:
+        ids = self._tokenizer.encode(text)
+        return ids[0].tolist() if hasattr(ids, "tolist") else list(ids[0])
+
+    def decode(self, token_ids: list[int]) -> str:
+        import torch as _torch
+
+        ids = _torch.tensor([token_ids])
+        return self._tokenizer.decode(ids)[0]

layer_scan/backends/transformers_backend.py

@@ -0,0 +1,254 @@
+"""HuggingFace Transformers backend — the default reference implementation.
+
+This backend hooks into the model's forward pass at the layer level,
+implementing layer duplication by re-executing specified layers without
+modifying model weights. It works with any CausalLM model.
+
+Note: This is the reference backend for correctness. For large models
+(70B+), use the ExLlamaV2 backend for quantized inference on consumer GPUs.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import torch
+
+from layer_scan.backends.base import Backend
+from layer_scan.config import DuplicationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class TransformersBackend(Backend):
+    """HuggingFace Transformers backend with runtime layer duplication."""
+
+    def __init__(self) -> None:
+        self._model: Any = None
+        self._tokenizer: Any = None
+        self._layers: list[Any] = []
+        self._total_layers: int = 0
+
+    def load(self, model_path: str, **kwargs) -> None:
+        """Load a CausalLM model from path or HuggingFace ID.
+
+        Kwargs:
+            dtype: Torch dtype string (default: "float16").
+            device_map: Device map for model parallelism (default: "auto").
+            trust_remote_code: Whether to trust remote code (default: False).
+        """
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        dtype_str = kwargs.get("dtype", "float16")
+        dtype_map = {
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+            "float32": torch.float32,
+        }
+        dtype = dtype_map.get(dtype_str, torch.float16)
+
+        logger.info("Loading tokenizer from %s", model_path)
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=kwargs.get("trust_remote_code", False),
+        )
+
+        logger.info("Loading model from %s (dtype=%s)", model_path, dtype_str)
+        self._model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=dtype,
+            device_map=kwargs.get("device_map", "auto"),
+            trust_remote_code=kwargs.get("trust_remote_code", False),
+        )
+        self._model.eval()
+
+        # Discover layer structure
+        self._layers = self._find_layers()
+        self._total_layers = len(self._layers)
+        logger.info("Model loaded: %d layers discovered", self._total_layers)
+
+    def _find_layers(self) -> list[Any]:
+        """Discover the decoder layers in the model.
+
+        Supports common architectures: LLaMA, Mistral, Qwen, GPT-NeoX, etc.
+        """
+        model = self._model
+
+        # Try common layer container names
+        for attr_path in [
+            "model.layers",              # LLaMA, Mistral, Qwen2
+            "transformer.h",             # GPT-2, GPT-Neo
+            "gpt_neox.layers",           # GPT-NeoX, Pythia
+            "transformer.blocks",        # MPT
+            "model.decoder.layers",      # OPT
+        ]:
+            obj = model
+            found = True
+            for part in attr_path.split("."):
+                if hasattr(obj, part):
+                    obj = getattr(obj, part)
+                else:
+                    found = False
+                    break
+            if found and hasattr(obj, "__len__"):
+                return list(obj)
+
+        raise RuntimeError(
+            "Could not find decoder layers. Supported architectures: "
+            "LLaMA, Mistral, Qwen2, GPT-2, GPT-NeoX, MPT, OPT. "
+            "For other architectures, use a custom backend."
+        )
+
+    def get_total_layers(self) -> int:
+        if self._total_layers == 0:
+            raise RuntimeError("Model not loaded. Call load() first.")
+        return self._total_layers
+
+    def get_tokenizer(self):
+        if self._tokenizer is None:
+            raise RuntimeError("Model not loaded. Call load() first.")
+        return self._tokenizer
+
+    @torch.no_grad()
+    def forward_with_duplication(
+        self,
+        text: str,
+        duplication_config: DuplicationConfig | None = None,
+    ) -> torch.Tensor:
+        """Run forward pass, optionally duplicating layers [i..j-1].
+
+        For the baseline (duplication_config=None), uses the standard
+        model forward pass. For duplicated configs, hooks into the
+        layer execution to replay specified layers.
+
+        Returns:
+            Logits tensor of shape (vocab_size,) for the last token.
+        """
+        inputs = self._tokenizer(text, return_tensors="pt")
+        input_ids = inputs["input_ids"].to(self._model.device)
+        attention_mask = inputs.get("attention_mask")
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(self._model.device)
+
+        if duplication_config is None:
+            # Standard forward pass (baseline)
+            outputs = self._model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+            return outputs.logits[0, -1, :]
+
+        # Layer-duplicated forward pass
+        return self._forward_duplicated(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            config=duplication_config,
+        )
+
+    def _forward_duplicated(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        config: DuplicationConfig,
+    ) -> torch.Tensor:
+        """Execute the forward pass with layer duplication.
+
+        Execution order: [0...j-1, i...N-1]
+        Layers [i...j-1] execute twice — this gives the model
+        additional "thinking time" on the reasoning circuits.
+        """
+        model = self._model
+
+        # Get the embedding output (before decoder layers)
+        if hasattr(model, "model"):
+            # LLaMA-style: model.model.embed_tokens
+            base = model.model
+        elif hasattr(model, "transformer"):
+            # GPT-style: model.transformer.wte
+            base = model.transformer
+        else:
+            raise RuntimeError("Cannot determine embedding layer location")
+
+        # Compute embeddings
+        if hasattr(base, "embed_tokens"):
+            hidden_states = base.embed_tokens(input_ids)
+        elif hasattr(base, "wte"):
+            hidden_states = base.wte(input_ids)
+        else:
+            raise RuntimeError("Cannot find embedding layer")
+
+        # Build execution order
+        exec_order = config.execution_order()
+
+        # Execute layers in the duplicated order
+        # Note: For models with RoPE, position IDs remain based on
+        # sequence length — we don't change them for duplicated layers.
+        # This matches the RYS methodology.
+        position_ids = torch.arange(
+            input_ids.shape[1], device=input_ids.device
+        ).unsqueeze(0)
+
+        for layer_idx in exec_order:
+            layer = self._layers[layer_idx]
+            # Most decoder layers accept (hidden_states, attention_mask, position_ids)
+            # but the exact signature varies by architecture
+            layer_output = layer(
+                hidden_states,
+                attention_mask=self._prepare_causal_mask(
+                    attention_mask, hidden_states
+                ),
+                position_ids=position_ids,
+            )
+            # Layer output is typically a tuple; first element is hidden_states
+            if isinstance(layer_output, tuple):
+                hidden_states = layer_output[0]
+            else:
+                hidden_states = layer_output
+
+        # Apply final norm
+        if hasattr(base, "norm"):
+            hidden_states = base.norm(hidden_states)
+        elif hasattr(base, "ln_f"):
+            hidden_states = base.ln_f(hidden_states)
+
+        # Project to logits
+        if hasattr(model, "lm_head"):
+            logits = model.lm_head(hidden_states)
+        elif hasattr(base, "lm_head"):
+            logits = base.lm_head(hidden_states)
+        else:
+            raise RuntimeError("Cannot find LM head for logit projection")
+
+        return logits[0, -1, :]
+
+    @staticmethod
+    def _prepare_causal_mask(
+        attention_mask: torch.Tensor | None,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Prepare 4D causal attention mask if needed.
+
+        Some models expect a pre-expanded mask. We return None and let
+        the layer handle masking internally when possible.
+        """
+        # Most modern architectures handle masking internally
+        # when attention_mask is None or 2D
+        return attention_mask
+
+    def cleanup(self) -> None:
+        """Free GPU memory."""
+        if self._model is not None:
+            del self._model
+            self._model = None
+        if self._tokenizer is not None:
+            del self._tokenizer
+            self._tokenizer = None
+        self._layers = []
+        self._total_layers = 0
+
+        import gc
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()