Fix: Remove unsupported attn_implementation parameter

run_cloud_training.py

@@ -0,0 +1,503 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Fine-tuning script for DeepSeek-R1-Distill-Qwen-14B-bnb-4bit using unsloth
+RESEARCH TRAINING PHASE ONLY - No output generation
+WORKS WITH PRE-TOKENIZED DATASET - No re-tokenization
+"""
+
+import os
+import json
+import logging
+import argparse
+import numpy as np
+from dotenv import load_dotenv
+import torch
+from datasets import load_dataset
+import transformers
+from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForCausalLM, AutoConfig
+from transformers.data.data_collator import DataCollatorMixin
+from peft import LoraConfig
+from unsloth import FastLanguageModel
+
+# Disable flash attention globally
+os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1"
+
+# Check if tensorboard is available
+try:
+    import tensorboard
+    TENSORBOARD_AVAILABLE = True
+except ImportError:
+    TENSORBOARD_AVAILABLE = False
+    print("Tensorboard not available. Will skip tensorboard logging.")
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler("training.log")
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# Default dataset path - use the correct path with username
+DEFAULT_DATASET = "George-API/phi4-cognitive-dataset"
+
+def load_config(config_path):
+    """Load the transformers config from JSON file"""
+    logger.info(f"Loading config from {config_path}")
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    return config
+
+def load_and_prepare_dataset(dataset_name, config):
+    """
+    Load and prepare the dataset for fine-tuning.
+    Sort entries by prompt_number as required.
+    NO TOKENIZATION - DATASET IS ALREADY TOKENIZED
+    """
+    # Use the default dataset path if no specific path is provided
+    if dataset_name == "phi4-cognitive-dataset":
+        dataset_name = DEFAULT_DATASET
+        
+    logger.info(f"Loading dataset: {dataset_name}")
+    
+    try:
+        # Load dataset
+        dataset = load_dataset(dataset_name)
+        
+        # Extract the split we want to use (usually 'train')
+        if 'train' in dataset:
+            dataset = dataset['train']
+        
+        # Get the dataset config
+        dataset_config = config.get("dataset_config", {})
+        sort_field = dataset_config.get("sort_by_field", "prompt_number")
+        sort_direction = dataset_config.get("sort_direction", "ascending")
+        
+        # Sort the dataset by prompt_number
+        logger.info(f"Sorting dataset by {sort_field} in {sort_direction} order")
+        if sort_direction == "ascending":
+            dataset = dataset.sort(sort_field)
+        else:
+            dataset = dataset.sort(sort_field, reverse=True)
+        
+        # Add shuffle with fixed seed if specified
+        if "shuffle_seed" in dataset_config:
+            shuffle_seed = dataset_config.get("shuffle_seed")
+            logger.info(f"Shuffling dataset with seed {shuffle_seed}")
+            dataset = dataset.shuffle(seed=shuffle_seed)
+        
+        # Print dataset structure for debugging
+        logger.info(f"Dataset loaded with {len(dataset)} entries")
+        logger.info(f"Dataset columns: {dataset.column_names}")
+        
+        # Print a sample entry to understand structure
+        if len(dataset) > 0:
+            sample = dataset[0]
+            logger.info(f"Sample entry structure: {list(sample.keys())}")
+            if 'conversations' in sample:
+                logger.info(f"Sample conversations structure: {sample['conversations'][:1]}")
+                
+        return dataset
+    
+    except Exception as e:
+        logger.error(f"Error loading dataset: {str(e)}")
+        logger.info("Available datasets in the Hub:")
+        # Print a more helpful error message
+        print(f"Failed to load dataset: {dataset_name}")
+        print(f"Make sure the dataset exists and is accessible.")
+        print(f"If it's a private dataset, ensure your HF_TOKEN has access to it.")
+        raise
+
+def tokenize_string(text, tokenizer):
+    """Tokenize a string using the provided tokenizer"""
+    if not text:
+        return []
+    
+    # Tokenize the text
+    tokens = tokenizer.encode(text, add_special_tokens=False)
+    return tokens
+
+# Data collator for pre-tokenized dataset
+class PreTokenizedCollator(DataCollatorMixin):
+    """
+    Data collator for pre-tokenized datasets.
+    Expects input_ids and labels already tokenized.
+    """
+    def __init__(self, pad_token_id=0, tokenizer=None):
+        self.pad_token_id = pad_token_id
+        self.tokenizer = tokenizer  # Keep a reference to the tokenizer for string conversion
+        
+    def __call__(self, features):
+        # Print a sample feature to understand structure
+        if len(features) > 0:
+            logger.info(f"Sample feature keys: {list(features[0].keys())}")
+            
+        # Extract input_ids from conversations if needed
+        processed_features = []
+        for feature in features:
+            # If input_ids is not directly available, try to extract from conversations
+            if 'input_ids' not in feature and 'conversations' in feature:
+                # Extract from conversations based on your dataset structure
+                conversations = feature['conversations']
+                
+                # Debug the conversations structure
+                logger.info(f"Conversations type: {type(conversations)}")
+                if isinstance(conversations, list) and len(conversations) > 0:
+                    logger.info(f"First conversation type: {type(conversations[0])}")
+                    logger.info(f"First conversation: {conversations[0]}")
+                
+                # Try different approaches to extract input_ids
+                if isinstance(conversations, list) and len(conversations) > 0:
+                    # Case 1: If conversations is a list of dicts with 'content' field
+                    if isinstance(conversations[0], dict) and 'content' in conversations[0]:
+                        content = conversations[0]['content']
+                        logger.info(f"Found content field: {type(content)}")
+                        
+                        # If content is a string, tokenize it
+                        if isinstance(content, str) and self.tokenizer:
+                            logger.info(f"Tokenizing string content: {content[:50]}...")
+                            feature['input_ids'] = self.tokenizer.encode(content, add_special_tokens=False)
+                        # If content is already a list of integers, use it directly
+                        elif isinstance(content, list) and all(isinstance(x, int) for x in content):
+                            feature['input_ids'] = content
+                        # If content is already tokenized in some other format
+                        else:
+                            logger.warning(f"Unexpected content format: {type(content)}")
+                            
+                    # Case 2: If conversations is a list of dicts with 'input_ids' field
+                    elif isinstance(conversations[0], dict) and 'input_ids' in conversations[0]:
+                        feature['input_ids'] = conversations[0]['input_ids']
+                    
+                    # Case 3: If conversations itself contains the input_ids
+                    elif all(isinstance(x, int) for x in conversations):
+                        feature['input_ids'] = conversations
+                    
+                    # Case 4: If conversations is a list of strings
+                    elif all(isinstance(x, str) for x in conversations) and self.tokenizer:
+                        # Join all strings and tokenize
+                        full_text = " ".join(conversations)
+                        feature['input_ids'] = self.tokenizer.encode(full_text, add_special_tokens=False)
+            
+            # Ensure input_ids is a list of integers
+            if 'input_ids' in feature:
+                # If input_ids is a string, tokenize it
+                if isinstance(feature['input_ids'], str) and self.tokenizer:
+                    logger.info(f"Converting string input_ids to tokens: {feature['input_ids'][:50]}...")
+                    feature['input_ids'] = self.tokenizer.encode(feature['input_ids'], add_special_tokens=False)
+                # If input_ids is not a list, convert it
+                elif not isinstance(feature['input_ids'], list):
+                    try:
+                        feature['input_ids'] = list(feature['input_ids'])
+                    except:
+                        logger.error(f"Could not convert input_ids to list: {type(feature['input_ids'])}")
+            
+            processed_features.append(feature)
+            
+        # If we still don't have input_ids, log an error
+        if len(processed_features) > 0 and 'input_ids' not in processed_features[0]:
+            logger.error(f"Could not find input_ids in features. Available keys: {list(processed_features[0].keys())}")
+            if 'conversations' in processed_features[0]:
+                logger.error(f"Conversations structure: {processed_features[0]['conversations'][:1]}")
+            raise ValueError("Could not find input_ids in dataset. Please check dataset structure.")
+            
+        # Determine max length in this batch
+        batch_max_len = max(len(x["input_ids"]) for x in processed_features)
+        
+        # Initialize batch tensors
+        batch = {
+            "input_ids": torch.ones((len(processed_features), batch_max_len), dtype=torch.long) * self.pad_token_id,
+            "attention_mask": torch.zeros((len(processed_features), batch_max_len), dtype=torch.long),
+            "labels": torch.ones((len(processed_features), batch_max_len), dtype=torch.long) * -100  # -100 is ignored in loss
+        }
+        
+        # Fill batch tensors
+        for i, feature in enumerate(processed_features):
+            input_ids = feature["input_ids"]
+            seq_len = len(input_ids)
+            
+            # Convert to tensor if it's a list
+            if isinstance(input_ids, list):
+                input_ids = torch.tensor(input_ids, dtype=torch.long)
+                
+            # Copy data to batch tensors
+            batch["input_ids"][i, :seq_len] = input_ids
+            batch["attention_mask"][i, :seq_len] = 1
+            
+            # If there are labels, use them, otherwise use input_ids
+            if "labels" in feature:
+                labels = feature["labels"]
+                if isinstance(labels, list):
+                    labels = torch.tensor(labels, dtype=torch.long)
+                batch["labels"][i, :len(labels)] = labels
+            else:
+                batch["labels"][i, :seq_len] = input_ids
+        
+        return batch
+
+def create_training_marker(output_dir):
+    """Create a marker file to indicate training is active"""
+    # Create in current directory for app.py to find
+    with open("TRAINING_ACTIVE", "w") as f:
+        f.write(f"Training active in {output_dir}")
+    
+    # Also create in output directory
+    os.makedirs(output_dir, exist_ok=True)
+    with open(os.path.join(output_dir, "RESEARCH_TRAINING_ONLY"), "w") as f:
+        f.write("This model is for research training only. No interactive outputs.")
+
+def remove_training_marker():
+    """Remove the training marker file"""
+    if os.path.exists("TRAINING_ACTIVE"):
+        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
+            logger.info("Loading model with flash attention DISABLED")
+            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
+                use_flash_attention=False  # Explicitly disable flash attention
+            )
+            logger.info("Model loaded successfully with unsloth with 4-bit quantization and flash attention disabled")
+            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,
+                    use_flash_attention=False,  # Explicitly disable flash attention
+                )
+                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...")
+        
+        # Disable flash attention in transformers config
+        config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+        if hasattr(config, "use_flash_attention"):
+            config.use_flash_attention = False
+            logger.info("Disabled flash attention in model config")
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            config=config,
+            device_map="auto",
+            torch_dtype=dtype or torch.float16,
+            load_in_4bit=True
+        )
+        logger.info("Model loaded successfully with standard HF loading and flash attention disabled")
+        return model, tokenizer
+
+def train(config_path, dataset_name, output_dir):
+    """Main training function - RESEARCH TRAINING PHASE ONLY"""
+    # Load environment variables
+    load_dotenv()
+    config = load_config(config_path)
+    
+    # Extract configs
+    model_config = config.get("model_config", {})
+    training_config = config.get("training_config", {})
+    hardware_config = config.get("hardware_config", {})
+    lora_config = config.get("lora_config", {})
+    dataset_config = config.get("dataset_config", {})
+    
+    # Override flash attention setting to disable it
+    hardware_config["use_flash_attention"] = False
+    logger.info("Flash attention has been DISABLED due to GPU compatibility issues")
+    
+    # Verify this is training phase only
+    training_phase_only = dataset_config.get("training_phase_only", True)
+    if not training_phase_only:
+        logger.warning("This script is meant for research training phase only")
+        logger.warning("Setting training_phase_only=True")
+    
+    # Verify dataset is pre-tokenized
+    logger.info("IMPORTANT: Using pre-tokenized dataset - No tokenization will be performed")
+    
+    # Set the output directory
+    output_dir = output_dir or training_config.get("output_dir", "fine_tuned_model")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Create training marker
+    create_training_marker(output_dir)
+    
+    try:
+        # Print configuration summary
+        logger.info("RESEARCH TRAINING PHASE ACTIVE - No output generation")
+        logger.info("Configuration Summary:")
+        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")
+        
+        # Load and prepare the dataset
+        dataset = load_and_prepare_dataset(dataset_name, config)
+        
+        # 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_name,
+            trust_remote_code=True
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+        
+        # Initialize model with unsloth
+        logger.info("Initializing model with unsloth (preserving 4-bit quantization)")
+        max_seq_length = training_config.get("max_seq_length", 2048)
+        
+        # Create LoRA config directly
+        logger.info("Creating LoRA configuration")
+        lora_config_obj = LoraConfig(
+            r=lora_config.get("r", 16),
+            lora_alpha=lora_config.get("lora_alpha", 32),
+            lora_dropout=lora_config.get("lora_dropout", 0.05),
+            bias=lora_config.get("bias", "none"),
+            target_modules=lora_config.get("target_modules", ["q_proj", "k_proj", "v_proj", "o_proj"])
+        )
+        
+        # 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)
+        
+        # Try different approaches to apply LoRA
+        logger.info("Applying LoRA to model")
+        
+        # Skip unsloth's method and go directly to PEFT
+        logger.info("Using standard PEFT method to apply LoRA")
+        from peft import get_peft_model
+        model = get_peft_model(model, lora_config_obj)
+        logger.info("Successfully applied LoRA with standard PEFT")
+            
+        # No need to format the dataset - it's already pre-tokenized
+        logger.info("Using pre-tokenized dataset - skipping tokenization step")
+        training_dataset = dataset
+        
+        # Configure reporting backends with fallbacks
+        reports = []
+        if TENSORBOARD_AVAILABLE:
+            reports.append("tensorboard")
+            logger.info("Tensorboard available and enabled for reporting")
+        else:
+            logger.warning("Tensorboard not available - metrics won't be logged to tensorboard")
+            
+        if os.getenv("WANDB_API_KEY"):
+            reports.append("wandb")
+            logger.info("Wandb API key found, enabling wandb reporting")
+        
+        # Default to "none" if no reporting backends are available
+        if not reports:
+            reports = ["none"]
+            logger.warning("No reporting backends available - training metrics won't be logged")
+        
+        # Set up training arguments with flash attention disabled
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=training_config.get("num_train_epochs", 3),
+            per_device_train_batch_size=training_config.get("per_device_train_batch_size", 2),
+            gradient_accumulation_steps=training_config.get("gradient_accumulation_steps", 4),
+            learning_rate=training_config.get("learning_rate", 2e-5),
+            lr_scheduler_type=training_config.get("lr_scheduler_type", "cosine"),
+            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"),
+            logging_steps=training_config.get("logging_steps", 10),
+            save_steps=training_config.get("save_steps", 200),
+            save_total_limit=training_config.get("save_total_limit", 3),
+            fp16=hardware_config.get("fp16", True),
+            bf16=hardware_config.get("bf16", False),
+            max_grad_norm=training_config.get("max_grad_norm", 0.3),
+            report_to=reports,
+            logging_first_step=training_config.get("logging_first_step", True),
+            disable_tqdm=training_config.get("disable_tqdm", False),
+            # Important: Don't remove columns that don't match model's forward method
+            remove_unused_columns=False
+        )
+        
+        # Create trainer with pre-tokenized collator
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=training_dataset,
+            data_collator=PreTokenizedCollator(pad_token_id=tokenizer.pad_token_id, tokenizer=tokenizer),
+        )
+        
+        # Start training
+        logger.info("Starting training - RESEARCH PHASE ONLY")
+        trainer.train()
+        
+        # Save the model
+        logger.info(f"Saving model to {output_dir}")
+        trainer.save_model(output_dir)
+        
+        # Save LoRA adapter separately for easier deployment
+        lora_output_dir = os.path.join(output_dir, "lora_adapter")
+        model.save_pretrained(lora_output_dir)
+        logger.info(f"Saved LoRA adapter to {lora_output_dir}")
+        
+        # Save tokenizer for completeness
+        tokenizer_output_dir = os.path.join(output_dir, "tokenizer")
+        tokenizer.save_pretrained(tokenizer_output_dir)
+        logger.info(f"Saved tokenizer to {tokenizer_output_dir}")
+        
+        # Copy config file for reference
+        with open(os.path.join(output_dir, "training_config.json"), "w") as f:
+            json.dump(config, f, indent=2)
+        
+        logger.info("Training complete - RESEARCH PHASE ONLY")
+        return output_dir
+    
+    finally:
+        # Always remove the training marker when done
+        remove_training_marker()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Fine-tune Unsloth/DeepSeek-R1-Distill-Qwen-14B-4bit model (RESEARCH ONLY)")
+    parser.add_argument("--config", type=str, default="transformers_config.json", 
+                        help="Path to the transformers config JSON file")
+    parser.add_argument("--dataset", type=str, default="phi4-cognitive-dataset", 
+                        help="Dataset name or path")
+    parser.add_argument("--output_dir", type=str, default=None, 
+                        help="Output directory for the fine-tuned model")
+    
+    args = parser.parse_args()
+    
+    # Run training - Research phase only
+    try:
+        output_path = train(args.config, args.dataset, args.output_dir)
+        print(f"Research training completed. Model saved to: {output_path}")
+    except Exception as e:
+        logger.error(f"Training failed: {str(e)}")
+        remove_training_marker()  # Clean up marker if training fails
+        raise