app.py

import os
import torch
import glob
import gc
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer,
    BitsAndBytesConfig, 
    TrainingArguments, 
    Trainer,
    DataCollatorForLanguageModeling,
    AutoTokenizer,
    LlamaConfig,
    AutoConfig
)
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from datasets import Dataset
from huggingface_hub import snapshot_download
from tqdm import tqdm
import gradio as gr
import math
from accelerate import Accelerator
import subprocess
import sys
import json
import shutil

# --- Configuration ---
YOUR_HF_USERNAME = "Twelve2five"
MODEL_REPO_NAME = "llama-3-8b-rvq-resized"
DATASET_REPO_NAME = "podcast-dialogue-rvq-pairs-3items"

hf_model_repo_id = f"{YOUR_HF_USERNAME}/{MODEL_REPO_NAME}"
hf_dataset_repo_id = f"{YOUR_HF_USERNAME}/{DATASET_REPO_NAME}"

# Output directories
OUTPUT_TRAINING_DIR = "./llama3-8b-rvq-qlora-finetuned-run"
LOGGING_DIR = "./llama3-8b-rvq-qlora-logs-run"
local_download_path = "./downloaded_dataset_files"

# Training parameters
NUM_EPOCHS = 1
BATCH_SIZE_PER_DEVICE = 1
GRAD_ACCUMULATION_STEPS = 64
LEARNING_RATE = 1e-4
WEIGHT_DECAY = 0.01
WARMUP_RATIO = 0.03
LR_SCHEDULER = "cosine"
OPTIMIZER = "paged_adamw_8bit"
MAX_SEQ_LENGTH = 256
MICRO_BATCH_SIZE = 1

# Multi-GPU configuration
accelerator = Accelerator()

# Configure environment for multi-GPU
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"

