vlora-dev 0.2.0__py3-none-any.whl

vlora/io.py

@@ -0,0 +1,191 @@
+"""Load and parse PEFT LoRA adapters from disk or HuggingFace Hub.
+
+Reads safetensors + adapter_config.json without requiring the PEFT
+library at runtime.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+import torch
+
+logger = logging.getLogger("vlora")
+from safetensors.torch import load_file, save_file
+from torch import Tensor
+
+
+@dataclass
+class LoRAWeights:
+    """Parsed LoRA adapter weights grouped by layer."""
+
+    layer_names: list[str]
+    lora_a: dict[str, Tensor]  # layer_name -> (rank, in_features)
+    lora_b: dict[str, Tensor]  # layer_name -> (out_features, rank)
+    rank: int
+    metadata: dict = field(default_factory=dict)
+
+
+# Pattern to extract layer name + side from PEFT state dict keys.
+# Handles: base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight
+#       or: model.layers.0.self_attn.q_proj.lora_A.weight
+_LORA_KEY_RE = re.compile(
+    r"(?:base_model\.model\.)?(.+)\.(lora_[AB])\.(?:weight|default\.weight)"
+)
+
+
+def parse_state_dict(
+    state_dict: dict[str, Tensor],
+) -> tuple[dict[str, Tensor], dict[str, Tensor], list[str]]:
+    """Parse a PEFT state dict into grouped A and B weight dicts.
+
+    Returns:
+        lora_a: {layer_name: tensor}
+        lora_b: {layer_name: tensor}
+        layer_names: sorted unique layer names
+    """
+    lora_a: dict[str, Tensor] = {}
+    lora_b: dict[str, Tensor] = {}
+
+    for key, tensor in state_dict.items():
+        match = _LORA_KEY_RE.search(key)
+        if match is None:
+            continue
+        layer_name = match.group(1)
+        side = match.group(2)
+        if side == "lora_A":
+            lora_a[layer_name] = tensor
+        else:
+            lora_b[layer_name] = tensor
+
+    layer_names = sorted(set(lora_a.keys()) & set(lora_b.keys()))
+    # Keep only layers where both A and B exist
+    lora_a = {k: lora_a[k] for k in layer_names}
+    lora_b = {k: lora_b[k] for k in layer_names}
+
+    return lora_a, lora_b, layer_names
+
+
+def load_adapter(path: str | Path) -> LoRAWeights:
+    """Load a PEFT LoRA adapter from a local directory.
+
+    Expects adapter_model.safetensors and adapter_config.json in the
+    given directory.
+    """
+    path = Path(path)
+
+    # Load safetensors weights
+    safetensors_path = path / "adapter_model.safetensors"
+    if not safetensors_path.exists():
+        raise FileNotFoundError(f"No adapter_model.safetensors in {path}")
+    state_dict = load_file(str(safetensors_path))
+
+    # Load config for metadata
+    config_path = path / "adapter_config.json"
+    metadata: dict = {}
+    rank = 0
+    if config_path.exists():
+        with open(config_path) as f:
+            config = json.load(f)
+        rank = config.get("r", 0)
+        metadata = config
+
+    lora_a, lora_b, layer_names = parse_state_dict(state_dict)
+    logger.debug("Loaded adapter from %s: %d layers", path, len(layer_names))
+
+    # Infer rank from weight shapes if not in config
+    if rank == 0 and layer_names:
+        rank = lora_a[layer_names[0]].shape[0]
+
+    return LoRAWeights(
+        layer_names=layer_names,
+        lora_a=lora_a,
+        lora_b=lora_b,
+        rank=rank,
+        metadata=metadata,
+    )
+
+
+def load_adapter_from_hub(repo_id: str, revision: str | None = None) -> LoRAWeights:
+    """Load a PEFT LoRA adapter from HuggingFace Hub.
+
+    Requires the `huggingface-hub` package (install with `pip install vlora-dev[hub]`).
+    """
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError:
+        raise ImportError(
+            "huggingface-hub is required to load from Hub. "
+            "Install with: pip install vlora-dev[hub]"
+        )
+
+    local_dir = snapshot_download(
+        repo_id,
+        revision=revision,
+        allow_patterns=["adapter_model.safetensors", "adapter_config.json"],
+    )
+    return load_adapter(local_dir)
+
+
+def stack_lora_weights(
+    adapters: list[LoRAWeights],
+    side: Literal["A", "B"],
+) -> dict[str, Tensor]:
+    """Stack LoRA weight matrices from multiple adapters per layer.
+
+    For side="A": each adapter's A matrix is (rank, in_features).
+    We flatten each to a row vector and stack N adapters into (N, rank*in_features).
+
+    This produces the "factor data matrix" the paper feeds into SVD.
+
+    Returns:
+        {layer_name: (N, flattened_dim)} stacked matrix.
+    """
+    if not adapters:
+        raise ValueError("Need at least one adapter to stack")
+
+    # Use intersection of all adapters' layer names
+    layer_set = set(adapters[0].layer_names)
+    for adapter in adapters[1:]:
+        layer_set &= set(adapter.layer_names)
+    layer_names = sorted(layer_set)
+
+    stacked: dict[str, Tensor] = {}
+    for layer in layer_names:
+        matrices = []
+        for adapter in adapters:
+            w = adapter.lora_a[layer] if side == "A" else adapter.lora_b[layer]
+            matrices.append(w.flatten())
+        stacked[layer] = torch.stack(matrices)
+
+    return stacked
+
+
+def save_adapter(weights: LoRAWeights, path: str | Path) -> None:
+    """Save LoRA weights back to PEFT-compatible format."""
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=True)
+
+    # Rebuild state dict with PEFT key format
+    state_dict = {}
+    for layer_name in weights.layer_names:
+        state_dict[f"base_model.model.{layer_name}.lora_A.weight"] = weights.lora_a[layer_name]
+        state_dict[f"base_model.model.{layer_name}.lora_B.weight"] = weights.lora_b[layer_name]
+
+    save_file(state_dict, str(path / "adapter_model.safetensors"))
+
+    # Save config — include defaults needed by vLLM/TGI
+    config = dict(weights.metadata) if weights.metadata else {}
+    config.setdefault("r", weights.rank)
+    config.setdefault("lora_alpha", weights.rank)  # alpha=rank → scaling=1.0
+    config.setdefault("peft_type", "LORA")
+    config.setdefault("task_type", "CAUSAL_LM")
+    config.setdefault("bias", "none")
+    config.setdefault("lora_dropout", 0.0)
+    with open(path / "adapter_config.json", "w") as f:
+        json.dump(config, f, indent=2)

vlora/merge.py

@@ -0,0 +1,229 @@
+"""Adapter merging — task arithmetic, TIES, and DARE.
+
+These techniques operate on LoRA weight matrices directly, producing
+a single merged adapter from multiple inputs. All three methods work
+on a per-layer basis and return a new LoRAWeights object.
+
+References:
+    - Task Arithmetic: Ilharco et al., "Editing Models with Task Arithmetic" (2023)
+    - TIES: Yadav et al., "TIES-Merging: Resolving Interference When Merging Models" (2023)
+    - DARE: Yu et al., "Language Models are Super Mario" (2024)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from vlora._validate import check_adapters_compatible
+from vlora.io import LoRAWeights
+
+logger = logging.getLogger("vlora")
+
+
+def task_arithmetic(
+    adapters: list[LoRAWeights],
+    weights: list[float] | None = None,
+) -> LoRAWeights:
+    """Merge adapters via weighted average of their weight matrices.
+
+    Each adapter's LoRA A and B matrices are averaged (optionally weighted)
+    per-layer to produce a single merged adapter.
+
+    Args:
+        adapters: List of adapters to merge (must share layers and rank).
+        weights: Per-adapter weights. Defaults to uniform (1/N each).
+
+    Returns:
+        Merged LoRAWeights.
+    """
+    check_adapters_compatible(adapters)
+    n = len(adapters)
+    if n == 0:
+        raise ValueError("Need at least one adapter to merge.")
+
+    if weights is None:
+        weights = [1.0 / n] * n
+    if len(weights) != n:
+        raise ValueError(f"weights length ({len(weights)}) must match adapters ({n})")
+
+    # Use intersection of layers
+    layer_names = sorted(set.intersection(*(set(a.layer_names) for a in adapters)))
+    if not layer_names:
+        raise ValueError("Adapters share no common layers.")
+
+    logger.info("Task arithmetic merge: %d adapters, %d layers", n, len(layer_names))
+
+    lora_a: dict[str, Tensor] = {}
+    lora_b: dict[str, Tensor] = {}
+
+    for layer in layer_names:
+        merged_a = sum(w * adapters[i].lora_a[layer] for i, w in enumerate(weights))
+        merged_b = sum(w * adapters[i].lora_b[layer] for i, w in enumerate(weights))
+        lora_a[layer] = merged_a
+        lora_b[layer] = merged_b
+
+    return LoRAWeights(
+        layer_names=layer_names,
+        lora_a=lora_a,
+        lora_b=lora_b,
+        rank=adapters[0].rank,
+    )
+
+
+def ties_merge(
+    adapters: list[LoRAWeights],
+    density: float = 0.5,
+    weights: list[float] | None = None,
+) -> LoRAWeights:
+    """Merge adapters using TIES: Trim, Elect sign, and merge.
+
+    1. Trim: zero out the smallest elements per adapter (keep top `density` fraction)
+    2. Elect sign: for each position, choose the majority sign across adapters
+    3. Merge: average only the values that agree with the elected sign
+
+    Args:
+        adapters: List of adapters to merge.
+        density: Fraction of elements to keep per adapter (0, 1]. Default 0.5.
+        weights: Per-adapter weights for the final average.
+
+    Returns:
+        Merged LoRAWeights.
+    """
+    check_adapters_compatible(adapters)
+    n = len(adapters)
+    if n == 0:
+        raise ValueError("Need at least one adapter to merge.")
+    if not 0 < density <= 1:
+        raise ValueError(f"density must be in (0, 1], got {density}")
+
+    if weights is None:
+        weights = [1.0 / n] * n
+    if len(weights) != n:
+        raise ValueError(f"weights length ({len(weights)}) must match adapters ({n})")
+
+    layer_names = sorted(set.intersection(*(set(a.layer_names) for a in adapters)))
+    if not layer_names:
+        raise ValueError("Adapters share no common layers.")
+
+    logger.info("TIES merge: %d adapters, density=%.2f, %d layers", n, density, len(layer_names))
+
+    lora_a: dict[str, Tensor] = {}
+    lora_b: dict[str, Tensor] = {}
+
+    for layer in layer_names:
+        for side, out_dict, attr in [("a", lora_a, "lora_a"), ("b", lora_b, "lora_b")]:
+            # Stack all adapters for this layer/side
+            tensors = [getattr(adapters[i], attr)[layer].clone() for i in range(n)]
+
+            # Step 1: Trim — zero out smallest elements per adapter
+            for t in tensors:
+                flat = t.flatten().abs()
+                k = max(1, int(density * flat.numel()))
+                threshold = flat.topk(k).values[-1]
+                t[t.abs() < threshold] = 0.0
+
+            stacked = torch.stack(tensors)  # (N, *shape)
+
+            # Step 2: Elect sign — majority vote at each position
+            sign_votes = (stacked > 0).float().sum(dim=0) - (stacked < 0).float().sum(dim=0)
+            elected_sign = sign_votes.sign()
+            # Ties go positive (convention)
+            elected_sign[elected_sign == 0] = 1.0
+
+            # Step 3: Merge — weighted average of values matching elected sign
+            mask = (stacked.sign() == elected_sign.unsqueeze(0))
+            # Apply weights
+            w = torch.tensor(weights, dtype=stacked.dtype).view(-1, *([1] * (stacked.dim() - 1)))
+            weighted = stacked * w
+            # Zero out values with wrong sign
+            weighted = weighted * mask.float()
+            # Sum and normalize by number of contributors (avoid division by zero)
+            contributor_count = mask.float().sum(dim=0).clamp(min=1)
+            merged = weighted.sum(dim=0) * (n / contributor_count)
+
+            out_dict[layer] = merged
+
+    return LoRAWeights(
+        layer_names=layer_names,
+        lora_a=lora_a,
+        lora_b=lora_b,
+        rank=adapters[0].rank,
+    )
+
+
+def dare_merge(
+    adapters: list[LoRAWeights],
+    drop_rate: float = 0.5,
+    weights: list[float] | None = None,
+    seed: int | None = None,
+) -> LoRAWeights:
+    """Merge adapters using DARE: Drop And REscale.
+
+    For each adapter, randomly drop elements with probability `drop_rate`
+    and rescale survivors by 1/(1-drop_rate). Then average the results.
+
+    Args:
+        adapters: List of adapters to merge.
+        drop_rate: Probability of dropping each element. Default 0.5.
+        weights: Per-adapter weights for the final average.
+        seed: Random seed for reproducibility.
+
+    Returns:
+        Merged LoRAWeights.
+    """
+    check_adapters_compatible(adapters)
+    n = len(adapters)
+    if n == 0:
+        raise ValueError("Need at least one adapter to merge.")
+    if not 0 <= drop_rate < 1:
+        raise ValueError(f"drop_rate must be in [0, 1), got {drop_rate}")
+
+    if weights is None:
+        weights = [1.0 / n] * n
+    if len(weights) != n:
+        raise ValueError(f"weights length ({len(weights)}) must match adapters ({n})")
+
+    layer_names = sorted(set.intersection(*(set(a.layer_names) for a in adapters)))
+    if not layer_names:
+        raise ValueError("Adapters share no common layers.")
+
+    logger.info("DARE merge: %d adapters, drop_rate=%.2f, %d layers", n, drop_rate, len(layer_names))
+
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    rescale = 1.0 / (1.0 - drop_rate) if drop_rate > 0 else 1.0
+
+    lora_a: dict[str, Tensor] = {}
+    lora_b: dict[str, Tensor] = {}
+
+    for layer in layer_names:
+        for side, out_dict, attr in [("a", lora_a, "lora_a"), ("b", lora_b, "lora_b")]:
+            merged = torch.zeros_like(getattr(adapters[0], attr)[layer])
+
+            for i, adapter in enumerate(adapters):
+                t = getattr(adapter, attr)[layer].clone()
+                if drop_rate > 0:
+                    mask = torch.bernoulli(torch.full_like(t, 1.0 - drop_rate))
+                    t = t * mask * rescale
+                merged = merged + weights[i] * t
+
+            out_dict[layer] = merged
+
+    return LoRAWeights(
+        layer_names=layer_names,
+        lora_a=lora_a,
+        lora_b=lora_b,
+        rank=adapters[0].rank,
+    )
+
+
+MERGE_METHODS = {
+    "average": task_arithmetic,
+    "ties": ties_merge,
+    "dare": dare_merge,
+}

