import math import traceback import random import time import json from pathlib import Path from contextlib import nullcontext from tqdm import tqdm import torch import torch.nn.functional as F from torch.amp import GradScaler, autocast from torch.utils.data import DataLoader from transformers import get_cosine_schedule_with_warmup from safetensors.torch import load_file from lmr.checkpointing import Checkpointing from lmr.ddp import setup_ddp, cleanup_ddp, initialize_model_ddp, unwrap_model, initialize_samplers_ddp from lmr.utils.logger import Logger class Trainer: def __init__(self, config, model, tokenizer, splits, checkpointing, samplers=None): """ 适配 main.py 的参数: config: 这里的 config 对应 main.py 里的 config.training splits: 数据集字典 """ self.training_config = config self.model = model self.tokenizer = tokenizer self.splits = splits self.checkpointing = checkpointing self.samplers = samplers # 默认设备设置（会在 _setup_training 中根据 DDP 更新） self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 识别 Token self.pad_token_id = getattr(self.tokenizer, 'pad_token_id', 0) if self.pad_token_id is None: self.pad_token_id = 0 self.eos_token_id = getattr(self.tokenizer, 'eos_token_id', None) # 精度设置 - 增加默认值以防 config 缺失 precision = getattr(self.training_config, "precision", "bfloat16") self.autocast_dtype = getattr(torch, precision) if hasattr(torch, precision) else torch.bfloat16 self.use_ddp = False self.rank = 0 self.world_size = 1 self.debug_prompts = [ "Question: What is 15 + 32?\nAnswer:", "Question: There are 5 birds on a tree. 2 fly away. How many are left?\nSolution:", ] # ========================================================================= # 权重加载逻辑 (保留你提供的鲁棒版本) # ========================================================================= def _is_lora_param(self, param_name): return any(indicator in param_name for indicator in ['lora_A', 'lora_B', 'lora_dropout']) def load_only_model_weights(self, checkpoint_path, map_location="cpu", strict=False, verbose=True): path_obj = Path(checkpoint_path) if not path_obj.exists(): return model_state = {} if path_obj.is_dir(): index_file = path_obj / "model.safetensors.index.json" if index_file.exists(): with open(index_file, 'r') as f: index_data = json.load(f) weight_map = index_data.get("weight_map", {}) for shard_name in set(weight_map.values()): model_state.update(load_file(str(path_obj / shard_name), device=str(map_location))) else: possible = list(path_obj.glob("*.safetensors")) + list(path_obj.glob("*.pt")) if possible: path_obj = possible[0] if not model_state and path_obj.is_file(): if path_obj.suffix == ".safetensors": model_state = load_file(str(path_obj), device=str(map_location)) else: ckpt = torch.load(path_obj, map_location=map_location) model_state = ckpt.get("model", ckpt.get("state_dict", ckpt)) ckpt_keys_map = {k.replace("module.", "").replace("_orig_mod.", "").replace("model.", ""): k for k in model_state.keys()} load_target = unwrap_model(self.model) target_state = load_target.state_dict() filtered_state = {} for k_target, v_target in target_state.items(): k_clean = k_target.replace("module.", "").replace("_orig_mod.", "").replace("model.", "") if k_clean in ckpt_keys_map: v_ckpt = model_state[ckpt_keys_map[k_clean]] if v_ckpt.shape == v_target.shape: filtered_state[k_target] = v_ckpt load_target.load_state_dict(filtered_state, strict=strict) if verbose and self.rank == 0: print(f"✅ Loaded {len(filtered_state)} parameters from {checkpoint_path}") # ========================================================================= # 训练设置 & 初始化 # ========================================================================= def _setup_training(self): # 1. Dataloaders self.train_dataloader = self._get_dataloader("train") # 适配不同数据集可能的 split 命名 val_key = "val" if "val" in self.splits else "validation" self.validation_dataloader = self._get_dataloader(val_key) if val_key in self.splits else None # 2. 梯度累积 if getattr(self.training_config, "use_grad_accum", False): if self.training_config.grad_accum_steps == "auto": tps = self.training_config.batch_size * 1024 * self.world_size self.grad_accum_steps = max(1, getattr(self.training_config, "tokens_per_step", 32768) // tps) else: self.grad_accum_steps = self.training_config.grad_accum_steps else: self.grad_accum_steps = 1 self.steps_per_epoch = len(self.train_dataloader) // self.grad_accum_steps self.tokens_per_batch = self.training_config.batch_size * 1024 self.tokens_per_step = self.grad_accum_steps * self.tokens_per_batch * self.world_size self.tokens_per_epoch = self.tokens_per_step * self.steps_per_epoch # 3. 基础恢复 try: self.checkpointing.load_model_states("recent") except: pass # 4. 设备 & 编译 self.device = torch.device(f"cuda:{self.rank}") # if getattr(self.training_config, "compile", False): # self.model = torch.compile(self.model, mode=getattr(self.training_config, "compile_mode", "default")) self.model.to(self.device) # 5. DDP 包装 if self.use_ddp: self.model = initialize_model_ddp(self.model, self.rank) # 6. 优化器 & 调度器 self._initialize_optimizer() self._initialize_scheduler() self._initialize_scaler() # 7. 训练状态恢复 try: self.checkpointing.load_training_states("recent") except: pass def _initialize_optimizer(self): # 自动判断 LoRA is_lora = any(getattr(unwrap_model(self.model).config, f"use_lora_{x}", False) for x in ["phi_attention", "icl_attention"]) if is_lora: params = [p for n, p in self.model.named_parameters() if self._is_lora_param(n)] for n, p in self.model.named_parameters(): p.requires_grad = self._is_lora_param(n) else: params = self.model.parameters() self.optimizer = torch.optim.AdamW( params, lr=self.training_config.lr, betas=getattr(self.training_config, "betas", (0.9, 0.95)), weight_decay=self.training_config.weight_decay ) self.checkpointing.optimizer = self.optimizer def _initialize_scheduler(self): warmup = getattr(self.training_config, "warmup_steps", 100) total = self.steps_per_epoch * self.training_config.max_epochs self.scheduler = get_cosine_schedule_with_warmup(self.optimizer, warmup, total) self.checkpointing.scheduler = self.scheduler def _initialize_scaler(self): self.scaler = GradScaler("cuda") if self.autocast_dtype == torch.float16 else None self.checkpointing.scaler = self.scaler def _get_dataloader(self, split_name): return DataLoader( self.splits[split_name], batch_size=self.training_config.batch_size, num_workers=getattr(self.training_config, "num_workers", 4), sampler=self.samplers[split_name] if self.samplers else None, shuffle=(split_name == "train" and self.samplers is None), pin_memory=True, drop_last=True ) # ========================================================================= # 核心循环 # ========================================================================= def _calculate_training_tokens(self, epoch, step): return epoch * self.tokens_per_epoch + step * self.tokens_per_step def _step_loss(self, batch): batch = batch.to(self.device, non_blocking=True) input_tokens = batch[:, :] vocab_limit = self.model.config.vocab_size if (input_tokens >= vocab_limit).any(): print(f"警告：发现越界 Token! 最大 ID: {input_tokens.max()}") input_tokens[input_tokens >= vocab_limit] = 0 # 临时替换为 0 防止崩溃 target_tokens = batch[:, 1:].clone() if self.pad_token_id is not None: target_tokens[target_tokens == self.pad_token_id] = 0 with autocast(device_type="cuda", dtype=self.autocast_dtype): logits = self.model(input_tokens) loss = unwrap_model(self.model).calculate_loss(logits[:, :-1], target_tokens) return loss def _train(self): self._setup_training() start_epoch = self.checkpointing.epoch start_step = self.checkpointing.step tokens_trained = self.checkpointing.tokens_trained if self.rank == 0: Logger.log(f"🚀 Starting training from Epoch {start_epoch}, Step {start_step}") for epoch in range(start_epoch, self.training_config.max_epochs): if self.use_ddp: self.train_dataloader.sampler.set_epoch(epoch) pbar = tqdm(total=self.steps_per_epoch, desc=f"Epoch {epoch}", disable=(self.rank != 0)) self.model.train() accum_loss = 0.0 for micro_step, batch in enumerate(self.train_dataloader): step = micro_step // self.grad_accum_steps is_update_step = ((micro_step + 1) % self.grad_accum_steps == 0) # 跳过已训练步骤 (Resume) if epoch == start_epoch and step < start_step: if is_update_step: pbar.update(1) continue sync_ctx = self.model.no_sync() if (self.use_ddp and not is_update_step) else nullcontext() with sync_ctx: loss = self._step_loss(batch) / self.grad_accum_steps if self.scaler: self.scaler.scale(loss).backward() else: loss.backward() accum_loss += loss.item() if is_update_step: if self.scaler: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.scaler.step(self.optimizer) self.scaler.update() else: torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.optimizer.step() self.optimizer.zero_grad(set_to_none=True) self.scheduler.step() if pbar: pbar.set_postfix(loss=f"{accum_loss:.4f}") pbar.update(1) accum_loss = 0.0 # Epoch 结束逻辑 if self.rank == 0: val_loss = self._validate() if self.validation_dataloader else 0.0 tokens_trained = self._calculate_training_tokens(epoch + 1, 0) self.checkpointing.save_checkpoint(epoch + 1, 0, 0, val_loss, tokens_trained) Logger.log(f"Epoch {epoch} Done. Val Loss: {val_loss:.4f}") @torch.no_grad() def _validate(self): self.model.eval() total_loss = 0 for batch in tqdm(self.validation_dataloader, desc="Validating", leave=False, disable=(self.rank != 0)): total_loss += self._step_loss(batch).item() return total_loss / len(self.validation_dataloader) def train(self): """外部唯一调用入口""" if getattr(self.training_config, "use_ddp", False): self.use_ddp = True self.rank, self.world_size = setup_ddp() try: self.samplers = initialize_samplers_ddp(self.splits, self.rank, self.world_size) self._train() finally: cleanup_ddp() else: self._train()