#!/usr/bin/env python3 """ SHOREKEEPER-4B Training Pipeline Runs on any CUDA device (RTX 3060, H100, etc.) """ import sys import json import torch import torch.nn as nn from pathlib import Path from tqdm import tqdm sys.path.insert(0, str(Path(__file__).parent.parent)) from src.shorekeeper import MemoryEfficientSHOREKEEPER from transformers import AutoTokenizer class SHOREKEEPERTrainer: """Simple training loop for SHOREKEEPER""" def __init__(self, model, tokenizer, config): self.model = model self.tokenizer = tokenizer self.device = next(model.parameters()).device self.learning_rate = config.get('learning_rate', 1e-4) self.epochs = config.get('epochs', 3) self.batch_size = config.get('batch_size', 2) self.gradient_accumulation = config.get('gradient_accumulation', 4) self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=self.learning_rate, weight_decay=0.01 ) self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( self.optimizer, T_max=1000, eta_min=1e-6 ) self.step = 0 def train_step(self, batch): """Single training step""" self.model.train() # Prepare batch texts = batch['text'] # Tokenize inputs = self.tokenizer( texts, return_tensors="pt", padding=True, truncation=True, max_length=512 ) input_ids = inputs['input_ids'].to(self.device) # Forward pass logits = self.model(input_ids) # Calculate loss (next token prediction) shift_logits = logits[..., :-1, :].contiguous() shift_labels = input_ids[..., 1:].contiguous() # Cross entropy loss - ignore padding tokens loss = nn.functional.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id else -100 ) # Backward loss.backward() # Gradient accumulation if (self.step + 1) % self.gradient_accumulation == 0: torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad() self.step += 1 return loss.item() def train(self, dataset, output_dir="./outputs"): """Full training loop""" print(f"\n{'='*60}") print("Starting Training") print(f"{'='*60}") print(f"Device: {self.device}") print(f"Training samples: {len(dataset)}") print(f"Batch size: {self.batch_size}") print(f"Learning rate: {self.learning_rate}") print(f"Epochs: {self.epochs}") print(f"{'='*60}\n") output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) for epoch in range(self.epochs): print(f"\nEpoch {epoch + 1}/{self.epochs}") print("-" * 40) total_loss = 0 steps = 0 # Create progress bar pbar = tqdm(dataset, desc=f"Training") for i, item in enumerate(pbar): # Format training text prompt = item.get('prompt', '') response = item.get('response', '') if not prompt or not response: continue # Create training text (prompt + response) text = f"{prompt}\n{response}" batch = {'text': [text]} try: loss = self.train_step(batch) total_loss += loss steps += 1 # Update progress bar pbar.set_postfix({'loss': f'{loss:.4f}'}) # Save checkpoint every 100 steps if steps % 100 == 0: checkpoint_path = output_dir / f"checkpoint_step_{steps}.pt" torch.save({ 'step': steps, 'model_state': self.model.state_dict(), 'optimizer_state': self.optimizer.state_dict(), 'loss': loss }, checkpoint_path) print(f"\n Saved checkpoint: {checkpoint_path}") except Exception as e: # Don't print every error to avoid spam if steps < 5: print(f"\n Error on step {steps}: {e}") continue avg_loss = total_loss / steps if steps > 0 else 0 print(f"\nEpoch {epoch + 1} complete: Avg Loss = {avg_loss:.4f}") # Save epoch checkpoint epoch_path = output_dir / f"epoch_{epoch + 1}.pt" torch.save({ 'epoch': epoch + 1, 'model_state': self.model.state_dict(), 'optimizer_state': self.optimizer.state_dict(), 'avg_loss': avg_loss }, epoch_path) print(f"Saved epoch checkpoint: {epoch_path}") # Save final model final_path = output_dir / "shorekeeper-4b-final.pt" torch.save(self.model.state_dict(), final_path) print(f"\n{'='*60}") print(f"✅ Training complete! Final model saved to: {final_path}") print(f"{'='*60}") return self.model def load_data(data_path, limit=None): """Load training data from JSONL file""" data = [] data_path = Path(data_path) if not data_path.exists(): print(f"Data file not found: {data_path}") return data with open(data_path, 'r') as f: for i, line in enumerate(f): if limit and i >= limit: break try: item = json.loads(line) data.append(item) except: continue print(f"Loaded {len(data)} examples from {data_path}") return data def main(): print("=" * 60) print("SHOREKEEPER-4B Training Pipeline") print("=" * 60) # Check device if torch.cuda.is_available(): device = torch.device("cuda") print(f"\n✓ CUDA available: {torch.cuda.get_device_name(0)}") print(f" Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") else: device = torch.device("cpu") print("\n⚠ No GPU detected, using CPU (will be slow)") # Load model print("\n1. Loading SHOREKEEPER model...") model = MemoryEfficientSHOREKEEPER(use_4bit=False) # Use full precision for training model = model.to(device) print(f" Model loaded on {device}") # Load tokenizer print("\n2. Loading tokenizer...") try: tokenizer = AutoTokenizer.from_pretrained("gpt2") tokenizer.pad_token = tokenizer.eos_token print(" ✓ Using GPT-2 tokenizer") except: print(" ⚠ Could not load GPT-2 tokenizer") return # Load training data print("\n3. Loading training data...") data_path = Path("./data/processed/train.jsonl") if not data_path.exists(): print(f"\n❌ No training data found at {data_path}") print(" Run: python3 scripts/01_download_data.py") print(" Then: python3 scripts/02_prepare_data.py") return print("\n Training options:") print(" [1] Quick test (50 examples, 1 epoch) - ~2 minutes") print(" [2] Small training (200 examples, 3 epochs) - ~10 minutes") print(" [3] Medium training (500 examples, 5 epochs) - ~30 minutes") print(" [4] Full training (all data, 10 epochs) - several hours") choice = input("\nChoose option (1/2/3/4): ").strip() if choice == "1": limit = 50 epochs = 1 learning_rate = 1e-4 elif choice == "2": limit = 200 epochs = 3 learning_rate = 5e-5 elif choice == "3": limit = 500 epochs = 5 learning_rate = 3e-5 else: limit = None epochs = 10 learning_rate = 1e-5 # Load data data = load_data(data_path, limit=limit) if not data: print("\n❌ No training data available!") return print(f"\n Training with {len(data)} examples, {epochs} epochs") print(f" Learning rate: {learning_rate}") # Training config config = { 'learning_rate': learning_rate, 'epochs': epochs, 'batch_size': 2, 'gradient_accumulation': 4 } # Create trainer print("\n4. Initializing trainer...") trainer = SHOREKEEPERTrainer(model, tokenizer, config) # Start training print("\n5. Starting training...") print(" Press Ctrl+C to stop early\n") try: trained_model = trainer.train(data, output_dir="./outputs") except KeyboardInterrupt: print("\n\n⚠ Training interrupted by user") print("Saving current model...") torch.save(model.state_dict(), "./outputs/shorekeeper-interrupted.pt") print("Model saved to: ./outputs/shorekeeper-interrupted.pt") except Exception as e: print(f"\n❌ Training failed: {e}") import traceback traceback.print_exc() print("\n" + "=" * 60) print("Next steps:") print(" 1. Run GRPO training: python3 scripts/05_grpo_train.py") print(" 2. Convert to 4-bit: python3 scripts/06_convert_to_4bit.py") print(" 3. Run SHOREKEEPER: python3 scripts/07_run_shorekeeper.py") print("=" * 60) if __name__ == "__main__": main()