# Print GPU information
print(f"Available GPUs: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
    print(f"GPU {i}: {torch.cuda.get_device_name(i)} with {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")

def seq2seq_causal_collator(features):
    """
    Collator that concatenates context (input_ids) and target (labels)
    for Causal LM sequence-to-sequence training.
    Masks the loss for the context part of the sequence.
    Pads sequences to the maximum length in the batch.
    """
    batch = {}
    concatenated_input_ids = []
    concatenated_labels = []
    max_len = 0

    # --- First pass: Concatenate, create masked labels, find max length ---
    for feature in features:
        # Dataset transform should provide tensors here
        input_ids = feature['input_ids']
        labels = feature['labels']

        # Ensure tensors are 1D (handle potential extra dims if any)
        if input_ids.dim() > 1: input_ids = input_ids.squeeze()
        if labels.dim() > 1: labels = labels.squeeze()

        context_len = input_ids.shape[0]
        target_len = labels.shape[0]

        # Concatenate context and target for input
        combined_ids = torch.cat([input_ids, labels], dim=0)
        concatenated_input_ids.append(combined_ids)

        # Create labels: -100 for context, actual labels for target
        masked_labels = torch.cat([
            torch.full((context_len,), -100, dtype=torch.long, device=input_ids.device),
            labels
        ], dim=0)
        concatenated_labels.append(masked_labels)

        # Track max length for padding
        if combined_ids.shape[0] > max_len:
            max_len = combined_ids.shape[0]

    # --- Second pass: Pad to max length ---
    padded_input_ids = []
    padded_labels = []
    input_pad_token_id = 0
    label_pad_token_id = -100

    for i in range(len(features)):
        ids = concatenated_input_ids[i]
        lbls = concatenated_labels[i]

        padding_len = max_len - ids.shape[0]

        # Pad on the right side
        padded_input_ids.append(torch.nn.functional.pad(
            ids, (0, padding_len), value=input_pad_token_id
        ))
        padded_labels.append(torch.nn.functional.pad(
            lbls, (0, padding_len), value=label_pad_token_id
        ))

    # --- Stack and create final batch ---
    batch['input_ids'] = torch.stack(padded_input_ids)
    batch['labels'] = torch.stack(padded_labels)

    # Create attention mask (1 for real tokens, 0 for padding)
    batch['attention_mask'] = batch['input_ids'].ne(input_pad_token_id).long()

    return batch

def prepare_for_dataset(batch):
    output = {'input_ids': [], 'labels': []}
    for item in batch:
        output['input_ids'].append(item['input_ids'].cpu().tolist())
        output['labels'].append(item['labels'].cpu().tolist())
    return output

def load_model():
    print(f"Loading base model architecture from: {hf_model_repo_id}")
    
    # Get information about GPU with most free memory
    gpu_id = 0  # Default to first GPU
    max_free_memory = 0
    
    for i in range(torch.cuda.device_count()):
        free_memory = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_allocated(i)
        if free_memory > max_free_memory:
            max_free_memory = free_memory
            gpu_id = i
    
    print(f"Loading model on GPU {gpu_id} with {max_free_memory / 1e9:.2f}GB free memory")
    
    # Configure quantization
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    
    # Load the model
    try:
        # First update transformers to make sure we have latest version
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers"])
        
        # Now try loading with explicit config class to avoid auto-detection issues
        from transformers import LlamaConfig
        
        # Load config first
        config = LlamaConfig.from_pretrained(
            hf_model_repo_id,
            trust_remote_code=True
        )
        
        # Then load model with explicit config
        model = AutoModelForCausalLM.from_pretrained(
            hf_model_repo_id,
            config=config,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True
        )
        log.append(f"Loaded model vocab size: {model.config.vocab_size}")
        log.append(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
    except Exception as e:
        error_msg = f"Error loading model from Hub: {e}"
        log.append(error_msg)
        # Try with a fallback method
        try:
            log.append("Attempting alternative loading method...")
            # Try loading without auto detection
            model = AutoModelForCausalLM.from_pretrained(
                hf_model_repo_id,
                quantization_config=bnb_config,
                device_map="auto",
                trust_remote_code=True,
                torch_dtype=torch.bfloat16,
                # Add these to help with the loading
                revision="main",
                low_cpu_mem_usage=True,
            )
            log.append("Alternative loading successful!")
            log.append(f"Loaded model vocab size: {model.config.vocab_size}")
        except Exception as e2:
            log.append(f"Alternative loading also failed: {e2}")
            return "\n".join(log)
    
    # Try to load the tokenizer from the model repository directly
    progress(0.3, desc="Loading tokenizer...")
    try:
        # First attempt: Try loading from local path
        tokenizer = AutoTokenizer.from_pretrained(
            local_model_path,
            padding_side="right",
            use_fast=True,
        )
        log.append("Tokenizer loaded from local files")
    except Exception as e:
        log.append(f"Could not load tokenizer from local files: {e}")
        
        # Second attempt: Try loading directly from HF repo
        try:
            log.append("Attempting to load tokenizer directly from Hugging Face...")
            tokenizer = AutoTokenizer.from_pretrained(
                hf_model_repo_id,
                padding_side="right",
                use_fast=True,
            )
            log.append("Tokenizer loaded from Hugging Face repository")
        except Exception as e2:
            # Third attempt: Try loading a compatible tokenizer
            log.append(f"Could not load tokenizer from repo: {e2}")
            log.append("Attempting to load a compatible LlamaTokenizer...")
            try:
                from transformers import LlamaTokenizer
                
                # Try Meta's standard Llama tokenizer
                tokenizer = LlamaTokenizer.from_pretrained(
                    "meta-llama/Llama-2-7b-hf",  # Standard Llama tokenizer
                    padding_side="right",
                    use_fast=False,  # Try the Python version
                )
                log.append("Loaded a compatible LlamaTokenizer as fallback")
            except Exception as e3:
                error_msg = f"Failed to load any compatible tokenizer: {e3}"
                log.append(error_msg)
                return "\n".join(log)

    # Set pad token if not already set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        log.append("Set pad_token to eos_token")
    
    print(f"Loaded tokenizer vocabulary size: {len(tokenizer)}")
    
    # Print information about input embeddings
    print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
    
    # Prepare model for k-bit training
    model = prepare_model_for_kbit_training(model)
    
    # Define LoRA configuration
    lora_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM
    )
    
    # Apply LoRA to model
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    return model, tokenizer  # Return both model and tokenizer

def load_dataset():
    # --- Download the dataset repository files ---
    try:
        os.makedirs(local_download_path, exist_ok=True)
        downloaded_repo_root = snapshot_download(
            repo_id=hf_dataset_repo_id,
            repo_type="dataset",
            local_dir=local_download_path,
            local_dir_use_symlinks=False
        )
        print(f"Dataset repository content downloaded to: {downloaded_repo_root}")
    except Exception as e:
        print(f"Error downloading dataset: {e}")
        return None
    
    # --- Load .pt files into a Hugging Face Dataset object ---
    pairs_dir = os.path.join(downloaded_repo_root, "final_rvq_pairs")
    all_pair_files = glob.glob(os.path.join(pairs_dir, "*_rvq_pairs.pt"))
    
    if not all_pair_files:
        all_pair_files = glob.glob(os.path.join(downloaded_repo_root, "*_rvq_pairs.pt"))
        if not all_pair_files:
            print("No RVQ pair files found!")
            return None
    
    print(f"Found {len(all_pair_files)} RVQ pair files.")
    
    # Load data from .pt files into memory
    all_data_pairs = []
    for file_path in tqdm(all_pair_files, desc="Loading pair files"):
        try:
            episode_pairs = torch.load(file_path, map_location='cpu')
            all_data_pairs.extend(episode_pairs)
        except Exception as e:
            print(f"Warning: Could not load file {file_path}: {e}")
    
    if not all_data_pairs:
        return None
    
    print(f"Loaded {len(all_data_pairs)} training pairs.")
    
    # Convert to Hugging Face Dataset
    chunk_size = 1000
    processed_data = {'input_ids': [], 'labels': []}
    for i in tqdm(range(0, len(all_data_pairs), chunk_size), desc="Preparing data"):
        batch = all_data_pairs[i:i + chunk_size]
        prepared_batch = prepare_for_dataset(batch)
        processed_data['input_ids'].extend(prepared_batch['input_ids'])
        processed_data['labels'].extend(prepared_batch['labels'])
    
    hf_dataset = Dataset.from_dict(processed_data)
    
    # Transform to get tensors back
    hf_dataset.set_transform(lambda batch: {
        'input_ids': [torch.tensor(ids, dtype=torch.long) for ids in batch['input_ids']],
        'labels': [torch.tensor(lbls, dtype=torch.long) for lbls in batch['labels']]
    })
    
    # Cleanup
    del all_data_pairs
    del processed_data
    gc.collect()
    
    return hf_dataset

# Memory cleaning function
def clean_memory():
    gc.collect()
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            with torch.cuda.device(f'cuda:{i}'):
                torch.cuda.empty_cache()
                torch.cuda.reset_peak_memory_stats()

def train_model(
    hf_username, 
    model_repo_name, 
    dataset_repo_name, 
    epochs=1, 
    batch_size=4,  # Increased for A100
    grad_accum_steps=4,
    learning_rate=2e-4,
    progress=gr.Progress()
):
    progress(0, desc="Setting up environment...")
    log = []
    
    # Clean up any existing model files to save space
    if os.path.exists("./model_files"):
        try:
            shutil.rmtree("./model_files")
        except Exception as e:
            log.append(f"Warning: Could not remove existing model files: {e}")
    
    if os.path.exists("./downloaded_dataset_files"):
        try:
            shutil.rmtree("./downloaded_dataset_files")
        except Exception as e:
            log.append(f"Warning: Could not remove existing dataset files: {e}")
            
    # Print GPU info - using imported torch, not a local variable
    if torch.cuda.is_available():
        log.append(f"Available GPUs: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_memory = torch.cuda.get_device_properties(i).total_memory / (1024**3)
            log.append(f"GPU {i}: {gpu_name} with {gpu_memory:.2f} GB")
    
    # Import required libraries
    try:
        from datasets import Dataset
        from huggingface_hub import snapshot_download
        # Don't import torch again, since it's already imported
        import transformers
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from transformers import BitsAndBytesConfig, TrainingArguments, Trainer
        from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
        
        log.append(f"Transformers version: {transformers.__version__}")
        log.append(f"PyTorch version: {torch.__version__}")
    except ImportError as e:
        log.append(f"Error importing libraries: {e}")
        return "\n".join(log)
    
    # --- Configuration ---
    progress(0.05, desc="Setting up configuration...")
    hf_model_repo_id = f"{hf_username}/{model_repo_name}"
    hf_dataset_repo_id = f"{hf_username}/{dataset_repo_name}"
    
    log.append(f"Model repo: {hf_model_repo_id}")
    log.append(f"Dataset repo: {hf_dataset_repo_id}")
    
    # Check if running on multiple GPUs
    n_gpus = torch.cuda.device_count()
    log.append(f"Number of GPUs available: {n_gpus}")
    
    # --- DeepSpeed Configuration ---
    # Create DeepSpeed config file
    progress(0.1, desc="Setting up DeepSpeed configuration...")
    
    # Create a simpler config since we have plenty of memory on A100
    ds_config = {
        "bf16": {
            "enabled": "auto"
        },
        "zero_optimization": {
            "stage": 1,  # Lower stage is fine for A100-80GB
            "contiguous_gradients": True,
            "overlap_comm": True
        },
        "gradient_accumulation_steps": grad_accum_steps,
        "gradient_clipping": 1.0,
        "train_batch_size": batch_size * grad_accum_steps * max(1, n_gpus)
    }
    
    ds_config_path = "ds_config.json"
    with open(ds_config_path, "w") as f:
        json.dump(ds_config, f, indent=4)
    
    log.append("DeepSpeed configuration created successfully")
    
    # --- Download and Load Model ---
    progress(0.15, desc="Downloading model...")
    
    try:
        # Download model files
        local_model_path = "./model_files"
        snapshot_download(
            repo_id=hf_model_repo_id,
            local_dir=local_model_path,
            use_auth_token=False,
            resume_download=True
        )
        log.append(f"Model files downloaded to {local_model_path}")
        
        # Check and fix the model config if needed
        config_path = os.path.join(local_model_path, "config.json")
        if os.path.exists(config_path):
            with open(config_path, 'r') as f:
                config_data = json.load(f)
            
            # Fix the rope_scaling configuration
            if 'rope_scaling' in config_data:
                if not isinstance(config_data['rope_scaling'], dict):
                    config_data['rope_scaling'] = {"type": "linear", "factor": 2.0}
                elif 'rope_type' in config_data['rope_scaling']:
                    # Convert complex rope_scaling to the simple format expected
                    rope_factor = config_data['rope_scaling'].get('factor', 2.0)
                    config_data['rope_scaling'] = {"type": "linear", "factor": rope_factor}
                
                # Write the updated config back
                with open(config_path, 'w') as f:
                    json.dump(config_data, f, indent=2)
                log.append("Updated model configuration for rope_scaling")
        
        # Create a bnb configuration for loading the model in 4-bit
        progress(0.25, desc="Loading model...")
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=False
        )
        
        # Load the model with fixed configuration
        model = AutoModelForCausalLM.from_pretrained(
            local_model_path,
            quantization_config=bnb_config,
            device_map="auto",
            use_cache=False,  # Needed for gradient checkpointing
            torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        )
        
        # Load the tokenizer
        tokenizer = AutoTokenizer.from_pretrained(
            local_model_path,
            padding_side="right",
            use_fast=True,
        )
        tokenizer.pad_token = tokenizer.eos_token
        
        # Find model's architecture type
        model_type = model.config.model_type
        log.append(f"Model architecture type: {model_type}")
        
        # PEFT Configuration (Smaller LoRA for faster iteration)
        model = prepare_model_for_kbit_training(model)
        log.append("Model prepared for k-bit training")
        
        lora_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=16,  # Keeping higher rank for A100
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]  # Fewer modules for faster training
        )
        peft_model = get_peft_model(model, lora_config)
        trainable_params = peft_model.print_trainable_parameters()
        log.append(f"LoRA applied to model")
        model_to_train = peft_model
    
    except Exception as e:
        error_msg = f"Error preparing model for training: {str(e)}"
        log.append(error_msg)
        return "\n".join(log)
    
    # --- Download and Process Dataset ---
    progress(0.4, desc="Downloading dataset...")
    
    try:
        dataset_path = "./downloaded_dataset_files"
        snapshot_download(
            repo_id=hf_dataset_repo_id,
            local_dir=dataset_path,
            use_auth_token=False,
            resume_download=True
        )
        log.append(f"Dataset repository content downloaded to: {dataset_path}")
        
        # Load dataset from PT files
        progress(0.5, desc="Processing dataset...")
        
        # Load RVQ pairs
        pair_files = glob.glob(f"{dataset_path}/*_rvq_pairs.pt")
        log.append(f"Found {len(pair_files)} RVQ pair files.")
        
        all_pairs = []
        for file in pair_files:
            pairs = torch.load(file)
            all_pairs.extend(pairs)
        
        log.append(f"Loaded a total of {len(all_pairs)} training pairs into memory.")
        
        # Process pairs into a format suitable for training
        all_texts = []
        for pair in all_pairs:
            # Create instruction format
            if isinstance(pair, dict):
                instruction = pair.get("instruction", "")
                input_text = pair.get("input", "")
                output = pair.get("output", "")
                
                # ALPACA format
                if instruction and input_text:
                    text = f"### Instruction: {instruction}\n### Input: {input_text}\n### Response: {output}"
                elif instruction:
                    text = f"### Instruction: {instruction}\n### Response: {output}"
                else:
                    text = output
            else:
                # Simple prompt-completion format
                if isinstance(pair, tuple) and len(pair) == 2:
                    prompt, completion = pair
                    text = f"{prompt}{completion}"
                else:
                    text = str(pair)
            
            all_texts.append({"text": text})
        
        # Create HF dataset
        train_dataset = Dataset.from_list(all_texts)
        
        # Function to tokenize the dataset
        def tokenize_function(examples):
            return tokenizer(
                examples["text"],
                padding=False,
                truncation=True,
                max_length=2048,
                return_tensors=None,
            )
        
        # Tokenize the dataset
        tokenized_dataset = train_dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=["text"],
            desc="Tokenizing dataset",
        )
        
        train_dataset = tokenized_dataset
        
        # Data collator
        from transformers import DataCollatorForLanguageModeling
        
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer, 
            mlm=False
        )
        
    except Exception as e:
        error_msg = f"Error loading dataset: {str(e)}"
        log.append(error_msg)
        return "\n".join(log)
    
    # --- Training Arguments ---
    progress(0.75, desc="Setting up training arguments...")
    output_dir = f"./results_{model_repo_name}"
    os.makedirs(output_dir, exist_ok=True)
    
    # Optimize settings for A100
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=float(epochs),
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=grad_accum_steps,
        learning_rate=learning_rate,
        weight_decay=0.01,
        logging_dir=f"{output_dir}/logs",
        logging_steps=10,
        save_steps=100,
        save_total_limit=3,
        remove_unused_columns=False,
        push_to_hub=False,
        disable_tqdm=False,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        report_to="tensorboard",
        bf16=True if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else False,
        gradient_checkpointing=True,  # Still useful for efficiency
        gradient_checkpointing_kwargs={'use_reentrant': False},
        ddp_find_unused_parameters=False,
        deepspeed=ds_config_path if n_gpus > 1 else None,  # Only use DeepSpeed for multi-GPU
    )
    
    # --- Initialize Trainer ---
    progress(0.8, desc="Initializing trainer...")
    trainer = Trainer(
        model=model_to_train,
        args=training_args,
        train_dataset=train_dataset,
        data_collator=data_collator,
    )
    
    log.append("Trainer initialized for training.")
    
    # --- Start Training ---
    # Clear cache before starting
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    try:
        progress(0.85, desc="Starting training...")
        log.append("Starting training...")
        train_result = trainer.train()
        progress(0.95, desc="Saving model...")
    
        # Save final model (adapter weights) and training state
        final_save_path = os.path.join(training_args.output_dir, "final_checkpoint")
        log.append(f"Saving final model checkpoint to {final_save_path}...")
        trainer.save_model(final_save_path)
        trainer.save_state()
    
        # Log metrics
        metrics = train_result.metrics
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        
        for key, value in metrics.items():
            log.append(f"{key}: {value}")
    
    except Exception as e:
        error_msg = f"An error occurred during training: {e}"
        log.append(error_msg)
        return "\n".join(log)
    
    progress(1.0, desc="Training complete!")
    log.append("Training process complete.")
    return "\n".join(log)

