Update app.py

app.py

@@ -330,49 +330,49 @@ def train_model(
     dataset_repo_name, 
     epochs=1, 
     batch_size=1, 
-    grad_accum_steps=4,
+    grad_accum_steps=16,  # Increased from 8 to 16
     learning_rate=1e-4,
     progress=gr.Progress()
 ):
     progress(0, desc="Setting up environment...")
     log = []
     
-    # Completely clean up transformers installation
-    log.append("Completely reinstalling transformers and dependencies...")
-    
-    # First uninstall any existing transformers
-    subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", "transformers"])
-    
-    # Clean any cached files that might be causing issues
-    cache_dirs = [
-        os.path.expanduser("~/.cache/huggingface"),
-        os.path.expanduser("~/.cache/pip")
-    ]
-    
-    for cache_dir in cache_dirs:
-        if os.path.exists(cache_dir):
-            log.append(f"Cleaning cache directory: {cache_dir}")
-            try:
-                shutil.rmtree(cache_dir)
-            except Exception as e:
-                log.append(f"Warning: Could not clean {cache_dir}: {e}")
+    # Aggressive memory cleanup
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        # Reset peak memory stats
+        torch.cuda.reset_peak_memory_stats()
     
-    # Install a stable version of transformers known to work with Llama models
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers==4.35.2", "sentencepiece"])
+    # 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}")
     
-    # Install other dependencies
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", 
-                          "accelerate", "bitsandbytes==0.41.1", "peft==0.6.1", 
-                          "datasets", "huggingface_hub", "deepspeed==0.12.3"])
+    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
+    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")
     
-    # Now import everything after installation to ensure we use the correct versions
+    # Import required libraries
     try:
         from datasets import Dataset
         from huggingface_hub import snapshot_download
         import torch
         import transformers
         from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
-        from transformers import BitsAndBytesConfig, TrainingArguments, Trainer
+        from transformers import BitsAndBytesConfig, TrainingArguments, Trainer, AutoTokenizer
         from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
         
         log.append(f"Transformers version: {transformers.__version__}")
@@ -393,59 +393,60 @@ def train_model(
     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,
-    )
-    
-    # --- Load Base Model (with quantization) ---
+    # --- Load Base Model (with extreme quantization) ---
     progress(0.1, desc="Loading base model...")
+    local_model_path = "./model_files"
     try:
-        # First try to download the repo without loading the model
-        local_model_path = "./model_files"
-        if os.path.exists(local_model_path):
-            shutil.rmtree(local_model_path)  # Clean up any previous files
-            
+        # Download the model files
         snapshot_download(
             repo_id=hf_model_repo_id,
             local_dir=local_model_path,
             local_dir_use_symlinks=False
         )
-        
         log.append(f"Model files downloaded to {local_model_path}")
         
-        # Check if this is a Llama model by looking at config.json
-        if os.path.exists(os.path.join(local_model_path, "config.json")):
-            with open(os.path.join(local_model_path, "config.json"), "r") as f:
-                config_data = json.load(f)
-                log.append(f"Model architecture type: {config_data.get('model_type', 'unknown')}")
-                
-                # Force model_type to llama
-                config_data["model_type"] = "llama"
-                if "architectures" in config_data:
-                    config_data["architectures"] = ["LlamaForCausalLM"]
-                
-                with open(os.path.join(local_model_path, "config.json"), "w") as f:
-                    json.dump(config_data, f)
-                log.append("Updated config.json to use llama model_type")
+        # Ensure model_type is set correctly in the config
+        config_path = os.path.join(local_model_path, "config.json")
+        with open(config_path, "r") as f:
+            config_data = json.load(f)
         
-        # Now try to load with explicit Llama classes
-        config = LlamaConfig.from_pretrained(
-            local_model_path,
-            trust_remote_code=False
-        )
+        model_type = config_data.get("model_type", "")
+        log.append(f"Model architecture type: {model_type}")
         
+        # Force model_type to be "llama" if it's not already
+        if model_type != "llama":
+            config_data["model_type"] = "llama"
+            # Also ensure architectures is set correctly
+            config_data["architectures"] = ["LlamaForCausalLM"]
+            with open(config_path, "w") as f:
+                json.dump(config_data, f, indent=2)
+            log.append("Updated config.json to use llama model_type")
+        
+        # Load the config first
+        config = LlamaConfig.from_pretrained(local_model_path)
         log.append(f"Successfully loaded config: {config.model_type}")
         
