import time import re import json import shutil from pathlib import Path from glob import glob from typing import Dict, Any, Tuple, List, Optional import torch from lmr.utils.logger import Logger from lmr.utils.parsing import int_to_formatted_string from lmr.ddp import unwrap_model def _find_checkpoint_file_in_dir(path_dir: Path) -> Path: if not path_dir.is_dir(): raise FileNotFoundError(f"{path_dir} is not a dir") idx = path_dir / "model.safetensors.index.json" if idx.exists(): return idx safes = list(path_dir.glob("*.safetensors")) if safes: for cand in safes: if cand.name == "model.safetensors": return cand safes.sort(key=lambda p: p.stat().st_mtime, reverse=True) return safes[0] bins = list(path_dir.glob("pytorch_model.bin")) + list(path_dir.glob("*.pt")) if bins: for cand in bins: if cand.name == "pytorch_model.bin": return cand bins.sort(key=lambda p: p.stat().st_mtime, reverse=True) return bins[0] others = [p for p in path_dir.iterdir() if p.is_file()] if others: others.sort(key=lambda p: p.stat().st_mtime, reverse=True) return others[0] raise FileNotFoundError(f"No checkpoint file found in directory: {path_dir}") def _load_safetensors_shards_from_index(index_path: Path, map_location="cpu") -> Dict[str, torch.Tensor]: from safetensors.torch import load_file base_dir = index_path.parent with open(index_path, "r") as f: index_data = json.load(f) weight_map = index_data.get("weight_map", {}) shards = sorted(set(weight_map.values())) merged = {} for shard_name in shards: shard_path = base_dir / shard_name if not shard_path.exists(): raise FileNotFoundError(f"Shard not found: {shard_path}") shard = load_file(str(shard_path), device=map_location) merged.update(shard) return merged def _safe_torch_load(path: Path, map_location="cpu", allow_unsafe_fallback: bool = True) -> Any: try: return torch.load(str(path), map_location=map_location, weights_only=True) except TypeError: return torch.load(str(path), map_location=map_location) except Exception: if not allow_unsafe_fallback: raise return torch.load(str(path), map_location=map_location, weights_only=False) class Checkpointing: def __init__(self, model, checkpoint_dir, optimizer=None, scheduler=None, scaler=None, map_device="cpu"): self.model = model self.optimizer = optimizer self.scheduler = scheduler self.scaler = scaler self.map_device = map_device # state self.epoch = 0 self.step = 0 self.train_loss = float("inf") self.val_loss = float("inf") self.tokens_trained = 0 # NEW: accuracies self.val_acc = None self.train_acc = None self.use_ddp = torch.distributed.is_available() and torch.distributed.is_initialized() self.is_main_process = (not self.use_ddp) or (torch.distributed.get_rank() == 0) self.checkpoint_dir = Path(checkpoint_dir) if not self.is_main_process: return self.checkpoint_dir.mkdir(parents=True, exist_ok=True) self.log_path = self.checkpoint_dir / "_checkpoint_log.tsv" self._create_log() self.best_val_loss = self._get_best_val_loss() def _barrier(self): if self.use_ddp: torch.distributed.barrier() # ---------- best loss scan ---------- def _get_best_val_loss(self): best_dirs = glob(str(self.checkpoint_dir / "best_epoch*_val=*")) best_val = float("inf") for dir_path in best_dirs: match = re.search(r"val=([0-9.]+)", dir_path) if match: try: best_val = min(best_val, float(match.group(1))) except ValueError: pass return best_val def _checkpoint_dirname(self, epoch, step=None, val_loss=None, tokens_trained=None, prefix=None): name = f"epoch_{epoch:03d}" if prefix is not None: name = f"{prefix}_{name}" if step is not None: name += f"_step_{step:09d}" if tokens_trained is not None: name += f"_tokens_{int_to_formatted_string(tokens_trained)}" if val_loss is not None: name += f"_val={val_loss:.4f}" return name # ---------- log ---------- def _create_log(self): if self.log_path.exists(): return header = "Time\tCheckpoint_Type\tEpoch\tStep\tTrain_Loss\tVal_Loss\tTrain_Acc\tVal_Acc\tTokens_Trained\tDirname\n" with open(self.log_path, "w", encoding="utf-8") as f: f.write(header) def _update_log(self, kind, dirname): ts = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) epoch_str = f"{self.epoch:03d}" step_str = f"{self.step:09d}" line = ( f"{ts}\t{kind}\t{epoch_str}\t{step_str}\t" f"{'' if self.train_loss is None else self.train_loss}\t" f"{'' if self.val_loss is None else self.val_loss}\t" f"{'' if self.train_acc is None else self.train_acc}\t" f"{'' if self.val_acc is None else self.val_acc}\t" f"{self.tokens_trained}\t" f"{dirname}\n" ) with open(self.log_path, "a", encoding="utf-8") as f: f.write(line) # ---------- remove old ---------- def _remove_old(self, pattern: str): """ Delete old checkpoint directories matching pattern, e.g. "recent_epoch*" Keeps only the most recent one if multiple exist. """ if not self.is_main_process: return dirs = sorted(glob(str(self.checkpoint_dir / pattern))) if len(dirs) <= 1: return # keep newest by mtime dirs_paths = [Path(d) for d in dirs] dirs_paths.sort(key=lambda p: p.stat().st_mtime, reverse=True) keep = dirs_paths[0] to_remove = dirs_paths[1:] # for p in to_remove: # try: # shutil.rmtree(p) # Logger.log(f"🧹 Removed old checkpoint dir: {p.name}") # except Exception as e: # Logger.log(f"⚠️ Failed to remove old checkpoint dir {p}: {e}") # ---------- save ---------- def _save_state(self, dirname, include_training_states=False): save_path = self.checkpoint_dir / dirname save_path.mkdir(parents=True, exist_ok=True) model = unwrap_model(self.model) if hasattr(model, "save_pretrained"): # Recommended: safetensors + shards model.save_pretrained(save_path, safe_serialization=True, max_shard_size="5GB") else: Logger.log("⚠️ Model does not have save_pretrained; saving state_dict.") from safetensors.torch import save_file state_dict = {k: v.cpu().contiguous() for k, v in model.state_dict().items()} save_file(state_dict, str(save_path / "model.safetensors")) # trainer metadata metadata = { "epoch": self.epoch, "step": self.step, "train_loss": self.train_loss, "val_loss": self.val_loss, "tokens_trained": self.tokens_trained, "train_acc": self.train_acc, "val_acc": self.val_acc, } with open(save_path / "trainer_state.json", "w") as f: json.dump(metadata, f, indent=4) if include_training_states: train_state = {} if self.optimizer is not None: train_state["optimizer"] = self.optimizer.state_dict() if self.scheduler is not None and hasattr(self.scheduler, "state_dict"): train_state["scheduler"] = self.scheduler.state_dict() if self.scaler is not None and hasattr(self.scaler, "state_dict"): train_state["scaler"] = self.scaler.state_dict() torch.save(train_state, save_path / "optimizer.pt") Logger.log(f"💾 Checkpoint saved to: {save_path}") def _update_state( self, epoch, step=0, train_loss=None, val_loss=None, tokens_trained=None, train_acc=None, val_acc=None, ): self.epoch = int(epoch) self.step = int(step) if step is not None else 0 if train_loss is not None: self.train_loss = train_loss if val_loss is not None: self.val_loss = val_loss if tokens_trained is not None: self.tokens_trained = int(tokens_trained) # NEW if train_acc is not None: self.train_acc = float(train_acc) if val_acc is not None: self.val_acc = float(val_acc) def _save_best(self): self._remove_old("best_epoch*") dirname = self._checkpoint_dirname(self.epoch, self.step, self.val_loss, self.tokens_trained, prefix="best") self._save_state(dirname, include_training_states=False) self._update_log("best", dirname) def _save_recent(self): self._remove_old("recent_epoch*") dirname = self._checkpoint_dirname(self.epoch, self.step, self.val_loss, self.tokens_trained, prefix="recent") self._save_state(dirname, include_training_states=True) self._update_log("recent", dirname) def _save_epoch(self): dirname = self._checkpoint_dirname(self.epoch, None, self.val_loss, self.tokens_trained) self._save_state(dirname, include_training_states=False) self._update_log("epoch", dirname) # ---------- public save ---------- def save_checkpoint( self, epoch, step=None, train_loss=None, val_loss=None, tokens_trained=None, val_acc=None, train_acc=None, ): if self.is_main_process: self._update_state( epoch, step=step, train_loss=train_loss, val_loss=val_loss, tokens_trained=tokens_trained, train_acc=train_acc, val_acc=val_acc, ) self._save_recent() if step is None: self.step = 0 self._save_epoch() if (self.val_loss is not None) and (self.val_loss < self.best_val_loss): self.best_val_loss = self.val_loss self._save_best() self._barrier() # ---------- find ckpt dir ---------- def _get_checkpoint_path(self, checkpoint_type): if checkpoint_type == "best": pattern = "best_epoch*" elif checkpoint_type == "recent": pattern = "recent_epoch*" elif checkpoint_type.startswith("epoch_"): pattern = f"{checkpoint_type}*" else: path = self.checkpoint_dir / checkpoint_type if path.exists(): return path return None candidates = sorted(glob(str(self.checkpoint_dir / pattern))) if not candidates: Logger.log(f"No checkpoint found for {checkpoint_type}") return None return Path(candidates[-1]) # ---------- load model weights + metadata ---------- def load_model_states(self, checkpoint_type="recent"): ckpt_dir = self._get_checkpoint_path(checkpoint_type) if not ckpt_dir: return None Logger.log(f"📂 Loading model from dir: {ckpt_dir}") meta_path = ckpt_dir / "trainer_state.json" if meta_path.exists(): with open(meta_path, "r") as f: state = json.load(f) self.epoch = int(state.get("epoch", 0)) self.step = int(state.get("step", 0)) self.train_loss = state.get("train_loss", None) self.val_loss = state.get("val_loss", None) self.tokens_trained = int(state.get("tokens_trained", 0)) self.train_acc = state.get("train_acc", None) self.val_acc = state.get("val_acc", None) try: self.load_only_model_weights(str(ckpt_dir), map_location=self.map_device, strict=False, verbose=True) except Exception as e: Logger.log(f"⚠️ Failed to load model weights from {ckpt_dir}: {e}") self._barrier() return ckpt_dir def load_only_model_weights( self, checkpoint_path: str, model: Optional[torch.nn.Module] = None, device: Optional[str] = None, map_location: Optional[str] = "cpu", save_filtered_to: Optional[str] = None, strict: bool = False, verbose: bool = True, allow_unsafe_fallback: bool = True, allow_continue_on_failure: bool = False, **kwargs ) -> Dict[str, Any]: load_model = model if model is not None else unwrap_model(self.model) p = Path(checkpoint_path) final_map = device or map_location # read raw ckpt try: if p.exists() and p.is_dir(): candidate = _find_checkpoint_file_in_dir(p) if candidate.name == "model.safetensors.index.json": if verbose: Logger.log(f"[ckpt_loader] Detected safetensors sharded index at {candidate}; loading shards...") raw_state = _load_safetensors_shards_from_index(candidate, map_location=final_map) else: if verbose: Logger.log(f"[ckpt_loader] Directory provided; selected checkpoint file: {candidate}") if candidate.suffix == ".safetensors": from safetensors.torch import load_file raw_state = load_file(str(candidate), device=final_map) else: raw_state = _safe_torch_load(candidate, map_location=final_map, allow_unsafe_fallback=allow_unsafe_fallback) elif p.exists() and p.is_file(): if p.suffix == ".safetensors": from safetensors.torch import load_file raw_state = load_file(str(p), device=final_map) else: raw_state = _safe_torch_load(p, map_location=final_map, allow_unsafe_fallback=allow_unsafe_fallback) else: raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}") except Exception as e_load: msg = f"[ckpt_loader] Failed to load checkpoint '{checkpoint_path}': {e_load}" if not allow_continue_on_failure: raise RuntimeError(msg) Logger.log(msg + " -- continuing (allow_continue_on_failure=True).") return {"matched": 0, "total_target": len(load_model.state_dict()), "error": str(e_load)} # extract state_dict def _extract_state_dict(raw): if isinstance(raw, dict): sample_keys = list(raw.keys())[:10] if any(("weight" in k or "bias" in k or "embed" in k) for k in sample_keys): return raw for candidate_key in ("model", "state_dict", "model_state_dict", "state"): if candidate_key in raw and isinstance(raw[candidate_key], dict): return raw[candidate_key] if isinstance(raw, dict): dicts = [v for v in raw.values() if isinstance(v, dict)] if dicts: return max(dicts, key=lambda x: len(x)) raise RuntimeError(f"Could not find a state_dict inside checkpoint (raw type={type(raw)})") state_dict = _extract_state_dict(raw_state) # normalize keys def _normalize_keys(state: Dict[str, Any], prefixes=("module.", "_orig_mod.", "model_state.")): normalized = {} for k, v in state.items(): new_k = k for pfx in prefixes: if new_k.startswith(pfx): new_k = new_k[len(pfx):] normalized[new_k] = v return normalized normalized = _normalize_keys(state_dict) if verbose: print(f"[ckpt_loader] extracted {len(normalized)} params (sample: {list(normalized.keys())[:10]})") # match by name+shape target_state = load_model.state_dict() filtered = {} size_mismatch: List[Tuple[str, Any, Any]] = [] missing: List[str] = [] matched = 0 for tk, tv in target_state.items(): if tk in normalized: ck = normalized[tk] if getattr(ck, "shape", None) == getattr(tv, "shape", None): try: filtered[tk] = ck.to(tv.device) if hasattr(ck, "to") else ck except Exception: filtered[tk] = ck matched += 1 else: size_mismatch.append((tk, getattr(ck, "shape", None), tv.shape)) else: missing.append(tk) if verbose: print(f"[ckpt_loader] Matched: {matched}/{len(target_state)}; size_mismatch: {len(size_mismatch)}; missing: {len(missing)}") load_msg = load_model.load_state_dict(filtered, strict=False) import pdb # pdb.set_trace() if verbose: print(f"[ckpt_loader] load_state_dict: {load_msg}") if save_filtered_to: try: torch.save(filtered, save_filtered_to) if verbose: print(f"[ckpt_loader] Saved filtered weights to {save_filtered_to}") except Exception as e_save: if verbose: print(f"[ckpt_loader] Warning: failed to save filtered weights: {e_save}") return { "matched": matched, "total_target": len(target_state), "size_mismatch": size_mismatch, "missing": missing, "load_msg": load_msg, } # ---------- load optimizer/scheduler/scaler ---------- def load_training_states(self, checkpoint_type="recent"): ckpt_dir = self._get_checkpoint_path(checkpoint_type) if not ckpt_dir: return opt_path = ckpt_dir / "optimizer.pt" if not opt_path.exists(): Logger.log(f"⚠️ Optimizer state not found in {ckpt_dir}") return state = torch.load(opt_path, map_location=self.map_device) if self.optimizer is not None and "optimizer" in state: self.optimizer.load_state_dict(state["optimizer"]) if self.scheduler is not None and "scheduler" in state: self.scheduler.load_state_dict(state["scheduler"]) if self.scaler is not None and "scaler" in state: self.scaler.load_state_dict(state["scaler"]) Logger.log(f"✅ Loaded training states from {opt_path}") self._barrier()