gptmed 0.0.1__py3-none-any.whl

gptmed/training/utils.py

@@ -0,0 +1,212 @@
+"""
+Training Utilities
+
+PURPOSE:
+Helper functions for training loop:
+- Gradient clipping
+- Learning rate scheduling
+- Training state management
+
+WHAT THIS FILE DOES:
+1. Clip gradients to prevent explosion
+2. Calculate learning rate with warmup + decay
+3. Compute gradient norms for monitoring
+
+WHY THESE ARE CRITICAL:
+- Gradient clipping: Prevents training collapse from exploding gradients
+- LR warmup: Stabilizes early training (large steps can destabilize)
+- LR decay: Helps model converge to better minima
+
+PACKAGES USED:
+- torch: Gradient operations
+- math: Cosine calculations
+
+FILES FROM THIS PROJECT:
+- None (utility functions)
+
+COMMON FAILURE MODES:
+- No gradient clipping → NaN loss
+- No warmup → unstable first epochs
+- Constant LR → suboptimal convergence
+"""
+
+import torch
+import torch.nn as nn
+import math
+from typing import Optional
+
+
+def clip_grad_norm(model: nn.Module, max_norm: float) -> float:
+    """
+    Clip gradient norms to prevent explosion.
+
+    Args:
+        model: Model with gradients
+        max_norm: Maximum gradient norm
+
+    Returns:
+        Total gradient norm (before clipping)
+
+    How it works:
+    - Compute total gradient norm across all parameters
+    - If norm > max_norm, scale all gradients down
+    - This prevents single large gradients from destroying training
+
+    Typical values:
+    - max_norm=1.0: Standard for transformers
+    - max_norm=5.0: More lenient
+    - max_norm=0.5: Very conservative
+    """
+    total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
+    return total_norm.item()
+
+
+def get_lr_with_warmup(
+    step: int,
+    warmup_steps: int,
+    max_lr: float,
+    min_lr: float,
+    max_steps: int,
+    decay_type: str = "cosine",
+) -> float:
+    """
+    Calculate learning rate with warmup and decay.
+
+    Args:
+        step: Current training step
+        warmup_steps: Number of warmup steps
+        max_lr: Peak learning rate
+        min_lr: Minimum learning rate
+        max_steps: Total training steps
+        decay_type: 'cosine', 'linear', or 'constant'
+
+    Returns:
+        Learning rate for this step
+
+    Schedule:
+    1. Warmup (0 to warmup_steps): Linear increase from 0 to max_lr
+    2. Decay (warmup_steps to max_steps): Cosine/linear decay to min_lr
+
+    Why warmup?
+    - Random initialization → large gradients early on
+    - Large LR + large gradients = explosion
+    - Warmup gives model time to stabilize
+    """
+    # Warmup phase
+    if step < warmup_steps:
+        # Linear warmup from 0 to max_lr
+        return max_lr * (step / warmup_steps)
+
+    # After warmup
+    if decay_type == "constant":
+        return max_lr
+
+    # Progress through decay phase
+    progress = (step - warmup_steps) / (max_steps - warmup_steps)
+    progress = min(progress, 1.0)  # Cap at 1.0
+
+    if decay_type == "cosine":
+        # Cosine decay: smooth curve to min_lr
+        lr = min_lr + (max_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))
+    elif decay_type == "linear":
+        # Linear decay: straight line to min_lr
+        lr = max_lr - (max_lr - min_lr) * progress
+    else:
+        raise ValueError(f"Unknown decay_type: {decay_type}")
+
+    return lr
+
+
+def set_learning_rate(optimizer: torch.optim.Optimizer, lr: float):
+    """
+    Set learning rate for optimizer.
+
+    Args:
+        optimizer: PyTorch optimizer
+        lr: New learning rate
+    """
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = lr
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Count trainable parameters."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def estimate_loss_single_batch(model: nn.Module, batch: tuple, device: str) -> float:
+    """
+    Compute loss on a single batch (for evaluation).
+
+    Args:
+        model: Model to evaluate
+        batch: (input_ids, target_ids) batch
+        device: Device to run on
+
+    Returns:
+        Loss value
+    """
+    input_ids, target_ids = batch
+    input_ids = input_ids.to(device)
+    target_ids = target_ids.to(device)
+
+    # Forward pass
+    logits = model(input_ids)
+
+    # Compute loss
+    # Reshape for CrossEntropyLoss:
+    # logits: [batch_size * seq_len, vocab_size]
+    # targets: [batch_size * seq_len]
+    batch_size, seq_len, vocab_size = logits.shape
+    logits = logits.view(batch_size * seq_len, vocab_size)
+    targets = target_ids.view(batch_size * seq_len)
+
+    loss = nn.functional.cross_entropy(logits, targets)
+
+    return loss.item()
+
+
+def estimate_loss_dataloader(
+    model: nn.Module, dataloader, device: str, max_batches: Optional[int] = None
+) -> float:
+    """
+    Estimate average loss over a dataloader.
+
+    Args:
+        model: Model to evaluate
+        dataloader: DataLoader to evaluate on
+        device: Device to run on
+        max_batches: Maximum number of batches to evaluate (None = all)
+
+    Returns:
+        Average loss
+    """
+    model.eval()
+    total_loss = 0.0
+    num_batches = 0
+
+    with torch.no_grad():
+        for i, batch in enumerate(dataloader):
+            if max_batches is not None and i >= max_batches:
+                break
+
+            loss = estimate_loss_single_batch(model, batch, device)
+            total_loss += loss
+            num_batches += 1
+
+    model.train()
+
+    return total_loss / num_batches if num_batches > 0 else 0.0
+
+
+def compute_perplexity(loss: float) -> float:
+    """
+    Compute perplexity from loss.
+
+    Perplexity = exp(loss)
+
+    Lower is better. Perplexity measures how "surprised" the model is.
+    - Perplexity of 10 = model is choosing between ~10 likely tokens
+    - Perplexity of 100 = model is very uncertain
+    """
+    return math.exp(loss)

