gptmed 0.0.1__py3-none-any.whl

gptmed/training/train.py

@@ -0,0 +1,272 @@
+"""
+Main Training Script
+
+PURPOSE:
+Entry point for training the GPT model.
+Ties everything together and starts training.
+
+WHAT THIS FILE DOES:
+1. Load configuration (model + training)
+2. Create model
+3. Load tokenized data
+4. Create optimizer
+5. Start training
+6. Handle command-line arguments
+
+USAGE:
+    python training/train.py                    # Use default config
+    python training/train.py --batch-size 32    # Override batch size
+    python training/train.py --resume          # Resume from checkpoint
+
+PACKAGES USED:
+- torch: PyTorch
+- argparse: Command-line arguments
+- pathlib: Path handling
+
+FILES FROM THIS PROJECT:
+- All model, training, and utility modules
+
+EXECUTION ORDER:
+1. Parse arguments
+2. Set random seeds (reproducibility)
+3. Create model
+4. Load data
+5. Create optimizer
+6. Initialize trainer
+7. Start training
+"""
+
+import torch
+import argparse
+from pathlib import Path
+import random
+import numpy as np
+import sys
+
+# Add parent directory to path for imports
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from llm_med.model.architecture import GPTTransformer
+from llm_med.model.configs.model_config import get_small_config, get_tiny_config
+from llm_med.configs.train_config import get_default_config, get_quick_test_config
+from llm_med.training.dataset import create_dataloaders
+from llm_med.training.trainer import Trainer
+
+
+def set_seed(seed: int):
+    """
+    Set random seeds for reproducibility.
+
+    Args:
+        seed: Random seed
+
+    Why this matters:
+    - Makes training reproducible
+    - Critical for debugging (can recreate issues)
+    - Scientific experiments need reproducibility
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+    # Make cudnn deterministic (slower but reproducible)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def count_parameters(model):
+    """Count trainable parameters."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train GPT model on MedQuAD")
+
+    # Model config
+    parser.add_argument(
+        "--model-size",
+        type=str,
+        default="small",
+        choices=["tiny", "small", "medium"],
+        help="Model size (tiny/small/medium)",
+    )
+
+    # Training config
+    parser.add_argument(
+        "--batch-size", type=int, default=None, help="Batch size (overrides config)"
+    )
+    parser.add_argument(
+        "--learning-rate", type=float, default=None, help="Learning rate (overrides config)"
+    )
+    parser.add_argument(
+        "--num-epochs", type=int, default=None, help="Number of epochs (overrides config)"
+    )
+    parser.add_argument(
+        "--quick-test", action="store_true", help="Use quick test config (small batches, few steps)"
+    )
+
+    # Paths
+    parser.add_argument(
+        "--train-data", type=str, default="./data/tokenized/train.npy", help="Path to training data"
+    )
+    parser.add_argument(
+        "--val-data", type=str, default="./data/tokenized/val.npy", help="Path to validation data"
+    )
+    parser.add_argument(
+        "--checkpoint-dir", type=str, default="./model/checkpoints", help="Checkpoint directory"
+    )
+
+    # Resume training
+    parser.add_argument("--resume", action="store_true", help="Resume from latest checkpoint")
+    parser.add_argument(
+        "--resume-from", type=str, default=None, help="Resume from specific checkpoint"
+    )
+
+    # Device
+    parser.add_argument(
+        "--device", type=str, default="cuda", choices=["cuda", "cpu"], help="Device to train on"
+    )
+
+    # Misc
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("GPT Training - MedQuAD")
+    print("=" * 60)
+
+    # Check CUDA availability
+    if args.device == "cuda" and not torch.cuda.is_available():
+        print("WARNING: CUDA not available, using CPU")
+        args.device = "cpu"
+
+    # Set random seed
+    print(f"\nSetting random seed: {args.seed}")
+    set_seed(args.seed)
+
+    # Load configurations
+    print(f"\nLoading configurations...")
+
+    # Model config
+    if args.model_size == "tiny":
+        model_config = get_tiny_config()
+    elif args.model_size == "small":
+        model_config = get_small_config()
+    else:
+        raise ValueError(f"Unknown model size: {args.model_size}")
+
+    print(f"Model config: {args.model_size}")
+    print(f"  d_model: {model_config.d_model}")
+    print(f"  n_layers: {model_config.n_layers}")
+    print(f"  n_heads: {model_config.n_heads}")
+
+    # Training config
+    if args.quick_test:
+        train_config = get_quick_test_config()
+        print("Using quick test config (fast debugging)")
+    else:
+        train_config = get_default_config()
+
+    # Override with command-line args
+    if args.batch_size is not None:
+        train_config.batch_size = args.batch_size
+    if args.learning_rate is not None:
+        train_config.learning_rate = args.learning_rate
+    if args.num_epochs is not None:
+        train_config.num_epochs = args.num_epochs
+    if args.train_data:
+        train_config.train_data_path = args.train_data
+    if args.val_data:
+        train_config.val_data_path = args.val_data
+    if args.checkpoint_dir:
+        train_config.checkpoint_dir = args.checkpoint_dir
+
+    train_config.device = args.device
+    train_config.seed = args.seed
+
+    print(f"\nTraining config:")
+    print(f"  Batch size: {train_config.batch_size}")
+    print(f"  Learning rate: {train_config.learning_rate}")
+    print(f"  Num epochs: {train_config.num_epochs}")
+    print(f"  Device: {train_config.device}")
+
+    # Create model
+    print(f"\nCreating model...")
+    model = GPTTransformer(model_config)
+    total_params = count_parameters(model)
+    print(f"Model created with {total_params:,} parameters")
+    print(f"Model size: ~{total_params * 4 / 1024 / 1024:.2f} MB")
+
+    # Load data
+    print(f"\nLoading data...")
+    train_loader, val_loader = create_dataloaders(
+        train_path=Path(train_config.train_data_path),
+        val_path=Path(train_config.val_data_path),
+        batch_size=train_config.batch_size,
+        num_workers=0,
+    )
+
+    # Create optimizer
+    print(f"\nCreating optimizer...")
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=train_config.learning_rate,
+        betas=train_config.betas,
+        eps=train_config.eps,
+        weight_decay=train_config.weight_decay,
+    )
+    print(f"Optimizer: AdamW")
+    print(f"  LR: {train_config.learning_rate}")
+    print(f"  Weight decay: {train_config.weight_decay}")
+
+    # Create trainer
+    print(f"\nInitializing trainer...")
+    trainer = Trainer(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        optimizer=optimizer,
+        config=train_config,
+        device=args.device,
+    )
+
+    # Resume if requested
+    if args.resume or args.resume_from:
+        print(f"\nResuming training...")
+        checkpoint_path = Path(args.resume_from) if args.resume_from else None
+        trainer.resume_from_checkpoint(checkpoint_path)
+
+    # Start training
+    print(f"\n{'='*60}")
+    print("Ready to train!")
+    print(f"{'='*60}\n")
+
+    try:
+        trainer.train()
+    except KeyboardInterrupt:
+        print("\n\nTraining interrupted by user")
+        print("Saving checkpoint...")
+        trainer.checkpoint_manager.save_checkpoint(
+            model=model,
+            optimizer=optimizer,
+            step=trainer.global_step,
+            epoch=trainer.current_epoch,
+            val_loss=trainer.best_val_loss,
+            model_config=model_config.to_dict(),
+            train_config=train_config.to_dict(),
+        )
+        print("Checkpoint saved. You can resume with --resume")
+
+    print("\n" + "=" * 60)
+    print("Training finished!")
+    print("=" * 60)
+    print(f"\nBest model saved in: {train_config.checkpoint_dir}/best_model.pt")
+    print(f"Logs saved in: {train_config.log_dir}")
+
+
+if __name__ == "__main__":
+    main()

