Update app.py

app.py

@@ -1,352 +1,431 @@
-# -*- coding: utf-8 -*-
-"""
-Script for fine-tuning Llama-3-8B with RVQ tokens on multiple GPUs
-"""
-
-# Basic setup and installations
-!pip install -q -U transformers accelerate bitsandbytes peft torch datasets huggingface_hub deepspeed
-
-# No need for notebook_login on Hugging Face platform
-# Authentication is handled automatically
-
+import gradio as gr
+import subprocess
+import sys
+import os
+import glob
+import json
+import math
 import torch
-from transformers import AutoModelForCausalLM, BitsAndBytesConfig
-from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
 import gc
-import os
+from tqdm import tqdm
 from datasets import Dataset
 from huggingface_hub import snapshot_download
-import glob
-from tqdm import tqdm
-
-# --- 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}"
-
-# Check if running on multiple GPUs
-n_gpus = torch.cuda.device_count()
-print(f"Number of GPUs available: {n_gpus}")
-
-# --- Quantization Configuration ---
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-)
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, Trainer
+from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
 
-# --- Load Base Model (with quantization) ---
-try:
-    # For multi-GPU QLoRA, we'll use device_map="auto" and let DeepSpeed handle distribution later
-    model = AutoModelForCausalLM.from_pretrained(
-        hf_model_repo_id,
-        quantization_config=bnb_config,
-        device_map="auto",  # Will be overridden by DeepSpeed
-        trust_remote_code=True
+# Function to run the training process
+def train_model(
+    hf_username, 
+    model_repo_name, 
+    dataset_repo_name, 
+    epochs=1, 
+    batch_size=1, 
+    grad_accum_steps=4,
+    learning_rate=1e-4,
+    progress=gr.Progress()
+):
+    progress(0, desc="Installing dependencies...")
+    # Install required packages if needed
+    try:
+        import transformers
+        import accelerate
+        import bitsandbytes
+        import peft
+        import deepspeed
+    except ImportError:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-U", 
+                          "transformers", "accelerate", "bitsandbytes", "peft", 
+                          "torch", "datasets", "huggingface_hub", "deepspeed"])
+    
+    # --- 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 = []
+    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}")
+    
+    # --- Quantization Configuration ---
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
     )
-    print(f"Loaded model vocab size: {model.config.vocab_size}")
-    print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
-except Exception as e:
-    print(f"Error loading model from Hub: {e}")
-    raise SystemExit("Model loading failed.")
-
-# --- Prepare for K-bit Training & Apply LoRA ---
-model = prepare_model_for_kbit_training(model)
-
-lora_config = LoraConfig(
-    task_type=TaskType.CAUSAL_LM,
-    r=16,
-    lora_alpha=32,
-    lora_dropout=0.05,
-    bias="none",
-    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
-)
-peft_model = get_peft_model(model, lora_config)
-peft_model.print_trainable_parameters()
-model_to_train = peft_model
-
-# Cleanup
-gc.collect()
-if torch.cuda.is_available():
-    torch.cuda.empty_cache()
-
-# --- Load Dataset from Hub ---
-local_download_path = "./downloaded_dataset_files"
-
-try:
-    downloaded_repo_root = snapshot_download(
-        repo_id=hf_dataset_repo_id,
-        repo_type="dataset",
-        local_dir=local_download_path,
-        local_dir_use_symlinks=False
+    
+    # --- Load Base Model (with quantization) ---
+    progress(0.1, desc="Loading base model...")
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            hf_model_repo_id,
+            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)
+        return "\n".join(log)
+    
+    # --- Prepare for K-bit Training & Apply LoRA ---
+    progress(0.15, desc="Preparing model for fine-tuning...")
+    model = prepare_model_for_kbit_training(model)
+    
+    lora_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        bias="none",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
     )
