Implement manual LoRA merging to fix PEFT key naming conflicts

app.py

@@ -3,6 +3,7 @@ import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel, PeftConfig
+from safetensors.torch import load_file
 import gc
 from huggingface_hub import login, snapshot_download
 import logging
@@ -61,6 +62,83 @@ class ModelMerger:
             logger.error(f"Login failed: {str(e)}")
             return f"❌ Login failed: {str(e)}"
     
+    def manual_merge_lora(self, model, adapter_path, progress=gr.Progress()):
+        """Manually merge LoRA weights into model to avoid PEFT key naming issues"""
+        import json
+        from tqdm import tqdm
+        
+        logger.info("Using manual LoRA merge to avoid key naming conflicts...")
+        progress(0.55, desc="Loading LoRA adapter weights...")
+        
+        # Load adapter weights
+        adapter_file = os.path.join(adapter_path, "adapter_model.safetensors")
+        adapter_weights = load_file(adapter_file)
+        logger.info(f"Loaded {len(adapter_weights)} adapter weight tensors")
+        
+        # Load adapter config
+        config_file = os.path.join(adapter_path, "adapter_config.json")
+        with open(config_file) as f:
+            adapter_config = json.load(f)
+        
+        lora_alpha = adapter_config["lora_alpha"]
+        r = adapter_config["r"]
+        scaling = lora_alpha / r
+        logger.info(f"LoRA scaling: {scaling} (alpha={lora_alpha}, r={r})")
+        
+        # Group LoRA A and B weights
+        lora_pairs = {}
+        for key in adapter_weights.keys():
+            if "lora_A" in key:
+                base_key = key.replace(".lora_A.weight", "")
+                lora_pairs[base_key] = {
+                    "A": adapter_weights[key],
+                    "B": adapter_weights.get(base_key + ".lora_B.weight")
+                }
+        
+        logger.info(f"Found {len(lora_pairs)} LoRA pairs to merge")
+        
+        progress(0.65, desc=f"Merging {len(lora_pairs)} LoRA layers...")
+        
+        # Get model state dict
+        model_state_dict = model.state_dict()
+        merged_count = 0
+        
+        for adapter_key, lora_weights in lora_pairs.items():
+            # adapter_key: base_model.model.model.layers.0.self_attn.q_proj
+            # Need to find corresponding key in model_state_dict
+            
+            # Remove 'base_model.model.' prefix
+            if adapter_key.startswith("base_model.model."):
+                search_key = adapter_key[len("base_model.model."):]
+            else:
+                search_key = adapter_key
+            
+            # Find matching key in model
+            model_key = None
+            for mk in model_state_dict.keys():
+                if search_key in mk or mk.endswith(search_key.split(".")[-4:][0]):
+                    # Match by layer structure
+                    if all(part in mk for part in search_key.split(".")[-4:]):
+                        model_key = mk
+                        break
+            
+            if model_key and model_key in model_state_dict:
+                lora_A = lora_weights["A"].to(model_state_dict[model_key].device)
+                lora_B = lora_weights["B"].to(model_state_dict[model_key].device)
+                
+                # Merge: W_new = W_old + (lora_B @ lora_A) * scaling
+                delta_W = (lora_B @ lora_A) * scaling
+                model_state_dict[model_key] = model_state_dict[model_key] + delta_W.to(model_state_dict[model_key].dtype)
+                merged_count += 1
+        
+        logger.info(f"Successfully merged {merged_count}/{len(lora_pairs)} LoRA weights")
+        
+        # Load merged weights back
+        progress(0.75, desc="Loading merged weights into model...")
+        model.load_state_dict(model_state_dict, strict=False)
+        
+        return model
+    
     def merge_models(self, hf_token, use_8bit=False, progress=gr.Progress()):
         """Merge LoRA adapters with base model"""
         try:
@@ -118,24 +196,16 @@ class ModelMerger:
                 precision_desc = "bfloat16"
             
             try:
-                # For merging, use sequential device map to avoid complex key nesting
-                # This ensures consistent key names between training and merging
+                # Try loading with balanced device map to distribute evenly
                 load_kwargs = {
                     "trust_remote_code": True,
                     "low_cpu_mem_usage": True,
-                    "device_map": "sequential",  # Changed from "auto" to avoid key nesting issues
+                    "device_map": "balanced",  # Distribute layers evenly across GPUs
                     "max_memory": max_memory,
+                    "torch_dtype": torch.bfloat16,
                 }
                 
