training/pretrain_v2/trainer.py

#!/usr/bin/env python3
"""
Self-contained pretrain trainer for Stage-A and Stage-B.
=========================================================
• Loads Qwen2.5-VL-3B-Instruct + LoRA (or resumes from Stage-A adapter)
• Causal LM loss with proper label masking (prompt tokens → -100)
• BF16, gradient accumulation, linear warmup + decay scheduler
• WandB logging, periodic eval, best-model checkpoint
"""

import json
import math
from contextlib import nullcontext
from datetime import datetime
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import AutoProcessor, AutoModelForVision2Seq, get_linear_schedule_with_warmup
from peft import LoraConfig, get_peft_model, PeftModel, TaskType
from tqdm import tqdm

try:
    import wandb
    _HAS_WANDB = True
except ImportError:
    _HAS_WANDB = False

from config import TrainConfig

# Qwen VL helper (handles dynamic resolution)
try:
    from qwen_vl_utils import process_vision_info as _qwen_process_vision_info
    _HAS_QWEN_UTILS = True
except ImportError:
    _HAS_QWEN_UTILS = False


class PretrainTrainer:
    """
    Training loop for pretrain_v2 Stage-A / Stage-B.

    Args:
        cfg:          TrainConfig dataclass
        train_loader: DataLoader (from dataset.py / collate_fn)
        val_loader:   DataLoader
        stage:        "A" or "B"
    """

    def __init__(
        self,
        cfg: TrainConfig,
        train_loader: DataLoader,
        val_loader: DataLoader,
        stage: str = "A",
    ):
        self.cfg          = cfg
        self.train_loader = train_loader
        self.val_loader   = val_loader
        self.stage        = stage

        self.device     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.global_step = 0
        self.best_val_loss = float("inf")

        self.output_dir = Path(cfg.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.train_log = self.output_dir / "train_metrics.jsonl"
        self.val_log   = self.output_dir / "val_metrics.jsonl"

        self._init_model()
        self._init_optimizer()

        if cfg.use_wandb and _HAS_WANDB:
            run_name = cfg.wandb_run_name or f"stage_{stage}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            wandb.init(
                project=cfg.wandb_project,
                name=run_name,
                config={
                    "stage": stage,
                    "lr": cfg.learning_rate,
                    "epochs": cfg.num_epochs,
                    "grad_acc": cfg.gradient_accumulation_steps,
                    "lora_r": cfg.lora.r,
                },
            )
        else:
            if cfg.use_wandb:
                print("⚠  wandb not available; skipping wandb logging.")
            cfg.use_wandb = False

    # ── model / optimizer ─────────────────────────────────────────────────────

    def _init_model(self):
        cfg = self.cfg
        print("=" * 60)
        print(f"Loading VLM backbone from {cfg.model_path}")

        self.processor = AutoProcessor.from_pretrained(
            cfg.model_path,
            trust_remote_code=True,
            min_pixels=4 * 28 * 28,
            max_pixels=cfg.max_pixels_single,
        )
        # Second processor with reduced pixel budget for multi-frame sequences
        self.processor_seq = AutoProcessor.from_pretrained(
            cfg.model_path,
            trust_remote_code=True,
            min_pixels=4 * 28 * 28,
            max_pixels=cfg.max_pixels_sequence,
        )

        for proc in (self.processor, self.processor_seq):
            if proc.tokenizer.pad_token is None:
                proc.tokenizer.pad_token    = proc.tokenizer.eos_token
                proc.tokenizer.pad_token_id = proc.tokenizer.eos_token_id

        model = AutoModelForVision2Seq.from_pretrained(
            cfg.model_path,
            torch_dtype=torch.bfloat16 if cfg.bf16 else torch.float32,
            trust_remote_code=True,
        )
        model.config.use_cache = False

        if cfg.pretrained_lora_path:
            # Stage-B: load Stage-A LoRA and continue training
            print(f"Loading Stage-A LoRA from {cfg.pretrained_lora_path}")
            model = PeftModel.from_pretrained(model, cfg.pretrained_lora_path, is_trainable=True)
            print("Stage-A LoRA loaded (trainable)")
        else:
            # Stage-A: fresh LoRA
            lora_cfg = LoraConfig(
                r=cfg.lora.r,
                lora_alpha=cfg.lora.alpha,
                target_modules=cfg.lora.target_modules,
                lora_dropout=cfg.lora.dropout,
                bias="none",
                task_type=TaskType.CAUSAL_LM,
            )
            model = get_peft_model(model, lora_cfg)
            model.print_trainable_parameters()

        try:
            model.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": False}
            )
        except TypeError:
            model.gradient_checkpointing_enable()

        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()

        model.to(self.device)
        self.model = model
        print(f"Model on {self.device}")
        print("=" * 60)

    def _init_optimizer(self):
        cfg     = self.cfg
        params  = [p for p in self.model.parameters() if p.requires_grad]
        if not params:
            raise RuntimeError("No trainable parameters found.")

        self.optimizer = torch.optim.AdamW(
            params,
            lr=cfg.learning_rate,
            weight_decay=cfg.weight_decay,
        )

        grad_acc           = max(1, cfg.gradient_accumulation_steps)
        updates_per_epoch  = math.ceil(len(self.train_loader) / grad_acc)
        total_steps        = updates_per_epoch * cfg.num_epochs
        warmup_steps       = int(total_steps * cfg.warmup_ratio)

        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_steps,
        )

        print(f"Optimizer: AdamW lr={cfg.learning_rate}")
        print(f"  batches/epoch={len(self.train_loader)}, "
              f"updates/epoch={updates_per_epoch}, "
              f"total={total_steps}, warmup={warmup_steps}")

    # ── label construction ────────────────────────────────────────────────────

    def _build_inputs_and_labels(self, batch: dict) -> dict:
        """
        Given a batch from collate_fn, build model inputs with masked labels.
        Frame format: batch['frames'] = List[List[PIL.Image]]
        """
        frames_list = batch["frames"]    # List[List[PIL]]
        prompts     = batch["prompts"]   # List[str]
        labels_text = batch["labels"]    # List[str]

        # Build chat messages per sample
        messages_batch = []
        for frames, prompt in zip(frames_list, prompts):
            content = [{"type": "image", "image": f} for f in frames]
            content.append({"type": "text", "text": prompt})
            messages_batch.append([{"role": "user", "content": content}])

        # Apply chat template → prompt texts only
        prompt_texts = [
            self.processor.apply_chat_template(
                msg, tokenize=False, add_generation_prompt=True
            )
            for msg in messages_batch
        ]

        # Full texts = prompt + label + eos
        eos = self.processor.tokenizer.eos_token or ""
        full_texts = [p + l + eos for p, l in zip(prompt_texts, labels_text)]

        # Build images_nested: list of list of PIL (required by Qwen processor)
        images_nested = frames_list  # already List[List[PIL]]

        # Use reduced-pixel processor for multi-frame to avoid OOM / truncation issues
        is_sequence = len(frames_list[0]) > 1
        proc = self.processor_seq if is_sequence else self.processor
        # For sequences, avoid hard truncation (image tokens alone can exceed 2048)
        max_len = None if is_sequence else 1024

        autocast_ctx = (
            torch.cuda.amp.autocast(dtype=torch.bfloat16)
            if self.cfg.bf16 else nullcontext()
        )

        with autocast_ctx:
            if max_len is not None:
                prompt_enc = proc(
                    text=prompt_texts, images=images_nested,
                    return_tensors="pt", padding=True,
                    truncation=True, max_length=max_len,
                )
                full_enc = proc(
                    text=full_texts, images=images_nested,
                    return_tensors="pt", padding=True,
                    truncation=True, max_length=max_len,
                )
            else:
                prompt_enc = proc(
                    text=prompt_texts, images=images_nested,
                    return_tensors="pt", padding=True,
                )
                full_enc = proc(
                    text=full_texts, images=images_nested,
                    return_tensors="pt", padding=True,
                )

        # Build labels tensor: mask prompt tokens with -100
        lbl = full_enc["input_ids"].clone()
        for i in range(lbl.shape[0]):
            prompt_len = int(prompt_enc["attention_mask"][i].sum().item())
            lbl[i, :prompt_len] = -100
        lbl[full_enc["attention_mask"] == 0] = -100
        full_enc["labels"] = lbl

        # Move to device; cast floats to model dtype
        model_dtype = next(self.model.parameters()).dtype
        inputs = {}
        for k, v in full_enc.items():
            if torch.is_tensor(v):
                inputs[k] = v.to(self.device, dtype=model_dtype if v.is_floating_point() else v.dtype)
            else:
                inputs[k] = v

        return inputs

    # ── forward / loss ────────────────────────────────────────────────────────

    def _compute_loss(self, batch: dict) -> torch.Tensor:
        inputs = self._build_inputs_and_labels(batch)
        autocast_ctx = (
            torch.cuda.amp.autocast(dtype=torch.bfloat16)
            if self.cfg.bf16 else nullcontext()
        )
        with autocast_ctx:
            outputs = self.model(**inputs)
        return outputs.loss

    # ── eval ─────────────────────────────────────────────────────────────────

    @torch.no_grad()
    def evaluate(self, epoch: int) -> float:
        self.model.eval()
        total_loss = 0.0
        n = 0
        for batch in tqdm(self.val_loader, desc="  Val"):
            try:
                loss = self._compute_loss(batch)
                total_loss += float(loss.detach())
                n += 1
            except Exception as e:
                print(f"  Val batch error: {e}")
                continue
        val_loss = total_loss / max(1, n)

        record = {"step": self.global_step, "epoch": epoch, "val/loss": val_loss}
        self._log_jsonl(self.val_log, record)
        if self.cfg.use_wandb and _HAS_WANDB:
            wandb.log(record, step=self.global_step)

        self.model.train()
        return val_loss

    # ── checkpoint ────────────────────────────────────────────────────────────

    def save(self, tag: str, is_best: bool = False):
        save_dir = self.output_dir / ("best_model" if is_best else tag)
        save_dir.mkdir(parents=True, exist_ok=True)
        self.model.save_pretrained(save_dir)
        self.processor.save_pretrained(save_dir)
        torch.save(
            {"global_step": self.global_step, "best_val_loss": self.best_val_loss},
            save_dir / "trainer_state.pt",
        )
        print(f"  ✓ Saved {'best model' if is_best else tag} → {save_dir}")
        if not is_best:
            self._rotate_checkpoints()

    def _rotate_checkpoints(self):
        limit = self.cfg.save_total_limit
        if limit <= 0:
            return
        ckpts = sorted(
            [p for p in self.output_dir.glob("checkpoint-*") if p.is_dir()],
            key=lambda p: int(p.name.split("-")[-1]) if p.name.split("-")[-1].isdigit() else 0,
        )
        for p in ckpts[:-limit]:
            import shutil
            shutil.rmtree(p, ignore_errors=True)

    # ── helpers ───────────────────────────────────────────────────────────────

    def _log_jsonl(self, path: Path, record: dict):
        record["time"] = datetime.now().isoformat(timespec="seconds")
        with open(path, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    # ── train loop ────────────────────────────────────────────────────────────

    def train(self):
        cfg      = self.cfg
        grad_acc = max(1, cfg.gradient_accumulation_steps)

        print("\n" + "=" * 60)
        print(f"Training Stage-{self.stage}  "
              f"epochs={cfg.num_epochs}  grad_acc={grad_acc}")
        print("=" * 60)

        for epoch in range(cfg.num_epochs):
            self.model.train()
            self.optimizer.zero_grad(set_to_none=True)

            win_loss, win_n = 0.0, 0
            pbar = tqdm(self.train_loader,
                        desc=f"Epoch {epoch+1}/{cfg.num_epochs}")

            for step, batch in enumerate(pbar):
                try:
                    loss = self._compute_loss(batch)
                except Exception as e:
                    print(f"\n  Batch {step} error: {e}")
                    self.optimizer.zero_grad(set_to_none=True)
                    continue

                scaled = loss / grad_acc
                scaled.backward()

                do_update = (
                    (step + 1) % grad_acc == 0
                    or (step + 1) == len(self.train_loader)
                )
                if not do_update:
                    win_loss += float(loss.detach())
                    win_n    += 1
                    continue

                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), cfg.max_grad_norm
                )
                self.optimizer.step()
                self.scheduler.step()
                self.optimizer.zero_grad(set_to_none=True)

                self.global_step += 1
                win_loss += float(loss.detach())
                win_n    += 1

                if self.global_step % cfg.logging_steps == 0:
                    avg = win_loss / max(1, win_n)
                    lr  = float(self.scheduler.get_last_lr()[0])
                    record = {
                        "step": self.global_step,
                        "epoch": epoch,
                        "train/loss": avg,
                        "train/lr": lr,
                    }
                    if torch.cuda.is_available():
                        record["gpu_mb"] = round(
                            torch.cuda.memory_allocated() / 1024 / 1024, 1
                        )
                    self._log_jsonl(self.train_log, record)
                    if cfg.use_wandb and _HAS_WANDB:
                        wandb.log(record, step=self.global_step)
                    pbar.set_postfix(loss=f"{avg:.4f}", lr=f"{lr:.2e}")
                    win_loss, win_n = 0.0, 0

                if cfg.save_steps > 0 and self.global_step % cfg.save_steps == 0:
                    self.save(f"checkpoint-{self.global_step}")

                if cfg.eval_steps > 0 and self.global_step % cfg.eval_steps == 0:
                    val_loss = self.evaluate(epoch)
                    print(f"\n  [step {self.global_step}] val_loss={val_loss:.4f}")
                    if val_loss < self.best_val_loss:
                        self.best_val_loss = val_loss
                        self.save("best_model", is_best=True)
                        print(f"  ★ New best! val_loss={val_loss:.4f}")

            # end-of-epoch eval
            val_loss = self.evaluate(epoch)
            print(f"\n[Epoch {epoch+1}] val_loss={val_loss:.4f}")
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                self.save("best_model", is_best=True)
                print(f"  ★ New best! val_loss={val_loss:.4f}")

        # final checkpoint
        self.save(f"checkpoint-{self.global_step}")

        print("\n" + "=" * 60)
        print(f"Stage-{self.stage} training complete!")
        print(f"Best val_loss: {self.best_val_loss:.4f}")
        print(f"Checkpoint dir: {self.output_dir}")
        print("=" * 60)

        if cfg.use_wandb and _HAS_WANDB:
            wandb.finish()