Upload run_cloud_training.py with huggingface_hub

run_cloud_training.py

@@ -16,7 +16,7 @@ from dotenv import load_dotenv
 import torch
 from datasets import load_dataset
 import transformers
-from transformers import AutoTokenizer, TrainingArguments, Trainer
+from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForCausalLM
 from transformers.data.data_collator import DataCollatorMixin
 from peft import LoraConfig
 from unsloth import FastLanguageModel
@@ -153,6 +153,58 @@ def remove_training_marker():
         os.remove("TRAINING_ACTIVE")
         logger.info("Removed training active marker")
 
+def load_model_safely(model_name, max_seq_length, dtype=None):
+    """
+    Load the model in a safe way that works with Qwen models
+    by trying different loading strategies.
+    """
+    try:
+        logger.info(f"Attempting to load model with unsloth optimizations: {model_name}")
+        # First try the standard unsloth loading
+        try:
+            # Try loading with unsloth but without the problematic parameter
+            model, tokenizer = FastLanguageModel.from_pretrained(
+                model_name=model_name,
+                max_seq_length=max_seq_length,
+                dtype=dtype,
+                load_in_4bit=True,  # This should work for already quantized models
+            )
+            logger.info("Model loaded successfully with unsloth with 4-bit quantization")
+            return model, tokenizer
+            
+        except TypeError as e:
+            # If we get a TypeError about unexpected keyword arguments
+            if "unexpected keyword argument" in str(e):
+                logger.warning(f"Unsloth loading error with 4-bit: {e}")
+                logger.info("Trying alternative loading method for Qwen model...")
+                
+                # Try loading with different parameters for Qwen model
+                model, tokenizer = FastLanguageModel.from_pretrained(
+                    model_name=model_name,
+                    max_seq_length=max_seq_length,
+                    dtype=dtype,
+                )
+                logger.info("Model loaded successfully with unsloth using alternative method")
+                return model, tokenizer
+            else:
+                # Re-raise if it's a different type error
+                raise
+                
+    except Exception as e:
+        # Fallback to standard loading if unsloth methods fail
+        logger.warning(f"Unsloth loading failed: {e}")
+        logger.info("Falling back to standard Hugging Face loading...")
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            torch_dtype=dtype or torch.float16,
+            load_in_4bit=True,
+        )
+        logger.info("Model loaded successfully with standard HF loading")
+        return model, tokenizer
+
 def train(config_path, dataset_name, output_dir):
     """Main training function - RESEARCH TRAINING PHASE ONLY"""
     # Load environment variables
@@ -186,7 +238,8 @@ def train(config_path, dataset_name, output_dir):
         # Print configuration summary
         logger.info("RESEARCH TRAINING PHASE ACTIVE - No output generation")
         logger.info("Configuration Summary:")
-        logger.info(f"Model: {model_config.get('model_name_or_path')}")
+        model_name = model_config.get("model_name_or_path")
+        logger.info(f"Model: {model_name}")
         logger.info(f"Dataset: {dataset_name if dataset_name != 'phi4-cognitive-dataset' else DEFAULT_DATASET}")
         logger.info(f"Output directory: {output_dir}")
         logger.info("IMPORTANT: Using already 4-bit quantized model - not re-quantizing")
@@ -197,7 +250,7 @@ def train(config_path, dataset_name, output_dir):
         # Initialize tokenizer (just for model initialization, not for tokenizing data)
         logger.info("Loading tokenizer (for model initialization only, not for tokenizing data)")
         tokenizer = AutoTokenizer.from_pretrained(
-            model_config.get("model_name_or_path"),
+            model_name,
             trust_remote_code=True
         )
         tokenizer.pad_token = tokenizer.eos_token
@@ -215,15 +268,13 @@ def train(config_path, dataset_name, output_dir):
             target_modules=lora_config.get("target_modules", ["q_proj", "k_proj", "v_proj", "o_proj"])
         )
         
-        # Initialize model with unsloth, preserving existing 4-bit quantization
-        logger.info("Loading pre-quantized model with unsloth")
-        model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=model_config.get("model_name_or_path"),
-            max_seq_length=max_seq_length,
-            dtype=torch.float16 if hardware_config.get("fp16", True) else None,
-            load_in_4bit=False,  # Don't re-quantize, model is already 4-bit
-            use_existing_bnb_quantization=True  # Use the existing quantization
-        )
+        # Initialize model with our safe loading function
+        logger.info("Loading pre-quantized model safely")
+        dtype = torch.float16 if hardware_config.get("fp16", True) else None
+        model, tokenizer = load_model_safely(model_name, max_seq_length, dtype)
+        
+        # Apply LoRA
+        logger.info("Applying LoRA to model")
         model = FastLanguageModel.get_peft_model(
             model,
             peft_config=peft_config,