-    print(f"Dataset repository content downloaded to: {downloaded_repo_root}")
-except Exception as e:
-    print(f"Error downloading dataset repository from Hub: {e}")
-    raise SystemExit("Dataset download failed.")
-
-# --- Find and load the .pt files ---
-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:
-        raise FileNotFoundError(f"No RVQ pair files found in expected directories")
-
-print(f"Found {len(all_pair_files)} RVQ pair files.")
-
-# --- Load data from .pt files ---
-all_data_pairs = []
-for file_path in tqdm(all_pair_files, desc="Loading pair files"):
+    peft_model = get_peft_model(model, lora_config)
+    trainable_params = peft_model.print_trainable_parameters()
+    log.append(f"Trainable parameters: {trainable_params}")
+    model_to_train = peft_model
+    
+    # Cleanup
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    # --- Load Dataset from Hub ---
+    progress(0.2, desc="Downloading dataset...")
+    local_download_path = "./downloaded_dataset_files"
+    
     try:
-        episode_pairs = torch.load(file_path, map_location='cpu')
-        all_data_pairs.extend(episode_pairs)
+        downloaded_repo_root = snapshot_download(
+            repo_id=hf_dataset_repo_id,
+            repo_type="dataset",
+            local_dir=local_download_path,
+            local_dir_use_symlinks=False
+        )
+        log.append(f"Dataset repository content downloaded to: {downloaded_repo_root}")
     except Exception as e:
