import math import traceback import random import time 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 from lmr.utils.parsing import int_to_formatted_string class Trainer: def __init__(self, training_config, model, tokenizer, splits, checkpointing, samplers=None, device=None): self.training_config = training_config self.model = model self.tokenizer = tokenizer self.splits = splits self.checkpointing = checkpointing self.samplers = samplers self.device = device # Identify special tokens for masking/loss 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) # Set precision self.autocast_dtype = getattr(torch, self.training_config.precision) self.use_ddp = False self.rank = 0 self.world_size = 1 # Prompts used for generation checks during validation 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:", ] # ========================================================================= # Visual & Debug Helpers # ========================================================================= def _log_batch_samples(self, batch, title="SAMPLE CHECK"): """Decodes and prints samples from a batch to verify data integrity.""" if self.rank != 0: return num_show = min(len(batch), 2) input_tokens = batch[:, :-1] target_tokens = batch[:, 1:] display_targets = target_tokens.clone() if self.pad_token_id is not None: display_targets[display_targets == self.pad_token_id] = -100 print(f"\n{'='*20} {title} (First {num_show} samples) {'='*20}") for i in range(num_show): inp_ids = input_tokens[i].tolist() # Filter out pads and -100 for clean decoding valid_inp = [x for x in inp_ids if x != self.pad_token_id and x != -100] try: text_inp = self.tokenizer.decode(valid_inp) except Exception as e: text_inp = f"[Decode Error: {e}]" tgt_ids = display_targets[i].tolist() valid_tgt = [x for x in tgt_ids if x != -100] try: text_tgt = self.tokenizer.decode(valid_tgt) except Exception as e: text_tgt = f"[Decode Error: {e}]" print(f"[Sample {i}]") print(f" Input: {text_inp[:100]} ...") print(f" Target: {text_tgt[:100]} ...") print("-" * 40) print(f"{'='*60}\n") def _generate_debug_samples(self, max_new_tokens=50): """Runs greedy generation to sanity check model output.""" if self.rank != 0: return print(f"\n{'='*20} GENERATION CHECK (Eval Mode) {'='*20}") self.model.eval() for prompt in self.debug_prompts: input_ids = self.tokenizer.encode(prompt) if isinstance(input_ids, list): input_tensor = torch.tensor(input_ids, dtype=torch.long, device=self.device).unsqueeze(0) else: input_tensor = input_ids.to(self.device).unsqueeze(0) if input_tensor.dim() == 1: input_tensor = input_tensor.unsqueeze(0) generated = input_tensor.clone() with torch.no_grad(): for _ in range(max_new_tokens): cond = generated if generated.shape[1] > self.model.config.max_seq_len: cond = generated[:, -self.model.config.max_seq_len:] with autocast(device_type="cuda", dtype=self.autocast_dtype): logits = self.model(cond) next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True) generated = torch.cat((generated, next_token), dim=1) if self.eos_token_id is not None and next_token.item() == self.eos_token_id: break full_text = self.tokenizer.decode(generated[0].tolist()) new_text = full_text[len(prompt):] print(f"📝 Prompt: {prompt.strip()}") print(f"🤖 Gen: {new_text.strip().replace(chr(10), ' ')}") print("-" * 40) self.model.train() print(f"{'='*60}\n") # ========================================================================= # State Loading & Safetensors Logic # ========================================================================= def _strip_prefixes(self, state_dict, prefixes=None): if prefixes is None: prefixes = ("module.", "model.", "_orig_mod.") new = {} for k, v in state_dict.items(): new_k = k for p in prefixes: if k.startswith(p): new_k = k[len(p):] break new[new_k] = v return new def _is_lora_param(self, param_name): lora_indicators = ['lora_A', 'lora_B', 'lora_dropout'] return any(indicator in param_name for indicator in lora_indicators) def load_only_model_weights(self, checkpoint_path, map_location="cpu", strict=True, verbose=True): path_obj = Path(checkpoint_path) if not path_obj.exists(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") model_state = {} is_sharded = False # --- 1. Load the Checkpoint (Sharded or Single) --- if path_obj.is_dir(): index_file = path_obj / "model.safetensors.index.json" if index_file.exists(): if verbose: print(f"🔹 Detected sharded safetensors folder: {path_obj}") is_sharded = True import json with open(index_file, 'r') as f: index_data = json.load(f) weight_map = index_data.get("weight_map", {}) shards = set(weight_map.values()) for shard_name in shards: shard_path = path_obj / shard_name shard_weights = load_file(str(shard_path), device=str(map_location)) model_state.update(shard_weights) else: # Fallback to single file in dir possible = list(path_obj.glob("*.safetensors")) + list(path_obj.glob("*.pt")) if not possible: raise FileNotFoundError(f"No weights in {path_obj}") path_obj = possible[0] if not is_sharded and path_obj.is_file(): if path_obj.suffix == ".safetensors": # if verbose: print(f"🔹 Attempting to load safetensors: {path_obj}") # try: model_state = load_file(str(path_obj), device=str(map_location)) # except Exception as e: # 关键修复：如果 safetensors 加载失败，尝试转为 torch.load # if verbose: print(f"⚠️ Safetensors load failed, trying torch.load fallback...") # model_state = torch.load(str(path_obj), map_location=map_location) else: if verbose: print(f"🔹 Loading pickle (.pt): {path_obj}") ckpt = torch.load(str(path_obj), map_location=map_location) model_state = ckpt.get("model", ckpt.get("state_dict", ckpt)) # --- 2. Normalize Keys for Matching --- # Strip prefixes from CHECKPOINT keys ckpt_keys_map = {} # cleaned_key -> original_ckpt_key for k in model_state.keys(): clean_k = k.replace("module.", "").replace("_orig_mod.", "").replace("model.", "") ckpt_keys_map[clean_k] = k load_target = unwrap_model(self.model) target_state = load_target.state_dict() filtered_state = {} missing_in_ckpt = [] size_mismatches = [] # Loop through TARGET keys (which might have _orig_mod.) for k_target, v_target in target_state.items(): # Clean target key to match checkpoint style # We treat 'model.' as optional because HF checkpoints usually start with 'model.', # but internal state_dicts might not if unwrapped differently. k_target_clean = k_target.replace("module.", "").replace("_orig_mod.", "").replace("model.", "") # Look up in checkpoint map if k_target_clean in ckpt_keys_map: real_ckpt_key = ckpt_keys_map[k_target_clean] v_ckpt = model_state[real_ckpt_key] if v_ckpt.shape == v_target.shape: filtered_state[k_target] = v_ckpt # Use TARGET key for loading else: size_mismatches.append(f"{k_target} (ckpt: {v_ckpt.shape}, target: {v_target.shape})") else: missing_in_ckpt.append(k_target) # --- 3. Load --- try: msg = load_target.load_state_dict(filtered_state, strict=False) if verbose: print(f"✅ Weights loaded.") print(f" - Matched keys: {len(filtered_state)}") print(f" - Missing keys: {len(missing_in_ckpt)}") if len(missing_in_ckpt) > 0: # Filter out LoRA keys from error reporting real_missing = [k for k in missing_in_ckpt if not self._is_lora_param(k)] if real_missing: print(f"⚠️ Real Missing (non-LoRA): {len(real_missing)} (e.g. {real_missing[:3]})") print(f" (Target clean key example: {k_target_clean})") print(f" (Ckpt clean key example: {list(ckpt_keys_map.keys())[0]})") return msg except Exception as e: raise RuntimeError(f"Failed to load model weights: {e}") # ========================================================================= # Setup & Initialization # ========================================================================= def _setup_training(self): # 1. Data self.train_dataloader = self._get_dataloader("train") self.validation_dataloader = self._get_dataloader("validation") # 2. Grad Accumulation Calculation if self.training_config.use_grad_accum and self.training_config.grad_accum_steps == "auto": tokens_per_model_step = self.training_config.batch_size * self.model.config.max_seq_len * self.world_size self.grad_accum_steps = max(1, self.training_config.tokens_per_step // tokens_per_model_step) elif self.training_config.use_grad_accum: 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 * self.model.config.max_seq_len 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. Resume Model Weights (Auto-detect Recent) try: self.checkpointing.load_model_states("recent") except Exception: pass # 4. Device & Compilation self.device = torch.device(f"cuda:{self.rank}") if self.training_config.compile: self.model = torch.compile(self.model, mode=self.training_config.compile_mode) self.model.to(self.device) # 5. Explicit Resume (CLI override) resume_path = getattr(self.training_config, "resume_checkpoint_path", None) if resume_path: Logger.log(f"🔄 Forcing resume from: {resume_path}") self.load_only_model_weights(resume_path, map_location="cpu", strict=False) # 6. DDP Wrapping if self.use_ddp: self.model = initialize_model_ddp(self.model, self.rank) self.model.train() # 7. Optimizers & Schedulers self._initialize_optimizer() self._initialize_scheduler() self._initialize_scaler() # 8. Resume Training State (Optimizer/Step counts) try: self.checkpointing.load_training_states("recent") except Exception: pass def _is_lora_finetuning_mode(self): model = unwrap_model(self.model) config = getattr(model, 'config', None) if config is None: return False lora_flags = [ getattr(config, 'use_lora_phi_attention', False), getattr(config, 'use_lora_phi_mlp', False), getattr(config, 'use_lora_icl_attention', False), getattr(config, 'use_lora_icl_mlp', False), ] return any(lora_flags) def _initialize_optimizer(self): lora_enabled = self._is_lora_finetuning_mode() if lora_enabled: lora_params = [] for name, param in self.model.named_parameters(): if self._is_lora_param(name): param.requires_grad = True lora_params.append({'name': name, 'param': param}) else: param.requires_grad = False if self.rank == 0: total = sum(p.numel() for p in self.model.parameters()) trainable = sum(p['param'].numel() for p in lora_params) print(f"🔧 LoRA Mode: {trainable:,} trainable params ({100 * trainable / total:.3f}%)") self.optimizer = torch.optim.AdamW( [p['param'] for p in lora_params], lr=self.training_config.lr, betas=self.training_config.betas, weight_decay=self.training_config.weight_decay ) else: if self.rank == 0: print(f"🔧 Full Fine-Tuning Mode") self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=self.training_config.lr, betas=self.training_config.betas, weight_decay=self.training_config.weight_decay ) self.checkpointing.optimizer = self.optimizer def _initialize_scheduler(self): self.scheduler = get_cosine_schedule_with_warmup( optimizer=self.optimizer, num_warmup_steps=self.training_config.warmup_steps, num_training_steps=self.steps_per_epoch * self.training_config.max_epochs ) self.checkpointing.scheduler = self.scheduler def _initialize_scaler(self): if self.autocast_dtype == torch.float16: self.scaler = GradScaler("cuda") else: self.scaler = None self.checkpointing.scaler = self.scaler def _get_dataloader(self, split_name): num_workers = 1 if split_name == "validation" else max(1, self.training_config.num_workers - 1) return DataLoader( self.splits[split_name], batch_size=self.training_config.batch_size, num_workers=num_workers, shuffle=(split_name == "train" and self.samplers is None), sampler=None if self.samplers is None else self.samplers[split_name], pin_memory=True, drop_last=True ) # ========================================================================= # Core Training Logic # ========================================================================= 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[:, :-1] target_tokens = batch[:, 1:].clone() # Mask padding in target if self.pad_token_id is not None: target_tokens[target_tokens == self.pad_token_id] = -100 with autocast(device_type="cuda", dtype=self.autocast_dtype): logits = self.model(input_tokens) loss = unwrap_model(self.model).calculate_loss(logits, target_tokens, l1_loss_lambda=self.training_config.l1_loss_lambda) return loss def _validate(self): self.model.eval() loss_sum = torch.tensor(0.0, device=self.device) count = torch.tensor(0, device=self.device, dtype=torch.long) with torch.no_grad(): for batch in tqdm(self.validation_dataloader, desc="Validating", leave=False): loss = self._step_loss(batch).detach() loss_sum += loss * self.tokens_per_batch count += self.tokens_per_batch self._reduce(loss_sum) self._reduce(count) # Optional: Run visual generation check # self._generate_debug_samples() self.model.train() return (loss_sum / count).item() def _log_training_msg(self, resume=False): model = unwrap_model(self.model) msg = f"{'Resuming' if resume else 'Starting'} training | Model: {model.full_name} | Device: {self.device} | DDP: {self.use_ddp}" Logger.log(msg) def _ddp_barrier(self): if self.use_ddp: torch.distributed.barrier() def _reduce(self, item): if self.use_ddp: torch.distributed.all_reduce(item, op=torch.distributed.ReduceOp.SUM) def _train(self): self._setup_training() # Debug Data Check try: first_batch = next(iter(self.train_dataloader)) self._log_batch_samples(first_batch, title="TRAINING START DATA CHECK") except StopIteration: Logger.log("⚠️ Train dataloader is empty!") start_epoch = self.checkpointing.epoch start_step = self.checkpointing.step tokens_trained = self.checkpointing.tokens_trained resume = start_step != 0 self._log_training_msg(resume=resume) mr_step_loss = self.checkpointing.train_loss mr_validation_loss = self.checkpointing.val_loss for epoch in range(start_epoch, self.training_config.max_epochs): if self.train_dataloader.sampler is not None and hasattr(self.train_dataloader.sampler, "set_epoch"): self.train_dataloader.sampler.set_epoch(epoch) pbar = tqdm(total=self.steps_per_epoch, desc=f"Epoch {epoch}") if self.rank == 0 else None step_loss_accum = 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) if step >= self.steps_per_epoch: break # Fast-forward if resuming mid-epoch if resume and step < start_step: if pbar is not None and is_update_step: pbar.update(1) self.scheduler.step() # Sync scheduler continue elif resume: resume = False sync_ctx = self.model.no_sync() if (self.use_ddp and not is_update_step) else nullcontext() # --- Forward & Backward --- with sync_ctx: loss = self._step_loss(batch) loss = loss / self.grad_accum_steps step_loss_accum += loss.item() if self.scaler is not None: self.scaler.scale(loss).backward() else: loss.backward() if not is_update_step: continue # --- Optimizer Step --- if self.scaler is not None: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.scaler.step(self.optimizer) self.scaler.update() else: torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) self.optimizer.step() self.optimizer.zero_grad(set_to_none=True) self.scheduler.step() # --- Logging & Checkpointing --- tokens_trained = self._calculate_training_tokens(epoch, step + 1) mr_step_loss = step_loss_accum step_loss_accum = 0.0 if self.training_config.validation_steps is not None and (step + 1) % self.training_config.validation_steps == 0: self._ddp_barrier() mr_validation_loss = self._validate() self.checkpointing.save_checkpoint(epoch, step + 1, mr_step_loss, mr_validation_loss, tokens_trained) self._ddp_barrier() if pbar is not None: pbar.set_postfix(loss=f"{mr_step_loss:.4f}", val_loss=f"{mr_validation_loss:.4f}") pbar.update(1) # End of Epoch tokens_trained = self._calculate_training_tokens(epoch + 1, 0) self._ddp_barrier() mr_validation_loss = self._validate() self.checkpointing.save_checkpoint(epoch + 1, None, mr_step_loss, mr_validation_loss, tokens_trained) self._ddp_barrier() Logger.log(f"Epoch {epoch + 1} Complete | Val Loss: {mr_validation_loss:.4f}") def _train_ddp(self): 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() except Exception: print(f"[Rank {self.rank}] Exception occurred:") traceback.print_exc() finally: cleanup_ddp() def train(self): if self.training_config.use_ddp: self._train_ddp() else: self._train()