Upload train3.py

train3.py

@@ -7,43 +7,66 @@ import sys
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import wandb
 from torch.utils.data import TensorDataset, DataLoader
 from transformers import AutoModelForCausalLM, get_linear_schedule_with_warmup
 from peft import PeftModel
 from torch.cuda.amp import GradScaler, autocast
 from tqdm.auto import tqdm
 from multiprocessing import freeze_support
+import shutil # Import shutil for removing old checkpoints
+import collections # Import collections for deque
 
 def main():
     # --- Config ---
     PRET_FILE      = "pretokenized_queries.pt"
     MODEL_NAME     = "google/gemma-3-1b-pt"
-    LORA_DIR       = "phase2_triplet_amp/final"
-    BATCH_SIZE     = 64
+    LORA_DIR       = "phase2_triplet_amp/final" # Adapters from previous stage
+    BATCH_SIZE     = 200 
     LR             = 1e-5
     WEIGHT_DECAY   = 0.01
-    NUM_EPOCHS     = 1
+    NUM_EPOCHS     = 1 # As per our discussion, 1 epoch is likely sufficient given fast convergence
     TEMP           = 0.05
-    OUTPUT_DIR     = "phase3_self_contrast"
+    OUTPUT_DIR     = "phase3_self_contrast_wandb"
     GRAD_CLIP_NORM = 1.0
     SEED           = 42
+    WANDB_PROJECT  = "query-encoder-phase3"
+
+    # --- Checkpointing Configuration ---
+    SAVE_INTERVAL = 1000 # Save a checkpoint every N steps
+    KEEP_LAST_CKPTS = 5  # Keep only the last N checkpoints (to save disk space)
 
     os.makedirs(OUTPUT_DIR, exist_ok=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     torch.manual_seed(SEED)
 
+    # --- Initialize WandB ---
+    wandb.init(
+        project=WANDB_PROJECT,
+        config={
+            "model_name": MODEL_NAME, "lora_dir": LORA_DIR, "batch_size": BATCH_SIZE, 
+            "lr": LR, "num_epochs": NUM_EPOCHS, "seed": SEED,
+            "save_interval_steps": SAVE_INTERVAL,
+            "keep_last_checkpoints": KEEP_LAST_CKPTS,
+        }
+    )
+
     # --- Load pretokenized queries safely ---
+    print(f"Loading pretokenized queries from {PRET_FILE}...")
     data = torch.load(PRET_FILE, weights_only=True)
     input_ids      = data["input_ids"]
     attention_mask = data["attention_mask"]
     dataset = TensorDataset(input_ids, attention_mask)
     loader  = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+    print(f"Loaded {len(dataset)} samples.")
 
-    # --- Load base model + LoRA adapters ---
+    # --- Load base model + LoRA adapters from previous stage ---
+    print(f"Loading base model '{MODEL_NAME}' and LoRA adapters from '{LORA_DIR}'...")
     base = AutoModelForCausalLM.from_pretrained(MODEL_NAME, attn_implementation="eager")
     peft = PeftModel.from_pretrained(base, LORA_DIR).to(device)
+    print("LoRA adapters loaded.")
 
-    # --- Projection head now outputs hidden_size instead of 256 ---
+    # --- Projection head now outputs hidden_size ---
     class GemmaSelfContrast(nn.Module):
         def __init__(self, peft_model):
             super().__init__()
@@ -69,6 +92,9 @@ def main():
             return z / norm
 
     model = GemmaSelfContrast(peft).to(device)
+    print("Encoder model (with projection head) initialized.")
+    # Watch the model with wandb (optional, can be slow, but good for tracking gradients)
+    # wandb.watch(model, log="all", log_freq=100) # Commented out due to potential slowdown
 
     # --- Optimizer, scheduler, AMP scaler ---
     optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
@@ -79,12 +105,18 @@ def main():
         num_training_steps=total_steps
     )
     scaler = GradScaler()
+    print(f"Training will run for {total_steps} steps.")
+
+    # Deque to manage checkpoint paths and enforce keeping only the last N
+    checkpoint_paths = collections.deque(maxlen=KEEP_LAST_CKPTS)
 
     # --- Training loop ---
     model.train()
