Update app.py

app.py

@@ -4,10 +4,12 @@ import glob
 import gc
 from transformers import (
     AutoModelForCausalLM, 
+    AutoTokenizer,
     BitsAndBytesConfig, 
     TrainingArguments, 
     Trainer,
-    DataCollatorForLanguageModeling
+    DataCollatorForLanguageModeling,
+    AutoTokenizer
 )
 from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
 from datasets import Dataset
@@ -130,71 +132,71 @@ def prepare_for_dataset(batch):
     return output
 
 def load_model():
-    clean_memory()  # Start with clean memory
-    
     print(f"Loading base model architecture from: {hf_model_repo_id}")
     
-    # Even more extreme quantization
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.float16,  # Use float16 instead of bfloat16
-        bnb_4bit_use_double_quant=True,
-    )
+    # Get information about GPU with most free memory
+    gpu_id = 0  # Default to first GPU
+    max_free_memory = 0
     
-    # For 4-bit training, we need to load on a single device
-    # Choose GPU with most available memory
-    free_memory = []
     for i in range(torch.cuda.device_count()):
-        total_memory = torch.cuda.get_device_properties(i).total_memory
-        reserved_memory = torch.cuda.memory_reserved(i)
-        free_memory.append((total_memory - reserved_memory) / 1e9)  # Convert to GB
+        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
     
-    # Choose the GPU with the most free memory
-    target_gpu = free_memory.index(max(free_memory))
-    print(f"Loading model on GPU {target_gpu} with {free_memory[target_gpu]:.2f}GB free memory")
+    print(f"Loading model on GPU {gpu_id} with {max_free_memory / 1e9:.2f}GB free memory")
     
-    # Use target GPU for model loading (crucial for 4-bit training)
-    device_map = {'': target_gpu}
+    # 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 model on the single target GPU
+    # Load the model
     model = AutoModelForCausalLM.from_pretrained(
         hf_model_repo_id,
         quantization_config=bnb_config,
-        device_map=device_map,  # Place entire model on one GPU
-        trust_remote_code=True,
-        use_cache=False,
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True,
+        device_map={"": gpu_id},
+        torch_dtype=torch.bfloat16,
     )
     
-    # Add print statement to check which device the model is on
-    print(f"Model loaded on device: {next(model.parameters()).device}")
+    print(f"Model loaded on device: cuda:{gpu_id}")
+    
+    # Load tokenizer as well
+    tokenizer = AutoTokenizer.from_pretrained(hf_model_repo_id)
+    print(f"Loaded model vocab size: {len(tokenizer)}")
     
-    # Continue with the LoRA config as before
-    print(f"Loaded model vocab size: {model.get_input_embeddings().weight.shape[0]}")
+    # Print information about input embeddings
     print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}")
     
-    # --- Configure PEFT/LoRA ---
+    # Prepare model for k-bit training
+    model = prepare_model_for_kbit_training(model)
+    
+    # Define LoRA configuration
     lora_config = LoraConfig(
-        r=16,  # rank
+        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,
-        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        task_type=TaskType.CAUSAL_LM
     )
     
-    # Prepare model for k-bit training
-    model = prepare_model_for_kbit_training(model)
-    
-    # Add LoRA adapters
+    # Apply LoRA to model
     model = get_peft_model(model, lora_config)
-    
-    # Log number of trainable parameters
     model.print_trainable_parameters()
     
-    return model
+    return model, tokenizer  # Return both model and tokenizer
 
 def load_dataset():
     # --- Download the dataset repository files ---
@@ -275,7 +277,7 @@ def train_model(progress=gr.Progress()):
     clean_memory()
     
     # Load model with optimized memory settings
-    model = load_model()
+    model, tokenizer = load_model()
     
     # Load and prepare dataset
     progress(0.1, desc="Loading dataset...")
@@ -302,13 +304,14 @@ def train_model(progress=gr.Progress()):
             logging_first_step=True,        # Force log on first step
         )
         
-        # Create a simple trainer
+        # Create a simple trainer with the tokenizer
         trainer = Trainer(
             model=model,
             args=training_args,
             train_dataset=train_dataset,
             data_collator=DataCollatorForLanguageModeling(
-                tokenizer=None, mlm=False
+                tokenizer=tokenizer,
+                mlm=False
             )
         )