temporarily diable cuda when loading lora

app.py

@@ -11,6 +11,7 @@ from copy import deepcopy
 import numpy as np
 import re
 from bs4 import BeautifulSoup
+import torch
 
 
 import logging
@@ -37,15 +38,38 @@ max_new_tokens = 5000
 obj_context_feature_type = "text"
 
 
-tokenizer, model, data_loader = load_model_and_dataloader(
-    model_path=model_path,
-    model_base=model_base,
-    load_8bit=load_8bit,
-    load_4bit=load_4bit,
-    load_bf16=load_bf16,
-    scene_to_obj_mapping=scene_to_obj_mapping,
-    device_map='cpu',
-)  # Huggingface Zero-GPU has to use .to(device) to set the device, otherwise it will fail
+class TemporaryDisableCUDA:
+    def __enter__(self):
+        # Store the current CUDA_VISIBLE_DEVICES environment variable
+        self.original_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
+        # Disable CUDA by setting CUDA_VISIBLE_DEVICES to an invalid value
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+        # Reinitialize CUDA state in PyTorch
+        torch.cuda.device_count()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Restore the original CUDA_VISIBLE_DEVICES environment variable
+        if self.original_cuda_visible_devices is not None:
+            os.environ["CUDA_VISIBLE_DEVICES"] = self.original_cuda_visible_devices
+        else:
+            del os.environ["CUDA_VISIBLE_DEVICES"]
+        # Reinitialize CUDA state in PyTorch
+        torch.cuda.device_count()
+
+
+# HACK: we need to temporarily disable CUDA b/c when loading model,
+# PEFT somehow always uses CUDA even if we set device_map to 'cpu'
+with TemporaryDisableCUDA():
+    tokenizer, model, data_loader = load_model_and_dataloader(
+        model_path=model_path,
+        model_base=model_base,
+        load_8bit=load_8bit,
+        load_4bit=load_4bit,
+        load_bf16=load_bf16,
+        scene_to_obj_mapping=scene_to_obj_mapping,
+        device_map='cpu',
+    )  # Huggingface Zero-GPU has to use .to(device) to set the device, otherwise it will fail
+
 model.to("cuda")  # Huggingface Zero-GPU requires explicit device placement
 
 def get_chatbot_response(user_chat_input, scene_id):