CUDA fix

model_handler.py

@@ -14,23 +14,47 @@ def load_model_and_tokenizer():
     offload_dir = "offload_dir"
     os.makedirs(offload_dir, exist_ok=True)
     
-    # Load base model with 8-bit quantization to reduce memory usage
-    base_model = AutoModelForCausalLM.from_pretrained(
-        base_model_name,
-        torch_dtype=torch.float32,  # Use float32 for CPU
-        device_map="auto",
-        offload_folder=offload_dir,  # Add offload directory
-        load_in_8bit=True,          # Use 8-bit quantization
-        low_cpu_mem_usage=True      # Optimize for low memory
-    )
+    # Check if CUDA is available
+    use_cuda = torch.cuda.is_available()
     
-    # Load adapter weights
-    model = PeftModel.from_pretrained(
-        base_model, 
-        "phi2-grpo-qlora-final",
-        device_map="auto",
-        offload_folder=offload_dir  # Add offload directory
-    )
+    try:
+        # First try loading with quantization if CUDA is available
+        if use_cuda:
+            base_model = AutoModelForCausalLM.from_pretrained(
+                base_model_name,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                offload_folder=offload_dir,
+                load_in_8bit=True,
+                low_cpu_mem_usage=True
+            )
+        else:
+            # CPU-only loading without quantization
+            base_model = AutoModelForCausalLM.from_pretrained(
+                base_model_name,
+                torch_dtype=torch.float32,
+                device_map="auto",
+                offload_folder=offload_dir,
+                low_cpu_mem_usage=True
+            )
+        
+        # Load adapter weights
+        model = PeftModel.from_pretrained(
+            base_model, 
+            "phi2-grpo-qlora-final",
+            device_map="auto",
+            offload_folder=offload_dir
+        )
+    except Exception as e:
+        print(f"Error loading with adapter: {e}")
+        print("Falling back to base model only...")
+        # Fallback to just the base model if adapter loading fails
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model_name,
+            torch_dtype=torch.float32,
+            device_map="auto",
+            low_cpu_mem_usage=True
+        )
     
     # Set to evaluation mode
     model.eval()