""" Training loop for SLM. Handles the complete training process including: - Mixed precision training - Gradient accumulation - Checkpointing - Logging """ import os import time import json from dataclasses import dataclass, asdict from typing import Optional, Dict, Any, Callable from pathlib import Path import torch import torch.nn as nn from torch.utils.data import DataLoader from torch.cuda.amp import autocast, GradScaler from tqdm import tqdm from .loss import LanguageModelingLoss, compute_perplexity, compute_accuracy from .optimizer import create_optimizer, create_scheduler, clip_grad_norm @dataclass class TrainingConfig: """Configuration for training.""" # Optimization learning_rate: float = 3e-4 weight_decay: float = 0.1 warmup_ratio: float = 0.1 min_lr_ratio: float = 0.1 max_grad_norm: float = 1.0 label_smoothing: float = 0.0 # Training num_epochs: int = 5 gradient_accumulation_steps: int = 4 fp16: bool = True # Checkpointing checkpoint_dir: str = "checkpoints" save_steps: int = 1000 save_total_limit: int = 3 # Evaluation eval_steps: int = 500 logging_steps: int = 10 # Early stopping early_stopping_patience: int = 5 # Stop after N evals without improvement early_stopping_threshold: float = 0.01 # Minimum improvement to reset patience # Device device: str = "auto" # Compile model (torch.compile) compile_model: bool = False def to_dict(self) -> Dict[str, Any]: return asdict(self) class Trainer: """Training loop for SLM model.""" def __init__( self, model: nn.Module, config: TrainingConfig, train_dataloader: DataLoader, val_dataloader: Optional[DataLoader] = None, wandb_project: Optional[str] = None, ): """Initialize trainer. Args: model: The model to train config: Training configuration train_dataloader: Training data loader val_dataloader: Optional validation data loader wandb_project: Optional W&B project name for logging """ self.config = config self.train_dataloader = train_dataloader self.val_dataloader = val_dataloader # Setup device if config.device == "auto": if torch.cuda.is_available(): self.device = torch.device("cuda") elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): self.device = torch.device("mps") else: self.device = torch.device("cpu") else: self.device = torch.device(config.device) print(f"Training on device: {self.device}") # Move model to device self.model = model.to(self.device) # Get vocab size from model if hasattr(model, "config"): self.vocab_size = model.config.vocab_size else: self.vocab_size = model.embed_tokens.num_embeddings # Setup loss function self.loss_fn = LanguageModelingLoss( vocab_size=self.vocab_size, label_smoothing=config.label_smoothing, ) # Calculate total steps self.steps_per_epoch = len(train_dataloader) self.total_steps = self.steps_per_epoch * config.num_epochs self.total_steps = self.total_steps // config.gradient_accumulation_steps # Setup optimizer and scheduler self.optimizer = create_optimizer( model, learning_rate=config.learning_rate, weight_decay=config.weight_decay, ) self.scheduler = create_scheduler( self.optimizer, num_training_steps=self.total_steps, warmup_ratio=config.warmup_ratio, min_lr_ratio=config.min_lr_ratio, ) # Setup mixed precision self.use_amp = config.fp16 and self.device.type == "cuda" self.scaler = GradScaler() if self.use_amp else None # Tracking self.global_step = 0 self.epoch = 0 self.best_val_loss = float("inf") # Early stopping tracking self.early_stopping_counter = 0 self.should_stop = False # Checkpoint directory os.makedirs(config.checkpoint_dir, exist_ok=True) # W&B logging self.wandb = None if wandb_project: try: import wandb wandb.init(project=wandb_project, config=config.to_dict()) self.wandb = wandb except ImportError: print("wandb not installed, skipping logging") def train(self) -> Dict[str, Any]: """Run the full training loop. Returns: Dictionary with training results """ print(f"\n{'='*60}") print("STARTING TRAINING") print(f"{'='*60}") print(f"Total epochs: {self.config.num_epochs}") print(f"Steps per epoch: {self.steps_per_epoch}") print(f"Total optimization steps: {self.total_steps}") print(f"Gradient accumulation: {self.config.gradient_accumulation_steps}") print(f"Mixed precision: {self.use_amp}") if self.config.early_stopping_patience > 0: print(f"Early stopping: patience={self.config.early_stopping_patience}") print(f"{'='*60}\n") training_start = time.time() # FIX: Start from loaded epoch (for resume), not always from 0 start_epoch = self.epoch if start_epoch > 0: print(f"Resuming from epoch {start_epoch + 1}") for epoch in range(start_epoch, self.config.num_epochs): self.epoch = epoch epoch_loss = self._train_epoch() print(f"\nEpoch {epoch + 1}/{self.config.num_epochs} - Loss: {epoch_loss:.4f}") # Validation if self.val_dataloader is not None: val_metrics = self.evaluate() print(f"Validation - Loss: {val_metrics['loss']:.4f}, PPL: {val_metrics['perplexity']:.2f}") # Early stopping check if val_metrics["loss"] < self.best_val_loss - self.config.early_stopping_threshold: self.best_val_loss = val_metrics["loss"] self.early_stopping_counter = 0 self.save_checkpoint("best") print(f" New best model saved!") else: self.early_stopping_counter += 1 print(f" No improvement. Early stopping: {self.early_stopping_counter}/{self.config.early_stopping_patience}") if self.config.early_stopping_patience > 0 and self.early_stopping_counter >= self.config.early_stopping_patience: print(f"\nEarly stopping triggered after {self.early_stopping_counter} evaluations without improvement.") self.should_stop = True # Save epoch checkpoint self.save_checkpoint(f"epoch_{epoch + 1}") # Check early stopping if self.should_stop: print("Stopping training early.") break training_time = time.time() - training_start print(f"\n{'='*60}") print(f"TRAINING COMPLETE") print(f"Total time: {training_time / 3600:.2f} hours") print(f"Best validation loss: {self.best_val_loss:.4f}") if self.should_stop: print(f"Stopped early at epoch {self.epoch + 1}") print(f"{'='*60}") return { "total_steps": self.global_step, "training_time": training_time, "best_val_loss": self.best_val_loss, } def _train_epoch(self) -> float: """Train for one epoch. Returns: Average training loss for the epoch """ self.model.train() total_loss = 0.0 num_batches = 0 accumulated_loss = 0.0 num_accumulated_batches = 0 # FIX: Track actual number of batches for correct averaging # Create progress bar pbar = tqdm( enumerate(self.train_dataloader), total=len(self.train_dataloader), desc=f"Epoch {self.epoch + 1}", ncols=100, ) for step, batch in pbar: # Move batch to device input_ids = batch["input_ids"].to(self.device) labels = batch["labels"].to(self.device) # Note: attention_mask from dataloader is padding mask (1/0) # The model creates its own causal mask internally # We handle padding via -100 labels in the loss function # Forward pass with optional mixed precision with autocast(enabled=self.use_amp): outputs = self.model(input_ids) # Handle different output types (tensor, tuple, or dataclass) if isinstance(outputs, torch.Tensor): logits = outputs elif hasattr(outputs, 'logits'): logits = outputs.logits else: logits = outputs[0] loss = self.loss_fn(logits, labels) loss = loss / self.config.gradient_accumulation_steps # Backward pass if self.use_amp: self.scaler.scale(loss).backward() else: loss.backward() # FIX: Track unscaled loss correctly unscaled_loss = loss.item() * self.config.gradient_accumulation_steps accumulated_loss += unscaled_loss num_accumulated_batches += 1 total_loss += unscaled_loss num_batches += 1 # Gradient accumulation if (step + 1) % self.config.gradient_accumulation_steps == 0: # Gradient clipping if self.use_amp: self.scaler.unscale_(self.optimizer) grad_norm = clip_grad_norm(self.model, self.config.max_grad_norm) # Optimizer step if self.use_amp: self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() self.scheduler.step() self.optimizer.zero_grad() self.global_step += 1 # Logging if self.global_step % self.config.logging_steps == 0: # FIX: Divide by actual number of accumulated batches, not gradient_accumulation_steps avg_loss = accumulated_loss / max(num_accumulated_batches, 1) lr = self.scheduler.get_last_lr()[0] # Update progress bar pbar.set_postfix({ 'loss': f'{avg_loss:.4f}', 'lr': f'{lr:.2e}', 'step': f'{self.global_step}/{self.total_steps}' }) tqdm.write( f"Step {self.global_step}/{self.total_steps} | " f"Loss: {avg_loss:.4f} | " f"LR: {lr:.2e} | " f"Grad: {grad_norm:.2f}" ) if self.wandb: self.wandb.log({ "train/loss": avg_loss, "train/learning_rate": lr, "train/grad_norm": grad_norm, "train/epoch": self.epoch, }, step=self.global_step) # Reset accumulators accumulated_loss = 0.0 num_accumulated_batches = 0 # Evaluation if self.config.eval_steps > 0 and self.global_step % self.config.eval_steps == 0: if self.val_dataloader is not None: val_metrics = self.evaluate() print(f" Eval - Loss: {val_metrics['loss']:.4f}, PPL: {val_metrics['perplexity']:.2f}") if self.wandb: self.wandb.log({ "eval/loss": val_metrics["loss"], "eval/perplexity": val_metrics["perplexity"], }, step=self.global_step) # Early stopping check during training if val_metrics["loss"] < self.best_val_loss - self.config.early_stopping_threshold: self.best_val_loss = val_metrics["loss"] self.early_stopping_counter = 0 self.save_checkpoint("best") print(f" New best model! Loss: {self.best_val_loss:.4f}") else: self.early_stopping_counter += 1 if self.config.early_stopping_patience > 0: print(f" No improvement ({self.early_stopping_counter}/{self.config.early_stopping_patience})") if self.early_stopping_counter >= self.config.early_stopping_patience: print(f"\n Early stopping triggered!") self.should_stop = True break # Exit the training loop # Checkpointing if self.config.save_steps > 0 and self.global_step % self.config.save_steps == 0: self.save_checkpoint(f"step_{self.global_step}") # Check if early stopping was triggered if self.should_stop: break return total_loss / max(num_batches, 1) @torch.no_grad() def evaluate(self) -> Dict[str, float]: """Evaluate the model on validation data. Returns: Dictionary with evaluation metrics """ self.model.eval() total_loss = 0.0 total_accuracy = 0.0 num_batches = 0 for batch in self.val_dataloader: input_ids = batch["input_ids"].to(self.device) labels = batch["labels"].to(self.device) with autocast(enabled=self.use_amp): outputs = self.model(input_ids) # Handle different output types (tensor, tuple, or dataclass) if isinstance(outputs, torch.Tensor): logits = outputs elif hasattr(outputs, 'logits'): logits = outputs.logits else: logits = outputs[0] loss = self.loss_fn(logits, labels) total_loss += loss.item() total_accuracy += compute_accuracy(logits, labels).item() num_batches += 1 self.model.train() avg_loss = total_loss / max(num_batches, 1) avg_accuracy = total_accuracy / max(num_batches, 1) return { "loss": avg_loss, "perplexity": compute_perplexity(torch.tensor(avg_loss)).item(), "accuracy": avg_accuracy, } def save_checkpoint(self, name: str): """Save a checkpoint. Args: name: Checkpoint name (e.g., "best", "epoch_1", "step_1000") """ checkpoint_path = os.path.join(self.config.checkpoint_dir, name) os.makedirs(checkpoint_path, exist_ok=True) # Save model model_path = os.path.join(checkpoint_path, "model.pt") torch.save(self.model.state_dict(), model_path) # Save optimizer and scheduler optimizer_path = os.path.join(checkpoint_path, "optimizer.pt") torch.save({ "optimizer": self.optimizer.state_dict(), "scheduler": self.scheduler.state_dict(), "global_step": self.global_step, "epoch": self.epoch, "best_val_loss": self.best_val_loss, "early_stopping_counter": self.early_stopping_counter, }, optimizer_path) # Save config config_path = os.path.join(checkpoint_path, "config.json") with open(config_path, "w") as f: json.dump(self.config.to_dict(), f, indent=2) print(f"Saved checkpoint: {checkpoint_path}") # Cleanup old checkpoints self._cleanup_checkpoints() def load_checkpoint(self, checkpoint_path: str): """Load a checkpoint. Args: checkpoint_path: Path to checkpoint directory """ # Load model model_path = os.path.join(checkpoint_path, "model.pt") state_dict = torch.load(model_path, map_location=self.device) # FIX: Handle torch.compile prefix (_orig_mod.) if present if any(k.startswith("_orig_mod.") for k in state_dict.keys()): print(" Detected compiled model checkpoint, removing _orig_mod. prefix...") state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()} self.model.load_state_dict(state_dict) # Load optimizer and scheduler optimizer_path = os.path.join(checkpoint_path, "optimizer.pt") if os.path.exists(optimizer_path): state = torch.load(optimizer_path, map_location=self.device) self.optimizer.load_state_dict(state["optimizer"]) self.scheduler.load_state_dict(state["scheduler"]) self.global_step = state["global_step"] self.epoch = state["epoch"] self.best_val_loss = state.get("best_val_loss", float("inf")) self.early_stopping_counter = state.get("early_stopping_counter", 0) # FIX: Increment epoch to start from next epoch (we saved after completing this epoch) # Only if checkpoint was saved at end of epoch (epoch_* checkpoints) if "epoch_" in checkpoint_path: self.epoch += 1 print(f" Checkpoint was end-of-epoch, will start from epoch {self.epoch + 1}") print(f"Loaded checkpoint: {checkpoint_path}") print(f" Resuming from step {self.global_step}, epoch {self.epoch}") print(f" Best val loss so far: {self.best_val_loss:.4f}") def _cleanup_checkpoints(self): """Remove old checkpoints to save disk space.""" if self.config.save_total_limit <= 0: return checkpoint_dir = Path(self.config.checkpoint_dir) checkpoints = sorted( [d for d in checkpoint_dir.iterdir() if d.is_dir() and d.name.startswith("step_")], key=lambda x: int(x.name.split("_")[1]), ) # Keep only the most recent checkpoints (plus "best" and "epoch_*") while len(checkpoints) > self.config.save_total_limit: old_checkpoint = checkpoints.pop(0) print(f"Removing old checkpoint: {old_checkpoint}") import shutil shutil.rmtree(old_checkpoint)