src/lmr/training/trainer_1_7.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 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()