import time import torch import re import json import shutil # 新增：用于删除文件夹 from pathlib import Path from glob import glob from lmr.utils.logger import Logger from lmr.utils.parsing import int_to_formatted_string from lmr.ddp import unwrap_model 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 fields self.epoch = 0 self.step = 0 self.train_loss = float("inf") self.val_loss = float("inf") self.tokens_trained = 0 # DDP flags 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() # ---------- ddp ---------- def _barrier(self): if self.use_ddp: torch.distributed.barrier() # ---------- logging + util ---------- def _remove_old(self, pattern): """ 删除旧的 Checkpoint 文件夹。 pattern 例如: "recent_epoch*_step*" """ if not self.is_main_process: return # 查找匹配的文件夹 candidates = sorted(glob(str(self.checkpoint_dir / pattern))) for path_str in candidates: path = Path(path_str) try: if path.is_dir(): shutil.rmtree(path) # 递归删除文件夹 print(f"🗑️ Removed old checkpoint dir: {path.name}") else: path.unlink() except Exception as e: print(f"⚠️ Failed to remove {path}: {e}") def _get_best_val_loss(self): # 扫描文件夹名称来获取 best loss # 文件夹格式通常是: best_epoch_010_step_5000_val=0.4500 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 def _create_log(self): if self.log_path.exists(): return header = "Time\tCheckpoint_Type\tEpoch\tStep\tTrain_Loss\tVal_Loss\tTokens_Trained\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"{dirname}\n" ) with open(self.log_path, "a", encoding="utf-8") as f: f.write(line) # ---------- state saving ---------- def _save_state(self, dirname, include_training_states=False): """ 创建文件夹并使用 save_pretrained 保存分片模型。 """ save_path = self.checkpoint_dir / dirname save_path.mkdir(parents=True, exist_ok=True) # 1. 保存模型权重 (使用 HF 原生分片逻辑) model = unwrap_model(self.model) if hasattr(model, "save_pretrained"): # 这是生成 model.safetensors.index.json 和 model-000xx.safetensors 的关键 # max_shard_size 可以控制切分大小，默认通常是 5GB 或 10GB model.save_pretrained(save_path, safe_serialization=False, max_shard_size="5GB") # 同时也保存 tokenizer (如果有的话，建议在 Trainer 里传入 tokenizer 并挂载到 self.model 上，或者手动调 tokenizer.save_pretrained) # if hasattr(self, 'tokenizer') and self.tokenizer: # self.tokenizer.save_pretrained(save_path) else: # 如果不是 HF 模型，回退到存单文件 Logger.log("⚠️ Model does not have save_pretrained method. Saving single file.") 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")) # 2. 保存自定义 Metadata metadata = { "epoch": self.epoch, "step": self.step, "train_loss": self.train_loss, "val_loss": self.val_loss, "tokens_trained": self.tokens_trained, } with open(save_path / "trainer_state.json", 'w') as f: json.dump(metadata, f, indent=4) # 3. 保存 Optimizer/Scheduler (放在同一个文件夹里) 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): 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 = tokens_trained def _save_best(self): # 移除旧的 best 文件夹 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) # Best 通常只存权重，不存 optimizer 以节省空间 self._update_log("best", dirname) def _save_recent(self): # 移除旧的 recent 文件夹 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) # Recent 必须存 optimizer 用于恢复 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) def _save_step(self): dirname = self._checkpoint_dirname(self.epoch, self.step, self.val_loss, self.tokens_trained) self._save_state(dirname, include_training_states=False) self._update_log("step", dirname) # ---------- checkpoint saving ---------- def save_checkpoint(self, epoch, step=None, train_loss=None, val_loss=None, tokens_trained=None): if self.is_main_process: self._update_state(epoch, step=step, train_loss=train_loss, val_loss=val_loss, tokens_trained=tokens_trained) # 保存 recent (用于断点续训) self._save_recent() # 保存 epoch 存档 (可选) if step is None: self.step = 0 self._save_epoch() # 只有当你想每隔多少步存一个永久档时才打开这个 # else: # self._save_step() # 保存 best (用于推理) 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() # ---------- loading (适配文件夹结构) ---------- 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]) def load_model_states(self, checkpoint_type="recent"): ckpt_dir = self._get_checkpoint_path(checkpoint_type) if not ckpt_dir: return Logger.log(f"📂 Loading model from dir: {ckpt_dir}") # 这里的加载逻辑其实在 Trainer.load_only_model_weights 里有更详细的实现 # 这里主要是为了让 Checkpointing 类也能独立运作 model = unwrap_model(self.model) # 优先尝试 HF 原生加载 (支持分片) if hasattr(model, "from_pretrained"): # 注意：from_pretrained 是类方法，但我们已经有实例了。 # 对于已有实例，通常没有直接的 "load_pretrained" 方法来处理分片。 # 所以我们还是依赖 Trainer 里的那个能够处理 index.json 的 load_only_model_weights 函数。 # 这里我们只负责加载 metadata。 pass # 读取 Metadata 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 = state.get("tokens_trained", 0) self._barrier() return ckpt_dir # 返回路径给 Trainer 用 def load_training_states(self, checkpoint_type="recent"): ckpt_dir = self._get_checkpoint_path(checkpoint_type) if not ckpt_dir: return # Optimizer 在文件夹里面 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()