| """Recover a LoRA adapter from a merged LoRA model: for every weight that differs from the base model, |
| delta = merged - base (fp32) is factorised by SVD and truncated to rank r (LoRA training with rank r makes the delta |
| rank <= r, up to bf16 rounding). Saved in PEFT format with lora_alpha = r (scaling 1), so base + adapter = merged. |
| |
| python scripts/extract_lora.py --base models/gemma-4-E4B-it-text --merged models/sft_gmmlu_lora/step-357 \ |
| --rank 16 --out models/adapter_A |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file, save_file |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
|
|
| def load_dir(d: Path) -> dict[str, torch.Tensor]: |
| out = {} |
| for f in sorted(d.glob("*.safetensors")): |
| out.update(load_file(f)) |
| return out |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base", required=True, type=Path) |
| ap.add_argument("--merged", required=True, type=Path) |
| ap.add_argument("--rank", type=int, default=16) |
| ap.add_argument("--out", required=True, type=Path) |
| args = ap.parse_args() |
| base, merged = load_dir(args.base), load_dir(args.merged) |
| assert base.keys() == merged.keys(), "base and merged checkpoints have different tensors" |
|
|
| tensors, targets, n_params, worst = {}, [], 0, 0.0 |
| for k in sorted(base): |
| if torch.equal(base[k], merged[k]): |
| continue |
| if base[k].dim() != 2 or not k.endswith(".weight"): |
| raise SystemExit(f"non-linear tensor changed ({k}); this is not a pure LoRA merge") |
| delta = merged[k].cuda().float() - base[k].cuda().float() |
| U, S, Vh = torch.linalg.svd(delta, full_matrices=False) |
| r = args.rank |
| B = (U[:, :r] * S[:r]).contiguous() |
| A = Vh[:r].contiguous() |
| err = ((delta - B @ A).norm() / delta.norm().clamp_min(1e-12)).item() |
| worst = max(worst, err) |
| module = k[: -len(".weight")] |
| tensors[f"base_model.model.{module}.lora_A.weight"] = A.to(torch.bfloat16).cpu() |
| tensors[f"base_model.model.{module}.lora_B.weight"] = B.to(torch.bfloat16).cpu() |
| targets.append(module) |
| n_params += A.numel() + B.numel() |
|
|
| args.out.mkdir(parents=True, exist_ok=True) |
| save_file(tensors, args.out / "adapter_model.safetensors") |
| cfg = {"peft_type": "LORA", "task_type": "CAUSAL_LM", "r": args.rank, "lora_alpha": args.rank, |
| "lora_dropout": 0.0, "bias": "none", "target_modules": targets, "fan_in_fan_out": False, |
| "use_rslora": False, "modules_to_save": None, "base_model_name_or_path": str(args.base), |
| "inference_mode": True} |
| (args.out / "adapter_config.json").write_text(json.dumps(cfg, indent=2)) |
| (args.out / "SOURCE.json").write_text(json.dumps({"extracted_from": str(args.merged), "base": str(args.base), |
| "rank": args.rank, "modules": len(targets), |
| "adapter_params": n_params, |
| "max_relative_reconstruction_error": worst}, indent=2)) |
| print(f"{len(targets)} modules, {n_params / 1e6:.1f}M adapter parameters, " |
| f"max relative reconstruction error {worst:.4f} -> {args.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
