Add safe_merge and better error handling for LoRA merge with MoE models

app.py

@@ -158,18 +158,64 @@ class ModelMerger:
             progress(0.50, desc="Loading LoRA adapters...")
             logger.info(f"Loading LoRA adapters from: {LORA_MODEL_NAME}")
             
-            # Merge LoRA weights
-            self.merged_model = PeftModel.from_pretrained(
-                self.base_model,
-                LORA_MODEL_NAME,
-                torch_dtype=torch.bfloat16,
-            )
-            logger.info("LoRA adapters loaded successfully")
+            # 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)}")
             
-            progress(0.70, desc="Merging LoRA weights with base model...")
-            logger.info("Merging LoRA weights...")
-            self.merged_model = self.merged_model.merge_and_unload()
-            logger.info("Models merged successfully")
+            # 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"
+                
+                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)
+                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...")