# Define the Gradio interface
def create_interface():
    with gr.Blocks(title="Llama 3.2 1B RVQ Fine-tuning") as demo:
        gr.Markdown("# Llama 3.2 1B RVQ LoRA Fine-tuning")
        gr.Markdown("Fine-tune a Llama 3.2 1B model with RVQ token embeddings using LoRA")
        
        with gr.Row():
            with gr.Column():
                hf_username = gr.Textbox(label="HuggingFace Username", value="Twelve2five")
                model_repo = gr.Textbox(label="Model Repository Name", value="llama-3.2-1b-rvq")
                dataset_repo = gr.Textbox(label="Dataset Repository Name", value="podcast-dialogue-rvq-pairs-3items")
            
            with gr.Column():
                epochs = gr.Number(label="Number of Epochs", value=3, minimum=1, maximum=10)
                batch_size = gr.Number(label="Batch Size per Device", value=4, minimum=1, maximum=16)
                grad_accum = gr.Number(label="Gradient Accumulation Steps", value=2, minimum=1, maximum=16)
                lr = gr.Number(label="Learning Rate", value=2e-4)
        
        start_btn = gr.Button("Start Training")
        output = gr.Textbox(label="Training Log", lines=20)
        
        start_btn.click(
            fn=train_model,
            inputs=[hf_username, model_repo, dataset_repo, epochs, batch_size, grad_accum, lr],
            outputs=output
        )
        
    return demo

# Create and launch the interface
demo = create_interface()
if __name__ == "__main__":
    demo.launch()