merge

app.py

@@ -1,18 +1,22 @@
 import os
 import torch
-import psutil
-from pathlib import Path
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
 from huggingface_hub import login, create_repo, HfApi
 import gradio as gr
-import queue
 import time
 import shutil
 from gradio_log import Log
 import logging
 
 
+MEMORY = int(os.getenv("MEMORY", 16)[:-2]) # 64Gi
+CPU_CORES = int(os.getenv("CPU_CORES", 4)) # 4
+SPACE_AUTHOR_NAME = os.getenv("SPACE_AUTHOR_NAME", "Steven10429") # str
+SPACE_REPO_NAME = os.getenv("SPACE_REPO_NAME", "apply_lora_and_quantize") # str
+SPACE_ID = os.getenv("SPACE_ID", "apply_lora_and_quantize") # str
+
+
 # 全局日志
 log = logging.getLogger("space_convert")
 log.setLevel(logging.INFO)
@@ -45,31 +49,37 @@ def get_model_size_in_gb(model_name):
 def check_system_resources(model_name):
     """检查系统资源，决定使用 CPU 或 GPU"""
     log.info("Checking system resources...")
-    system_memory = psutil.virtual_memory()
-    total_memory_gb = system_memory.total / (1024 ** 3)
-    log.info(f"Total system memory: {total_memory_gb:.1f}GB")
+    log.info(f"Total CPU cores: {CPU_CORES}")
+    log.info(f"Total system memory: {MEMORY}GB")
     
     model_size_gb = get_model_size_in_gb(model_name)
-    required_memory_gb = model_size_gb * 2.5  # 预留额外内存
+    required_memory_gb_16bit = model_size_gb * 1.5
+    required_memory_gb = required_memory_gb_16bit
+    
+    
     log.info(f"Estimated required memory for model: {required_memory_gb:.1f}GB")
     
-    if torch.cuda.is_available():
-        gpu_name = torch.cuda.get_device_name(0)
-        gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
-        log.info(f"Detected GPU: {gpu_name} with {gpu_memory_gb:.1f}GB memory")
-        if gpu_memory_gb >= required_memory_gb:
-            log.info("✅ Sufficient GPU memory available; using GPU.")
-            return "cuda", gpu_memory_gb
-        else:
-            log.warning(f"⚠️ Insufficient GPU memory (requires {required_memory_gb:.1f}GB, found {gpu_memory_gb:.1f}GB).")
-    else:
-        log.error("❌ No GPU detected.")
+    # if torch.cuda.is_available(): # failed with torch complie without GPU FLAG
+    #     gpu_name = torch.cuda.get_device_name(0)
+    #     gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
+    #     log.info(f"Detected GPU: {gpu_name} with {gpu_memory_gb:.1f}GB memory")
+    #     if gpu_memory_gb >= required_memory_gb:
+    #         log.info("✅ Sufficient GPU memory available; using GPU.")
+    #         return "cuda", gpu_memory_gb
+    #     else:
+    #         log.warning(f"⚠️ Insufficient GPU memory (requires {required_memory_gb:.1f}GB, found {gpu_memory_gb:.1f}GB).")
+    # else:
+    #     log.error("❌ No GPU detected.")
+    # just use CPU, it's enough for merge and quantize
         
-    if total_memory_gb >= required_memory_gb:
+    if MEMORY >= required_memory_gb:
         log.info("✅ Sufficient CPU memory available; using CPU.")
-        return "cpu", total_memory_gb
+        return "cpu", MEMORY
     else:
-        raise MemoryError(f"❌ Insufficient system memory (requires {required_memory_gb:.1f}GB, available {available_memory_gb:.1f}GB).")
+        log.warning(f"⚠️ Insufficient CPU memory (requires {required_memory_gb:.1f}GB, found {MEMORY}GB).")
+        log.error("❌ No CPU detected.")
+        log.error("Will try low memory mode, but it may fail.")
+        return "cpu", MEMORY
 
 @timeit
 def setup_environment(model_name):
@@ -114,19 +124,20 @@ def download_and_merge_model(base_model_name, lora_model_name, output_dir, devic
     """
     os.makedirs("temp", exist_ok=True)
     log.info("Loading base model...")
-    model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True, device_map="auto", trust_remote_code=True, offload_folder="temp")
+    model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True, device_map="auto", force_download=True, trust_remote_code=True, torch_dtype=torch.float16)
     log.info("Loading adapter tokenizer...")
-    adapter_tokenizer = AutoTokenizer.from_pretrained(lora_model_name, trust_remote_code=True, device_map="auto", offload_folder="temp")
+    adapter_tokenizer = AutoTokenizer.from_pretrained(lora_model_name, trust_remote_code=True, device_map="auto", force_download=True, trust_remote_code=True, torch_dtype=torch.float16)
     log.info("Resizing token embeddings...")
     added_tokens_decoder = adapter_tokenizer.added_tokens_decoder
     model.resize_token_embeddings(adapter_tokenizer.vocab_size + len(added_tokens_decoder))
     log.info("Loading LoRA adapter...")
-    peft_model = PeftModel.from_pretrained(model, lora_model_name, low_cpu_mem_usage=True, device_map="auto", trust_remote_code=True, offload_folder="temp")
+    peft_model = PeftModel.from_pretrained(model, lora_model_name, low_cpu_mem_usage=True, device_map="auto", force_download=True, trust_remote_code=True, torch_dtype=torch.float16)
     log.info("Merging and unloading model...")
     model = peft_model.merge_and_unload()
     log.info("Saving model...")
     model.save_pretrained(output_dir)
     adapter_tokenizer.save_pretrained(output_dir)
+    del model, peft_model
     return output_dir
 
 @timeit