vlora/model.py

@@ -0,0 +1,148 @@
+"""VLoRAModel — inference wrapper that applies reconstructed LoRA deltas."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from vlora._validate import check_task_exists
+from vlora.subspace import SharedSubspace
+
+
+class VLoRAModel(nn.Module):
+    """Wraps a base model with a shared subspace for multi-task LoRA inference.
+
+    Reconstructs task-specific LoRA deltas on demand and applies them to
+    the base model's linear layers during forward pass.
+
+    Usage:
+        subspace = SharedSubspace.load("shared_subspace/")
+        base_model = AutoModelForCausalLM.from_pretrained("model-name")
+        model = VLoRAModel(base_model, subspace)
+
+        model.set_task("task_0")
+        output = model(input_ids)
+
+        model.set_task("task_1")  # switches adapter, cached if same task
+        output = model(input_ids)
+    """
+
+    def __init__(
+        self,
+        base_model: nn.Module,
+        subspace: SharedSubspace,
+        scaling: float | None = None,
+        lora_alpha: float | None = None,
+    ):
+        super().__init__()
+        self.base_model = base_model
+        self.subspace = subspace
+
+        # Resolve scaling: explicit scaling > lora_alpha/rank > 1.0
+        if scaling is not None:
+            self.scaling = scaling
+        elif lora_alpha is not None:
+            self.scaling = lora_alpha / subspace.rank
+        else:
+            self.scaling = 1.0
+        self._active_task: str | None = None
+        self._cached_deltas: dict[str, Tensor] | None = None
+        self._hooks: list[torch.utils.hooks.RemovableHook] = []
+
+    def set_task(self, task_id: str) -> None:
+        """Set the active task adapter. Reconstructs and caches if changed."""
+        if task_id == self._active_task:
+            return
+
+        check_task_exists(self.subspace, task_id)
+
+        # Reconstruct and cache the LoRA deltas
+        weights = self.subspace.reconstruct(task_id)
+        self._cached_deltas = {}
+        for layer_name in weights.layer_names:
+            # delta_W = B @ A
+            delta = weights.lora_b[layer_name] @ weights.lora_a[layer_name]
+            self._cached_deltas[layer_name] = delta
+
+        self._active_task = task_id
+        self._apply_hooks()
+
+    def clear_task(self) -> None:
+        """Remove the active task adapter."""
+        self._remove_hooks()
+        self._active_task = None
+        self._cached_deltas = None
+
+    def _apply_hooks(self) -> None:
+        """Register forward hooks on matching linear layers."""
+        self._remove_hooks()
+
+        if self._cached_deltas is None:
+            return
+
+        for name, module in self.base_model.named_modules():
+            if name in self._cached_deltas and isinstance(module, nn.Linear):
+                delta = self._cached_deltas[name]
+                hook = module.register_forward_hook(
+                    self._make_lora_hook(delta)
+                )
+                self._hooks.append(hook)
+
+    def _remove_hooks(self) -> None:
+        """Remove all registered forward hooks."""
+        for hook in self._hooks:
+            hook.remove()
+        self._hooks.clear()
+
+    def _make_lora_hook(self, delta: Tensor):
+        """Create a forward hook that adds LoRA delta to the output."""
+        scaling = self.scaling
+
+        def hook(module: nn.Module, input: Any, output: Tensor) -> Tensor:
+            # input[0] is the input tensor to the linear layer
+            x = input[0] if isinstance(input, tuple) else input
+            lora_out = x @ delta.T.to(x.device, x.dtype)
+            return output + scaling * lora_out
+
+        return hook
+
+    def forward(self, *args, **kwargs):
+        """Forward pass through the base model with active LoRA adapter."""
+        return self.base_model(*args, **kwargs)
+
+    @property
+    def active_task(self) -> str | None:
+        """Currently active task ID, or None."""
+        return self._active_task
+
+    @property
+    def available_tasks(self) -> list[str]:
+        """List of available task IDs."""
+        return sorted(self.subspace.tasks.keys())
+
+    def reconstruct_state_dict(self, task_id: str) -> dict[str, Tensor]:
+        """Get the LoRA delta weight dict for a task without applying hooks.
+
+        Returns dict of {layer_name: delta_W} where delta_W = B @ A.
+        Useful for manual integration with custom model architectures.
+        """
+        weights = self.subspace.reconstruct(task_id)
+        deltas = {}
+        for layer_name in weights.layer_names:
+            deltas[layer_name] = weights.lora_b[layer_name] @ weights.lora_a[layer_name]
+        return deltas
+
+    def compile(self, **kwargs) -> VLoRAModel:
+        """Compile the base model with torch.compile for faster inference.
+
+        Passes all kwargs to torch.compile(). The LoRA hooks remain
+        uncompiled (they're lightweight matmuls) while the base model
+        benefits from fusion and kernel optimization.
+
+        Returns self for chaining.
+        """
+        self.base_model = torch.compile(self.base_model, **kwargs)
+        return self