#!/usr/bin/env python3 """ MR-JEPA Phase 3 Training — Enriched Evidence + Generative Decoder Loads the best Phase 2 checkpoint and: 1. Enables OCR token injection (from TextVQA ocr_tokens or simple extraction) 2. Trains the generative head on open-ended benchmarks (DocVQA, ChartQA, TextVQA) 3. Continues JEPA + discriminative training on ScienceQA 4. Full end-to-end fine-tuning of all components Training data: - ScienceQA train (MC, JEPA + task loss) - DocVQA validation (open-ended, generative loss) - ChartQA test (open-ended, generative loss) - TextVQA train (open-ended, generative loss, OCR tokens available) Eval: - ScienceQA test (accuracy) - DocVQA validation (ANLS) - ChartQA test (relaxed accuracy) - TextVQA validation (VQA accuracy) Phase 3 hyperparameters (from ARCHITECTURE.md): LR: 5e-5 (core), 5e-6 (backbone) Batch: 16, grad_accum: 8 Epochs: 10 Cosine schedule + warmup (10%) Usage: python train_phase3.py python train_phase3.py --epochs 10 --core_lr 5e-5 """ import os import sys import json import math import copy import logging import argparse from collections import defaultdict import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.optim import AdamW from torch.utils.data import Dataset, DataLoader from PIL import Image logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", datefmt="%H:%M:%S", ) log = logging.getLogger("mrjepa-p3") # ══════════════════════════════════════════════════════════════════════════ # OPEN-ENDED DATASET (DocVQA / ChartQA / TextVQA) # ══════════════════════════════════════════════════════════════════════════ class OpenEndedDataset(Dataset): """Dataset for open-ended VQA benchmarks (Phase 3 generative training).""" def __init__(self, benchmark, split, max_samples=0, transform=None, tokenizer=None, max_len=192, max_gen_len=64): from datasets import load_dataset self.benchmark = benchmark self.transform = transform self.tokenizer = tokenizer self.max_len = max_len self.max_gen_len = max_gen_len log.info(f"Loading {benchmark} {split}...") if benchmark == "docvqa": ds = load_dataset("lmms-lab/DocVQA", "DocVQA", split=split) elif benchmark == "chartqa": ds = load_dataset("lmms-lab/ChartQA", split=split) elif benchmark == "textvqa": ds = load_dataset("lmms-lab/textvqa", split=split) else: raise ValueError(f"Unknown benchmark: {benchmark}") if max_samples > 0: ds = ds.select(range(min(max_samples, len(ds)))) self.data = ds log.info(f"Loaded {len(ds)} samples from {benchmark} {split}") def __len__(self): return len(self.data) def __getitem__(self, idx): row = self.data[idx] # Image img = row.get("image") if img is None: img = Image.new("RGB", (256, 256), "white") else: img = img.convert("RGB") # Question question = row["question"] # Answer (target for generative head) if self.benchmark == "docvqa": answers = row.get("answers", [""]) answer = answers[0] if answers else "" all_answers = answers elif self.benchmark == "chartqa": answer = str(row.get("answer", "")) all_answers = [answer] elif self.benchmark == "textvqa": answers = row.get("answers", [""]) # Use most common answer from collections import Counter answer_counts = Counter(a.lower().strip() for a in answers) answer = answer_counts.most_common(1)[0][0] if answer_counts else "" all_answers = answers else: answer = "" all_answers = [""] # OCR tokens (TextVQA provides them; others we skip for now) ocr_tokens = row.get("ocr_tokens", []) ocr_text = " ".join(ocr_tokens[:50]) if ocr_tokens else "" # Build text: question + optional OCR context text = question if ocr_text: text += f" [OCR: {ocr_text}]" return { "image": img, "text": text, "answer": answer, "all_answers": all_answers, "benchmark": self.benchmark, "ocr_text": ocr_text, "question_type": row.get("type", row.get("question_types", [""])), } def collate_open_ended(batch, transform, tokenizer, max_len, max_gen_len): """Collate function for open-ended VQA batches.""" images = [s["image"] for s in batch] texts = [s["text"] for s in batch] answers = [s["answer"] for s in batch] # Process images if hasattr(transform, '__call__') and not hasattr(transform, 'feature_extractor'): pixel_values = torch.stack([transform(img) for img in images]) else: pixel_values = transform(images=images, return_tensors="pt")["pixel_values"] # Tokenize questions tok = tokenizer( texts, padding="max_length", truncation=True, max_length=max_len, return_tensors="pt" ) # Tokenize answers (generative targets) # Add BOS/EOS tokens if available answer_texts = [a if a else " " for a in answers] gen_tok = tokenizer( answer_texts, padding="max_length", truncation=True, max_length=max_gen_len, return_tensors="pt" ) return { "pixel_values": pixel_values, "input_ids": tok["input_ids"], "attention_mask": tok["attention_mask"], "gen_target_ids": gen_tok["input_ids"], "gen_attention_mask": gen_tok["attention_mask"], "batch_size": len(batch), "benchmarks": [s["benchmark"] for s in batch], "all_answers": [s["all_answers"] for s in batch], "question_types": [s.get("question_type", "") for s in batch], } # ══════════════════════════════════════════════════════════════════════════ # GENERATIVE HEAD (Lightweight Transformer Decoder) # ══════════════════════════════════════════════════════════════════════════ class GenerativeDecoderLayer(nn.Module): """Transformer decoder layer with cross-attention to latent state and evidence.""" def __init__(self, hidden_dim, num_heads, dropout=0.1): super().__init__() # Causal self-attention self.self_attn = nn.MultiheadAttention( embed_dim=hidden_dim, num_heads=num_heads, dropout=dropout, batch_first=True, ) self.self_attn_norm = nn.LayerNorm(hidden_dim) # Cross-attention to latent state z_K self.state_cross_attn = nn.MultiheadAttention( embed_dim=hidden_dim, num_heads=num_heads, dropout=dropout, batch_first=True, ) self.state_cross_norm = nn.LayerNorm(hidden_dim) # Cross-attention to evidence memory self.evidence_cross_attn = nn.MultiheadAttention( embed_dim=hidden_dim, num_heads=num_heads, dropout=dropout, batch_first=True, ) self.evidence_cross_norm = nn.LayerNorm(hidden_dim) # FFN self.ffn = nn.Sequential( nn.Linear(hidden_dim, hidden_dim * 4), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim * 4, hidden_dim), nn.Dropout(dropout), ) self.ffn_norm = nn.LayerNorm(hidden_dim) def forward(self, x, z_final, evidence, causal_mask=None): # Causal self-attention r = x x2 = self.self_attn_norm(x) x2, _ = self.self_attn(x2, x2, x2, attn_mask=causal_mask) x = r + x2 # Cross-attention to latent state r = x x2 = self.state_cross_norm(x) x2, _ = self.state_cross_attn(x2, z_final, z_final) x = r + x2 # Cross-attention to evidence r = x x2 = self.evidence_cross_norm(x) x2, _ = self.evidence_cross_attn(x2, evidence, evidence) x = r + x2 # FFN r = x x = r + self.ffn(self.ffn_norm(x)) return x class GenerativeHead(nn.Module): """ Lightweight generative decoder for Phase 3. Cross-attends to z_K and evidence memory to generate short answers. Uses the text encoder's tokenizer vocabulary. """ def __init__(self, hidden_dim, vocab_size, num_layers=4, num_heads=12, max_gen_len=64, dropout=0.1): super().__init__() self.hidden_dim = hidden_dim self.vocab_size = vocab_size self.max_gen_len = max_gen_len # Token embedding + positional encoding self.token_embedding = nn.Embedding(vocab_size, hidden_dim) self.pos_embedding = nn.Embedding(max_gen_len, hidden_dim) # Decoder layers self.layers = nn.ModuleList([ GenerativeDecoderLayer(hidden_dim, num_heads, dropout) for _ in range(num_layers) ]) # Output self.output_norm = nn.LayerNorm(hidden_dim) self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False) # Tie weights self.lm_head.weight = self.token_embedding.weight def forward(self, z_final, evidence, target_ids, pad_token_id=0): """Teacher-forced forward pass.""" B, seq_len = target_ids.shape device = target_ids.device positions = torch.arange(seq_len, device=device).unsqueeze(0) x = self.token_embedding(target_ids) + self.pos_embedding(positions) # Causal mask causal_mask = torch.triu( torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), diagonal=1 ) for layer in self.layers: x = layer(x, z_final, evidence, causal_mask) logits = self.lm_head(self.output_norm(x)) # Loss: next-token prediction shift_logits = logits[:, :-1].contiguous() shift_labels = target_ids[:, 1:].contiguous() # Mask padding tokens loss = F.cross_entropy( shift_logits.view(-1, self.vocab_size), shift_labels.view(-1), ignore_index=pad_token_id, ) return logits, loss @torch.no_grad() def generate(self, z_final, evidence, start_token_id, max_length=64, eos_token_id=None): """Autoregressive generation.""" B = z_final.size(0) device = z_final.device generated = torch.full((B, 1), start_token_id, dtype=torch.long, device=device) for step in range(max_length - 1): seq_len = generated.size(1) positions = torch.arange(seq_len, device=device).unsqueeze(0) x = self.token_embedding(generated) + self.pos_embedding(positions) causal_mask = torch.triu( torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), diagonal=1 ) for layer in self.layers: x = layer(x, z_final, evidence, causal_mask) logits = self.lm_head(self.output_norm(x[:, -1:])) next_token = logits.argmax(dim=-1) generated = torch.cat([generated, next_token], dim=1) if eos_token_id is not None and (next_token == eos_token_id).all(): break return generated # ══════════════════════════════════════════════════════════════════════════ # EVALUATION FUNCTIONS (Phase 3 Metrics) # ══════════════════════════════════════════════════════════════════════════ def normalized_levenshtein(s1, s2): """Normalized Levenshtein distance.""" s1 = s1.lower().strip() s2 = s2.lower().strip() if s1 == s2: return 0.0 len1, len2 = len(s1), len(s2) if len1 == 0 or len2 == 0: return 1.0 matrix = [[0] * (len2 + 1) for _ in range(len1 + 1)] for i in range(len1 + 1): matrix[i][0] = i for j in range(len2 + 1): matrix[0][j] = j for i in range(1, len1 + 1): for j in range(1, len2 + 1): cost = 0 if s1[i-1] == s2[j-1] else 1 matrix[i][j] = min(matrix[i-1][j]+1, matrix[i][j-1]+1, matrix[i-1][j-1]+cost) return matrix[len1][len2] / max(len1, len2) def compute_anls(predictions, ground_truths, threshold=0.5): """ANLS metric for DocVQA.""" scores = [] for pred, gts in zip(predictions, ground_truths): max_score = 0.0 for gt in gts: nl_dist = normalized_levenshtein(str(pred), str(gt)) score = 1.0 - nl_dist if nl_dist < threshold else 0.0 max_score = max(max_score, score) scores.append(max_score) return np.mean(scores) * 100 if scores else 0.0 def compute_vqa_accuracy(predictions, ground_truths): """VQA accuracy for TextVQA.""" scores = [] for pred, gts in zip(predictions, ground_truths): pred_norm = str(pred).lower().strip() matching = sum(1 for gt in gts if str(gt).lower().strip() == pred_norm) scores.append(min(matching / 3.0, 1.0)) return np.mean(scores) * 100 if scores else 0.0 def compute_relaxed_accuracy(predictions, ground_truths, tolerance=0.05): """Relaxed accuracy for ChartQA.""" correct = [] for pred, gt in zip(predictions, ground_truths): pred_str = str(pred).strip().lower() gt_str = str(gt).strip().lower() try: gt_val = float(gt_str.replace(',', '').replace('%', '')) pred_val = float(pred_str.replace(',', '').replace('%', '')) if gt_val == 0: is_correct = abs(pred_val) <= tolerance else: is_correct = abs(pred_val - gt_val) / abs(gt_val) <= tolerance except (ValueError, ZeroDivisionError): is_correct = pred_str == gt_str correct.append(is_correct) return np.mean(correct) * 100 if correct else 0.0 # ══════════════════════════════════════════════════════════════════════════ # PHASE 3 MAIN # ══════════════════════════════════════════════════════════════════════════ def download_phase2_checkpoint(hub_model_id, run_name="hybrid_main"): from huggingface_hub import hf_hub_download path = hf_hub_download( repo_id=hub_model_id, filename=f"checkpoints/{run_name}_best.pt", repo_type="model" ) log.info(f"Downloaded Phase 2 checkpoint: {path}") return path def main(): parser = argparse.ArgumentParser(description="MR-JEPA Phase 3 Training") parser.add_argument("--checkpoint", type=str, default=None) parser.add_argument("--hub_model_id", default="JorgeAV/MR-JEPA") parser.add_argument("--run_name", default="hybrid_main_phase3") parser.add_argument("--phase2_run", default="hybrid_main") parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--grad_accum", type=int, default=8) parser.add_argument("--core_lr", type=float, default=5e-5) parser.add_argument("--backbone_lr", type=float, default=5e-6) parser.add_argument("--text_lr", type=float, default=5e-6) parser.add_argument("--gen_weight", type=float, default=0.5, help="Weight for generative loss relative to task loss") parser.add_argument("--max_eval_samples", type=int, default=500) parser.add_argument("--max_gen_len", type=int, default=64) parser.add_argument("--max_train_samples", type=int, default=0, help="0 = all samples") parser.add_argument("--output_dir", default="./outputs/mrjepa_phase3") parser.add_argument("--trackio_space", default="JorgeAV/MR-JEPA-Trackio") args = parser.parse_args() # ── Download Phase 1 training script (has all model definitions) ── log.info("Downloading Phase 1 training script for model definitions...") from huggingface_hub import hf_hub_download p1_script = hf_hub_download( repo_id=args.hub_model_id, filename="train_mrjepa.py", repo_type="model" ) import importlib.util spec = importlib.util.spec_from_file_location("train_mrjepa", p1_script) p1 = importlib.util.module_from_spec(spec) spec.loader.exec_module(p1) # ── Load Phase 2 checkpoint ── if args.checkpoint and os.path.exists(args.checkpoint): ckpt_path = args.checkpoint else: ckpt_path = download_phase2_checkpoint(args.hub_model_id, args.phase2_run) log.info(f"Loading Phase 2 checkpoint: {ckpt_path}") ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) saved_cfg = ckpt["config"] cfg = p1.Config() for k, v in saved_cfg.items(): if hasattr(cfg, k): setattr(cfg, k, v) cfg.phase = 3 cfg.epochs = args.epochs cfg.batch_size = args.batch_size cfg.grad_accum = args.grad_accum cfg.lr = args.core_lr cfg.backbone_lr = args.backbone_lr cfg.output_dir = args.output_dir cfg.run_name = args.run_name cfg.freeze_backbone = True # Will unfreeze manually below cfg.freeze_text = True cfg.max_eval_samples = args.max_eval_samples cfg.resolve() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") log.info(f"Device: {device}") os.makedirs(cfg.output_dir, exist_ok=True) # ── Initialize Trackio ── import trackio trackio.init( name=args.run_name, project="MR-JEPA", space_id=args.trackio_space, config={ "phase": 3, "epochs": args.epochs, "core_lr": args.core_lr, "backbone_lr": args.backbone_lr, "text_lr": args.text_lr, "gen_weight": args.gen_weight, "batch_size": args.batch_size, "grad_accum": args.grad_accum, "backbone": cfg.backbone, "K": cfg.K, "use_jepa": cfg.use_jepa, "loss_fn": cfg.loss_fn, "max_gen_len": args.max_gen_len, "phase2_best_acc": ckpt.get("eval_acc", "unknown"), } ) log.info(f"Trackio initialized → https://huggingface.co/spaces/{args.trackio_space}") # ── Build model ── log.info("Building model...") model = p1.MRJEPAModel(cfg) model.evidence.load_state_dict(ckpt["evidence"]) model.rollout.load_state_dict(ckpt["rollout"]) model.disc.load_state_dict(ckpt["disc"]) model.target.t_ev.load_state_dict(ckpt["target_ev"]) model.target.t_ro.load_state_dict(ckpt["target_ro"]) log.info(f"Loaded Phase 2 weights (epoch={ckpt.get('epoch','?')}, " f"eval_acc={ckpt.get('eval_acc','?')}%)") # ── Add generative head ── tokenizer = model.txt.tokenizer # Use len(tokenizer) not tokenizer.vocab_size — Qwen3 has special tokens # beyond vocab_size (pad_token_id=151643 >= vocab_size=151643) actual_vocab_size = len(tokenizer) log.info(f"Adding generative head: actual_vocab_size={actual_vocab_size}, " f"hidden_dim={cfg.rollout_dim}, layers=4") gen_head = GenerativeHead( hidden_dim=cfg.rollout_dim, vocab_size=actual_vocab_size, num_layers=4, num_heads=cfg.predictor_heads, max_gen_len=args.max_gen_len, dropout=0.1, ) model.gen_head = gen_head # ── Unfreeze backbone layers (same as Phase 2 — keep them unfrozen) ── log.info("Unfreezing last 6 visual layers, last 4 text layers") model.vis.unfreeze_last(6) model.txt.unfreeze_last(4) model = model.to(device) total_p = sum(p.numel() for p in model.parameters()) train_p = sum(p.numel() for p in model.parameters() if p.requires_grad) log.info(f"Total: {total_p:,} | Trainable: {train_p:,} ({100*train_p/total_p:.1f}%)") trackio.log({ "model/total_params": total_p, "model/trainable_params": train_p, "model/trainable_pct": 100 * train_p / total_p }) # ── Build datasets ── transform = model.vis.get_transform() # MC dataset (ScienceQA) — keep JEPA + task loss mc_max = args.max_train_samples if args.max_train_samples > 0 else 0 train_mc_ds = p1.ScienceQADataset( "train", max_samples=mc_max, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_opts=cfg.max_options ) eval_mc_ds = p1.ScienceQADataset( "test", max_samples=cfg.max_eval_samples, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_opts=cfg.max_options ) mc_coll = lambda batch: p1.collate_fn( batch, transform, tokenizer, cfg.max_text_len, cfg.max_options ) train_mc_dl = DataLoader( train_mc_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=2, collate_fn=mc_coll, pin_memory=True, drop_last=True ) eval_mc_dl = DataLoader( eval_mc_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, collate_fn=mc_coll, pin_memory=True ) # Open-ended datasets max_open_train = args.max_train_samples if args.max_train_samples > 0 else 5000 # DocVQA — use validation as training (5349 samples, no explicit train split available) train_docvqa_ds = OpenEndedDataset( "docvqa", "validation", max_samples=max_open_train, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_gen_len=args.max_gen_len ) # ChartQA — use test (2500 samples) train_chartqa_ds = OpenEndedDataset( "chartqa", "test", max_samples=max_open_train, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_gen_len=args.max_gen_len ) # TextVQA — use train split (34602 samples, has OCR tokens) train_textvqa_ds = OpenEndedDataset( "textvqa", "train", max_samples=max_open_train, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_gen_len=args.max_gen_len ) # Evaluation splits eval_docvqa_ds = OpenEndedDataset( "docvqa", "validation", max_samples=args.max_eval_samples, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_gen_len=args.max_gen_len ) eval_chartqa_ds = OpenEndedDataset( "chartqa", "test", max_samples=args.max_eval_samples, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_gen_len=args.max_gen_len ) eval_textvqa_ds = OpenEndedDataset( "textvqa", "validation", max_samples=args.max_eval_samples, transform=transform, tokenizer=tokenizer, max_len=cfg.max_text_len, max_gen_len=args.max_gen_len ) open_coll = lambda batch: collate_open_ended( batch, transform, tokenizer, cfg.max_text_len, args.max_gen_len ) train_open_dls = { "docvqa": DataLoader( train_docvqa_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=2, collate_fn=open_coll, pin_memory=True, drop_last=True ), "chartqa": DataLoader( train_chartqa_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=2, collate_fn=open_coll, pin_memory=True, drop_last=True ), "textvqa": DataLoader( train_textvqa_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=2, collate_fn=open_coll, pin_memory=True, drop_last=True ), } eval_open_dls = { "docvqa": DataLoader( eval_docvqa_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, collate_fn=open_coll, pin_memory=True ), "chartqa": DataLoader( eval_chartqa_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, collate_fn=open_coll, pin_memory=True ), "textvqa": DataLoader( eval_textvqa_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, collate_fn=open_coll, pin_memory=True ), } # ── Optimizer with parameter groups ── backbone_params = [p for p in model.vis.parameters() if p.requires_grad] text_params = [p for p in model.txt.parameters() if p.requires_grad] bb_txt_ids = {id(p) for p in backbone_params + text_params} core_params = [p for p in model.parameters() if p.requires_grad and id(p) not in bb_txt_ids] param_groups = [ {"params": core_params, "lr": args.core_lr}, {"params": backbone_params, "lr": args.backbone_lr}, {"params": text_params, "lr": args.text_lr}, ] log.info(f"Optimizer: core={len(core_params)} @ {args.core_lr}, " f"backbone={len(backbone_params)} @ {args.backbone_lr}, " f"text={len(text_params)} @ {args.text_lr}") optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay) # Estimate total steps across all dataloaders mc_steps_per_epoch = len(train_mc_dl) open_steps_per_epoch = sum(len(dl) for dl in train_open_dls.values()) total_batches_per_epoch = mc_steps_per_epoch + open_steps_per_epoch total_steps = cfg.epochs * total_batches_per_epoch // cfg.grad_accum warmup_steps = int(total_steps * 0.1) # Phase 3: 10% warmup def lr_lambda(step): if step < warmup_steps: return step / max(warmup_steps, 1) progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) return 0.01 + 0.99 * 0.5 * (1 + math.cos(math.pi * progress)) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) # ── Pad token ID for generative loss masking ── pad_token_id = tokenizer.pad_token_id if pad_token_id is None: pad_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0 log.info(f"Pad token ID for gen loss: {pad_token_id}") log.info(f"Phase 3: {cfg.epochs} epochs") log.info(f" MC batches/epoch: {mc_steps_per_epoch}") log.info(f" Open batches/epoch: {open_steps_per_epoch}") log.info(f" Total opt steps: ~{total_steps}, warmup: {warmup_steps}") global_step = 0 best_composite = 0.0 amp_dtype = torch.bfloat16 if cfg.bf16 else torch.float32 trainable = [p for p in model.parameters() if p.requires_grad] try: for epoch in range(cfg.epochs): model.train() epoch_losses = defaultdict(list) epoch_mc_correct = 0 epoch_mc_total = 0 optimizer.zero_grad() batch_count = 0 # ── Phase 3A: MC training (ScienceQA) — JEPA + task loss ── log.info(f"Phase 3 Epoch {epoch}: MC training on ScienceQA...") for batch_idx, batch in enumerate(train_mc_dl): batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=cfg.bf16 and device.type == "cuda"): losses, preds = model(**batch) loss = losses["total"] / cfg.grad_accum loss.backward() batch_count += 1 if batch_count % cfg.grad_accum == 0: nn.utils.clip_grad_norm_(trainable, cfg.max_grad_norm) optimizer.step(); scheduler.step(); optimizer.zero_grad() model.update_target(global_step, total_steps) global_step += 1 for k, v in losses.items(): if isinstance(v, torch.Tensor): epoch_losses[f"mc_{k}"].append(v.item()) epoch_mc_correct += (preds == batch["labels"]).sum().item() epoch_mc_total += batch["batch_size"] if batch_idx % 100 == 0: avg = {k: np.mean(v[-100:]) for k, v in epoch_losses.items() if k.startswith("mc_")} mc_acc = epoch_mc_correct / max(epoch_mc_total, 1) * 100 log.info(f"P3 E{epoch} MC B{batch_idx}/{mc_steps_per_epoch} | " f"loss={avg.get('mc_total',0):.4f} | acc={mc_acc:.1f}%") trackio.log({ "train/mc_loss": avg.get("mc_total", 0), "train/mc_jepa": avg.get("mc_jepa", 0), "train/mc_task": avg.get("mc_task", 0), "train/mc_accuracy": mc_acc, "train/lr": scheduler.get_last_lr()[0], "train/epoch": epoch, "train/step": global_step, }) # ── Phase 3B: Open-ended training (DocVQA, ChartQA, TextVQA) ── log.info(f"Phase 3 Epoch {epoch}: Open-ended training...") epoch_gen_losses = defaultdict(list) # Interleave open-ended datasets open_iters = {name: iter(dl) for name, dl in train_open_dls.items()} open_active = set(open_iters.keys()) open_batch_idx = 0 while open_active: for name in list(open_active): try: batch = next(open_iters[name]) except StopIteration: open_active.discard(name) continue batch_t = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=cfg.bf16 and device.type == "cuda"): # Forward through perception + reasoning vis_tok = model.vis(batch_t["pixel_values"]).float() txt_tok = model.txt(batch_t["input_ids"], batch_t["attention_mask"]).float() evidence, _, ev_mask = model.evidence(vis_tok, txt_tok, batch_t["attention_mask"]) if model._use_rollout: traj, z_final, z_proj = model.rollout(evidence) else: B = batch_t["batch_size"] z0 = model.rollout.init_tokens.expand(B, -1, -1) + \ model.rollout.z0_proj(F.adaptive_avg_pool1d( evidence.permute(0,2,1), model.rollout.num_tokens ).permute(0,2,1)) z_final = z0 z_proj = model.rollout.out_proj(z0).unsqueeze(1) # JEPA loss (still active in Phase 3) jepa_loss_val = torch.tensor(0.0, device=device) if model._use_jepa: target_proj = model.target( vis_tok.detach(), txt_tok.detach(), batch_t["attention_mask"].detach() ) jepa_losses = model.jepa_loss( z_proj, target_proj, torch.tensor(0.0, device=device) ) jepa_loss_val = jepa_losses["jepa"] + jepa_losses["reg"] # Generative loss gen_logits, gen_loss = model.gen_head( z_final, evidence, batch_t["gen_target_ids"], pad_token_id=pad_token_id ) # Total loss for open-ended: JEPA + generative total_loss = (cfg.jepa_weight * jepa_loss_val + args.gen_weight * gen_loss) loss = total_loss / cfg.grad_accum loss.backward() batch_count += 1 if batch_count % cfg.grad_accum == 0: nn.utils.clip_grad_norm_(trainable, cfg.max_grad_norm) optimizer.step(); scheduler.step(); optimizer.zero_grad() model.update_target(global_step, total_steps) global_step += 1 epoch_gen_losses[f"{name}_gen"].append(gen_loss.item()) epoch_gen_losses[f"{name}_total"].append(total_loss.item()) epoch_losses["gen_total"].append(total_loss.item()) open_batch_idx += 1 if open_batch_idx % 100 == 0: avg_gen = {k: np.mean(v[-100:]) for k, v in epoch_gen_losses.items()} log.info(f"P3 E{epoch} OPEN B{open_batch_idx} | " + " | ".join(f"{k}={v:.4f}" for k, v in avg_gen.items())) trackio.log({ f"train/{k}": v for k, v in avg_gen.items() }) # ── Epoch-end evaluation ── log.info(f"Phase 3 Epoch {epoch}: Evaluating...") # MC eval (ScienceQA) mc_eval_acc = p1.evaluate(model, eval_mc_dl, device, cfg) log.info(f" ScienceQA eval accuracy: {mc_eval_acc:.1f}%") # Open-ended eval eval_results = evaluate_generative( model, eval_open_dls, device, cfg, tokenizer, pad_token_id, args.max_gen_len, amp_dtype ) for bm, metrics in eval_results.items(): for mk, mv in metrics.items(): log.info(f" {bm} {mk}: {mv:.2f}") # Composite score (average of all metrics) all_scores = [mc_eval_acc] for bm, metrics in eval_results.items(): all_scores.extend(metrics.values()) composite = np.mean(all_scores) log.info(f"=== Phase 3 Epoch {epoch} | MC: {mc_eval_acc:.1f}% | " f"Composite: {composite:.1f} ===") trackio.log({ "eval/scienceqa_accuracy": mc_eval_acc, "eval/composite_score": composite, "eval/epoch": epoch, **{f"eval/{bm}_{mk}": mv for bm, metrics in eval_results.items() for mk, mv in metrics.items()}, }) # Save best if composite > best_composite: best_composite = composite save_phase3_checkpoint( model, cfg, epoch, mc_eval_acc, eval_results, composite, is_best=True ) log.info(f"New best composite: {best_composite:.1f}") log.info(f"Phase 3 complete. Best composite score: {best_composite:.1f}") finally: trackio.log({ "final/best_composite": best_composite, "final/phase": 3, "final/total_steps": global_step }) log.info("Finishing Trackio...") trackio.finish() # ── Push results to Hub ── if cfg.push_to_hub: push_phase3_results(cfg, args, best_composite, eval_results) @torch.no_grad() def evaluate_generative(model, eval_dls, device, cfg, tokenizer, pad_token_id, max_gen_len, amp_dtype): """Evaluate on open-ended benchmarks via generation.""" model.eval() results = {} # Get start token ID start_token_id = tokenizer.bos_token_id if start_token_id is None: start_token_id = tokenizer.cls_token_id or 1 eos_token_id = tokenizer.eos_token_id for benchmark, dl in eval_dls.items(): predictions = [] ground_truths = [] for batch in dl: batch_t = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=cfg.bf16 and device.type == "cuda"): vis_tok = model.vis(batch_t["pixel_values"]).float() txt_tok = model.txt(batch_t["input_ids"], batch_t["attention_mask"]).float() evidence, _, _ = model.evidence(vis_tok, txt_tok, batch_t["attention_mask"]) if model._use_rollout: _, z_final, _ = model.rollout(evidence) else: B = batch_t["batch_size"] z_final = model.rollout.init_tokens.expand(B, -1, -1) + \ model.rollout.z0_proj(F.adaptive_avg_pool1d( evidence.permute(0,2,1), model.rollout.num_tokens ).permute(0,2,1)) # Generate answers gen_ids = model.gen_head.generate( z_final, evidence, start_token_id, max_length=max_gen_len, eos_token_id=eos_token_id ) # Decode for i in range(gen_ids.size(0)): pred_text = tokenizer.decode( gen_ids[i], skip_special_tokens=True ).strip() predictions.append(pred_text) ground_truths.extend(batch["all_answers"]) # Compute metrics if benchmark == "docvqa": score = compute_anls(predictions, ground_truths) results[benchmark] = {"anls": score} elif benchmark == "chartqa": # Ground truths are single strings wrapped in lists gt_flat = [gt[0] if isinstance(gt, list) else gt for gt in ground_truths] score = compute_relaxed_accuracy(predictions, gt_flat) results[benchmark] = {"relaxed_accuracy": score} elif benchmark == "textvqa": score = compute_vqa_accuracy(predictions, ground_truths) results[benchmark] = {"vqa_accuracy": score} log.info(f" {benchmark}: {results[benchmark]}") model.train() return results def save_phase3_checkpoint(model, cfg, epoch, mc_acc, open_results, composite, is_best=False): """Save Phase 3 checkpoint.""" tag = "best" if is_best else f"epoch{epoch}" path = os.path.join(cfg.output_dir, f"checkpoint_{tag}.pt") state = { "evidence": model.evidence.state_dict(), "rollout": model.rollout.state_dict(), "disc": model.disc.state_dict(), "gen_head": model.gen_head.state_dict(), "target_ev": model.target.t_ev.state_dict(), "target_ro": model.target.t_ro.state_dict(), "config": cfg.__dict__, "epoch": epoch, "mc_eval_acc": mc_acc, "open_results": open_results, "composite_score": composite, "phase": 3, } torch.save(state, path) log.info(f"Saved Phase 3 checkpoint: {path} (composite={composite:.1f})") def push_phase3_results(cfg, args, best_composite, eval_results): """Push Phase 3 results and checkpoint to Hub.""" try: from huggingface_hub import HfApi api = HfApi() results = { "run_name": cfg.run_name, "phase": 3, "backbone": cfg.backbone, "K": cfg.K, "use_jepa": cfg.use_jepa, "loss_fn": cfg.loss_fn, "best_composite_score": best_composite, "epochs": cfg.epochs, "core_lr": args.core_lr, "backbone_lr": args.backbone_lr, "text_lr": args.text_lr, "gen_weight": args.gen_weight, "batch_size": cfg.batch_size, "grad_accum": cfg.grad_accum, "open_results": {k: v for k, v in (eval_results or {}).items()}, } result_path = os.path.join(cfg.output_dir, f"results_{cfg.run_name}.json") with open(result_path, "w") as f: json.dump(results, f, indent=2) api.upload_file( path_or_fileobj=result_path, path_in_repo=f"results/{cfg.run_name}.json", repo_id=cfg.hub_model_id, repo_type="model", ) best_ckpt = os.path.join(cfg.output_dir, "checkpoint_best.pt") if os.path.exists(best_ckpt): api.upload_file( path_or_fileobj=best_ckpt, path_in_repo=f"checkpoints/{cfg.run_name}_best.pt", repo_id=cfg.hub_model_id, repo_type="model", ) log.info(f"Pushed Phase 3 results to {cfg.hub_model_id}") except Exception as e: log.error(f"Push failed: {e}") if __name__ == "__main__": main()