Add LoRA adapter merge utility

scripts/merge_lora_adapters.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+"""Merge PEFT/LoRA adapters with linear model soup or SLERP.
+
+This is intended for adapters that share the same base model and LoRA config,
+for example B2, DPO, and PPO adapters trained from the same LLaVA-Med base.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import shutil
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+
+COPY_FILES = [
+    "adapter_config.json",
+    "tokenizer_config.json",
+    "tokenizer.json",
+    "processor_config.json",
+    "chat_template.jinja",
+]
+
+
+def parse_weighted_adapter(raw: str) -> tuple[Path, float]:
+    if "=" not in raw:
+        return Path(raw), 1.0
+    path, weight = raw.rsplit("=", 1)
+    return Path(path), float(weight)
+
+
+def normalize_weights(weights: list[float]) -> list[float]:
+    total = sum(weights)
+    if not math.isfinite(total) or total <= 0:
+        raise ValueError("Adapter weights must sum to a positive finite value.")
+    return [weight / total for weight in weights]
+
+
+def adapter_model_path(adapter_dir: Path) -> Path:
+    path = adapter_dir / "adapter_model.safetensors"
+    if not path.exists():
+        raise FileNotFoundError(f"Missing adapter weights: {path}")
+    return path
+
+
+def load_adapter_state(adapter_dir: Path) -> dict[str, torch.Tensor]:
+    return load_file(str(adapter_model_path(adapter_dir)), device="cpu")
+
+
+def validate_states(states: list[dict[str, torch.Tensor]]) -> list[str]:
+    keys = sorted(states[0].keys())
+    expected = set(keys)
+    for idx, state in enumerate(states[1:], start=2):
+        if set(state.keys()) != expected:
+            missing = sorted(expected - set(state.keys()))[:5]
+            extra = sorted(set(state.keys()) - expected)[:5]
+            raise ValueError(f"Adapter {idx} has different keys. Missing={missing}, extra={extra}")
+        for key in keys:
+            if state[key].shape != states[0][key].shape:
+                raise ValueError(f"Shape mismatch for {key}: {state[key].shape} != {states[0][key].shape}")
+    return keys
+
+
+def linear_soup(states: list[dict[str, torch.Tensor]], weights: list[float], keys: list[str]) -> dict[str, torch.Tensor]:
+    merged: dict[str, torch.Tensor] = {}
+    for key in keys:
+        ref = states[0][key]
+        tensor = torch.zeros_like(ref, dtype=torch.float32)
+        for state, weight in zip(states, weights):
+            tensor += state[key].float() * weight
+        merged[key] = tensor.to(dtype=ref.dtype)
+    return merged
+
+
+def flatten_state(state: dict[str, torch.Tensor], keys: list[str]) -> torch.Tensor:
+    return torch.cat([state[key].float().reshape(-1) for key in keys])
+
+
+def unflatten_state(vector: torch.Tensor, reference: dict[str, torch.Tensor], keys: list[str]) -> dict[str, torch.Tensor]:
+    merged: dict[str, torch.Tensor] = {}
+    offset = 0
+    for key in keys:
+        ref = reference[key]
+        size = ref.numel()
+        merged[key] = vector[offset : offset + size].reshape(ref.shape).to(dtype=ref.dtype)
+        offset += size
+    return merged
+
+
+def slerp_vectors(a: torch.Tensor, b: torch.Tensor, t: float, eps: float = 1e-8) -> torch.Tensor:
+    a_norm = torch.linalg.vector_norm(a)
+    b_norm = torch.linalg.vector_norm(b)
+    if a_norm < eps or b_norm < eps:
+        return (1.0 - t) * a + t * b
+
+    a_unit = a / a_norm
+    b_unit = b / b_norm
+    dot = torch.clamp(torch.dot(a_unit, b_unit), -1.0, 1.0)
+    omega = torch.acos(dot)
+    sin_omega = torch.sin(omega)
+    if torch.abs(sin_omega) < eps:
+        return (1.0 - t) * a + t * b
+
+    direction = (torch.sin((1.0 - t) * omega) / sin_omega) * a + (torch.sin(t * omega) / sin_omega) * b
+    target_norm = (1.0 - t) * a_norm + t * b_norm
+    direction_norm = torch.linalg.vector_norm(direction)
+    if direction_norm < eps:
+        return direction
+    return direction / direction_norm * target_norm
+
+
+def iterative_slerp(states: list[dict[str, torch.Tensor]], weights: list[float], keys: list[str]) -> dict[str, torch.Tensor]:
+    vectors = [flatten_state(state, keys) for state in states]
+    merged = vectors[0]
+    accumulated = weights[0]
+    for vector, weight in zip(vectors[1:], weights[1:]):
+        t = weight / (accumulated + weight)
+        merged = slerp_vectors(merged, vector, t)
+        accumulated += weight
+    return unflatten_state(merged, states[0], keys)
+
+
+def copy_adapter_files(reference_dir: Path, output_dir: Path) -> None:
+    for filename in COPY_FILES:
+        src = reference_dir / filename
+        if src.exists():
+            shutil.copy2(src, output_dir / filename)
+
+
+def write_metadata(
+    output_dir: Path,
+    method: str,
+    adapters: list[Path],
+    weights: list[float],
+    base_model: str,
+) -> None:
+    metadata = {
+        "method": method,
+        "base_model": base_model,
+        "adapters": [str(path) for path in adapters],
+        "weights": weights,
+    }
+    (output_dir / "merge_config.json").write_text(json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8")
+    readme = [
+        "---",
+        "library_name: peft",
+        "tags:",
+        "- peft",
+        "- lora",
+        "- visual-question-answering",
+        f"base_model: {base_model}",
+        "---",
+        "",
+        "# Medical VQA Merged Adapter",
+        "",
+        f"Merge method: `{method}`",
+        "",
+        "| Adapter | Weight |",
+        "|---|---:|",
+    ]
+    for path, weight in zip(adapters, weights):
+        readme.append(f"| `{path}` | {weight:.4f} |")
+    readme.extend(
+        [
+            "",
+            "This adapter was created by merging fine-tuned LoRA adapters without additional training.",
+        ]
+    )
+    (output_dir / "README.md").write_text("\n".join(readme) + "\n", encoding="utf-8")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Merge PEFT LoRA adapters via model soup or SLERP.")
+    parser.add_argument(
+        "--adapter",
+        action="append",
+        required=True,
+        help="Adapter directory, optionally with weight: path=0.4. Repeat for each adapter.",
+    )
+    parser.add_argument("--output", required=True, help="Output adapter directory.")
+    parser.add_argument("--method", choices=["linear", "slerp"], default="linear")
+    parser.add_argument("--base-model", default="chaoyinshe/llava-med-v1.5-mistral-7b-hf")
+    parser.add_argument(
+        "--reference-index",
+        type=int,
+        default=0,
+        help="Adapter index used as source for adapter_config/tokenizer files.",
+    )
+    args = parser.parse_args()
+
+    adapters, raw_weights = zip(*(parse_weighted_adapter(raw) for raw in args.adapter))
+    adapters = [path.expanduser().resolve() for path in adapters]
+    weights = normalize_weights(list(raw_weights))
+    output_dir = Path(args.output).expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if not 0 <= args.reference_index < len(adapters):
+        raise ValueError("--reference-index is out of range.")
+
+    states = [load_adapter_state(path) for path in adapters]
+    keys = validate_states(states)
+    merged = linear_soup(states, weights, keys) if args.method == "linear" else iterative_slerp(states, weights, keys)
+
+    save_file(merged, str(output_dir / "adapter_model.safetensors"))
+    copy_adapter_files(adapters[args.reference_index], output_dir)
+    write_metadata(output_dir, args.method, adapters, weights, args.base_model)
+    print(f"Saved merged adapter to: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()