Update app.py

app.py

@@ -643,56 +643,16 @@ def train_model(
     
     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
+    # --- 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": "auto",
-            "loss_scale": 0,
-            "loss_scale_window": 1000,
-            "initial_scale_power": 16,
-            "hysteresis": 2,
-            "min_loss_scale": 1
+            "enabled": False
         },
         "bf16": {
-            "enabled": "auto"
+            "enabled": True
         },
         "zero_optimization": {
             "stage": 3,
@@ -708,60 +668,88 @@ def train_model(
             "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
+            "stage3_param_persistence_threshold": "auto"
         },
-        "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
+        "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)
     
-    # Configure for multi-GPU training using DeepSpeed
+    log.append("DeepSpeed configuration created successfully")
+
+    # --- Training Arguments ---
     progress(0.75, 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
     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,
+        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,
         gradient_checkpointing_kwargs={'use_reentrant': False},
         
-        # Multi-GPU specific settings
-        deepspeed="ds_config.json",
+        # For multi-GPU - use a different approach for DeepSpeed
         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.")
+    # 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
+            }
+            
+            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")
     
     # --- Start Training ---
     # Clear cache before starting