#!/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