Upload extract_lora.py

extract_lora.py

@@ -0,0 +1,139 @@
+import torch
+import argparse
+from safetensors.torch import save_file, safe_open
+from tqdm import tqdm
+import sys
+
+
+def get_torch_dtype(dtype_str: str):
+    """Converts a string to a torch.dtype object."""
+    if dtype_str == "fp32":
+        return torch.float32
+    if dtype_str == "fp16":
+        return torch.float16
+    if dtype_str == "bf16":
+        return torch.bfloat16
+    raise ValueError(f"Unsupported dtype: {dtype_str}")
+
+
+def extract_and_svd_lora(model_a_path: str, model_b_path: str, output_path: str, rank: int, device: str, alpha: float,
+                         dtype: torch.dtype):
+    """
+    Extracts the difference between two models, applies SVD to reduce the rank,
+    and saves the result as a LoRA file.
+    """
+    print(f"Loading base model A: {model_a_path}")
+    print(f"Loading finetuned model B: {model_b_path}")
+
+    lora_tensors = {}
+
+    with safe_open(model_a_path, framework="pt", device="cpu") as f_a, \
+            safe_open(model_b_path, framework="pt", device="cpu") as f_b:
+
+        keys_a = set(f_a.keys())
+        keys_b = set(f_b.keys())
+        common_keys = keys_a.intersection(keys_b)
+
+        # Filter for processable layers (typically linear and conv weights)
+        # We exclude biases and non-weight tensors.
+        weight_keys = {k for k in common_keys if k.endswith('.weight') and 'lora_' not in k}
+
+        if not weight_keys:
+            print("No common weight keys found between the two models. Exiting.")
+            sys.exit(1)
+
+        print(f"Found {len(weight_keys)} common weight keys to process.")
+
+        # Main processing loop with progress bar
+        for key in tqdm(sorted(list(weight_keys)), desc="Processing Layers"):
+            try:
+                # Load tensors and move to the selected device and dtype
+                tensor_a = f_a.get_tensor(key).to(device=device, dtype=dtype)
+                tensor_b = f_b.get_tensor(key).to(device=device, dtype=dtype)
+
+                if tensor_a.shape != tensor_b.shape:
+                    print(f"Skipping key {key} due to shape mismatch: A={tensor_a.shape}, B={tensor_b.shape}")
+                    continue
+
+                # Calculate the difference (delta weight)
+                delta_w = tensor_b - tensor_a
+
+                # SVD works on 2D matrices. Reshape conv layers and other ND tensors.
+                original_shape = delta_w.shape
+                if delta_w.dim() > 2:
+                    delta_w = delta_w.view(original_shape[0], -1)
+
+                # --- Core SVD Logic ---
+                # ΔW ≈ U * S * Vh
+                # U: Left singular vectors
+                # S: Singular values (a 1D vector)
+                # Vh: Right singular vectors (transposed)
+                U, S, Vh = torch.linalg.svd(delta_w, full_matrices=False)
+
+                # Truncate to the desired rank
+                current_rank = min(rank, S.size(0))  # Ensure rank is not > possible rank
+                U = U[:, :current_rank]
+                S = S[:current_rank]
+                Vh = Vh[:current_rank, :]
+
+                # --- Decompose into LoRA A and B matrices ---
+                # LoRA A (lora_down) is Vh
+                # LoRA B (lora_up) is U * S
+                # We scale lora_up by the singular values to retain the magnitude
+                lora_down = Vh
+                lora_up = U @ torch.diag(S)
+
+                # Reshape back to original conv format if necessary
+                if len(original_shape) > 2:
+                    # For Conv2D, lora_down is (rank, in_channels * k_h * k_w)
+                    # and lora_up is (out_channels, rank). No reshape needed for up.
+                    pass  # The matrix form is standard for LoRA conv layers
+
+                # Create LoRA tensor names
+                base_name = key.replace('.weight', '')
+                lora_down_name = f"{base_name}.lora_down.weight"
+                lora_up_name = f"{base_name}.lora_up.weight"
+                alpha_name = f"{base_name}.alpha"
+
+                # Store tensors, moving them to CPU for saving
+                lora_tensors[lora_down_name] = lora_down.contiguous().cpu().to(torch.float32)
+                lora_tensors[lora_up_name] = lora_up.contiguous().cpu().to(torch.float32)
+                lora_tensors[alpha_name] = torch.tensor(alpha).to(torch.float32)
+
+            except Exception as e:
+                print(f"Failed to process key {key}: {e}")
+
+    # Save the final LoRA file
+    if not lora_tensors:
+        print("No tensors were processed. Output file will not be created.")
+        return
+
+    print(f"\nSaving {len(lora_tensors)} tensors to {output_path}...")
+    save_file(lora_tensors, output_path)
+    print("✅ Done!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Extract and SVD a LoRA from two SafeTensors checkpoints.")
+
+    parser.add_argument("model_a", type=str, help="Path to the base model (A) checkpoint in .safetensors format.")
+    parser.add_argument("model_b", type=str, help="Path to the finetuned model (B) checkpoint in .safetensors format.")
+    parser.add_argument("output", type=str, help="Path to save the output LoRA file in .safetensors format.")
+
+    parser.add_argument("--rank", type=int, required=True, help="The target rank for the SVD.")
+    parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"],
+                        help="Device to use for computation ('cuda' or 'cpu').")
+    parser.add_argument("--alpha", type=float, default=1.0, help="The alpha (scaling) factor for the LoRA.")
+    parser.add_argument("--precision", type=str, default="fp32", choices=["fp32", "fp16", "bf16"],
+                        help="Precision to use for calculations.")
+
+    args = parser.parse_args()
+
+    # Device check
+    if args.device == "cuda" and not torch.cuda.is_available():
+        print("CUDA is not available. Falling back to CPU.")
+        args.device = "cpu"
+
+    dtype = get_torch_dtype(args.precision)
+
+    extract_and_svd_lora(args.model_a, args.model_b, args.output, args.rank, args.device, args.alpha, dtype)