#!/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()