""" Checkpoint management for training. Single Responsibility: Only handles saving and loading checkpoints. """ import os import torch import threading from typing import Dict, Any, Optional, List from pathlib import Path from dataclasses import dataclass from .utils import log @dataclass class TrainingState: """Immutable training state for checkpointing.""" step: int epoch: int loss: float text_ratio: float best_loss: float = float("inf") def to_dict(self) -> Dict[str, Any]: return { "step": self.step, "epoch": self.epoch, "loss": self.loss, "text_ratio": self.text_ratio, "best_loss": self.best_loss, } @classmethod def from_dict(cls, d: Dict[str, Any]) -> 'TrainingState': return cls( step=d.get("step", 0), epoch=d.get("epoch", 0), loss=d.get("loss", 0.0), text_ratio=d.get("text_ratio", 0.9), best_loss=d.get("best_loss", float("inf")), ) class CheckpointManager: """ Manages checkpoint saving and loading with async support. Single Responsibility: Only handles checkpoint I/O. Open/Closed: Can extend with new checkpoint formats without modification. """ def __init__( self, output_dir: str, prefix: str = "checkpoint", verbose: bool = True, max_checkpoints: Optional[int] = None ): """ Initialize checkpoint manager. Args: output_dir: Directory for saving checkpoints prefix: Prefix for checkpoint filenames verbose: Whether to log operations max_checkpoints: Maximum checkpoints to keep (None = keep all) """ self.output_dir = Path(output_dir) self.prefix = prefix self.verbose = verbose self.max_checkpoints = max_checkpoints self._save_threads: List[threading.Thread] = [] # Create output directory self.output_dir.mkdir(parents=True, exist_ok=True) def save( self, state_dict: Dict[str, Any], filename: str, async_save: bool = True ) -> str: """ Save checkpoint. Args: state_dict: State dictionary to save filename: Checkpoint filename async_save: Whether to save asynchronously Returns: Path to saved checkpoint """ path = self.output_dir / filename if async_save: self._save_async(state_dict, path) else: self._save_sync(state_dict, path) return str(path) def save_step( self, adapter_state: Dict[str, Any], optimizer_state: Dict[str, Any], training_state: TrainingState, async_save: bool = True ) -> str: """Save step checkpoint.""" state_dict = { "adapter": adapter_state, "optimizer": optimizer_state, **training_state.to_dict() } filename = f"{self.prefix}_step{training_state.step}.pt" return self.save(state_dict, filename, async_save) def save_epoch( self, adapter_state: Dict[str, Any], optimizer_state: Dict[str, Any], training_state: TrainingState, async_save: bool = True ) -> str: """Save epoch checkpoint.""" state_dict = { "adapter": adapter_state, "optimizer": optimizer_state, **training_state.to_dict() } filename = f"{self.prefix}_epoch{training_state.epoch}.pt" return self.save(state_dict, filename, async_save) def save_best( self, adapter_state: Dict[str, Any], training_state: TrainingState, lora_state: Optional[Dict[str, Any]] = None, async_save: bool = True ) -> str: """Save best model checkpoint.""" state_dict = { "adapter": adapter_state, **training_state.to_dict() } if lora_state is not None: state_dict["lora"] = lora_state filename = f"{self.prefix}_best.pt" return self.save(state_dict, filename, async_save) def load(self, path: str) -> Dict[str, Any]: """ Load checkpoint. Args: path: Path to checkpoint Returns: Loaded state dictionary """ if self.verbose: log(f"Loading checkpoint: {path}") return torch.load(path, map_location="cpu", weights_only=False) def load_latest(self) -> Optional[Dict[str, Any]]: """Load the most recent checkpoint.""" checkpoints = self._get_checkpoints() if not checkpoints: return None return self.load(str(checkpoints[-1])) def wait_for_saves(self): """Wait for all async saves to complete.""" for t in self._save_threads: t.join() self._save_threads = [] def _save_sync(self, state_dict: Dict[str, Any], path: Path): """Synchronous save.""" # Copy tensors to CPU state_copy = self._copy_to_cpu(state_dict) torch.save(state_copy, path) if self.verbose: log(f"[Checkpoint] Saved: {path.name}") def _save_async(self, state_dict: Dict[str, Any], path: Path): """Asynchronous save.""" # Clean up completed threads self._save_threads = [t for t in self._save_threads if t.is_alive()] # Copy tensors to CPU state_copy = self._copy_to_cpu(state_dict) def _save(): try: torch.save(state_copy, path) if self.verbose: log(f"[Checkpoint] Saved: {path.name}") except Exception as e: if self.verbose: log(f"[Checkpoint] Error saving {path.name}: {e}") thread = threading.Thread(target=_save, daemon=True) thread.start() self._save_threads.append(thread) # Cleanup old checkpoints if needed if self.max_checkpoints: self._cleanup_old_checkpoints() def _copy_to_cpu(self, obj: Any) -> Any: """Recursively copy tensors to CPU.""" if isinstance(obj, torch.Tensor): return obj.detach().cpu().clone() elif isinstance(obj, dict): return {k: self._copy_to_cpu(v) for k, v in obj.items()} elif isinstance(obj, list): return [self._copy_to_cpu(v) for v in obj] return obj def _get_checkpoints(self) -> List[Path]: """Get sorted list of checkpoint files.""" pattern = f"{self.prefix}_step*.pt" checkpoints = sorted( self.output_dir.glob(pattern), key=lambda p: int(p.stem.split("step")[-1]) ) return checkpoints def _cleanup_old_checkpoints(self): """Remove old checkpoints beyond max_checkpoints.""" if not self.max_checkpoints: return checkpoints = self._get_checkpoints() while len(checkpoints) > self.max_checkpoints: oldest = checkpoints.pop(0) try: oldest.unlink() if self.verbose: log(f"[Checkpoint] Removed old: {oldest.name}") except Exception: pass class Stage1CheckpointManager(CheckpointManager): """Checkpoint manager for Stage 1 training.""" def __init__(self, output_dir: str, **kwargs): super().__init__(output_dir, prefix="stage1", **kwargs) class Stage2CheckpointManager(CheckpointManager): """Checkpoint manager for Stage 2 training.""" def __init__(self, output_dir: str, **kwargs): super().__init__(output_dir, prefix="stage2", **kwargs) def save_best( self, adapter_state: Dict[str, Any], training_state: TrainingState, lora_state: Dict[str, Any], async_save: bool = True ) -> str: """Save best model with LoRA weights.""" return super().save_best( adapter_state, training_state, lora_state=lora_state, async_save=async_save )