+    global_step = 0
     for epoch in range(1, NUM_EPOCHS + 1):
         total_loss = 0.0
-        for ids, mask in tqdm(loader, desc=f"Epoch {epoch}", unit="batch"):
+        pbar = tqdm(loader, desc=f"Epoch {epoch}", unit="batch")
+        for ids, mask in pbar:
             ids, mask = ids.to(device), mask.to(device)
 
             with autocast():
@@ -105,23 +137,72 @@ def main():
 
             optimizer.zero_grad()
             scaler.scale(loss).backward()
-            scaler.unscale_(optimizer)
+            scaler.unscale_(optimizer) # Unscale gradients before clipping
             torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP_NORM)
             scaler.step(optimizer)
             scaler.update()
             scheduler.step()
-
+            
+            # --- Log metrics to WandB at every step ---
+            wandb.log({
+                "train/loss": loss.item(),
+                "train/lr": scheduler.get_last_lr()[0],
+                "train/epoch": epoch,
+                "train/global_step": global_step
+            }, step=global_step)
+            
+            pbar.set_postfix({"loss": f"{loss.item():.4f}"})
+
+            # --- PERIODIC SAVING BLOCK ---
+            # Save checkpoint every SAVE_INTERVAL steps
+            if (global_step + 1) % SAVE_INTERVAL == 0: 
+                # Create a unique directory for this checkpoint
+                ckpt_dir = os.path.join(OUTPUT_DIR, f"checkpoint-step-{global_step + 1}")
+                os.makedirs(ckpt_dir, exist_ok=True)
+                
+                print(f"\nSaving checkpoint to {ckpt_dir}...")
+                # Save the PEFT adapters
+                peft.save_pretrained(ckpt_dir)
+                # Save the trained projection head's state dictionary
+                torch.save(model.proj.state_dict(), os.path.join(ckpt_dir, "encoder_proj.pth"))
+                
+                # Manage old checkpoints
+                if len(checkpoint_paths) == KEEP_LAST_CKPTS:
+                    oldest_ckpt = checkpoint_paths.popleft() # Remove the oldest path from deque
+                    if os.path.isdir(oldest_ckpt):
+                        print(f"Removing old checkpoint: {oldest_ckpt}")
+                        shutil.rmtree(oldest_ckpt, ignore_errors=True) # Delete the directory
+                checkpoint_paths.append(ckpt_dir) # Add new checkpoint path
+                print("Checkpoint saved and old ones managed.")
+            # --- END PERIODIC SAVING ---
+
+            global_step += 1
             total_loss += loss.item()
 
         avg_loss = total_loss / len(loader)
-        print(f"Epoch {epoch} avg loss: {avg_loss:.6f}")
-
-    # --- Save only LoRA adapters ---
+        print(f"Epoch {epoch} training complete. Avg loss: {avg_loss:.6f}")
+        # Log average epoch loss as well
+        wandb.log({"train/epoch_avg_loss": avg_loss, "epoch": epoch}, step=global_step)
+
+    # --- Final Save for the "final" directory ---
+    # This ensures that even if you stop mid-epoch (after a checkpoint)
+    # or don't stop, there's always a clear 'final' model.
+    print("\nTraining finished. Saving final model to 'final' directory...")
     final_dir = os.path.join(OUTPUT_DIR, "final")
     os.makedirs(final_dir, exist_ok=True)
+    
+    # Save the LoRA adapters
     peft.save_pretrained(final_dir)
-    print("Phase 3 complete. LoRA adapters saved to", final_dir)
+    
+    # Save the trained projection head's state dictionary
+    torch.save(model.proj.state_dict(), os.path.join(final_dir, "encoder_proj.pth"))
+    
+    print(f"Phase 3 complete. LoRA adapters and projection head saved to {final_dir}")
+    
+    # --- Finalize WandB run ---
+    wandb.finish()
+
 
 if __name__ == "__main__":
     freeze_support()
-    main()
+    main()