extract_svd.py

import os
import argparse
import torch
try:
    from safetensors.torch import load_file
except ImportError:
    load_file = None

def load_weights(adapter_dir):
    """
    Attempts to load LoRA weights from adapter_model.bin or adapter_model.safetensors.
    """
    bin_path = os.path.join(adapter_dir, "adapter_model.bin")
    st_path = os.path.join(adapter_dir, "adapter_model.safetensors")
    
    if os.path.exists(bin_path):
        return torch.load(bin_path, map_location="cpu")
    elif os.path.exists(st_path) and load_file is not None:
        return load_file(st_path)
    else:
        return None

def compute_svd_for_dataset(base_dir, r=8, output_file="lora_svd_dataset.pt"):
    dataset = []
    
    for entry in os.scandir(base_dir):
        if not entry.is_dir():
            continue
            
        adapter_dir = entry.path
        state_dict = load_weights(adapter_dir)
        if state_dict is None:
            print(f"Warning: No adapter weights found in {adapter_dir}")
            continue
            
        modules = {}
        for k, v in state_dict.items():
            if "lora_A" in k or "lora_B" in k:
                # Group by base module name
                base_name = k.replace(".lora_A.default.weight", "").replace(".lora_B.default.weight", "")
                base_name = base_name.replace(".lora_A.weight", "").replace(".lora_B.weight", "")
                
                if base_name not in modules:
                    modules[base_name] = {}
                    
                if "lora_A" in k:
                    modules[base_name]["A"] = v
                elif "lora_B" in k:
                    modules[base_name]["B"] = v
                    
        U_dict = {}
        log_Sigma_dict = {}
        V_dict = {}
        
        valid_modules = 0
        for name, parts in modules.items():
            if "A" in parts and "B" in parts:
                A = parts["A"]
                B = parts["B"]
                
                # PEFT standard usage: forward pass uses output = input @ W + input @ A.T @ B.T
                # Meaning A is shaped (r, d_in) and B is shaped (d_out, r).
                # To match spec (Delta_W = A x B), let's ensure inner dimensions match.
                if A.shape[0] == B.shape[1]: 
                    delta_W = B @ A  # (d_out, r) @ (r, d_in) = (d_out, d_in)
                else:
                    delta_W = A @ B
                    
                # Perform SVD natively
                U, S, Vh = torch.linalg.svd(delta_W, full_matrices=False)
                
                # Truncate to rank r
                U_r = U[:, :r]
                S_r = S[:r]
                Vh_r = Vh[:r, :]
                
                # Convert Sigma to log(Sigma) to enforce strict positivity on recreation
                log_S_r = torch.log(S_r + 1e-8)
                
                U_dict[name] = U_r
                log_Sigma_dict[name] = log_S_r
                V_dict[name] = Vh_r
                valid_modules += 1
                
        if valid_modules > 0:
            # Fallback for task description if the text file isn't found
            desc_path = os.path.join(adapter_dir, "task_description.txt")
            if os.path.exists(desc_path):
                with open(desc_path, "r", encoding="utf-8") as f:
                    task_desc = f.read().strip()
            else:
                task_desc = f"Instruction tuning task for adapter {os.path.basename(adapter_dir)}"
                
            dataset.append({
                "task_description_string": task_desc,
                "tensors": {
                    "U_dict": U_dict,
                    "log_Sigma_dict": log_Sigma_dict,
                    "V_dict": V_dict
                }
            })
            print(f"Processed 1 adapter: '{task_desc[:30]}...' with {valid_modules} matrices.")
            
    torch.save(dataset, output_file)
    print(f"Extraction complete. Saved {len(dataset)} adapters to {output_file}.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Extract SVD components from pre-trained LoRA models")
    parser.add_argument("--base_dir", type=str, required=True, help="Directory containing the 50 LoRA subdirectories")
    parser.add_argument("--r", type=int, default=8, help="Truncation rank")
    parser.add_argument("--output_file", type=str, default="lora_svd_dataset.pt", help="Path to output .pt file")
    
    args = parser.parse_args()
    compute_svd_for_dataset(args.base_dir, args.r, args.output_file)