gptmed/utils/__init__.py

@@ -0,0 +1 @@
+"""Utils package for training utilities."""

gptmed/utils/checkpoints.py

@@ -0,0 +1,224 @@
+"""
+Model Checkpointing Utilities
+
+PURPOSE:
+Save and load model checkpoints during training. This is essential for:
+- Resuming interrupted training
+- Saving best model based on validation loss
+- Preventing loss of work from crashes
+
+WHAT THIS FILE DOES:
+1. Save model state_dict + optimizer + training state
+2. Load checkpoints to resume training
+3. Manage checkpoint files (keep only best/recent)
+4. Save configuration alongside weights
+
+PACKAGES USED:
+- torch: Save/load model state
+- pathlib: File management
+- json: Save metadata
+
+FILES FROM THIS PROJECT:
+- model/checkpoints/ (checkpoint directory)
+
+CHECKPOINT CONTENTS:
+- model_state_dict: Model weights
+- optimizer_state_dict: Optimizer state (for resuming)
+- step: Current training step
+- epoch: Current epoch
+- best_val_loss: Best validation loss so far
+- config: Model and training config
+
+COMMON ISSUES:
+- Not saving optimizer → can't resume training properly
+- Not saving RNG state → non-reproducible results
+- Checkpoints too large → disk space issues
+- Not testing loading → corrupt checkpoints discovered too late
+"""
+
+import torch
+from pathlib import Path
+import json
+from typing import Dict, Optional
+import shutil
+
+
+class CheckpointManager:
+    """
+    Manages model checkpoints during training.
+
+    Handles saving, loading, and cleanup of checkpoint files.
+    """
+
+    def __init__(self, checkpoint_dir: Path, keep_last_n: int = 3):
+        """
+        Args:
+            checkpoint_dir: Directory to save checkpoints
+            keep_last_n: Number of recent checkpoints to keep
+        """
+        self.checkpoint_dir = Path(checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        self.keep_last_n = keep_last_n
+
+        # Track best validation loss
+        self.best_val_loss = float("inf")
+
+        print(f"Checkpoint directory: {self.checkpoint_dir}")
+
+    def save_checkpoint(
+        self,
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        step: int,
+        epoch: int,
+        val_loss: float,
+        model_config: dict,
+        train_config: dict,
+        is_best: bool = False,
+    ):
+        """
+        Save a checkpoint.
+
+        Args:
+            model: Model to save
+            optimizer: Optimizer to save
+            step: Current training step
+            epoch: Current epoch
+            val_loss: Validation loss
+            model_config: Model configuration
+            train_config: Training configuration
+            is_best: Whether this is the best model so far
+        """
+        checkpoint = {
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "step": step,
+            "epoch": epoch,
+            "val_loss": val_loss,
+            "best_val_loss": self.best_val_loss,
+            "model_config": model_config,
+            "train_config": train_config,
+        }
+
+        # Save regular checkpoint
+        checkpoint_path = self.checkpoint_dir / f"checkpoint_step_{step}.pt"
+        torch.save(checkpoint, checkpoint_path)
+        print(f"Checkpoint saved: {checkpoint_path}")
+
+        # Save as best if applicable
+        if is_best or val_loss < self.best_val_loss:
+            self.best_val_loss = val_loss
+            best_path = self.checkpoint_dir / "best_model.pt"
+            torch.save(checkpoint, best_path)
+            print(f"Best model saved: {best_path} (val_loss: {val_loss:.4f})")
+
+        # Save as latest (for easy resuming)
+        latest_path = self.checkpoint_dir / "latest_checkpoint.pt"
+        torch.save(checkpoint, latest_path)
+
+        # Cleanup old checkpoints
+        self._cleanup_old_checkpoints()
+
+    def _cleanup_old_checkpoints(self):
+        """Remove old checkpoints, keeping only last N."""
+        # Get all checkpoint files (except best and latest)
+        checkpoints = sorted(
+            self.checkpoint_dir.glob("checkpoint_step_*.pt"),
+            key=lambda p: int(p.stem.split("_")[-1]),
+        )
+
+        # Remove old ones
+        if len(checkpoints) > self.keep_last_n:
+            for ckpt in checkpoints[: -self.keep_last_n]:
+                ckpt.unlink()
+                print(f"Removed old checkpoint: {ckpt.name}")
+
+    def load_checkpoint(
+        self,
+        model: torch.nn.Module,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+        checkpoint_path: Optional[Path] = None,
+        device: str = "cuda",
+    ) -> Dict:
+        """
+        Load a checkpoint.
+
+        Args:
+            model: Model to load weights into
+            optimizer: Optimizer to load state into (optional)
+            checkpoint_path: Specific checkpoint to load (or None for latest)
+            device: Device to load to
+
+        Returns:
+            Checkpoint dictionary with metadata
+        """
+        # Use latest if no path specified
+        if checkpoint_path is None:
+            checkpoint_path = self.checkpoint_dir / "latest_checkpoint.pt"
+        else:
+            checkpoint_path = Path(checkpoint_path)
+
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+        print(f"Loading checkpoint: {checkpoint_path}")
+
+        # Load checkpoint
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+
+        # Load model weights
+        model.load_state_dict(checkpoint["model_state_dict"])
+
+        # Load optimizer state if provided
+        if optimizer is not None and "optimizer_state_dict" in checkpoint:
+            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+
+        print(f"Loaded checkpoint from step {checkpoint['step']}, epoch {checkpoint['epoch']}")
+        print(f"Validation loss: {checkpoint['val_loss']:.4f}")
+
+        return checkpoint
+
+    def has_checkpoint(self) -> bool:
+        """Check if a checkpoint exists."""
+        latest_path = self.checkpoint_dir / "latest_checkpoint.pt"
+        return latest_path.exists()
+
+
+def save_model_for_inference(
+    model: torch.nn.Module, tokenizer_path: Path, save_path: Path, model_config: dict
+):
+    """
+    Save model for inference (weights only, no optimizer).
+
+    Args:
+        model: Trained model
+        tokenizer_path: Path to tokenizer model
+        save_path: Where to save
+        model_config: Model configuration
+    """
+    save_dict = {
+        "model_state_dict": model.state_dict(),
+        "model_config": model_config,
+        "tokenizer_path": str(tokenizer_path),
+    }
+
+    torch.save(save_dict, save_path)
+    print(f"Model saved for inference: {save_path}")
+
+
+def load_model_for_inference(model: torch.nn.Module, checkpoint_path: Path, device: str = "cuda"):
+    """
+    Load model for inference.
+
+    Args:
+        model: Model architecture (will be populated with weights)
+        checkpoint_path: Path to saved model
+        device: Device to load to
+    """
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    print(f"Model loaded for inference from: {checkpoint_path}")
+
+    return checkpoint.get("model_config")

gptmed/utils/logging.py

@@ -0,0 +1,189 @@
+"""
+Training Logging Utilities
+
+PURPOSE:
+Track and log training metrics (loss, learning rate, gradient norms).
+This is CRITICAL for debugging training issues.
+
+WHAT THIS FILE DOES:
+1. Log training metrics to console
+2. Track loss history
+3. Compute moving averages
+4. Save metrics to file
+
+WHY LOGGING IS CRITICAL:
+- Detect gradient explosions (loss → NaN)
+- Spot overfitting (train loss ↓, val loss ↑)
+- Monitor learning rate schedule
+- Debug slow convergence
+
+PACKAGES USED:
+- json: Save metrics
+- time: Track training speed
+- statistics: Compute averages
+
+FILES FROM THIS PROJECT:
+- None (utility module)
+
+COMMON ISSUES TO DETECT:
+- Loss = NaN → exploding gradients, LR too high
+- Loss stuck → LR too low, bad initialization
+- Loss oscillating → LR too high, reduce it
+- Val loss increasing → overfitting, add regularization
+"""
+
+import json
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+from collections import defaultdict
+
+
+class MetricsLogger:
+    """
+    Simple metrics logger for training.
+
+    Tracks loss, learning rate, and other metrics over time.
+    """
+
+    def __init__(self, log_dir: Path, experiment_name: str = "training"):
+        """
+        Args:
+            log_dir: Directory to save logs
+            experiment_name: Name for this training run
+        """
+        self.log_dir = Path(log_dir)
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+
+        self.experiment_name = experiment_name
+        self.log_file = self.log_dir / f"{experiment_name}_metrics.jsonl"
+
+        # Metrics storage
+        self.metrics = defaultdict(list)
+
+        # Timing
+        self.start_time = time.time()
+        self.step_times = []
+
+        print(f"Logging to: {self.log_file}")
+
+    def log(self, step: int, metrics: Dict[str, float]):
+        """
+        Log metrics for a training step.
+
+        Args:
+            step: Training step number
+            metrics: Dictionary of metric_name -> value
+        """
+        # Add timestamp and step
+        log_entry = {"step": step, "timestamp": time.time() - self.start_time, **metrics}
+
+        # Store in memory
+        for key, value in metrics.items():
+            self.metrics[key].append(value)
+
+        # Append to file (JSONL format)
+        with open(self.log_file, "a") as f:
+            f.write(json.dumps(log_entry) + "\n")
+
+    def log_epoch(self, epoch: int, train_loss: float, val_loss: float):
+        """Log epoch-level metrics."""
+        print(f"\nEpoch {epoch}:")
+        print(f"  Train loss: {train_loss:.4f}")
+        print(f"  Val loss: {val_loss:.4f}")
+
+        self.log(epoch, {"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss})
+
+    def get_last(self, metric_name: str, n: int = 1) -> Optional[float]:
+        """Get last n values of a metric."""
+        if metric_name not in self.metrics:
+            return None
+        values = self.metrics[metric_name]
+        if len(values) < n:
+            return None
+        return sum(values[-n:]) / n
+
+    def get_average(self, metric_name: str, window: int = 100) -> Optional[float]:
+        """Get moving average of a metric."""
+        if metric_name not in self.metrics:
+            return None
+        values = self.metrics[metric_name]
+        if len(values) == 0:
+            return None
+        window = min(window, len(values))
+        return sum(values[-window:]) / window
+
+
+def log_training_step(
+    step: int,
+    loss: float,
+    lr: float,
+    grad_norm: float,
+    tokens_per_sec: float,
+    print_output: bool = True,
+):
+    """
+    Log a single training step to console.
+
+    Args:
+        step: Training step
+        loss: Current loss
+        lr: Current learning rate
+        grad_norm: Gradient norm (before clipping)
+        tokens_per_sec: Processing speed
+        print_output: Whether to print to console
+    """
+    if not print_output:
+        return
+
+    # Format output
+    msg = f"Step {step:5d} | "
+    msg += f"Loss: {loss:.4f} | "
+    msg += f"LR: {lr:.2e} | "
+    msg += f"Grad: {grad_norm:.3f} | "
+    msg += f"Tok/s: {tokens_per_sec:.0f}"
+
+    # Check for issues
+    if loss != loss:  # NaN check
+        msg += " [WARNING: NaN loss!]"
+    elif loss > 100:
+        msg += " [WARNING: High loss!]"
+
+    if grad_norm > 10:
+        msg += " [WARNING: Large gradients!]"
+
+    print(msg)
+
+
+def log_validation(step: int, val_loss: float, val_perplexity: float):
+    """
+    Log validation metrics.
+
+    Args:
+        step: Training step
+        val_loss: Validation loss
+        val_perplexity: Validation perplexity (exp(loss))
+    """
+    print(f"\n{'='*60}")
+    print(f"Validation at step {step}")
+    print(f"  Loss: {val_loss:.4f}")
+    print(f"  Perplexity: {val_perplexity:.2f}")
+    print(f"{'='*60}\n")
+
+
+def save_training_summary(log_dir: Path, config: dict, final_metrics: dict):
+    """
+    Save final training summary.
+
+    Args:
+        log_dir: Directory to save summary
+        config: Training configuration
+        final_metrics: Final metrics (best val loss, etc.)
+    """
+    summary = {"config": config, "final_metrics": final_metrics, "timestamp": time.time()}
+
+    summary_file = log_dir / "training_summary.json"
+    with open(summary_file, "w") as f:
+        json.dump(summary, f, indent=2)
+
+    print(f"\nTraining summary saved to: {summary_file}")