-                if use_8bit:
-                    # Use 8-bit quantization for tighter memory constraints
-                    load_kwargs["load_in_8bit"] = True
-                    load_kwargs["llm_int8_enable_fp32_cpu_offload"] = True
-                    load_kwargs["llm_int8_threshold"] = 6.0
-                    logger.info("Enabling CPU offload for 8-bit quantization")
-                else:
-                    # Use bfloat16 for best quality when memory allows
-                    load_kwargs["torch_dtype"] = torch.bfloat16
+                logger.info("Loading base model with balanced device map...")
                 
                 self.base_model = AutoModelForCausalLM.from_pretrained(
                     BASE_MODEL_NAME,
@@ -155,68 +225,34 @@ class ModelMerger:
                     error_msg += "\n💡 **Try enabling 8-bit quantization** to reduce memory usage by ~50%."
                 raise Exception(error_msg)
             
-            # Load LoRA configuration
-            progress(0.50, desc="Loading LoRA adapters...")
-            logger.info(f"Loading LoRA adapters from: {LORA_MODEL_NAME}")
+            # Download LoRA adapters
+            progress(0.50, desc="Downloading LoRA adapters...")
+            logger.info(f"Downloading LoRA adapters from: {LORA_MODEL_NAME}")
             
-            # Check if LoRA model exists and is accessible
-            try:
-                from huggingface_hub import repo_info
-                info = repo_info(LORA_MODEL_NAME, token=hf_token)
-                logger.info(f"LoRA model found: {info}")
-            except Exception as e:
-                logger.warning(f"Could not verify LoRA model: {str(e)}")
+            # Download entire adapter folder
+            adapter_path = snapshot_download(
+                repo_id=LORA_MODEL_NAME,
+                token=hf_token,
+                allow_patterns=["adapter_*", "*.json"]
+            )
+            logger.info(f"LoRA adapters downloaded to: {adapter_path}")
+            
+            # Use manual merge to avoid PEFT key naming issues
+            progress(0.55, desc="Merging LoRA weights (manual merge)...")
+            logger.info("Using manual LoRA merge to avoid key naming conflicts with PEFT")
             
-            # Load LoRA adapters with additional parameters
             try:
-                logger.info("Attempting to load LoRA adapters...")
-                logger.info(f"LoRA targets attention layers: q_proj, k_proj, v_proj, o_proj")
-                
-                # Load PEFT model - this wraps the base model
-                peft_model = PeftModel.from_pretrained(
-                    self.base_model,
-                    LORA_MODEL_NAME,
-                    torch_dtype=torch.bfloat16 if not use_8bit else None,
-                    is_trainable=False,
-                )
-                logger.info("LoRA adapters loaded successfully")
-                
-                progress(0.70, desc="Merging LoRA weights with base model...")
-                logger.info("Merging LoRA weights into base model...")
-                
-                # Use merge_and_unload with explicit safe merge
-                try:
-                    self.merged_model = peft_model.merge_and_unload(safe_merge=True)
-                    logger.info("Models merged successfully with safe_merge=True")
-                except Exception as merge_error:
-                    logger.warning(f"safe_merge=True failed, trying without: {str(merge_error)}")
-                    # Fallback to regular merge
-                    self.merged_model = peft_model.merge_and_unload()
-                    logger.info("Models merged successfully")
-                
-            except KeyError as e:
-                # Handle missing keys - might be an architecture mismatch
-                error_key = str(e)
-                error_msg = f"Key error when loading LoRA adapters: {error_key}\n\n"
-                
-                if "block_sparse_moe" in error_key or "experts" in error_key:
-                    error_msg += "⚠️ This error is related to MoE (Mixture of Experts) layers.\n\n"
-                    error_msg += "The LoRA adapters only target attention layers (q/k/v/o_proj),\n"
-                    error_msg += "but there seems to be a key naming mismatch with the base model.\n\n"
-                    error_msg += "Possible causes:\n"
-                    error_msg += "1. The base model version has changed since training\n"
-                    error_msg += "2. Different transformers/peft library versions\n"
-                    error_msg += "3. Model was saved with different device_map than loading\n\n"
+                self.merged_model = self.manual_merge_lora(self.base_model, adapter_path, progress)
+                logger.info("✅ LoRA weights merged successfully using manual method")
                 
-                error_msg += "Please verify:\n"
-                error_msg += f"- Base model: {BASE_MODEL_NAME}\n"
-                error_msg += f"- LoRA model: {LORA_MODEL_NAME}\n"
-                error_msg += "- Both use the same transformers version\n"
-                logger.error(error_msg)
+            except Exception as merge_error:
+                logger.error(f"Manual merge failed: {str(merge_error)}", exc_info=True)
+                error_msg = f"Failed to merge LoRA adapters: {str(merge_error)}\n\n"
+                error_msg += "This could be due to:\n"
+                error_msg += "1. Incompatible model architectures\n"
+                error_msg += "2. Corrupted adapter files\n"
+                error_msg += "3. Memory issues during merge\n"
                 raise Exception(error_msg)
-            except Exception as e:
-                logger.error(f"Unexpected error during merge: {str(e)}", exc_info=True)
-                raise
             
             # Save merged model
             progress(0.85, desc="Saving merged model...")

merge_script.py

@@ -0,0 +1,140 @@
+"""
+Manual LoRA merging script that handles key naming issues
+"""
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from safetensors.torch import load_file, save_file
+import os
+import argparse
+from tqdm import tqdm
+
+def merge_lora_weights(
+    base_model_name,
+    adapter_path,
+    output_path,
+    device_map="auto"
+):
+    """Manually merge LoRA weights into base model"""
+    
+    print(f"Loading base model: {base_model_name}")
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.bfloat16,
+        device_map=device_map,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+    
+    print(f"Loading LoRA adapters from: {adapter_path}")
+    adapter_weights = load_file(os.path.join(adapter_path, "adapter_model.safetensors"))
+    
+    print(f"Loaded {len(adapter_weights)} adapter weights")
+    
+    # Load adapter config to get scaling factor
+    import json
+    with open(os.path.join(adapter_path, "adapter_config.json")) as f:
+        adapter_config = json.load(f)
+    
+    lora_alpha = adapter_config["lora_alpha"]
+    r = adapter_config["r"]
+    scaling = lora_alpha / r
+    
+    print(f"LoRA scaling factor: {scaling} (alpha={lora_alpha}, r={r})")
+    
+    # Group LoRA weights by layer
+    lora_pairs = {}
+    for key in adapter_weights.keys():
+        if "lora_A" in key:
+            base_key = key.replace(".lora_A.weight", "")
+            lora_pairs[base_key] = {
+                "A": adapter_weights[key],
+                "B": adapter_weights.get(base_key + ".lora_B.weight")
+            }
+    
+    print(f"Found {len(lora_pairs)} LoRA pairs to merge")
+    
+    # Get model state dict
+    model_state_dict = model.state_dict()
+    
+    # Map adapter keys to model keys
+    # Adapter keys: base_model.model.model.layers.X.self_attn.q_proj
+    # Model keys might be: model.layers.X.self_attn.q_proj (depending on device_map)
+    
+    print("\nMerging LoRA weights...")
+    merged_count = 0
+    
+    for adapter_key, lora_weights in tqdm(lora_pairs.items()):
+        # Remove 'base_model.model.' prefix from adapter key
+        # adapter_key looks like: base_model.model.model.layers.0.self_attn.q_proj
+        if adapter_key.startswith("base_model.model."):
+            model_key = adapter_key[len("base_model.model."):]
+        else:
+            model_key = adapter_key
+        
+        # Try to find the matching key in model
+        found = False
+        for mk in model_state_dict.keys():
+            if model_key in mk or mk.endswith(model_key):
+                model_key = mk
+                found = True
+                break
+        
+        if not found:
+            # Try alternative key formats
+            alternatives = [
+                model_key,
+                "model." + model_key,
+                model_key.replace("model.", ""),
+            ]
+            
+            for alt_key in alternatives:
+                if alt_key in model_state_dict:
+                    model_key = alt_key
+                    found = True
+                    break
+        
+        if found and model_key in model_state_dict:
+            # Merge: W' = W + (B @ A) * scaling
+            lora_A = lora_weights["A"]
+            lora_B = lora_weights["B"]
+            
+            # Move to same device as model weight
+            device = model_state_dict[model_key].device
+            lora_A = lora_A.to(device)
+            lora_B = lora_B.to(device)
+            
+            # Compute delta_W = (lora_B @ lora_A) * scaling
+            delta_W = (lora_B @ lora_A) * scaling
+            
+            # Add to original weight
+            model_state_dict[model_key] = model_state_dict[model_key] + delta_W.to(model_state_dict[model_key].dtype)
+            merged_count += 1
+        else:
+            print(f"Warning: Could not find model key for {adapter_key}")
+    
+    print(f"\nSuccessfully merged {merged_count}/{len(lora_pairs)} LoRA weights")
+    
+    # Load merged weights back into model
+    model.load_state_dict(model_state_dict, strict=False)
+    
+    # Save merged model
+    print(f"\nSaving merged model to: {output_path}")
+    os.makedirs(output_path, exist_ok=True)
+    model.save_pretrained(output_path, safe_serialization=True, max_shard_size="5GB")
+    
+    # Also save tokenizer
+    print("Saving tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+    tokenizer.save_pretrained(output_path)
+    
+    print("\n✅ Merge complete!")
+    return model
+
+if __name__ == "__main__":
+    # For use in the Space
+    BASE_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
+    ADAPTER_PATH = "/app/lora_adapters"  # We'll download here
+    OUTPUT_PATH = "/app/merged_model"
+    
+    merge_lora_weights(BASE_MODEL, ADAPTER_PATH, OUTPUT_PATH)
+