-        print(f"Warning: Could not load file {file_path}: {e}")
-
-if not all_data_pairs:
-    raise ValueError("No valid data pairs were loaded")
-
-print(f"Loaded a total of {len(all_data_pairs)} training pairs into memory.")
-
-# --- Convert to HF Dataset ---
-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
-
-chunk_size = 1000
-processed_data = {'input_ids': [], 'labels': []}
-for i in tqdm(range(0, len(all_data_pairs), chunk_size), desc="Preparing data for Dataset"):
-    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']]
-})
-
-train_dataset = hf_dataset
-
-# Cleanup
-del all_data_pairs
-del processed_data
-gc.collect()
-
-# --- Define Data Collator ---
-def seq2seq_causal_collator(features):
-    batch = {}
-    concatenated_input_ids = []
-    concatenated_labels = []
-    max_len = 0
-
-    # First pass: Concatenate, create masked labels, find max length
-    for feature in features:
-        input_ids = feature['input_ids']
-        labels = feature['labels']
-
-        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]
-
-        combined_ids = torch.cat([input_ids, labels], dim=0)
-        concatenated_input_ids.append(combined_ids)
-
-        masked_labels = torch.cat([
-            torch.full((context_len,), -100, dtype=torch.long, device=input_ids.device),
-            labels
-        ], dim=0)
-        concatenated_labels.append(masked_labels)
-
-        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]
-
-        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)
-    batch['attention_mask'] = batch['input_ids'].ne(input_pad_token_id).long()
-
-    return batch
-
-data_collator = seq2seq_causal_collator
-
-# --- Define Training Arguments and Initialize Trainer ---
-from transformers import TrainingArguments, Trainer
-import math
-
-# Output directories
-OUTPUT_TRAINING_DIR = "./llama3-8b-rvq-qlora-finetuned-run"
-LOGGING_DIR = "./llama3-8b-rvq-qlora-logs-run"
-
-# Training parameters - adjusted for 4x T4 GPUs
-NUM_EPOCHS = 1
-# Scale down per-device batch size since we have multiple GPUs now
-BATCH_SIZE_PER_DEVICE = 1  # Smaller per-device batch size to avoid OOM
-GRAD_ACCUMULATION_STEPS = 4
-LEARNING_RATE = 1e-4
-WEIGHT_DECAY = 0.01
-WARMUP_RATIO = 0.03
-LR_SCHEDULER = "cosine"
-OPTIMIZER = "paged_adamw_8bit"
-
-# Calculate total steps and warmup steps
-# Total batch size is now batch_size × num_gpus × grad_accum_steps
-total_train_batch_size = BATCH_SIZE_PER_DEVICE * n_gpus * GRAD_ACCUMULATION_STEPS
-num_training_steps = math.ceil((len(train_dataset) * NUM_EPOCHS) / total_train_batch_size)
-num_warmup_steps = int(num_training_steps * WARMUP_RATIO)
-
-# Logging/Saving frequency
-steps_per_epoch = math.ceil(len(train_dataset) / total_train_batch_size)
-LOGGING_STEPS = max(10, steps_per_epoch // 15)
-SAVE_STEPS = max(50, steps_per_epoch // 10)
-
-print(f"Dataset size: {len(train_dataset)}")
-print(f"Number of GPUs: {n_gpus}")
-print(f"Batch size per device: {BATCH_SIZE_PER_DEVICE}")
-print(f"Gradient Accumulation steps: {GRAD_ACCUMULATION_STEPS}")
-print(f"Total train batch size (effective): {total_train_batch_size}")
-print(f"Total optimization steps: {num_training_steps}")
-print(f"Warmup steps: {num_warmup_steps}")
-
-# Configure for multi-GPU training using DeepSpeed
-training_args = TrainingArguments(
-    output_dir=OUTPUT_TRAINING_DIR,
-    num_train_epochs=NUM_EPOCHS,
-    per_device_train_batch_size=BATCH_SIZE_PER_DEVICE,
-    gradient_accumulation_steps=GRAD_ACCUMULATION_STEPS,
-    optim=OPTIMIZER,
-    logging_dir=LOGGING_DIR,
-    logging_strategy="steps",
-    logging_steps=LOGGING_STEPS,
-    save_strategy="steps",
-    save_steps=SAVE_STEPS,
-    save_total_limit=2,
-    learning_rate=LEARNING_RATE,
-    weight_decay=WEIGHT_DECAY,
-    warmup_steps=num_warmup_steps,
-    lr_scheduler_type=LR_SCHEDULER,
-    report_to="tensorboard",
-    bf16=True if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else False,
-    gradient_checkpointing=True,
-    gradient_checkpointing_kwargs={'use_reentrant': False},
-    
-    # Multi-GPU specific settings
-    deepspeed="ds_config.json",  # We'll create this file below
-    ddp_find_unused_parameters=False,
-)
-
-# --- Create DeepSpeed configuration file ---
-import json
-
-# DeepSpeed ZeRO-3 config optimized for T4 GPUs
-ds_config = {
-    "fp16": {
-        "enabled": "auto",
-        "loss_scale": 0,
-        "loss_scale_window": 1000,
-        "initial_scale_power": 16,
-        "hysteresis": 2,
-        "min_loss_scale": 1
-    },
-    "bf16": {
-        "enabled": "auto"
-    },
-    "zero_optimization": {
-        "stage": 3,
-        "offload_optimizer": {
-            "device": "cpu",
-            "pin_memory": True
+        error_msg = f"Error downloading dataset repository from Hub: {e}"
+        log.append(error_msg)
+        return "\n".join(log)
+    
+    # --- Find and load the .pt files ---
+    progress(0.25, desc="Finding dataset files...")
+    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:
+            error_msg = "No RVQ pair files found in expected directories"
+            log.append(error_msg)
+            return "\n".join(log)
+    
+    log.append(f"Found {len(all_pair_files)} RVQ pair files.")
+    
+    # --- Load data from .pt files ---
+    progress(0.3, desc="Loading dataset files...")
+    all_data_pairs = []
+    for i, file_path in enumerate(all_pair_files):
+        progress(0.3 + (0.1 * i / len(all_pair_files)), desc=f"Loading file {i+1}/{len(all_pair_files)}")
+        try:
+            episode_pairs = torch.load(file_path, map_location='cpu')
+            all_data_pairs.extend(episode_pairs)
+        except Exception as e:
+            log.append(f"Warning: Could not load file {file_path}: {e}")
+    
+    if not all_data_pairs:
+        error_msg = "No valid data pairs were loaded"
+        log.append(error_msg)
+        return "\n".join(log)
+    
+    log.append(f"Loaded a total of {len(all_data_pairs)} training pairs into memory.")
+    
+    # --- Convert to HF Dataset ---
+    progress(0.45, desc="Converting to Hugging Face Dataset...")
+    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
+    
+    chunk_size = 1000
+    processed_data = {'input_ids': [], 'labels': []}
+    
+    total_chunks = len(range(0, len(all_data_pairs), chunk_size))
+    for i in range(0, len(all_data_pairs), chunk_size):
+        chunk_idx = i // chunk_size
+        progress(0.45 + (0.1 * chunk_idx / total_chunks), 
+                desc=f"Processing chunk {chunk_idx+1}/{total_chunks}")
+        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']]
+    })
+    
+    train_dataset = hf_dataset
+    
+    # Cleanup
+    del all_data_pairs
+    del processed_data
+    gc.collect()
+    
+    # --- Define Data Collator ---
+    progress(0.55, desc="Defining data collator...")
+    def seq2seq_causal_collator(features):
+        batch = {}
+        concatenated_input_ids = []
+        concatenated_labels = []
+        max_len = 0
+    
+        # First pass: Concatenate, create masked labels, find max length
+        for feature in features:
+            input_ids = feature['input_ids']
+            labels = feature['labels']
+    
+            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]
+    
+            combined_ids = torch.cat([input_ids, labels], dim=0)
+            concatenated_input_ids.append(combined_ids)
+    
+            masked_labels = torch.cat([
+                torch.full((context_len,), -100, dtype=torch.long, device=input_ids.device),
+                labels
+            ], dim=0)
+            concatenated_labels.append(masked_labels)
+    
+            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]
+    
+            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)
+        batch['attention_mask'] = batch['input_ids'].ne(input_pad_token_id).long()
+    
+        return batch
+    
+    data_collator = seq2seq_causal_collator
+    
+    # --- Define Training Arguments and Initialize Trainer ---
+    progress(0.65, desc="Setting up training configuration...")
+    
+    # Output directories
+    OUTPUT_TRAINING_DIR = "./llama3-8b-rvq-qlora-finetuned-run"
+    LOGGING_DIR = "./llama3-8b-rvq-qlora-logs-run"
+    
+    # Training parameters - adjusted for 4x T4 GPUs
+    NUM_EPOCHS = int(epochs)
+    BATCH_SIZE_PER_DEVICE = int(batch_size)  # Smaller per-device batch size to avoid OOM
+    GRAD_ACCUMULATION_STEPS = int(grad_accum_steps)
+    LEARNING_RATE = float(learning_rate)
+    WEIGHT_DECAY = 0.01
+    WARMUP_RATIO = 0.03
+    LR_SCHEDULER = "cosine"
+    OPTIMIZER = "paged_adamw_8bit"
+    
+    # Calculate total steps and warmup steps
+    # Total batch size is now batch_size × num_gpus × grad_accum_steps
+    total_train_batch_size = BATCH_SIZE_PER_DEVICE * n_gpus * GRAD_ACCUMULATION_STEPS
+    num_training_steps = math.ceil((len(train_dataset) * NUM_EPOCHS) / total_train_batch_size)
+    num_warmup_steps = int(num_training_steps * WARMUP_RATIO)
+    
+    # Logging/Saving frequency
+    steps_per_epoch = math.ceil(len(train_dataset) / total_train_batch_size)
+    LOGGING_STEPS = max(10, steps_per_epoch // 15)
+    SAVE_STEPS = max(50, steps_per_epoch // 10)
+    
+    log.append(f"Dataset size: {len(train_dataset)}")
+    log.append(f"Number of GPUs: {n_gpus}")
+    log.append(f"Batch size per device: {BATCH_SIZE_PER_DEVICE}")
+    log.append(f"Gradient Accumulation steps: {GRAD_ACCUMULATION_STEPS}")
+    log.append(f"Total train batch size (effective): {total_train_batch_size}")
+    log.append(f"Total optimization steps: {num_training_steps}")
+    log.append(f"Warmup steps: {num_warmup_steps}")
+    
+    # --- Create DeepSpeed configuration file ---
+    progress(0.7, desc="Creating DeepSpeed configuration...")
+    # DeepSpeed ZeRO-3 config optimized for T4 GPUs
+    ds_config = {
+        "fp16": {
+            "enabled": "auto",
+            "loss_scale": 0,
+            "loss_scale_window": 1000,
+            "initial_scale_power": 16,
+            "hysteresis": 2,
+            "min_loss_scale": 1
         },
-        "offload_param": {
-            "device": "cpu",
-            "pin_memory": True
+        "bf16": {
+            "enabled": "auto"
         },
-        "overlap_comm": True,
-        "contiguous_gradients": True,
-        "reduce_bucket_size": "auto",
-        "stage3_prefetch_bucket_size": "auto",
-        "stage3_param_persistence_threshold": "auto",
-        "gather_16bit_weights_on_model_save": True,
-        "stage3_max_live_parameters": 1e9,
-        "stage3_max_reuse_distance": 1e9
-    },
-    "gradient_accumulation_steps": GRAD_ACCUMULATION_STEPS,
-    "gradient_clipping": "auto",
-    "steps_per_print": 10,
-    "train_batch_size": "auto",
-    "train_micro_batch_size_per_gpu": "auto",
-    "wall_clock_breakdown": False
-}
-
-with open("ds_config.json", "w") as f:
-    json.dump(ds_config, f, indent=4)
-
-# --- Initialize Trainer ---
-trainer = Trainer(
-    model=model_to_train,
-    args=training_args,
-    train_dataset=train_dataset,
-    data_collator=data_collator,
-)
-
-print("Trainer initialized with DeepSpeed for multi-GPU training.")
-
-# --- Start Training ---
-# Clear cache before starting
-gc.collect()
-if torch.cuda.is_available():
-    torch.cuda.empty_cache()
-
-try:
-    print("Starting distributed training on multiple GPUs...")
-    train_result = trainer.train()
-
-    # Save final model (adapter weights) and training state
-    final_save_path = os.path.join(training_args.output_dir, "final_checkpoint")
-    print(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)
-
-except Exception as e:
-    print(f"An error occurred during training: {e}")
-    raise e
-
-print("Multi-GPU training process complete.")
+        "zero_optimization": {
+            "stage": 3,
+            "offload_optimizer": {
+                "device": "cpu",
+                "pin_memory": True
+            },
+            "offload_param": {
+                "device": "cpu",
+                "pin_memory": True
+            },
+            "overlap_comm": True,
+            "contiguous_gradients": True,
+            "reduce_bucket_size": "auto",
+            "stage3_prefetch_bucket_size": "auto",
+            "stage3_param_persistence_threshold": "auto",
+            "gather_16bit_weights_on_model_save": True,
+            "stage3_max_live_parameters": 1e9,
+            "stage3_max_reuse_distance": 1e9
+        },
+        "gradient_accumulation_steps": GRAD_ACCUMULATION_STEPS,
+        "gradient_clipping": "auto",
+        "steps_per_print": 10,
+        "train_batch_size": "auto",
+        "train_micro_batch_size_per_gpu": "auto",
+        "wall_clock_breakdown": False
+    }
+    
+    with open("ds_config.json", "w") as f:
+        json.dump(ds_config, f, indent=4)
+    
+    # Configure for multi-GPU training using DeepSpeed
+    progress(0.75, desc="Setting up training arguments...")
+    training_args = TrainingArguments(
+        output_dir=OUTPUT_TRAINING_DIR,
+        num_train_epochs=NUM_EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE_PER_DEVICE,
+        gradient_accumulation_steps=GRAD_ACCUMULATION_STEPS,
+        optim=OPTIMIZER,
+        logging_dir=LOGGING_DIR,
+        logging_strategy="steps",
+        logging_steps=LOGGING_STEPS,
+        save_strategy="steps",
+        save_steps=SAVE_STEPS,
+        save_total_limit=2,
+        learning_rate=LEARNING_RATE,
+        weight_decay=WEIGHT_DECAY,
+        warmup_steps=num_warmup_steps,
+        lr_scheduler_type=LR_SCHEDULER,
+        report_to="tensorboard",
+        bf16=True if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else False,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={'use_reentrant': False},
+        
+        # Multi-GPU specific settings
+        deepspeed="ds_config.json",
+        ddp_find_unused_parameters=False,
+    )
+    
+    # --- 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 with DeepSpeed for multi-GPU 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 distributed training on multiple GPUs...")
+        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("Multi-GPU training process complete.")
+    return "\n".join(log)
+
+# Define the Gradio interface
+def create_interface():
+    with gr.Blocks(title="Llama 3 8B RVQ Fine-tuning") as demo:
+        gr.Markdown("# Llama 3 8B RVQ LoRA Fine-tuning")
+        gr.Markdown("Fine-tune a Llama 3 8B model with RVQ token embeddings using LoRA on multiple GPUs")
+        
+        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-8b-rvq-resized")
+                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=1, minimum=1, maximum=10)
+                batch_size = gr.Number(label="Batch Size per Device", value=1, minimum=1, maximum=8)
+                grad_accum = gr.Number(label="Gradient Accumulation Steps", value=4, minimum=1, maximum=16)
+                lr = gr.Number(label="Learning Rate", value=1e-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()