src/lmr/training/trainer_bert_1_17.py

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()