#!/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 import sys 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 all attention optimizations that might cause issues os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1" os.environ["CUDA_LAUNCH_BLOCKING"] = "1" os.environ["XFORMERS_DISABLED"] = "1" # Completely disable xformers by removing it from sys.modules if it's loaded if 'xformers' in sys.modules: del sys.modules['xformers'] if 'xformers.ops' in sys.modules: del sys.modules['xformers.ops'] # Patch transformers to prevent xformers import def prevent_xformers_import(name, *args, **kwargs): if 'xformers' in name: raise ImportError(f"Import of {name} prevented") return original_import(name, *args, **kwargs) original_import = __import__ __builtins__['__import__'] = prevent_xformers_import # Configure PyTorch memory allocator for better memory management os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # Configure logging first 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__) # Make sure torch is installed and available before proceeding try: logger.info("Importing torch...") import torch logger.info(f"PyTorch version: {torch.__version__}") logger.info(f"CUDA available: {torch.cuda.is_available()}") if torch.cuda.is_available(): logger.info(f"CUDA version: {torch.version.cuda}") logger.info(f"GPU: {torch.cuda.get_device_name(0)}") except ImportError: logger.error("PyTorch not found. Installing torch first...") try: import subprocess import sys subprocess.check_call([sys.executable, "-m", "pip", "install", "torch"]) logger.info("PyTorch installed successfully. Importing...") import torch logger.info(f"PyTorch version: {torch.__version__}") except Exception as e: logger.error(f"Failed to install PyTorch: {e}") logger.error("Cannot proceed without PyTorch. Exiting.") raise # Now try to install flash-attention (for systems that support it) try: import subprocess import sys # Make sure torch is installed before attempting flash-attn try: logger.info("Ensuring PyTorch is installed before flash-attention...") subprocess.check_call([sys.executable, "-m", "pip", "install", "torch", "--quiet"]) logger.info("PyTorch installation verified") except Exception as torch_error: logger.warning(f"PyTorch installation check failed: {torch_error}") logger.info("Will continue with flash-attention installation anyway") logger.info("Attempting to install flash-attention...") # Try multiple installation approaches for flash-attention try: # First try with pip install logger.info("Trying standard pip install for flash-attn") subprocess.check_call([sys.executable, "-m", "pip", "install", "flash-attn"]) except Exception as pip_error: logger.warning(f"Standard installation failed: {pip_error}") logger.info("Trying alternative installation approach...") # Try the PIP_EXTRA_INDEX_URL approach env = os.environ.copy() if "PIP_EXTRA_INDEX_URL" not in env: env["PIP_EXTRA_INDEX_URL"] = "https://download.pytorch.org/whl/cu118" subprocess.check_call( [sys.executable, "-m", "pip", "install", "flash-attn"], env=env ) logger.info("Successfully installed flash-attention") except Exception as e: logger.warning(f"Failed to install flash-attention: {e}") logger.info("Continuing without flash-attention") # Check if flash attention was successfully installed flash_attention_available = False try: import flash_attn flash_attention_available = True logger.info(f"Flash Attention will be used (version: {flash_attn.__version__})") # We'll handle flash attention configuration during model loading except ImportError: logger.info("Flash Attention not available, will use standard attention mechanism") # Check if tensorboard is available try: import tensorboard TENSORBOARD_AVAILABLE = True except ImportError: TENSORBOARD_AVAILABLE = False print("Tensorboard not available. Will skip tensorboard logging.") # 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. Handles both pre-tokenized and string content. """ # 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") # Always sort in ascending order by prompt_number logger.info(f"Sorting dataset by {sort_field} in ascending order") dataset = dataset.sort(sort_field) # Verify sorting if len(dataset) > 1: first_prompt = dataset[0].get(sort_field, None) last_prompt = dataset[-1].get(sort_field, None) logger.info(f"Dataset sorted: first {sort_field}={first_prompt}, last {sort_field}={last_prompt}") # Additional verification of a few samples sample_indices = [0, len(dataset)//2, len(dataset)-1] sample_prompts = [dataset[i].get(sort_field, None) for i in sample_indices] logger.info(f"Sample prompt numbers: {sample_prompts}") # Verify order is ascending if not all(sample_prompts[i] <= sample_prompts[i+1] for i in range(len(sample_prompts)-1)): logger.warning("Dataset may not be properly sorted! Please check the ordering.") # 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())}") # Check if dataset is pre-tokenized or contains string content is_pre_tokenized = False if 'input_ids' in sample and isinstance(sample['input_ids'], list) and all(isinstance(x, int) for x in sample['input_ids']): logger.info("Dataset appears to be pre-tokenized with input_ids field") is_pre_tokenized = True elif 'conversations' in sample: logger.info(f"Sample conversations structure: {sample['conversations'][:1]}") # Check if conversations contain pre-tokenized data if isinstance(sample['conversations'], list) and len(sample['conversations']) > 0: conv = sample['conversations'][0] if isinstance(conv, dict) and 'input_ids' in conv and isinstance(conv['input_ids'], list): logger.info("Dataset appears to be pre-tokenized in conversations.input_ids") is_pre_tokenized = True elif isinstance(conv, dict) and 'content' in conv: content = conv['content'] if isinstance(content, list) and all(isinstance(x, int) for x in content): logger.info("Dataset appears to be pre-tokenized in conversations.content") is_pre_tokenized = True else: logger.info("Dataset appears to contain string content that will need tokenization") if is_pre_tokenized: logger.info("Using pre-tokenized dataset - tokenizer will only be used as fallback") else: logger.info("Dataset contains string content - tokenizer will be used") 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 that can handle both pre-tokenized datasets and string content. Will tokenize strings if necessary, but logs warnings. """ 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 fallback tokenization 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 directly available, use it without tokenization if 'input_ids' in feature and isinstance(feature['input_ids'], list): # Already tokenized, no processing needed processed_features.append(feature) continue # 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 (only for first batch) if len(processed_features) == 0: logger.info(f"Conversations type: {type(conversations)}") if isinstance(conversations, list) and len(conversations) > 0: logger.info(f"First conversation type: {type(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 'input_ids' field (pre-tokenized) if isinstance(conversations[0], dict) and 'input_ids' in conversations[0]: feature['input_ids'] = conversations[0]['input_ids'] # Case 2: If conversations itself contains the input_ids (pre-tokenized) elif all(isinstance(x, int) for x in conversations): feature['input_ids'] = conversations # Case 3: If conversations is a list of dicts with 'content' field elif isinstance(conversations[0], dict) and 'content' in conversations[0]: content = conversations[0]['content'] # If content is already a list of integers, use it directly if isinstance(content, list) and all(isinstance(x, int) for x in content): feature['input_ids'] = content # If content is a string, tokenize it with a warning elif isinstance(content, str) and self.tokenizer: logger.warning("Found string content in dataset. Tokenizing as fallback.") feature['input_ids'] = self.tokenizer.encode(content, add_special_tokens=False) else: logger.warning(f"Unexpected content format: {type(content)}") continue # 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 logger.warning("Found string conversations in dataset. Tokenizing as fallback.") 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.warning("Found string input_ids in dataset. Tokenizing as fallback.") 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'])}") continue else: logger.warning("No input_ids found in this example. Skipping.") continue processed_features.append(feature) # If we still don't have input_ids, log an error if len(processed_features) == 0: logger.error("No valid examples found in batch. Check dataset format.") raise ValueError("No valid examples found. Please check dataset structure.") if '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. """ global flash_attention_available # Force disable flash attention and xformers flash_attention_available = False os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1" os.environ["XFORMERS_DISABLED"] = "1" # Patch transformers attention implementation try: # Try to patch transformers attention implementation to avoid xformers import transformers.models.llama.modeling_llama as llama_modeling # Store original attention implementation if not hasattr(llama_modeling, '_original_forward'): # Only patch if not already patched logger.info("Patching LLaMA attention implementation to avoid xformers") # Store original implementation if hasattr(llama_modeling.LlamaAttention, 'forward'): llama_modeling._original_forward = llama_modeling.LlamaAttention.forward # Define a new forward method that doesn't use xformers def safe_attention_forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, output_attentions=False, use_cache=False): logger.info("Using safe attention implementation (no xformers)") # Force use_flash_attention to False self._attn_implementation = "eager" if hasattr(self, 'use_flash_attention'): self.use_flash_attention = False if hasattr(self, 'use_flash_attention_2'): self.use_flash_attention_2 = False # Call original implementation with flash attention disabled return llama_modeling._original_forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache) # Replace the forward method llama_modeling.LlamaAttention.forward = safe_attention_forward logger.info("Successfully patched LLaMA attention implementation") except Exception as e: logger.warning(f"Failed to patch attention implementation: {e}") logger.info("Will try to proceed with standard loading") try: logger.info(f"Attempting to load model with unsloth optimizations: {model_name}") # Create BitsAndBytesConfig for 4-bit quantization from transformers import BitsAndBytesConfig bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True ) # First try loading with unsloth but without flash attention try: logger.info("Loading model with unsloth optimizations") # Don't pass any flash attention parameters to unsloth model, tokenizer = FastLanguageModel.from_pretrained( model_name=model_name, max_seq_length=max_seq_length, dtype=dtype, quantization_config=bnb_config, attn_implementation="eager", # Force eager attention use_flash_attention=False, # Explicitly disable flash attention use_xformers_attention=False # Explicitly disable xformers ) logger.info("Model loaded successfully with unsloth") # Explicitly disable flash attention in model config if hasattr(model, 'config'): if hasattr(model.config, 'attn_implementation'): model.config.attn_implementation = "eager" if hasattr(model.config, 'use_flash_attention'): model.config.use_flash_attention = False if hasattr(model.config, 'use_flash_attention_2'): model.config.use_flash_attention_2 = False if hasattr(model.config, 'use_xformers_attention'): model.config.use_xformers_attention = False return model, tokenizer except Exception as e: logger.warning(f"Unsloth loading failed: {e}") logger.info("Falling back to standard Hugging Face loading...") # We'll try with HF loading attn_params = { "attn_implementation": "eager", # Always use eager "use_flash_attention": False, # Explicitly disable flash attention "use_xformers_attention": False # Explicitly disable xformers } # Approach 1: Using attn_implementation parameter (newer method) try: logger.info(f"Trying HF loading with attention parameters: {attn_params}") config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) # Disable flash attention in config if hasattr(config, 'attn_implementation'): config.attn_implementation = "eager" if hasattr(config, 'use_flash_attention'): config.use_flash_attention = False if hasattr(config, 'use_flash_attention_2'): config.use_flash_attention_2 = False if hasattr(config, 'use_xformers_attention'): config.use_xformers_attention = False tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) # The proper way to set attention implementation in newer transformers model = AutoModelForCausalLM.from_pretrained( model_name, config=config, device_map="auto", torch_dtype=dtype or torch.float16, quantization_config=bnb_config, trust_remote_code=True, **attn_params ) logger.info(f"Model loaded successfully with HF using attention parameters: {attn_params}") return model, tokenizer except Exception as e: logger.warning(f"HF loading with attn_implementation failed: {e}") logger.info("Trying fallback method...") # Approach 2: Complete fallback with minimal parameters config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) # Disable flash attention in config if hasattr(config, 'attn_implementation'): config.attn_implementation = "eager" if hasattr(config, 'use_flash_attention'): config.use_flash_attention = False if hasattr(config, 'use_flash_attention_2'): config.use_flash_attention_2 = False if hasattr(config, 'use_xformers_attention'): config.use_xformers_attention = False tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) # Most basic loading without any attention parameters model = AutoModelForCausalLM.from_pretrained( model_name, config=config, device_map="auto", torch_dtype=dtype or torch.float16, quantization_config=bnb_config, trust_remote_code=True, attn_implementation="eager", use_flash_attention=False, use_xformers_attention=False ) logger.info("Model loaded successfully with basic HF loading") return model, tokenizer except Exception as e: logger.error(f"All model loading attempts failed: {e}") raise 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", {}) # Force disable flash attention and xformers os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1" os.environ["XFORMERS_DISABLED"] = "1" os.environ["CUDA_LAUNCH_BLOCKING"] = "1" # Monkey patch torch.nn.functional to disable memory_efficient_attention try: import torch.nn.functional as F if hasattr(F, 'scaled_dot_product_attention'): logger.info("Monkey patching torch.nn.functional.scaled_dot_product_attention") original_sdpa = F.scaled_dot_product_attention def safe_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None): # Force disable memory efficient attention logger.info("Using safe scaled_dot_product_attention (no xformers)") return original_sdpa(query, key, value, attn_mask, dropout_p, is_causal, scale) F.scaled_dot_product_attention = safe_sdpa except Exception as e: logger.warning(f"Failed to patch scaled_dot_product_attention: {e}") # Completely remove xformers from sys.modules if it's loaded for module_name in list(sys.modules.keys()): if 'xformers' in module_name: logger.info(f"Removing {module_name} from sys.modules") del sys.modules[module_name] # Update flash attention setting to always use eager global flash_attention_available flash_attention_available = False logger.info("Flash Attention has been DISABLED globally") # Update hardware config to ensure eager attention hardware_config["attn_implementation"] = "eager" hardware_config["use_flash_attention"] = False hardware_config["use_xformers_attention"] = False # 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 # Force eager attention implementation os.environ["TRANSFORMERS_NO_FLASH_ATTENTION"] = "1" logger.info("Flash attention has been DISABLED globally via environment variable") # Update hardware config to ensure eager attention hardware_config["attn_implementation"] = "eager" hardware_config["use_flash_attention"] = False model, tokenizer = load_model_safely(model_name, max_seq_length, dtype) # Disable generation capabilities for research training logger.info("Disabling generation capabilities - Research training only") model.config.is_decoder = False model.config.task_specific_params = None # 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 dataset with flexible tokenization handling") logger.info("Will use pre-tokenized data if available, or tokenize strings as fallback") 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 correct parameters # Extract only the valid parameters from hardware_config training_args_dict = { "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), "remove_unused_columns": False, "seed": 42 } # Create TrainingArguments with validated parameters training_args = TrainingArguments(**training_args_dict) # 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