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 Bert_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', None) if self.pad_token_id is None: self.pad_token_id = 0 # print(self.tokenizer) self.mask_token_id = getattr(self.tokenizer, 'mask_token_id', None) if self.mask_token_id is None: raise ValueError("Tokenizer must have a mask_token_id for MLM (e.g., '[MASK]').") 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) # batch is expected shape (B, L) inputs = batch print(f"\n{'='*20} {title} (First {num_show} samples) {'='*20}") for i in range(num_show): seq_ids = inputs[i].tolist() # filter pads for display display_ids = [x for x in seq_ids if x != self.pad_token_id and x != -100] try: text = self.tokenizer.decode(display_ids) except Exception as e: text = f"[Decode Error: {e}]" print(f"[Sample {i}]") print(f" Tokens: {display_ids[:40]} ...") print(f" Text: {text[:200]} ...") print("-" * 40) print(f"{'='*60}\n") def _generate_debug_samples(self, max_new_tokens=50): """Runs greedy generation to sanity check model output. (Optional for MLM models)""" 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) # For non-autoregressive MLM, we can attempt to fill masks if any, otherwise just skip print("Skipping generation check for MLM trainer (not strictly causal generation).") self.model.train() print(f"{'='*60}\n") # ========================================================================= # MLM Masking Utility # ========================================================================= def _mask_tokens(self, inputs: torch.Tensor): """ Prepare masked tokens inputs/labels for masked language modeling: - 15% probability to select tokens for prediction. - of selected tokens: 80% -> [MASK], 10% -> random token, 10% -> keep same. Returns: inputs_masked: input ids with replacements applied labels: -100 for non-masked positions, otherwise original token id to compute loss """ device = inputs.device labels = inputs.clone() # Build a special token mask to avoid masking special tokens (if tokenizer supports) special_tokens_mask = None try: # tokenizer.get_special_tokens_mask expects list of ints for each example special_tokens_mask = [ self.tokenizer.get_special_tokens_mask(x, already_has_special_tokens=True) for x in inputs.tolist() ] special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool, device=device) except Exception: # fallback: only treat pad tokens as special special_tokens_mask = inputs.eq(self.pad_token_id) # Do not consider pad/special tokens for masking probability_matrix = torch.full(labels.shape, 0.15, device=device) probability_matrix.masked_fill_(special_tokens_mask, value=0.0) # Decide which tokens will be masked (boolean mask) masked_mask = torch.bernoulli(probability_matrix).bool() # Labels: only tokens we want predict keep original id; others set -100 labels[~masked_mask] = -100 # Prepare masked inputs inputs_masked = inputs.clone() # For each masked position, choose among 3 actions according to 80/10/10 rand_for_each = torch.rand(labels.shape, device=device) mask_token_mask = masked_mask & (rand_for_each < 0.8) random_token_mask = masked_mask & (rand_for_each >= 0.8) & (rand_for_each < 0.9) # the remaining masked_mask & (rand >=0.9) -> keep original token (no change to inputs_masked) # 80% -> replace with [MASK] if mask_token_mask.any(): inputs_masked[mask_token_mask] = self.mask_token_id # 10% -> replace with random token id from vocab if random_token_mask.any(): # vocab size detection try: vocab_size = self.tokenizer.vocab_size except Exception: try: vocab_size = len(self.tokenizer.get_vocab()) except Exception: raise RuntimeError("Unable to determine tokenizer vocab size for random token replacement.") rand_tokens = torch.randint(low=0, high=vocab_size, size=labels.shape, device=device) inputs_masked[random_token_mask] = rand_tokens[random_token_mask] return inputs_masked, labels # ========================================================================= # State Loading & Safetensors Logic (unchanged) # ========================================================================= 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": model_state = load_file(str(path_obj), device=str(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 --- 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 = [] for k_target, v_target in target_state.items(): k_target_clean = k_target.replace("module.", "").replace("_orig_mod.", "").replace("model.", "") 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 else: size_mismatches.append(f"{k_target} (ckpt: {v_ckpt.shape}, target: {v_target.shape})") else: missing_in_ckpt.append(k_target) 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: 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 (mostly unchanged) # ========================================================================= 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 (MLM-aware) # ========================================================================= def _calculate_training_tokens(self, epoch, step): return epoch * self.tokens_per_epoch + step * self.tokens_per_step def _step_loss(self, batch): """ For MLM we expect batch shape (B, L) with full sequences (no shifting). We will: - Move to device - Create masked inputs & labels via _mask_tokens() - Forward through model to get logits - Compute cross-entropy on masked positions (labels != -100) """ inputs = batch.to(self.device, non_blocking=True) # Build masked inputs and labels inputs_masked, labels = self._mask_tokens(inputs) # Forward with autocast(device_type="cuda", dtype=self.autocast_dtype): outputs = unwrap_model(self.model)(input_ids=inputs_masked, attention_mask=(inputs_masked != self.pad_token_id).long()) # outputs may be a ModelOutput with .logits or a tuple logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] # Compute loss (ignore_index = -100) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-100) 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() # accumulate by batch (we'll weight by batch size * seq_len) bsz = batch.size(0) seq_len = batch.size(1) tokens = bsz * seq_len loss_sum += loss * tokens count += tokens self._reduce(loss_sum) self._reduce(count) 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: {getattr(model, 'full_name', getattr(model, '__class__', 'model'))} | 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()