-        # Load model with specific Llama class
+        # Use 4-bit quantization for extreme memory savings
+        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 tokenizer first (needed for dataset preparation)
+        tokenizer = AutoTokenizer.from_pretrained(local_model_path)
+        
+        # Explicit device map to enable CPU offloading
+        max_memory = {0: "40GB", "cpu": "64GB"}
+        
+        # Load the model with extreme memory optimization
         model = LlamaForCausalLM.from_pretrained(
             local_model_path,
             config=config,
             quantization_config=bnb_config,
             device_map="auto",
+            max_memory=max_memory,
             torch_dtype=torch.bfloat16,
             low_cpu_mem_usage=True
         )
@@ -455,20 +456,7 @@ def train_model(
     except Exception as e:
         error_msg = f"Error loading model: {str(e)}"
         log.append(error_msg)
-        
-        # Try a fallback approach
-        try:
-            log.append("Trying fallback approach with AutoModelForCausalLM...")
-            model = AutoModelForCausalLM.from_pretrained(
-                local_model_path,
-                device_map="auto",
-                torch_dtype=torch.bfloat16,
-                low_cpu_mem_usage=True
-            )
-            log.append(f"Fallback model loaded successfully")
-        except Exception as e2:
-            log.append(f"Fallback approach also failed: {str(e2)}")
-            return "\n".join(log)
+        return "\n".join(log)
     
     # --- Prepare for K-bit Training & Apply LoRA ---
     progress(0.15, desc="Preparing model for fine-tuning...")
@@ -476,218 +464,138 @@ def train_model(
         model = prepare_model_for_kbit_training(model)
         log.append("Model prepared for k-bit training")
         
+        # Use minimal LoRA configuration with fewer parameters
         lora_config = LoraConfig(
             task_type=TaskType.CAUSAL_LM,
-            r=16,
-            lora_alpha=32,
+            r=8,  # Reduced from 16 to 8
+            lora_alpha=16,  # Reduced from 32 to 16
             lora_dropout=0.05,
             bias="none",
-            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+            # Target only key modules to reduce memory usage
+            target_modules=["q_proj", "v_proj"]  # Reduced target modules
         )
+        
+        # Apply LoRA
         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
+        log.append("LoRA applied to model")
+        
+        # Free memory
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
     except Exception as e:
         error_msg = f"Error preparing model for training: {str(e)}"
         log.append(error_msg)
         return "\n".join(log)
     
-    # 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"
-    
+    # --- Download and Process Dataset ---
+    progress(0.2, desc="Loading dataset...")
     try:
-        downloaded_repo_root = snapshot_download(
+        # Download the dataset files
+        dataset_dir = os.path.join(os.getcwd(), "downloaded_dataset_files")
+        snapshot_download(
             repo_id=hf_dataset_repo_id,
-            repo_type="dataset",
-            local_dir=local_download_path,
+            local_dir=dataset_dir,
             local_dir_use_symlinks=False
         )
-        log.append(f"Dataset repository content downloaded to: {downloaded_repo_root}")
+        log.append(f"Dataset repository content downloaded to: {dataset_dir}")
+        
+        # Find all RVQ pair files
+        rvq_pair_files = glob.glob(os.path.join(dataset_dir, "*_rvq_pairs.pt"))
+        log.append(f"Found {len(rvq_pair_files)} RVQ pair files.")
+        
+        # Load training pairs from the dataset
+        training_pairs = []
+        
+        # For memory conservation, use only half the dataset for now
+        max_file_count = min(12, len(rvq_pair_files))
+        
+        for i, pair_file in enumerate(rvq_pair_files[:max_file_count]):
+            try:
+                pairs = torch.load(pair_file)
+                training_pairs.extend(pairs)
+            except Exception as e:
+                log.append(f"Warning: Could not load {pair_file}: {e}")
+        
+        log.append(f"Loaded a total of {len(training_pairs)} training pairs into memory.")
+        
+        # Prepare dataset
+        dataset = Dataset.from_dict({
+            "input_ids": [pair[0].tolist() for pair in training_pairs],
+            "labels": [pair[1].tolist() for pair in training_pairs]
+        })
+        
+        # Clear the training_pairs to free memory
+        training_pairs = None
+        gc.collect()
+        torch.cuda.empty_cache()
+        
+        # Use a smaller max_length to reduce memory pressure
+        max_length = 512  # Reduced max sequence length
+        
+        # Create data collator that handles padding
+        def data_collator(examples):
+            # Convert lists back to tensors
+            for i in range(len(examples)):
+                examples[i]["input_ids"] = torch.tensor(examples[i]["input_ids"], dtype=torch.long)
+                examples[i]["labels"] = torch.tensor(examples[i]["labels"], dtype=torch.long)
+            
+            # Get max length in this batch
+            batch_max_length = min(
+                max(len(example["input_ids"]) for example in examples),
+                max_length
+            )
+            
+            batch = {
+                "input_ids": [],
+                "attention_mask": [],
+                "labels": []
+            }
+            
+            # Prepare sequences
+            for example in examples:
+                input_ids = example["input_ids"][:batch_max_length]
+                labels = example["labels"][:batch_max_length]
+                
+                # Pad sequences
+                padding_length = batch_max_length - len(input_ids)
+                attention_mask = torch.ones_like(input_ids)
+                
+                if padding_length > 0:
+                    padding = torch.ones(padding_length, dtype=input_ids.dtype) * tokenizer.pad_token_id
+                    input_ids = torch.cat([input_ids, padding])
+                    labels = torch.cat([labels, padding * -100])  # -100 to ignore in loss computation
+                    attention_mask = torch.cat([attention_mask, torch.zeros(padding_length, dtype=attention_mask.dtype)])
+                
+                batch["input_ids"].append(input_ids)
+                batch["attention_mask"].append(attention_mask)
+                batch["labels"].append(labels)
+            
+            # Convert lists to tensors
+            for key in batch:
+                batch[key] = torch.stack(batch[key])
+            
+            return batch
+        
+        # Convert to training dataset
+        train_dataset = dataset
+        
+        # Free memory
+        del dataset
+        gc.collect()
+        torch.cuda.empty_cache()
     except Exception as e:
-        error_msg = f"Error downloading dataset repository from Hub: {e}"
+        error_msg = f"Error loading dataset: {str(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
-    
-    # --- DeepSpeed Configuration ---
-    # Create DeepSpeed config file directly in Python instead of loading from a file
-    progress(0.15, desc="Setting up DeepSpeed configuration...")
-    
-    ds_config = {
-        "fp16": {
-            "enabled": False
-        },
-        "bf16": {
-            "enabled": True
-        },
-        "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"
-        },
-        "gradient_accumulation_steps": grad_accum_steps,
-        "train_micro_batch_size_per_gpu": batch_size,
-        "gradient_clipping": 1.0,
-        "steps_per_print": 10
-    }
-    
-    # Save the config to a file
-    with open("ds_config.json", "w") as f:
-        json.dump(ds_config, f, indent=4)
-    
-    log.append("DeepSpeed configuration created successfully")
-
     # --- Training Arguments ---
-    progress(0.75, desc="Setting up training arguments...")
+    progress(0.3, desc="Setting up training arguments...")
     output_dir = f"./results_{model_repo_name}"
     os.makedirs(output_dir, exist_ok=True)
     
-    # Create training arguments without DeepSpeed first
+    # Super-aggressive memory conservation
     training_args = TrainingArguments(
         output_dir=output_dir,
         num_train_epochs=float(epochs),
@@ -696,71 +604,90 @@ def train_model(
         learning_rate=learning_rate,
         weight_decay=0.01,
         logging_dir=f"{output_dir}/logs",
-        logging_steps=10,
-        save_steps=100,
-        save_total_limit=3,
+        logging_steps=1,  # Log frequently to see progress
+        save_steps=25,  # Save checkpoints more frequently
+        save_total_limit=1,  # Keep only one checkpoint to save space
         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,
+        bf16=True,
+        fp16=False,
+        
+        # Memory optimization
         gradient_checkpointing=True,
         gradient_checkpointing_kwargs={'use_reentrant': False},
+        max_grad_norm=0.3,  # Reduced from default 1.0
+        dataloader_pin_memory=False,  # Reduce memory pressure
+        
+        # Optimizer settings for memory efficiency
+        optim="adamw_torch",
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        adam_epsilon=1e-8,
         
-        # For multi-GPU - use a different approach for DeepSpeed
-        ddp_find_unused_parameters=False,
+        # Evaluation settings
+        do_eval=False,
+        evaluation_strategy="no",
+        
+        # Set this for smaller chunks of data processing
+        dataloader_num_workers=1,
+        
+        # For memory efficiency when loading datasets
+        dataloader_drop_last=True,
     )
     
-    # Now initialize DeepSpeed separately
-    if n_gpus > 1:
-        log.append("Setting up DeepSpeed for multi-GPU training")
-        try:
-            import deepspeed
-            from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
-            
-            # Modify the trainer to use DeepSpeed
-            trainer_kwargs = {
-                "model": model_to_train,
-                "args": training_args,
-                "train_dataset": train_dataset,
-                "data_collator": data_collator,
-                "deepspeed": ds_config,  # Pass the config as a dict
-            }
+    # --- Initialize Trainer ---
+    progress(0.4, desc="Initializing trainer...")
+    
+    # Use optimizer that requires less memory
+    class MemoryEfficientTrainer(Trainer):
+        def create_optimizer(self):
+            # Create optimizer with reduced memory footprint
+            optimizer = super().create_optimizer()
+            # Force optimizer to use CPU offloading for states
+            for param_group in optimizer.param_groups:
+                for param in param_group['params']:
+                    if param.requires_grad:
+                        param.data = param.data.to("cpu")
+                        if param.grad is not None:
+                            param.grad.data = param.grad.data.to("cpu")
+            return optimizer
             
-            trainer = Trainer(**trainer_kwargs)
-            log.append("Trainer initialized with DeepSpeed for multi-GPU training")
-        except Exception as e:
-            log.append(f"Warning: Could not initialize DeepSpeed: {e}")
-            # Fallback to standard distributed training
-            trainer = Trainer(
-                model=model_to_train,
-                args=training_args,
-                train_dataset=train_dataset,
-                data_collator=data_collator,
-            )
-            log.append("Falling back to standard distributed training")
-    else:
-        # Single GPU setup
-        trainer = Trainer(
-            model=model_to_train,
-            args=training_args,
-            train_dataset=train_dataset,
-            data_collator=data_collator,
-        )
-        log.append("Trainer initialized for single GPU training")
+        def training_step(self, *args, **kwargs):
+            # Memory cleanup before each training step
+            gc.collect()
+            torch.cuda.empty_cache()
+            return super().training_step(*args, **kwargs)
+    
+    trainer = MemoryEfficientTrainer(
+        model=model_to_train,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=data_collator,
+    )
     
-    # --- Start Training ---
-    # Clear cache before starting
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
+    log.append("Trainer initialized with memory-efficient settings")
     
+    # --- Start Training ---
     try:
-        progress(0.85, desc="Starting training...")
-        log.append("Starting distributed training on multiple GPUs...")
+        # Final memory cleanup before training
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            
+        progress(0.5, desc="Starting training...")
+        log.append("Starting training with extreme memory optimization...")
+        
+        # Train in smaller chunks to manage memory better
+        total_steps = len(train_dataset) // (batch_size * grad_accum_steps)
+        log.append(f"Total training steps: {total_steps}")
+        
+        # Train the model
         train_result = trainer.train()
+        
         progress(0.95, desc="Saving model...")
     
         # Save final model (adapter weights) and training state
@@ -776,21 +703,37 @@ def train_model(
         
         for key, value in metrics.items():
             log.append(f"{key}: {value}")
+            
+        # Print peak memory usage
+        if torch.cuda.is_available():
+            peak_memory = torch.cuda.max_memory_allocated() / (1024**3)
+            log.append(f"Peak GPU memory usage: {peak_memory:.2f} GB")
     
     except Exception as e:
-        error_msg = f"An error occurred during training: {e}"
+        error_msg = f"An error occurred during training: {str(e)}"
         log.append(error_msg)
+        
+        # Try to save checkpoint even if training failed
+        try:
+            # Save whatever we have
+            log.append("Attempting to save partial checkpoint...")
+            emergency_save_path = os.path.join(training_args.output_dir, "emergency_checkpoint")
+            trainer.save_model(emergency_save_path)
+            log.append(f"Saved emergency checkpoint to {emergency_save_path}")
+        except Exception as save_error:
+            log.append(f"Could not save emergency checkpoint: {save_error}")
+        
         return "\n".join(log)
     
     progress(1.0, desc="Training complete!")
-    log.append("Multi-GPU training process complete.")
+    log.append("Training process complete successfully.")
     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")
+        gr.Markdown("Fine-tune a Llama 3 8B model with RVQ token embeddings using LoRA with extreme memory optimization")
         
         with gr.Row():
             with gr.Column():
@@ -801,7 +744,7 @@ def create_interface():
             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)
+                grad_accum = gr.Number(label="Gradient Accumulation Steps", value=16, minimum=8, maximum=32)
                 lr = gr.Number(label="Learning Rate", value=1e-4)
         
         start_btn = gr.Button("Start Training")