""" SFT Trainer for TTA (Time-to-Accident) Regression - Loads Qwen2.5-VL backbone + LoRA - Trains belief_aggregator + TTA head - Supports resuming from SFT checkpoints (LoRA + heads + optional optimizer state) - Robust LoRA grad/update checks (no false-positive with grad accumulation / bf16 tiny updates) NEW in this version (for your request): 1) Reset best_val_loss when resuming (default: ON) 2) Optionally run a fresh evaluation on the NEW val dataset immediately after resume (default: ON) so "best" is re-defined under the new val split. """ from __future__ import annotations import os import json import time import math import random import logging import argparse from pathlib import Path from typing import Dict, List, Optional, Tuple, Any from collections import defaultdict import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torch.amp import GradScaler, autocast from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR from tqdm import tqdm # Optional deps try: import wandb HAS_WANDB = True except Exception: HAS_WANDB = False wandb = None try: from transformers import AutoProcessor, AutoModelForVision2Seq from peft import PeftModel, LoraConfig, get_peft_model HAS_TRANSFORMERS = True except Exception: HAS_TRANSFORMERS = False # Local imports from .dataset import SFTDataset, sft_collate_fn # ---------------- Logging ---------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logger = logging.getLogger("SFT.trainer") # ============================================================================ # Model Components # ============================================================================ class HazardHead(nn.Module): """Binary hazard head: outputs hazard_prob ∈ (0, 1). Initialized to be slightly below 0.5 (lean toward safe at start). """ def __init__(self, hidden_dim: int): super().__init__() self.fc = nn.Linear(hidden_dim, 1) nn.init.zeros_(self.fc.weight) self.fc.bias.data = torch.tensor([-1.0]) # sigmoid(-1) ≈ 0.27 def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: """Returns hazard_logit [B] (raw, pre-sigmoid).""" return self.fc(hidden_state).squeeze(-1) class TTAHead(nn.Module): """TTA Regression Head: outputs (tta_mean, tta_logvar).""" def __init__(self, hidden_dim: int, intermediate_dim: int = 512, dropout: float = 0.1): super().__init__() self.hidden_dim = hidden_dim self.intermediate_dim = intermediate_dim self.net = nn.Sequential( nn.Linear(hidden_dim, intermediate_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(intermediate_dim, intermediate_dim // 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(intermediate_dim // 2, 2), ) self._init_weights() def _init_weights(self): nn.init.zeros_(self.net[-1].weight) # bias: mean=5, logvar=0 self.net[-1].bias.data = torch.tensor([5.0, 0.0]) def forward(self, hidden_state: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: out = self.net(hidden_state) tta_mean = F.softplus(out[:, 0]) tta_logvar = out[:, 1] return tta_mean, tta_logvar class BeliefAggregator(nn.Module): """Aggregate token hidden states to a single belief vector. Strategies: - mean_pool : masked mean over all tokens -> [B, D] - last_token : hidden at last real token -> [B, D] - attention_pool : learned-query attention pool -> [B, D] - dual_pool : [mean(image_tokens) || mean(text_tokens)] -> [B, 2D] Requires image_token_id (and optionally video_token_id). This is P0.2 L1 "dual-modality pooling" — prevents the language prompt from being diluted 10× by image tokens. """ def __init__( self, hidden_dim: int, strategy: str = "mean_pool", image_token_id: Optional[int] = None, video_token_id: Optional[int] = None, ): super().__init__() self.hidden_dim = hidden_dim self.strategy = strategy self.image_token_id = image_token_id self.video_token_id = video_token_id if strategy == "attention_pool": self.query = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02) self.key_proj = nn.Linear(hidden_dim, hidden_dim) if strategy == "dual_pool" and image_token_id is None and video_token_id is None: raise ValueError("dual_pool requires image_token_id and/or video_token_id.") @property def output_dim(self) -> int: return 2 * self.hidden_dim if self.strategy == "dual_pool" else self.hidden_dim def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: if self.strategy == "mean_pool": return self._mean_pool(hidden_states, attention_mask) if self.strategy == "last_token": return self._last_token(hidden_states, attention_mask) if self.strategy == "attention_pool": return self._attention_pool(hidden_states, attention_mask) if self.strategy == "dual_pool": if input_ids is None: raise RuntimeError("dual_pool requires input_ids to separate image vs text tokens.") return self._dual_pool(hidden_states, attention_mask, input_ids) raise ValueError(f"Unknown strategy: {self.strategy}") def _mean_pool(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]) -> torch.Tensor: if attention_mask is None: return hidden_states.mean(dim=1) mask = attention_mask.unsqueeze(-1).float() masked = hidden_states * mask denom = mask.sum(dim=1).clamp(min=1e-9) return masked.sum(dim=1) / denom def _last_token(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]) -> torch.Tensor: if attention_mask is None: return hidden_states[:, -1, :] seq_lens = attention_mask.sum(dim=1).long() - 1 b = torch.arange(hidden_states.size(0), device=hidden_states.device) return hidden_states[b, seq_lens, :] def _attention_pool(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]) -> torch.Tensor: B, L, D = hidden_states.shape q = self.query.expand(B, -1, -1) k = self.key_proj(hidden_states) scores = torch.bmm(q, k.transpose(1, 2)) / math.sqrt(D) if attention_mask is not None: scores = scores.masked_fill(attention_mask.unsqueeze(1) == 0, -1e9) w = F.softmax(scores, dim=-1) return torch.bmm(w, hidden_states).squeeze(1) def _dual_pool( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor], input_ids: torch.Tensor, ) -> torch.Tensor: """Separately mean-pool image tokens and text tokens, concat -> [B, 2D].""" is_img = torch.zeros_like(input_ids, dtype=torch.bool) if self.image_token_id is not None: is_img = is_img | (input_ids == self.image_token_id) if self.video_token_id is not None: is_img = is_img | (input_ids == self.video_token_id) if attention_mask is not None: valid = attention_mask > 0 is_img = is_img & valid is_text = (~is_img) & valid else: is_text = ~is_img def _masked_mean(mask_bool: torch.Tensor) -> torch.Tensor: m = mask_bool.unsqueeze(-1).to(hidden_states.dtype) s = (hidden_states * m).sum(dim=1) denom = m.sum(dim=1).clamp(min=1e-6) return s / denom img_pool = _masked_mean(is_img) text_pool = _masked_mean(is_text) return torch.cat([img_pool, text_pool], dim=-1) # ============================================================================ # SFT Model # ============================================================================ class SFTModel(nn.Module): """VLM + LoRA + belief aggregator + HazardHead + TTAHead (dual head).""" def __init__( self, model_name: str = "Qwen/Qwen2.5-VL-3B-Instruct", pretrained_lora_path: Optional[str] = None, belief_strategy: str = "mean_pool", tta_intermediate_dim: int = 512, use_lora: bool = True, lora_r: int = 32, lora_alpha: int = 64, lora_dropout: float = 0.1, lora_target_modules: Optional[List[str]] = None, use_bf16: bool = True, device: str = "auto", max_pixels: Optional[int] = None, # None → 768*28*28 default # P0.3 PEFT upgrade flags use_dora: bool = False, use_rslora: bool = False, lora_init: str = "default", # "default" | "pissa" | "pissa_niter_16" | "olora" | "gaussian" attn_implementation: str = "flash_attention_2", ): super().__init__() if not HAS_TRANSFORMERS: raise RuntimeError("transformers/peft not available in this env.") self.model_name = model_name self.use_lora = use_lora self.use_bf16 = use_bf16 if lora_target_modules is None: lora_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] dtype = torch.bfloat16 if use_bf16 else torch.float32 logger.info(f"📦 Loading VLM: {model_name} (attn={attn_implementation})") self.vlm = AutoModelForVision2Seq.from_pretrained( model_name, torch_dtype=dtype, device_map="cuda:0", trust_remote_code=True, attn_implementation=attn_implementation, ) if hasattr(self.vlm, "config"): self.vlm.config.use_cache = False if hasattr(self.vlm, "gradient_checkpointing_enable"): try: self.vlm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) except TypeError: self.vlm.gradient_checkpointing_enable() if hasattr(self.vlm, "enable_input_require_grads"): try: self.vlm.enable_input_require_grads() except Exception: pass _min_pixels = 256 * 28 * 28 _max_pixels = max_pixels if max_pixels is not None else (768 * 28 * 28) logger.info(f" max_pixels: {_max_pixels} ({_max_pixels // (28*28)} tokens/frame max)") self.processor = AutoProcessor.from_pretrained( model_name, trust_remote_code=True, min_pixels=_min_pixels, max_pixels=_max_pixels, ) self.hidden_dim = getattr(self.vlm.config, "hidden_size", None) if self.hidden_dim is None: raise RuntimeError("Cannot infer hidden_size from model config.") logger.info(f" Hidden dim: {self.hidden_dim}") if use_lora: if pretrained_lora_path is not None: p = Path(pretrained_lora_path) if (p / "adapter_config.json").exists() and (p / "adapter_model.safetensors").exists(): logger.info(f" Loading pretrained LoRA via PeftModel.from_pretrained: {p}") self.vlm = PeftModel.from_pretrained(self.vlm, str(p), is_trainable=True) else: logger.warning(f"⚠️ pretrained_lora_path exists but missing adapter files: {p}. Creating new LoRA.") pretrained_lora_path = None if pretrained_lora_path is None: logger.info( f" Creating new LoRA (r={lora_r}, alpha={lora_alpha}, dropout={lora_dropout}, " f"use_dora={use_dora}, use_rslora={use_rslora}, init={lora_init})" ) lora_kwargs = dict( r=lora_r, lora_alpha=lora_alpha, target_modules=lora_target_modules, lora_dropout=lora_dropout, bias="none", task_type="CAUSAL_LM", ) if use_dora: lora_kwargs["use_dora"] = True if use_rslora: lora_kwargs["use_rslora"] = True if lora_init and lora_init != "default": # peft accepts: True | False | "gaussian" | "olora" | "pissa" | "pissa_niter_[N]" | "loftq" lora_kwargs["init_lora_weights"] = lora_init lora_config = LoraConfig(**lora_kwargs) self.vlm = get_peft_model(self.vlm, lora_config) base = self.get_base_model() if hasattr(base, "config"): base.config.use_cache = False if hasattr(base, "gradient_checkpointing_enable"): try: base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) except TypeError: base.gradient_checkpointing_enable() if hasattr(base, "enable_input_require_grads"): try: base.enable_input_require_grads() except Exception: pass try: self.vlm.print_trainable_parameters() except Exception: pass self._register_requires_grad_hooks() self.device = next(self.vlm.parameters()).device self.dtype = next(self.vlm.parameters()).dtype # Grab image / video token ids from the VLM config (Qwen2.5-VL: 151655 / 151656). _cfg = getattr(self.vlm, "config", None) img_tok_id = getattr(_cfg, "image_token_id", None) vid_tok_id = getattr(_cfg, "video_token_id", None) if img_tok_id is None: img_tok_id = 151655 # Qwen2.5-VL fallback if vid_tok_id is None: vid_tok_id = 151656 # Qwen2.5-VL fallback self.belief_aggregator = BeliefAggregator( self.hidden_dim, strategy=belief_strategy, image_token_id=img_tok_id, video_token_id=vid_tok_id, ).to(self.device, dtype=self.dtype) belief_dim = self.belief_aggregator.output_dim self.belief_dim = belief_dim self.hazard_head = HazardHead(belief_dim).to(self.device, dtype=self.dtype) self.tta_head = TTAHead(belief_dim, intermediate_dim=tta_intermediate_dim).to(self.device, dtype=self.dtype) trainable = [(n, p) for n, p in self.vlm.named_parameters() if p.requires_grad] lora_trainable = [(n, p) for n, p in trainable if "lora_" in n.lower()] logger.info(f" Trainable tensors: {len(trainable)}; LoRA trainable tensors: {len(lora_trainable)}") logger.info("✅ SFTModel initialized") logger.info(f" Device: {self.device}") logger.info(f" Dtype: {self.dtype}") logger.info(f" Belief strategy: {belief_strategy}") def get_base_model(self): if hasattr(self.vlm, "get_base_model"): try: return self.vlm.get_base_model() except Exception: pass return getattr(self.vlm, "model", self.vlm) def _register_requires_grad_hooks(self): def _force_requires_grad_hook(_module, _inp, out): try: if torch.is_tensor(out) and out.is_floating_point(): out.requires_grad_(True) elif isinstance(out, (tuple, list)): for t in out: if torch.is_tensor(t) and t.is_floating_point(): t.requires_grad_(True) except Exception: return base_model = self.get_base_model() try: emb = base_model.get_input_embeddings() if hasattr(base_model, "get_input_embeddings") else None if emb is not None: emb.register_forward_hook(_force_requires_grad_hook) logger.info("✅ Registered requires_grad hook on TEXT embeddings") except Exception as e: logger.warning(f"⚠️ Failed to hook TEXT embeddings: {e}") try: hooked = False for name in ["visual", "vision_tower", "vision_model", "vision_encoder"]: if hasattr(base_model, name): getattr(base_model, name).register_forward_hook(_force_requires_grad_hook) logger.info(f"✅ Registered requires_grad hook on VISION module: {name}") hooked = True break if not hooked: for n, m in base_model.named_modules(): nl = n.lower() if any(k in nl for k in ["visual", "vision", "patch_embed", "patch_embedding", "img_embed"]): m.register_forward_hook(_force_requires_grad_hook) logger.info(f"✅ Registered requires_grad hook on VISION submodule: {n}") break except Exception as e: logger.warning(f"⚠️ Failed to hook VISION module: {e}") def encode_observation(self, batch_inputs: Dict[str, torch.Tensor]) -> torch.Tensor: moved: Dict[str, Any] = {} for k, v in batch_inputs.items(): if not isinstance(v, torch.Tensor): moved[k] = v continue if k == "pixel_values": moved[k] = v.to(self.device, dtype=self.dtype, non_blocking=True) else: moved[k] = v.to(self.device, non_blocking=True) base = self.get_base_model() hidden_states = None core = getattr(base, "model", None) if core is not None: try: out = core( input_ids=moved["input_ids"], attention_mask=moved.get("attention_mask"), pixel_values=moved.get("pixel_values"), image_grid_thw=moved.get("image_grid_thw"), use_cache=False, return_dict=True, ) hidden_states = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] except TypeError: hidden_states = None if hidden_states is None: out = base( input_ids=moved["input_ids"], attention_mask=moved.get("attention_mask"), pixel_values=moved.get("pixel_values"), image_grid_thw=moved.get("image_grid_thw"), use_cache=False, return_dict=True, output_hidden_states=True, ) if not hasattr(out, "hidden_states") or out.hidden_states is None: raise RuntimeError("Model output has no hidden_states; cannot build belief.") hidden_states = out.hidden_states[-1] if hidden_states.dim() != 3 or hidden_states.size(-1) != self.hidden_dim: raise RuntimeError(f"Unexpected hidden_states shape {tuple(hidden_states.shape)}, expected [B,L,{self.hidden_dim}]") return self.belief_aggregator( hidden_states, moved.get("attention_mask"), moved.get("input_ids"), ) def forward(self, batch_inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: belief = self.encode_observation(batch_inputs) hazard_logit = self.hazard_head(belief) # raw logit [B] hazard_prob = torch.sigmoid(hazard_logit) # probability [B] tta_mean, tta_logvar = self.tta_head(belief) return { "hazard_logit": hazard_logit, "hazard_prob": hazard_prob, "tta_mean": tta_mean, "tta_logvar": tta_logvar, "belief": belief.detach(), } def save_checkpoint(self, save_dir: str, epoch: int = 0, step: int = 0): save_dir = Path(save_dir) save_dir.mkdir(parents=True, exist_ok=True) if self.use_lora: lora_dir = save_dir / "vlm_lora" self.vlm.save_pretrained(lora_dir) logger.info(f" Saved LoRA to {lora_dir}") torch.save(self.belief_aggregator.state_dict(), save_dir / "belief_aggregator.pt") torch.save(self.hazard_head.state_dict(), save_dir / "hazard_head.pt") torch.save(self.tta_head.state_dict(), save_dir / "tta_head.pt") cfg = { "model_name": self.model_name, "hidden_dim": self.hidden_dim, "belief_strategy": self.belief_aggregator.strategy, "belief_dim": self.belief_aggregator.output_dim, "image_token_id": self.belief_aggregator.image_token_id, "video_token_id": self.belief_aggregator.video_token_id, "tta_intermediate_dim": self.tta_head.intermediate_dim, "epoch": epoch, "step": step, } with open(save_dir / "config.json", "w") as f: json.dump(cfg, f, indent=2) logger.info(f"✅ Checkpoint saved to {save_dir}") # ============================================================================ # Loss # ============================================================================ def compute_sft_loss( hazard_logit: torch.Tensor, tta_mean: torch.Tensor, tta_logvar: torch.Tensor, hazard_label: torch.Tensor, hazard_weight: torch.Tensor, is_ego_positive: torch.Tensor, is_censored: torch.Tensor, tta_label: torch.Tensor, tta_cap: float = 10.0, nll_weight: float = 0.5, tta_obs_weight: float = 1.0, tta_cens_weight: float = 0.5, # legacy kwarg kept for callers that still pass hazard_prob hazard_prob: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Dict[str, float]]: """ Dual-head SFT loss. Hazard head (all samples): L_hazard = weighted_BCE_with_logits(hazard_logit, hazard_label) weights: ego_pos=1.0, safe_neg=1.0, non_ego=0.35, pre_risky=0.8 TTA head (ego_positive only): Observed (TTA ≤ 10s): MSE + nll_weight * NLL Censored (TTA > 10s): relu(tta_cap - tta_mean)² Non-ego and safe_neg: NO TTA gradient. """ hl = hazard_label.float() hw = hazard_weight.float() tm = tta_mean.float() tlv = tta_logvar.float() tl = tta_label.float().clamp(0.1, tta_cap) var = torch.exp(tlv).clamp(min=1e-6) zero = tta_mean.new_zeros(()) # ── hazard loss (logits → safe with autocast) ───────────────────────── hl_logit = hazard_logit.float() bce_unreduced = F.binary_cross_entropy_with_logits(hl_logit, hl, reduction="none") L_hazard = (bce_unreduced * hw).mean() # keep hp for metrics hp = torch.sigmoid(hl_logit).detach() # ── TTA loss: ego_positive observed ────────────────────────────────────── obs_mask = is_ego_positive & (~is_censored) cens_mask = is_ego_positive & is_censored if obs_mask.any(): m = tm[obs_mask]; l = tl[obs_mask] lv = tlv[obs_mask]; v = var[obs_mask] mse = F.mse_loss(m, l) nll = 0.5 * (lv + (m - l).pow(2) / v).mean() L_tta_obs = mse + nll_weight * nll else: mse = zero; nll = zero; L_tta_obs = zero # ── TTA loss: ego_positive censored ────────────────────────────────────── if cens_mask.any(): cm = tm[cens_mask] L_tta_cens = F.relu(tta_cap - cm).pow(2).mean() else: L_tta_cens = zero loss = L_hazard + tta_obs_weight * L_tta_obs + tta_cens_weight * L_tta_cens # ── metrics ────────────────────────────────────────────────────────────── n_obs = int(obs_mask.sum()) n_cens = int(cens_mask.sum()) n_pos = n_obs + n_cens n_noneego = int((~is_ego_positive & (hazard_label == 0) & (hazard_weight < 0.5)).sum()) hazard_pred_bin = (hp > 0.5).float() hazard_correct = (hazard_pred_bin == hl).float().mean() metrics: Dict[str, float] = { "loss": float(loss.detach()), "loss_hazard": float(L_hazard.detach()), "loss_tta_obs": float(L_tta_obs.detach()), "loss_tta_cens": float(L_tta_cens.detach()), "hazard_acc": float(hazard_correct), "n_obs": n_obs, "n_cens": n_cens, "n_pos": n_pos, "n_noneego": n_noneego, } if obs_mask.any(): metrics["tta_mae"] = float((tm[obs_mask] - tl[obs_mask]).abs().mean().detach()) metrics["tta_rmse"] = float((tm[obs_mask] - tl[obs_mask]).pow(2).mean().sqrt().detach()) metrics["mse_loss"] = float(mse.detach()) metrics["nll_loss"] = float(nll.detach()) else: metrics.update({"tta_mae": 0.0, "tta_rmse": 0.0, "mse_loss": 0.0, "nll_loss": 0.0}) return loss, metrics # ============================================================================ # Resume Helpers # ============================================================================ def _is_sft_ckpt_dir(d: Path) -> bool: return ( d.is_dir() and (d / "tta_head.pt").exists() and (d / "belief_aggregator.pt").exists() and (d / "config.json").exists() and (d / "vlm_lora" / "adapter_config.json").exists() and (d / "vlm_lora" / "adapter_model.safetensors").exists() ) def _parse_step(name: str) -> int: if name.startswith("step_"): try: return int(name.split("_", 1)[1]) except Exception: return -1 return -1 def find_auto_resume_checkpoint(output_dir: Path, experiment_name: str) -> Optional[Path]: candidates: List[Path] = [] exp_dir = output_dir / experiment_name if exp_dir.exists(): for child in exp_dir.iterdir(): if _is_sft_ckpt_dir(child): candidates.append(child) if not candidates: for d1 in output_dir.iterdir(): if not d1.is_dir(): continue for d2 in d1.iterdir(): if _is_sft_ckpt_dir(d2): candidates.append(d2) if not candidates: return None step_cands = [(c, _parse_step(c.name)) for c in candidates] step_cands = [x for x in step_cands if x[1] >= 0] if step_cands: step_cands.sort(key=lambda x: x[1], reverse=True) return step_cands[0][0] epoch_cands = [] for c in candidates: if c.name.startswith("epoch_"): try: epoch_cands.append((c, int(c.name.split("_", 1)[1]))) except Exception: pass if epoch_cands: epoch_cands.sort(key=lambda x: x[1], reverse=True) return epoch_cands[0][0] for c in candidates: if c.name == "best": return c candidates.sort(key=lambda p: p.stat().st_mtime, reverse=True) return candidates[0] def load_sft_heads(model: SFTModel, ckpt_dir: Path): b_path = ckpt_dir / "belief_aggregator.pt" h_path = ckpt_dir / "hazard_head.pt" t_path = ckpt_dir / "tta_head.pt" model.belief_aggregator.load_state_dict(torch.load(b_path, map_location=model.device), strict=True) if h_path.exists(): model.hazard_head.load_state_dict(torch.load(h_path, map_location=model.device), strict=True) logger.info(f" Loaded hazard_head from {h_path}") else: logger.warning("⚠️ hazard_head.pt not found in checkpoint; using fresh init.") model.tta_head.load_state_dict(torch.load(t_path, map_location=model.device), strict=True) logger.info(f"✅ Loaded heads from {ckpt_dir}") try: last = model.tta_head.net[-1] if hasattr(last, "bias") and last.bias is not None: logger.info(f" TTAHead last-layer bias(after load) = {last.bias.detach().float().cpu().tolist()}") except Exception: pass def compute_calibration_error( predictions: np.ndarray, uncertainties: np.ndarray, labels: np.ndarray, num_bins: int = 10 ) -> Tuple[float, np.ndarray, np.ndarray]: """ Compute Expected Calibration Error (ECE) for regression. Returns: ece: scalar observed_freq: per-bin observed frequencies expected_freq: per-bin expected frequencies """ predictions = np.asarray(predictions) uncertainties = np.asarray(uncertainties) labels = np.asarray(labels) # normalized error = |pred-label| / std errors = np.abs(predictions - labels) normalized_errors = errors / np.maximum(uncertainties, 1e-6) if normalized_errors.size == 0: return 0.0, np.array([]), np.array([]) # bin edges over normalized error max_ne = float(np.max(normalized_errors)) if not np.isfinite(max_ne) or max_ne <= 0: return 0.0, np.array([]), np.array([]) bin_edges = np.linspace(0.0, max_ne, num_bins + 1) observed_freq = [] expected_freq = [] # expected coverage for Gaussian within z std: erf(z/sqrt(2)) # (note: this is a simple reference curve; you can replace later with your preferred ECE definition) sqrt2 = math.sqrt(2.0) for i in range(num_bins): lo, hi = bin_edges[i], bin_edges[i + 1] mask = (normalized_errors >= lo) & (normalized_errors < hi) if mask.sum() == 0: continue z = 0.5 * (lo + hi) expected = math.erf(z / sqrt2) # in [0,1] observed = float(mask.mean()) observed_freq.append(observed) expected_freq.append(expected) observed_freq = np.asarray(observed_freq, dtype=np.float32) expected_freq = np.asarray(expected_freq, dtype=np.float32) ece = float(np.abs(observed_freq - expected_freq).mean()) if observed_freq.size > 0 else 0.0 return ece, observed_freq, expected_freq # ============================================================================ # Trainer # ============================================================================ class SFTTrainer: def __init__( self, model: SFTModel, train_dataset: SFTDataset, val_dataset: Optional[SFTDataset], num_epochs: int = 10, batch_size: int = 4, gradient_accumulation_steps: int = 4, learning_rate: float = 1e-4, tta_head_lr: float = 1e-3, vlm_lr_multiplier: float = 0.1, weight_decay: float = 0.01, max_grad_norm: float = 1.0, mse_weight: float = 1.0, nll_weight: float = 0.5, use_curriculum: bool = True, scheduler_type: str = "cosine", warmup_ratio: float = 0.1, output_dir: str = "./checkpoints/sft", experiment_name: str = "sft_default", logging_steps: int = 1250, eval_steps: int = 2500, save_steps: int = 5000, save_total_limit: int = 3, use_amp: bool = True, use_wandb: bool = True, wandb_project: str = "lkalert-sft", lora_update_patience: int = 30, disable_metadata_prompt: bool = False, # P0.1: drop weather/road/time context ): self.model = model self.train_dataset = train_dataset self.val_dataset = val_dataset self.num_epochs = num_epochs self.batch_size = batch_size self.gradient_accumulation_steps = gradient_accumulation_steps self.learning_rate = learning_rate self.tta_head_lr = tta_head_lr self.vlm_lr_multiplier = vlm_lr_multiplier self.weight_decay = weight_decay self.max_grad_norm = max_grad_norm self.mse_weight = mse_weight self.nll_weight = nll_weight self.use_curriculum = use_curriculum self.scheduler_type = scheduler_type self.warmup_ratio = warmup_ratio self.output_dir = Path(output_dir) / experiment_name self.output_dir.mkdir(parents=True, exist_ok=True) self.experiment_name = experiment_name self.logging_steps = logging_steps self.eval_steps = eval_steps self.save_steps = save_steps self.save_total_limit = save_total_limit # AMP self.use_amp = use_amp self.amp_dtype = torch.bfloat16 if self.model.dtype == torch.bfloat16 else torch.float16 self.use_scaler = self.use_amp and (self.amp_dtype == torch.float16) self.scaler = GradScaler("cuda", enabled=self.use_scaler) if self.use_amp else None logger.info(f"AMP enabled={self.use_amp}, amp_dtype={self.amp_dtype}, scaler_enabled={self.use_scaler}") # wandb self.use_wandb = bool(use_wandb and HAS_WANDB) if self.use_wandb: wandb.init( project=wandb_project, name=experiment_name, config={ "num_epochs": num_epochs, "batch_size": batch_size, "grad_accum": gradient_accumulation_steps, "learning_rate": learning_rate, "tta_head_lr": tta_head_lr, "vlm_lr_multiplier": vlm_lr_multiplier, "use_curriculum": use_curriculum, }, ) # loaders/optim/sched self._create_dataloaders() self._create_optimizer() self._create_scheduler() # training state self.global_step = 0 self.current_epoch = 0 self.best_ckpt_score = float("-inf") # higher is better (0.6*f1 - 0.4*mae/10) self.saved_checkpoints: List[Path] = [] # LoRA checks self._lora_grad_verified = False self._lora_update_verified = False self._lora_update_zero_steps = 0 self.lora_update_patience = int(lora_update_patience) # P0.1: if True, drop weather/road_type/time_of_day from the prompt. self.disable_metadata_prompt = bool(disable_metadata_prompt) logger.info("✅ SFTTrainer initialized") if self.disable_metadata_prompt: logger.info(" [P0.1] disable_metadata_prompt=True → metadata context stripped from prompt") logger.info(f" Output dir: {self.output_dir}") logger.info(f" Total steps: {self.total_steps}") logger.info(f" Effective batch size: {batch_size * gradient_accumulation_steps}") def _create_dataloaders(self): self.train_loader = DataLoader( self.train_dataset, batch_size=self.batch_size, shuffle=True, collate_fn=sft_collate_fn, num_workers=4, pin_memory=True, ) self.train_sampler = None self.val_loader = None if self.val_dataset is not None: self.val_loader = DataLoader( self.val_dataset, batch_size=self.batch_size * 2, shuffle=False, collate_fn=sft_collate_fn, num_workers=4, pin_memory=True, ) steps_per_epoch = max(1, len(self.train_loader) // self.gradient_accumulation_steps) self.total_steps = steps_per_epoch * self.num_epochs def _create_optimizer(self): vlm_params = [] for _, p in self.model.vlm.named_parameters(): if p.requires_grad: vlm_params.append(p) head_params = ( list(self.model.belief_aggregator.parameters()) + list(self.model.hazard_head.parameters()) + list(self.model.tta_head.parameters()) ) self.optimizer = AdamW( [ {"params": vlm_params, "lr": self.learning_rate * self.vlm_lr_multiplier}, {"params": head_params, "lr": self.tta_head_lr}, ], weight_decay=self.weight_decay, ) logger.info(f" VLM params: {len(vlm_params)} (lr={self.learning_rate * self.vlm_lr_multiplier})") logger.info(f" Head params: {len(head_params)} (lr={self.tta_head_lr})") def _create_scheduler(self): warmup_steps = int(self.total_steps * self.warmup_ratio) if self.scheduler_type == "cosine": warmup = LinearLR(self.optimizer, start_factor=0.1, end_factor=1.0, total_iters=max(1, warmup_steps)) cosine = CosineAnnealingLR(self.optimizer, T_max=max(1, self.total_steps - warmup_steps), eta_min=1e-6) self.scheduler = SequentialLR(self.optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps]) elif self.scheduler_type == "linear": self.scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=0.1, total_iters=max(1, self.total_steps)) else: self.scheduler = None def _build_prompt(self, batch: Dict, idx: int) -> str: metadata = batch["metadata"][idx] window_type = batch["window_types"][idx] window_str = f"{2.0 if window_type == 'standard' else 3.0}s" # P0.1: prompt ablation — drop metadata context entirely. # Tests whether weather/road_type/time_of_day in the prompt contributes any # real signal vs. merely being boilerplate for the VLM. if getattr(self, "disable_metadata_prompt", False): return ( f"Analyze this driving sequence ({window_str} window).\n" f"Estimate the time to potential collision. Output a single number in seconds." ) context_parts = [] if metadata.get("weather"): context_parts.append(f"Weather: {metadata['weather']}") if metadata.get("road_type"): context_parts.append(f"Road: {metadata['road_type']}") if metadata.get("time_of_day"): context_parts.append(f"Time: {metadata['time_of_day']}") context = ", ".join(context_parts) if context_parts else "Urban driving" return ( f"Analyze this driving sequence ({window_str} window).\n" f"Context: {context}\n" f"Estimate the time to potential collision. Output a single number in seconds." ) def _prepare_batch(self, batch: Dict) -> Dict[str, torch.Tensor]: system_prompt = "You are a driving safety AI analyzing dashcam footage for collision risk." texts = [] images = batch["images"] # List[List[PIL.Image]]: B x K proc = self.model.processor apply_chat = proc.apply_chat_template if hasattr(proc, "apply_chat_template") else proc.tokenizer.apply_chat_template for i in range(len(batch["video_ids"])): user_text = self._build_prompt(batch, i) frames = images[i] content = [{"type": "image"} for _ in range(len(frames))] content.append({"type": "text", "text": user_text}) messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": content}] texts.append(apply_chat(messages, tokenize=False, add_generation_prompt=False)) processed = proc( text=texts, images=images, return_tensors="pt", padding=True, truncation=True, ) return processed # -------- LoRA checks -------- def _verify_lora_grads_once(self): if self._lora_grad_verified: return lora = [(n, p) for n, p in self.model.vlm.named_parameters() if p.requires_grad and "lora_" in n.lower()] if not lora: logger.warning("⚠️ No trainable LoRA parameters found.") self._lora_grad_verified = True return non_none = 0 non_zero = 0 for _, p in lora: if p.grad is not None: non_none += 1 if float(p.grad.detach().abs().sum().item()) > 0: non_zero += 1 logger.info(f"🔎 LoRA grad check: total={len(lora)}, grad_non_none={non_none}, grad_non_zero={non_zero}") if non_none == 0 or non_zero == 0: logger.warning("⚠️ LoRA grads are missing/zero at this moment (may be before first real update).") else: logger.info("✅ LoRA gradient flow verified.") self._lora_grad_verified = True def _pick_probe_lora_param(self) -> Optional[Tuple[str, torch.nn.Parameter]]: candidates = [] for n, p in self.model.vlm.named_parameters(): if not p.requires_grad: continue if "lora_" not in n.lower(): continue if p.grad is None: continue if float(p.grad.detach().abs().sum().item()) == 0.0: continue candidates.append((n, p)) if candidates: return random.choice(candidates) fallback = [(n, p) for n, p in self.model.vlm.named_parameters() if p.requires_grad and "lora_" in n.lower()] if not fallback: return None return random.choice(fallback) def _post_step_lora_update_check(self, probe_name: Optional[str], before_fp32: Optional[torch.Tensor]): if probe_name is None or before_fp32 is None: return probe_param = None for n, p in self.model.vlm.named_parameters(): if n == probe_name: probe_param = p break if probe_param is None: return after_fp32 = probe_param.detach().float() delta = float((after_fp32 - before_fp32).abs().mean().item()) if delta == 0.0: self._lora_update_zero_steps += 1 lr0 = self.optimizer.param_groups[0]["lr"] logger.warning( f"⚠️ LoRA update probe delta==0 (name='{probe_name}'), " f"consecutive_zero_steps={self._lora_update_zero_steps}, lr={lr0:.2e}. " f"Will only raise after {self.lora_update_patience} consecutive steps." ) if self._lora_update_zero_steps >= self.lora_update_patience: raise RuntimeError( "LoRA probe parameter did not change for many optimizer steps. " "Likely lr too small for bf16 rounding, or LoRA params not in optimizer, or training graph bypassing LoRA." ) else: if not self._lora_update_verified: logger.info(f"✅ LoRA update verified: probe='{probe_name}', mean_abs_delta={delta:.6e}") self._lora_update_verified = True self._lora_update_zero_steps = 0 # -------- train/eval -------- def _batch_to_device(self, batch: Dict, keys) -> Dict: return {k: batch[k].to(self.model.device) for k in keys if k in batch} def train_step(self, batch: Dict) -> Dict[str, float]: self.model.train() inputs = self._prepare_batch(batch) t = self._batch_to_device(batch, [ "tta_labels", "hazard_labels", "hazard_weights", "is_ego_positive", "is_censored", ]) with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp): out = self.model(inputs) loss, metrics = compute_sft_loss( hazard_logit = out["hazard_logit"], tta_mean = out["tta_mean"], tta_logvar = out["tta_logvar"], hazard_label = t["hazard_labels"], hazard_weight = t["hazard_weights"], is_ego_positive = t["is_ego_positive"], is_censored = t["is_censored"], tta_label = t["tta_labels"], nll_weight = self.nll_weight, ) loss = loss / self.gradient_accumulation_steps if self.use_scaler: self.scaler.scale(loss).backward() else: loss.backward() if not self._lora_grad_verified: self._verify_lora_grads_once() return metrics def _optimizer_step(self): probe = self._pick_probe_lora_param() probe_name = probe[0] if probe else None before_fp32 = probe[1].detach().float().clone() if probe else None if self.use_scaler: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) if self.use_scaler: self.scaler.step(self.optimizer) self.scaler.update() else: self.optimizer.step() self.optimizer.zero_grad(set_to_none=True) if self.scheduler is not None: self.scheduler.step() self._post_step_lora_update_check(probe_name, before_fp32) self.global_step += 1 @torch.no_grad() def evaluate(self) -> Dict[str, float]: if self.val_loader is None: return {} self.model.eval() total_loss = 0.0 n = 0 preds, labels_all, stds = [], [], [] all_hazard_prob: List[np.ndarray] = [] all_hazard_label: List[np.ndarray] = [] all_is_noneego: List[np.ndarray] = [] all_is_ego_pos: List[np.ndarray] = [] for batch in tqdm(self.val_loader, desc="Evaluating", leave=False, ncols=60): inputs = self._prepare_batch(batch) t = self._batch_to_device(batch, [ "tta_labels", "hazard_labels", "hazard_weights", "is_ego_positive", "is_censored", ]) is_noneego_b = batch.get("is_non_ego", torch.zeros(len(batch["video_ids"]), dtype=torch.bool)) with autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp): out = self.model(inputs) loss, _ = compute_sft_loss( hazard_logit = out["hazard_logit"], tta_mean = out["tta_mean"], tta_logvar = out["tta_logvar"], hazard_label = t["hazard_labels"], hazard_weight = t["hazard_weights"], is_ego_positive = t["is_ego_positive"], is_censored = t["is_censored"], tta_label = t["tta_labels"], nll_weight = self.nll_weight, ) total_loss += float(loss.item()) n += 1 tta_mean = out["tta_mean"].detach().float().cpu().numpy() tta_label_np = t["tta_labels"].detach().float().cpu().numpy() tta_std = torch.exp(0.5 * out["tta_logvar"].detach().float()).cpu().numpy() preds.append(tta_mean) labels_all.append(tta_label_np) stds.append(tta_std) all_hazard_prob.append(out["hazard_prob"].detach().float().cpu().numpy()) all_hazard_label.append(t["hazard_labels"].detach().float().cpu().numpy()) all_is_noneego.append(is_noneego_b.cpu().numpy()) all_is_ego_pos.append(t["is_ego_positive"].cpu().numpy()) preds = np.concatenate(preds) if preds else np.zeros(0, np.float32) labels_all = np.concatenate(labels_all) if labels_all else np.zeros(0, np.float32) stds = np.concatenate(stds) if stds else np.zeros(0, np.float32) hp_all = np.concatenate(all_hazard_prob) if all_hazard_prob else np.zeros(0, np.float32) hl_all = np.concatenate(all_hazard_label) if all_hazard_label else np.zeros(0, np.float32) ne_all = np.concatenate(all_is_noneego) if all_is_noneego else np.zeros(0, bool) ep_all = np.concatenate(all_is_ego_pos) if all_is_ego_pos else np.zeros(0, bool) if preds.size == 0: self.model.train() return {"loss": float("inf"), "hazard_f1": 0.0, "pos_tta_mae": float("inf"), "ckpt_score": float("-inf")} # ── hazard metrics ─────────────────────────────────────────────────── hp_bin = (hp_all > 0.5).astype(np.float32) tp = float(((hp_bin == 1) & (hl_all == 1)).sum()) fp = float(((hp_bin == 1) & (hl_all == 0)).sum()) fn = float(((hp_bin == 0) & (hl_all == 1)).sum()) prec = tp / max(1, tp + fp) recall = tp / max(1, tp + fn) f1 = 2 * prec * recall / max(1e-9, prec + recall) ne_mask = ne_all.astype(bool) safe_neg_mask = (~ep_all) & (~ne_mask) ne_far = float((hp_bin[ne_mask] == 1).mean()) if ne_mask.any() else 0.0 sneg_fa = float((hp_bin[safe_neg_mask] == 1).mean()) if safe_neg_mask.any() else 0.0 # ── TTA metrics (positive-observed only) ───────────────────────────── obs_mask = ep_all & (labels_all < 9.9) if obs_mask.any(): pos_preds = preds[obs_mask]; pos_labels = labels_all[obs_mask] pos_mae = float(np.abs(pos_preds - pos_labels).mean()) pos_rmse = float(np.sqrt(((pos_preds - pos_labels)**2).mean())) low_mask = pos_labels <= 3.0 low_mae = float(np.abs(pos_preds[low_mask] - pos_labels[low_mask]).mean()) if low_mask.any() else 0.0 denom = float(((pos_labels - pos_labels.mean())**2).sum()) + 1e-12 pos_r2 = float(1.0 - ((pos_preds - pos_labels)**2).sum() / denom) else: pos_mae = pos_rmse = low_mae = 10.0; pos_r2 = 0.0 # ── checkpoint selection score ──────────────────────────────────────── # Higher is better: maximize hazard_f1, minimize normalized pos_tta_mae ckpt_score = 0.6 * f1 - 0.4 * (pos_mae / 10.0) metrics = { "loss": total_loss / max(1, n), "hazard_f1": f1, "hazard_precision": prec, "hazard_recall": recall, "pos_tta_mae": pos_mae, "pos_tta_rmse": pos_rmse, "pos_tta_r2": pos_r2, "low_tta_mae": low_mae, "non_ego_false_alert": ne_far, "safe_neg_false_alert": sneg_fa, "uncertainty_mean": float(stds.mean()), "ckpt_score": ckpt_score, } logger.info( f"Val: loss={metrics['loss']:.4f} hazard_f1={f1:.3f} " f"pos_tta_mae={pos_mae:.3f} ckpt_score={ckpt_score:.4f} " f"non_ego_fa={ne_far:.3f} safe_neg_fa={sneg_fa:.3f}" ) self.model.train() return metrics def save_checkpoint(self, name: str): ckpt_dir = self.output_dir / name self.model.save_checkpoint(str(ckpt_dir), epoch=self.current_epoch, step=self.global_step) torch.save( { "optimizer": self.optimizer.state_dict(), "scheduler": self.scheduler.state_dict() if self.scheduler else None, "scaler": self.scaler.state_dict() if self.scaler else None, "epoch": self.current_epoch, "global_step": self.global_step, "best_ckpt_score": self.best_ckpt_score, }, ckpt_dir / "training_state.pt", ) self.saved_checkpoints.append(ckpt_dir) if len(self.saved_checkpoints) > self.save_total_limit + 1: oldest = self.saved_checkpoints.pop(0) if oldest.name != "best" and oldest.exists(): import shutil shutil.rmtree(oldest, ignore_errors=True) def load_training_state(self, ckpt_dir: Path, reset_best_val_loss: bool = False): """ Loads optimizer/scheduler/scaler + epoch/global_step. If reset_best_val_loss=True: best_ckpt_score is forcibly reset to -inf so your NEW val split can define a NEW best from scratch. """ ts = ckpt_dir / "training_state.pt" if not ts.exists(): logger.warning(f"⚠️ No training_state.pt in {ckpt_dir}, resume weights only.") if reset_best_val_loss: self.best_ckpt_score = float("-inf") logger.info("✅ best_ckpt_score reset to -inf (weights-only path).") return obj = torch.load(ts, map_location="cpu") try: self.optimizer.load_state_dict(obj["optimizer"]) except Exception as e: logger.warning(f"⚠️ Failed to load optimizer state: {e}") if self.scheduler is not None and obj.get("scheduler") is not None: try: self.scheduler.load_state_dict(obj["scheduler"]) except Exception as e: logger.warning(f"⚠️ Failed to load scheduler state: {e}") if self.scaler is not None and obj.get("scaler") is not None: try: self.scaler.load_state_dict(obj["scaler"]) except Exception as e: logger.warning(f"⚠️ Failed to load scaler state: {e}") self.current_epoch = int(obj.get("epoch", 0)) self.global_step = int(obj.get("global_step", 0)) if reset_best_val_loss: self.best_ckpt_score = float("-inf") logger.info( f"✅ Loaded training_state from {ckpt_dir}: epoch={self.current_epoch}, global_step={self.global_step}, " f"best_ckpt_score RESET to -inf for NEW val." ) else: self.best_ckpt_score = float(obj.get("best_ckpt_score", obj.get("best_val_loss", float("-inf")))) logger.info( f"✅ Loaded training_state from {ckpt_dir}: epoch={self.current_epoch}, global_step={self.global_step}, " f"best_ckpt_score={self.best_ckpt_score:.4f}" ) def maybe_eval_and_set_new_best(self, force_save_best: bool = True): """ Evaluate once immediately (useful after resume + reset_best_val_loss). If best_ckpt_score is -inf, this will always become the new best. """ if self.val_loader is None: return val = self.evaluate() if self.use_wandb and val: wandb.log({"val/" + k: v for k, v in val.items()}, step=self.global_step) if not val: return score = val.get("ckpt_score", float("-inf")) improved = score > self.best_ckpt_score if improved: self.best_ckpt_score = score if force_save_best: self.save_checkpoint("best") logger.info( f"[InitEval] ckpt_score={score:.4f}, " f"best_ckpt_score={self.best_ckpt_score:.4f}, improved={improved}" ) def train(self): logger.info("=" * 60) logger.info(f"Starting SFT training: {self.experiment_name}") logger.info("=" * 60) start = time.time() for epoch in range(self.current_epoch, self.num_epochs): self.current_epoch = epoch progress = epoch / max(1, self.num_epochs) _ = progress # curriculum not used in new dataset pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs}", ncols=60) metrics_hist = defaultdict(list) accum = 0 for batch in pbar: m = self.train_step(batch) for k, v in m.items(): metrics_hist[k].append(v) accum += 1 if accum >= self.gradient_accumulation_steps: self._optimizer_step() accum = 0 if self.global_step % self.logging_steps == 0: avg = {k: float(np.mean(v[-self.logging_steps:])) for k, v in metrics_hist.items()} lr = self.optimizer.param_groups[0]["lr"] pbar.set_postfix({"loss": f"{avg['loss']:.4f}", "mae": f"{avg['tta_mae']:.3f}", "lr": f"{lr:.2e}"}) if self.use_wandb: wandb.log({"train/" + k: v for k, v in avg.items()} | {"train/lr": lr}, step=self.global_step) if self.val_loader and (self.global_step % self.eval_steps == 0): val = self.evaluate() if self.use_wandb and val: wandb.log({"val/" + k: v for k, v in val.items()}, step=self.global_step) if val: score = val.get("ckpt_score", float("-inf")) if score > self.best_ckpt_score: self.best_ckpt_score = score self.save_checkpoint("best") if self.global_step % self.save_steps == 0: self.save_checkpoint(f"step_{self.global_step}") if self.val_loader: val = self.evaluate() if val: score = val.get("ckpt_score", float("-inf")) if score > self.best_ckpt_score: self.best_ckpt_score = score self.save_checkpoint("best") self.save_checkpoint(f"epoch_{epoch+1}") logger.info("=" * 60) logger.info(f"Training completed in {(time.time()-start)/3600:.2f} hours") logger.info(f"Best ckpt_score: {self.best_ckpt_score:.4f}") logger.info(f"Checkpoints saved to: {self.output_dir}") logger.info("=" * 60) if self.use_wandb: wandb.finish() # ============================================================================ # Main # ============================================================================ def main(): parser = argparse.ArgumentParser("SFT Training for TTA Regression") # data — manifest-based parser.add_argument( "--manifest_dir", type=str, default="PROJECT_ROOT/data/sft_manifests", help="Directory containing split manifest JSONs from make_split_manifest.py", ) # legacy aliases (ignored if manifest_dir is provided via manifests) parser.add_argument("--nexar_root", type=str, default=None, help="(unused; kept for back-compat)") parser.add_argument("--dada_root", type=str, default=None, help="(unused; kept for back-compat)") # model parser.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct") parser.add_argument("--pretrained_lora", type=str, default=None) parser.add_argument( "--attn_implementation", type=str, default="flash_attention_2", choices=["flash_attention_2", "sdpa", "eager"], help="VLM attention backend. sdpa is safe fallback for Blackwell/new backbones.", ) parser.add_argument( "--belief_strategy", type=str, default="mean_pool", choices=["mean_pool", "last_token", "attention_pool", "dual_pool"], help="dual_pool = [mean(image_tokens) || mean(text_tokens)] (P0.2 L1)", ) # P0.1 — prompt ablation (drop weather/road_type/time_of_day) parser.add_argument( "--disable_metadata_prompt", action="store_true", default=False, help="P0.1: remove weather/road_type/time_of_day from the SFT prompt.", ) # P0.3 — PEFT upgrade flags parser.add_argument("--use_dora", action="store_true", default=False, help="P0.3: enable DoRA (Weight-Decomposed Low-Rank Adaptation).") parser.add_argument("--use_rslora", action="store_true", default=False, help="P0.3: enable rsLoRA (rank-stabilised scaling alpha/sqrt(r)).") parser.add_argument( "--lora_init", type=str, default="default", choices=["default", "gaussian", "pissa", "pissa_niter_16", "olora"], help="P0.3: initialisation scheme for fresh LoRA (ignored when resuming).", ) # resume parser.add_argument("--resume_from", type=str, default=None, help="Path to an SFT checkpoint dir that contains tta_head.pt, belief_aggregator.pt, and vlm_lora/") parser.add_argument("--resume_weights_only", action="store_true", help="If set, do not load optimizer/scheduler/scaler states (start new training state).") parser.add_argument("--auto_resume", action="store_true", default=True, help="Auto search a previous SFT checkpoint under output_dir (default: True).") parser.add_argument("--no_auto_resume", action="store_false", dest="auto_resume") # NEW: reset best + optional eval at start parser.add_argument("--reset_best_val_loss", action="store_true", default=True, help="Reset best_val_loss to +inf when resuming so NEW val can redefine best. (default: True)") parser.add_argument("--no_reset_best_val_loss", action="store_false", dest="reset_best_val_loss") parser.add_argument("--eval_on_start", action="store_true", default=True, help="Run a val evaluation immediately after resume (useful with reset_best_val_loss). (default: True)") parser.add_argument("--no_eval_on_start", action="store_false", dest="eval_on_start") # training parser.add_argument("--num_epochs", type=int, default=10) parser.add_argument("--batch_size", type=int, default=1) parser.add_argument("--gradient_accumulation_steps", type=int, default=4) parser.add_argument("--learning_rate", type=float, default=1e-4) parser.add_argument("--tta_head_lr", type=float, default=1e-3) parser.add_argument("--vlm_lr_multiplier", type=float, default=0.1) parser.add_argument("--weight_decay", type=float, default=0.01) parser.add_argument("--max_grad_norm", type=float, default=1.0) parser.add_argument("--mse_weight", type=float, default=1.0) parser.add_argument("--nll_weight", type=float, default=0.5) parser.add_argument("--max_pixels", type=int, default=None, help="Max pixels per frame for vision encoder. Default: 768*28*28=602112. " "Lower (e.g. 512*28*28=401408) reduces VRAM → allows larger batch.") parser.add_argument("--use_curriculum", action="store_true", default=True) parser.add_argument("--no_curriculum", action="store_false", dest="use_curriculum") # output/log parser.add_argument("--output_dir", type=str, required=True) parser.add_argument("--experiment_name", type=str, required=True) parser.add_argument("--use_wandb", action="store_true", default=True) parser.add_argument("--no_wandb", action="store_false", dest="use_wandb") # debug parser.add_argument("--debug", action="store_true") parser.add_argument("--debug_samples", type=int, default=100) args = parser.parse_args() # datasets — manifest-based logger.info("📊 Loading datasets from manifests...") manifest_dir = Path(args.manifest_dir) train_manifests = [ manifest_dir / "nexar_train.json", manifest_dir / "dada_pos_train.json", manifest_dir / "dada_noneego_train.json", manifest_dir / "dada_neg_train.json", ] val_manifests = [ manifest_dir / "nexar_val.json", manifest_dir / "dada_pos_val.json", manifest_dir / "dada_noneego_val.json", ] # Filter to existing manifests (graceful in case some sources are absent) train_manifests = [m for m in train_manifests if m.exists()] val_manifests = [m for m in val_manifests if m.exists()] if not train_manifests: raise RuntimeError(f"No train manifests found in {manifest_dir}. Run make_split_manifest.py first.") logger.info(f" Train manifests: {[m.name for m in train_manifests]}") logger.info(f" Val manifests: {[m.name for m in val_manifests]}") train_dataset = SFTDataset( manifests=train_manifests, split="train", debug=args.debug, debug_samples=args.debug_samples, ) val_dataset = SFTDataset( manifests=val_manifests, split="val", debug=args.debug, debug_samples=max(1, args.debug_samples // 2), ) if val_manifests else None # Decide resume checkpoint output_root = Path(args.output_dir) resume_dir: Optional[Path] = None if args.resume_from: resume_dir = Path(args.resume_from) if not _is_sft_ckpt_dir(resume_dir): raise RuntimeError(f"--resume_from is not a valid SFT checkpoint dir: {resume_dir}") elif args.auto_resume: resume_dir = find_auto_resume_checkpoint(output_root, args.experiment_name) if resume_dir is not None: logger.info(f"🔁 Auto-resume selected checkpoint: {resume_dir}") # If resume: load LoRA from ckpt/vlm_lora lora_path_for_init = args.pretrained_lora if resume_dir is not None: lora_path_for_init = str(resume_dir / "vlm_lora") # Create model logger.info("📦 Creating model...") model = SFTModel( model_name=args.model_name, pretrained_lora_path=lora_path_for_init, belief_strategy=args.belief_strategy, use_lora=True, use_bf16=True, device="auto", max_pixels=args.max_pixels, use_dora=args.use_dora, use_rslora=args.use_rslora, lora_init=args.lora_init, attn_implementation=args.attn_implementation, ) # If resume: load heads if resume_dir is not None: load_sft_heads(model, resume_dir) # Trainer trainer = SFTTrainer( model=model, train_dataset=train_dataset, val_dataset=val_dataset, num_epochs=args.num_epochs, batch_size=args.batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, learning_rate=args.learning_rate, tta_head_lr=args.tta_head_lr, vlm_lr_multiplier=args.vlm_lr_multiplier, weight_decay=args.weight_decay, max_grad_norm=args.max_grad_norm, mse_weight=args.mse_weight, nll_weight=args.nll_weight, use_curriculum=args.use_curriculum, output_dir=args.output_dir, experiment_name=args.experiment_name, use_wandb=args.use_wandb and HAS_WANDB, disable_metadata_prompt=args.disable_metadata_prompt, ) # Load training state if requested (but allow best reset) if resume_dir is not None and (not args.resume_weights_only): trainer.load_training_state(resume_dir, reset_best_val_loss=args.reset_best_val_loss) else: # weights-only path: still respect reset_best_val_loss if resume_dir is not None and args.reset_best_val_loss: trainer.best_ckpt_score = float("-inf") logger.info("✅ best_ckpt_score reset to -inf (resume_weights_only path).") # IMPORTANT: Evaluate immediately on the NEW val, so best is re-defined from scratch. if resume_dir is not None and args.eval_on_start and trainer.val_loader is not None: # If best is inf, this will always save "best" corresponding to the NEW val baseline. trainer.maybe_eval_and_set_new_best(force_save_best=True) trainer.train() if __name__ == "__main__": main()