gptmed/training/trainer.py

@@ -0,0 +1,331 @@
+"""
+Trainer Class
+
+PURPOSE:
+Core training loop logic. This is the heart of Phase 3.
+Handles forward pass, backward pass, optimization, and evaluation.
+
+WHAT THIS FILE DOES:
+1. Training loop: iterate over batches, compute loss, backprop
+2. Evaluation: compute validation loss
+3. Checkpointing: save model periodically
+4. Logging: track metrics
+
+TRAINING ALGORITHM:
+For each batch:
+  1. Forward pass: model(input) → logits
+  2. Compute loss: CrossEntropyLoss(logits, targets)
+  3. Backward pass: loss.backward()
+  4. Clip gradients: prevent explosion
+  5. Optimizer step: update weights
+  6. Update learning rate: warmup + decay
+
+PACKAGES USED:
+- torch: PyTorch training
+- time: Track speed
+
+FILES FROM THIS PROJECT:
+- model/architecture/transformer.py: Model
+- training/dataset.py: DataLoader
+- training/utils.py: Helper functions
+- utils/logging.py: Metrics logging
+- utils/checkpoints.py: Save/load
+
+TENSOR SHAPES:
+- Input: [batch_size, seq_len]
+- Logits: [batch_size, seq_len, vocab_size]
+- Targets: [batch_size, seq_len]
+- Loss: scalar
+
+COMMON TRAINING ISSUES:
+- Loss = NaN → gradient explosion (reduce LR, check grad clipping)
+- Loss stuck → LR too low, bad initialization
+- Slow convergence → LR too low, increase it
+- Overfitting → add dropout, weight decay
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import time
+from pathlib import Path
+from typing import Optional
+
+from llm_med.model.architecture import GPTTransformer
+from llm_med.training.utils import (
+    clip_grad_norm,
+    get_lr_with_warmup,
+    set_learning_rate,
+    estimate_loss_dataloader,
+    compute_perplexity,
+)
+from llm_med.utils.logging import MetricsLogger, log_training_step, log_validation
+from llm_med.utils.checkpoints import CheckpointManager
+
+
+class Trainer:
+    """
+    Training orchestrator for GPT model.
+
+    Handles the full training loop including:
+    - Forward/backward passes
+    - Optimization
+    - Evaluation
+    - Checkpointing
+    - Logging
+    """
+
+    def __init__(
+        self,
+        model: GPTTransformer,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        optimizer: torch.optim.Optimizer,
+        config,  # TrainingConfig
+        device: str = "cuda",
+    ):
+        """
+        Args:
+            model: GPT model to train
+            train_loader: Training data loader
+            val_loader: Validation data loader
+            optimizer: Optimizer (e.g., AdamW)
+            config: TrainingConfig object
+            device: Device to train on
+        """
+        self.model = model.to(device)
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.optimizer = optimizer
+        self.config = config
+        self.device = device
+
+        # Initialize utilities
+        self.logger = MetricsLogger(log_dir=config.log_dir, experiment_name="gpt_training")
+
+        self.checkpoint_manager = CheckpointManager(
+            checkpoint_dir=config.checkpoint_dir, keep_last_n=config.keep_last_n
+        )
+
+        # Training state
+        self.global_step = 0
+        self.current_epoch = 0
+        self.best_val_loss = float("inf")
+
+        # Calculate total steps
+        steps_per_epoch = len(train_loader)
+        if config.max_steps > 0:
+            self.total_steps = config.max_steps
+        else:
+            self.total_steps = steps_per_epoch * config.num_epochs
+
+        print(f"\nTrainer initialized:")
+        print(f"  Device: {device}")
+        print(f"  Total steps: {self.total_steps}")
+        print(f"  Steps per epoch: {steps_per_epoch}")
+        print(f"  Num epochs: {config.num_epochs}")
+
+    def train_step(self, batch: tuple) -> dict:
+        """
+        Single training step.
+
+        Args:
+            batch: (input_ids, target_ids) tuple
+
+        Returns:
+            Dictionary with step metrics
+        """
+        # Move batch to device
+        input_ids, target_ids = batch
+        input_ids = input_ids.to(self.device)
+        target_ids = target_ids.to(self.device)
+
+        # Forward pass
+        logits = self.model(input_ids)
+
+        # Compute loss
+        # CrossEntropyLoss expects:
+        #   - Input: [N, C] where N = batch_size * seq_len, C = vocab_size
+        #   - Target: [N] with class indices
+        batch_size, seq_len, vocab_size = logits.shape
+        logits_flat = logits.view(batch_size * seq_len, vocab_size)
+        targets_flat = target_ids.view(batch_size * seq_len)
+
+        loss = nn.functional.cross_entropy(logits_flat, targets_flat)
+
+        # Backward pass
+        self.optimizer.zero_grad()
+        loss.backward()
+
+        # Clip gradients (CRITICAL for stability)
+        grad_norm = clip_grad_norm(self.model, self.config.grad_clip)
+
+        # Optimizer step
+        self.optimizer.step()
+
+        # Return metrics
+        return {
+            "loss": loss.item(),
+            "grad_norm": grad_norm,
+            "batch_size": batch_size,
+            "seq_len": seq_len,
+        }
+
+    def evaluate(self) -> dict:
+        """
+        Evaluate on validation set.
+
+        Returns:
+            Dictionary with validation metrics
+        """
+        print("\nRunning validation...")
+
+        val_loss = estimate_loss_dataloader(
+            self.model, self.val_loader, self.device, max_batches=self.config.eval_iters
+        )
+
+        val_perplexity = compute_perplexity(val_loss)
+
+        log_validation(self.global_step, val_loss, val_perplexity)
+
+        return {"val_loss": val_loss, "val_perplexity": val_perplexity}
+
+    def train(self):
+        """
+        Main training loop.
+
+        This is where everything comes together.
+        """
+        print("\n" + "=" * 60)
+        print("Starting Training")
+        print("=" * 60)
+
+        self.model.train()
+
+        # Training loop
+        for epoch in range(self.config.num_epochs):
+            self.current_epoch = epoch
+
+            print(f"\n{'='*60}")
+            print(f"Epoch {epoch + 1}/{self.config.num_epochs}")
+            print(f"{'='*60}")
+
+            epoch_start_time = time.time()
+
+            for batch_idx, batch in enumerate(self.train_loader):
+                step_start_time = time.time()
+
+                # Update learning rate (warmup + decay)
+                lr = get_lr_with_warmup(
+                    step=self.global_step,
+                    warmup_steps=self.config.warmup_steps,
+                    max_lr=self.config.learning_rate,
+                    min_lr=self.config.min_lr,
+                    max_steps=self.total_steps,
+                    decay_type=self.config.lr_decay,
+                )
+                set_learning_rate(self.optimizer, lr)
+
+                # Training step
+                metrics = self.train_step(batch)
+
+                # Calculate tokens per second
+                step_time = time.time() - step_start_time
+                tokens_per_sec = (metrics["batch_size"] * metrics["seq_len"]) / step_time
+
+                # Log to console
+                if self.global_step % self.config.log_interval == 0:
+                    log_training_step(
+                        step=self.global_step,
+                        loss=metrics["loss"],
+                        lr=lr,
+                        grad_norm=metrics["grad_norm"],
+                        tokens_per_sec=tokens_per_sec,
+                    )
+
+                # Log metrics
+                self.logger.log(
+                    self.global_step,
+                    {
+                        "train_loss": metrics["loss"],
+                        "learning_rate": lr,
+                        "grad_norm": metrics["grad_norm"],
+                        "tokens_per_sec": tokens_per_sec,
+                    },
+                )
+
+                # Evaluate
+                if self.global_step % self.config.eval_interval == 0 and self.global_step > 0:
+                    val_metrics = self.evaluate()
+
+                    # Log validation metrics
+                    self.logger.log(self.global_step, val_metrics)
+
+                    # Check if best model
+                    if val_metrics["val_loss"] < self.best_val_loss:
+                        self.best_val_loss = val_metrics["val_loss"]
+                        is_best = True
+                    else:
+                        is_best = False
+
+                    # Save checkpoint
+                    self.checkpoint_manager.save_checkpoint(
+                        model=self.model,
+                        optimizer=self.optimizer,
+                        step=self.global_step,
+                        epoch=epoch,
+                        val_loss=val_metrics["val_loss"],
+                        model_config=self.model.config.to_dict(),
+                        train_config=self.config.to_dict(),
+                        is_best=is_best,
+                    )
+
+                    self.model.train()  # Back to training mode
+
+                # Save checkpoint periodically
+                if self.global_step % self.config.save_interval == 0 and self.global_step > 0:
+                    self.checkpoint_manager.save_checkpoint(
+                        model=self.model,
+                        optimizer=self.optimizer,
+                        step=self.global_step,
+                        epoch=epoch,
+                        val_loss=self.best_val_loss,
+                        model_config=self.model.config.to_dict(),
+                        train_config=self.config.to_dict(),
+                    )
+
+                self.global_step += 1
+
+                # Check if reached max steps
+                if self.config.max_steps > 0 and self.global_step >= self.config.max_steps:
+                    print(f"\nReached max_steps ({self.config.max_steps}). Stopping training.")
+                    return
+
+            # End of epoch
+            epoch_time = time.time() - epoch_start_time
+            print(f"\nEpoch {epoch + 1} completed in {epoch_time:.2f}s")
+
+        print("\n" + "=" * 60)
+        print("Training Complete!")
+        print("=" * 60)
+        print(f"Best validation loss: {self.best_val_loss:.4f}")
+
+    def resume_from_checkpoint(self, checkpoint_path: Optional[Path] = None):
+        """
+        Resume training from a checkpoint.
+
+        Args:
+            checkpoint_path: Path to checkpoint (or None for latest)
+        """
+        checkpoint = self.checkpoint_manager.load_checkpoint(
+            model=self.model,
+            optimizer=self.optimizer,
+            checkpoint_path=checkpoint_path,
+            device=self.device,
+        )
+
+        self.global_step = checkpoint["step"]
+        self.current_epoch = checkpoint["epoch"]
+        self.best_val_loss = checkpoint.get("best_val_loss", float("inf"))
+
+        print(f"Resumed from step {self.global_step}, epoch {self.current_epoch}")