Upload run_cloud_training.py with huggingface_hub

run_cloud_training.py

@@ -5,6 +5,7 @@ Simplified fine-tuning script for DeepSeek-R1-Distill-Qwen-14B-unsloth-bnb-4bit
 - Optimized for L40S GPU
 - Works with pre-tokenized datasets
 - Research training only (no inference)
+- Added CPU fallback support for Hugging Face Spaces
 """
 
 import os
@@ -30,6 +31,27 @@ DEFAULT_DATASET = "George-API/phi4-cognitive-dataset"
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# Check if CUDA is available for bitsandbytes
+def is_bnb_available():
+    """Check if bitsandbytes with CUDA is available"""
+    try:
+        import bitsandbytes as bnb
+        if torch.cuda.is_available():
+            # Try to create a dummy 4-bit tensor to see if it works
+            try:
+                _ = torch.zeros(1, dtype=torch.float16, device="cuda").to(bnb.nn.Linear4bit)
+                logger.info("BitsAndBytes with CUDA support is available")
+                return True
+            except Exception as e:
+                logger.warning(f"CUDA available but bitsandbytes test failed: {e}")
+                return False
+        else:
+            logger.warning("CUDA not available for bitsandbytes")
+            return False
+    except (ImportError, RuntimeError) as e:
+        logger.warning(f"Error checking bitsandbytes: {e}")
+        return False
+
 # Create a marker file to indicate training is active
 def create_training_marker(output_dir):
     os.makedirs(output_dir, exist_ok=True)
@@ -300,26 +322,57 @@ def train(config_path, dataset_name, output_dir, upload_to_hub=False, hub_repo_n
         )
         tokenizer.pad_token = tokenizer.eos_token
         
-        # Create quantization config
+        # Get quantization config
         quant_config = config.get("quantization_config", {})
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=quant_config.get("load_in_4bit", True),
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_quant_type=quant_config.get("bnb_4bit_quant_type", "nf4"),
-            bnb_4bit_use_double_quant=quant_config.get("bnb_4bit_use_double_quant", True)
-        )
+        
+        # Check if bitsandbytes with CUDA is available
+        use_4bit = is_bnb_available() and quant_config.get("load_in_4bit", True)
         
         # Create model with proper configuration
-        logger.info("Loading pre-quantized model")
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            quantization_config=bnb_config,
-            device_map="auto",
-            torch_dtype=torch.float16,
-            trust_remote_code=True,
-            use_cache=model_config.get("use_cache", False),
-            attn_implementation=hardware_config.get("attn_implementation", "eager")
-        )
+        logger.info(f"Loading model (4-bit quantization: {use_4bit})")
+        
+        if use_4bit:
+            # Create quantization config for GPU
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type=quant_config.get("bnb_4bit_quant_type", "nf4"),
+                bnb_4bit_use_double_quant=quant_config.get("bnb_4bit_use_double_quant", True)
+            )
+            
+            # Load 4-bit quantized model for GPU
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                quantization_config=bnb_config,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                use_cache=model_config.get("use_cache", False),
+                attn_implementation=hardware_config.get("attn_implementation", "eager")
+            )
+        else:
+            # CPU fallback (or non-quantized GPU) mode
+            logger.warning("Loading model in CPU fallback mode (no 4-bit quantization)")
+            
+            # Determine best dtype based on available hardware
+            if torch.cuda.is_available():
+                dtype = torch.float16
+                device_map = "auto"
+                logger.info("Using GPU with fp16")
+            else:
+                dtype = torch.float32
+                device_map = "cpu"
+                logger.info("Using CPU with fp32")
+            
+            # Load model without quantization
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                device_map=device_map,
+                torch_dtype=dtype,
+                trust_remote_code=True,
+                use_cache=model_config.get("use_cache", False),
+                low_cpu_mem_usage=True
+            )
         
         # Apply rope scaling if configured
         if "rope_scaling" in model_config:
@@ -342,7 +395,7 @@ def train(config_path, dataset_name, output_dir, upload_to_hub=False, hub_repo_n
         model = get_peft_model(model, lora_config_obj)
         logger.info("Successfully applied LoRA")
         
-        # Check for L40S GPU and optimize batch size
+        # Determine batch size based on available hardware
         if torch.cuda.is_available():
             gpu_info = torch.cuda.get_device_properties(0)
             logger.info(f"GPU: {gpu_info.name}, VRAM: {gpu_info.total_memory / 1e9:.2f} GB")
@@ -356,8 +409,9 @@ def train(config_path, dataset_name, output_dir, upload_to_hub=False, hub_repo_n
                 per_device_train_batch_size = 2
                 logger.info(f"Using conservative batch size for non-L40S GPU: {per_device_train_batch_size}")
         else:
+            # Use minimal batch size for CPU
             per_device_train_batch_size = 1
-            logger.warning("No GPU detected - using minimal batch size")
+            logger.warning("No GPU detected - using minimal batch size for CPU training")
         
         # Configure reporting backends
         reports = training_config.get("report_to", ["tensorboard"])
@@ -374,8 +428,8 @@ def train(config_path, dataset_name, output_dir, upload_to_hub=False, hub_repo_n
             warmup_ratio=training_config.get("warmup_ratio", 0.03),
             weight_decay=training_config.get("weight_decay", 0.01),
             optim=training_config.get("optim", "adamw_torch"),
-            fp16=hardware_config.get("fp16", True),
-            bf16=hardware_config.get("bf16", False),
+            fp16=torch.cuda.is_available() and hardware_config.get("fp16", True),
+            bf16=torch.cuda.is_available() and hardware_config.get("bf16", False),
             max_grad_norm=training_config.get("max_grad_norm", 0.3),
             logging_steps=training_config.get("logging_steps", 10),
             save_steps=training_config.get("save_steps", 200),
@@ -387,7 +441,7 @@ def train(config_path, dataset_name, output_dir, upload_to_hub=False, hub_repo_n
             logging_first_step=training_config.get("logging_first_step", True),
             disable_tqdm=training_config.get("disable_tqdm", False),
             remove_unused_columns=False,
-            gradient_checkpointing=hardware_config.get("gradient_checkpointing", True),
+            gradient_checkpointing=torch.cuda.is_available() and hardware_config.get("gradient_checkpointing", True),
             dataloader_num_workers=training_config.get("dataloader_num_workers", 4)
         )