# ============================================================================ # TRAINER: MEMORY-CLIP-SEQ — Sequence Reconstruction # # Extends the v2 trainer with: # - Teacher full sequence capture (ModernBERT last_hidden_state) # - Sequence reconstruction loss (reconstructed 77 vs teacher projected 77) # - Two-phase training: # Phase 1: freeze v1 memory weights, train only seq head # Phase 2: unfreeze all, joint fine-tune # - v1 checkpoint loading at startup # # Core training loop (InfoNCE + Procrustes + CV) UNCHANGED from v2. # Sequence loss is ADDED alongside existing losses. # ============================================================================ import gc import math import os import json import time from dataclasses import dataclass, asdict import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm from safetensors.torch import save_file as safetensors_save, load_file # ══════════════════════════════════════════════════════════════════ # CONFIG # ══════════════════════════════════════════════════════════════════ @dataclass class TrainSeqConfig: # Data max_train_samples: int = 50000 max_val_samples: int = 2000 min_caption_length: int = 100 # Training — phase 1 (seq head only) phase1_epochs: int = 5 phase1_lr_seq: float = 2e-3 phase1_lr_proj: float = 1e-3 # Training — phase 2 (joint fine-tune) phase2_epochs: int = 5 phase2_lr_bank: float = 5e-4 # reduced from v2's 2e-3 phase2_lr_output: float = 2e-4 # reduced phase2_lr_proj: float = 5e-4 phase2_lr_seq: float = 1e-3 # Shared batch_size: int = 64 min_lr: float = 1e-6 weight_decay: float = 0.01 grad_clip: float = 1.0 warmup_steps: int = 200 # Loss weights — existing (unchanged from v2) modern_weight: float = 1.0 procrustes_weight: float = 0.3 cv_weight: float = 0.05 temperature: float = 0.07 # Loss weights — sequence (NEW) sequence_weight: float = 1.0 # MSE between reconstructed and teacher seq sequence_cosine_weight: float = 0.5 # per-position cosine similarity # Teacher modern_max_len: int = 4096 procrustes_n_samples: int = 300 # v1 checkpoint — local path or HuggingFace URL v1_checkpoint: str = "" v1_repo_id: str = "AbstractPhil/geolip-clip-vit-large-patch14-ctx576" v1_filename: str = "model.safetensors" # Logging checkpoint_dir: str = "/home/claude/memory_clip_seq_checkpoints" tensorboard_dir: str = "/home/claude/memory_clip_seq_tb" metrics_file: str = "/home/claude/memory_clip_seq_checkpoints/metrics.json" log_every: int = 20 eval_every: int = 200 TCFG = TrainSeqConfig() # ══════════════════════════════════════════════════════════════════ # GEOMETRIC UTILITIES — IDENTICAL to v2 # ══════════════════════════════════════════════════════════════════ def cayley_menger_vol2(pts): with torch.amp.autocast("cuda", enabled=False): pts = pts.float() diff = pts.unsqueeze(-2) - pts.unsqueeze(-3) d2 = (diff * diff).sum(-1) B, V, _ = d2.shape cm = torch.zeros(B, V+1, V+1, device=d2.device, dtype=torch.float32) cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2 s = (-1.0)**V; f = math.factorial(V-1) return s / ((2.0**(V-1)) * f*f) * torch.linalg.det(cm) def pentachoron_cv(embeddings, n_samples=16): B = embeddings.shape[0] if B < 5: return torch.tensor(0.0, device=embeddings.device) vols = [] for _ in range(n_samples): idx = torch.randperm(B, device=embeddings.device)[:5] v2 = cayley_menger_vol2(embeddings[idx].unsqueeze(0)) vols.append(torch.sqrt(F.relu(v2[0]) + 1e-12)) stacked = torch.stack(vols) return stacked.std() / (stacked.mean() + 1e-8) def procrustes_alignment_loss(emb_a, emb_b): with torch.amp.autocast("cuda", enabled=False): A = F.normalize(emb_a.float(), dim=-1) B_e = F.normalize(emb_b.float(), dim=-1) A = A - A.mean(0, keepdim=True) B_e = B_e - B_e.mean(0, keepdim=True) S = torch.linalg.svdvals(A.T @ B_e) N, D = A.shape return 1.0 - S.sum() / (math.sqrt(N) * D) # ══════════════════════════════════════════════════════════════════ # LOSSES — v2 existing + NEW sequence loss # ══════════════════════════════════════════════════════════════════ def infonce_loss(emb_a, emb_b, temperature=0.07): a = F.normalize(emb_a, dim=-1) b = F.normalize(emb_b, dim=-1) logits = (a @ b.T) / temperature labels = torch.arange(logits.shape[0], device=logits.device) loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) / 2 with torch.no_grad(): acc = (logits.argmax(-1) == labels).float().mean().item() top5 = logits.topk(min(5, logits.shape[1]), dim=-1).indices acc5 = (top5 == labels.unsqueeze(-1)).any(-1).float().mean().item() return loss, acc, acc5 def batch_cv_loss(all_anchors, n_reals, cv_target=0.20): device = all_anchors.device B = all_anchors.shape[0] total_loss = torch.tensor(0.0, device=device) total_cv = 0.0; n_valid = 0 per_sample_cv = [] for b in range(B): n = n_reals[b].item() if isinstance(n_reals[b], torch.Tensor) else n_reals[b] if n < 5: continue cv_val = pentachoron_cv(all_anchors[b, :n], n_samples=16) total_loss = total_loss + (cv_val - cv_target).abs() total_cv += cv_val.item() per_sample_cv.append(cv_val.item()) n_valid += 1 stats = { "cv_raw": total_cv / max(n_valid, 1), "cv_std": float(np.std(per_sample_cv)) if per_sample_cv else 0.0, "cv_n_valid": n_valid, } return total_loss / max(n_valid, 1), stats def sequence_reconstruction_loss(pred_seq, target_seq): """ pred_seq: (B, 77, 768) — reconstructed sequence target_seq: (B, 77, 768) — teacher projected sequence Returns: mse_loss: mean squared error cos_loss: 1 - mean per-position cosine similarity mean_cos: scalar metric (not differentiable) """ mse = F.mse_loss(pred_seq, target_seq) # Per-position cosine similarity pred_norm = F.normalize(pred_seq, dim=-1) tgt_norm = F.normalize(target_seq, dim=-1) cos_sim = (pred_norm * tgt_norm).sum(-1) # (B, 77) cos_loss = 1.0 - cos_sim.mean() with torch.no_grad(): mean_cos = cos_sim.mean().item() return mse, cos_loss, mean_cos # ══════════════════════════════════════════════════════════════════ # TEACHER — returns BOTH pooled AND full sequence # ══════════════════════════════════════════════════════════════════ @torch.no_grad() def teacher_forward(model, tokenizer, texts, device, max_len): """Returns pooled (B, 1024) from ModernBERT. Sequence target comes from CLIP.""" inputs = tokenizer(texts, max_length=max_len, padding=True, truncation=True, return_tensors="pt").to(device) out = model(**inputs) mask = inputs.attention_mask.unsqueeze(-1).float() pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1) return pooled # ══════════════════════════════════════════════════════════════════ # PARAM GROUPS # ══════════════════════════════════════════════════════════════════ def make_param_groups_phase1(model): """Phase 1: Only train sequence head + teacher seq projector.""" seq_params = [] for name, param in model.named_parameters(): param.requires_grad = False # freeze everything first for name, param in model.named_parameters(): if "sequence_reconstructor" in name: param.requires_grad = True seq_params.append(param) # Also keep proj_modern trainable (it's the pooled projector) for name, param in model.named_parameters(): if "proj_modern" in name: param.requires_grad = True proj_params = [p for n, p in model.named_parameters() if "proj_modern" in n and p.requires_grad] groups = [ {"params": seq_params, "lr": TCFG.phase1_lr_seq, "name": "seq_head", "weight_decay": TCFG.weight_decay}, {"params": proj_params, "lr": TCFG.phase1_lr_proj, "name": "proj", "weight_decay": TCFG.weight_decay}, ] for g in groups: n = sum(p.numel() for p in g["params"]) print(f" {g['name']:12s}: {n:>10,} params @ lr={g['lr']}") return groups def make_param_groups_phase2(model): """Phase 2: Unfreeze everything, differential LRs.""" # Unfreeze all trainable (non-CLIP) params for name, param in model.named_parameters(): if "clip_text" not in name and "_clip_text" not in name: param.requires_grad = True bank_names = {"bank.", "clip_cross_attn", "clip_cross_norms", "clip_cross_ffns", "clip_cross_ffn_norms"} seq_names = {"sequence_reconstructor"} proj_names = {"proj_modern"} bank_p, seq_p, proj_p, output_p = [], [], [], [] for name, param in model.named_parameters(): if not param.requires_grad: continue if any(name.startswith(s) or s in name for s in seq_names): seq_p.append(param) elif any(name.startswith(s) or s in name for s in proj_names): proj_p.append(param) elif any(name.startswith(s) or s in name for s in bank_names): bank_p.append(param) else: output_p.append(param) groups = [ {"params": bank_p, "lr": TCFG.phase2_lr_bank, "name": "bank", "weight_decay": TCFG.weight_decay}, {"params": seq_p, "lr": TCFG.phase2_lr_seq, "name": "seq_head", "weight_decay": TCFG.weight_decay}, {"params": proj_p, "lr": TCFG.phase2_lr_proj, "name": "proj", "weight_decay": TCFG.weight_decay}, {"params": output_p, "lr": TCFG.phase2_lr_output, "name": "output", "weight_decay": TCFG.weight_decay}, ] for g in groups: n = sum(p.numel() for p in g["params"]) print(f" {g['name']:12s}: {n:>10,} params @ lr={g['lr']}") return groups # ══════════════════════════════════════════════════════════════════ # PROCRUSTES INIT — IDENTICAL to v2 # ══════════════════════════════════════════════════════════════════ @torch.no_grad() def compute_and_init_procrustes(student_model, modern_model, modern_tok, captions, device): print(f"\n Computing static Procrustes on {len(captions)} captions...") student_embs, modern_embs = [], [] clip_tok = student_model.clip_tokenizer for i in range(0, len(captions), 16): batch = captions[i:i+16] tokens = clip_tok(batch, max_length=77, padding=True, truncation=True, return_tensors="pt").to(device) clip_out = student_model.clip_text( input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=False) student_embs.append(clip_out.pooler_output.cpu()) pooled = teacher_forward(modern_model, modern_tok, batch, device, TCFG.modern_max_len) modern_embs.append(pooled.cpu()) student_all = torch.cat(student_embs) modern_all = torch.cat(modern_embs) print(f" Student: {student_all.shape}, Teacher: {modern_all.shape}") X = student_all.float(); Y = modern_all.float() mu_x, mu_y = X.mean(0), Y.mean(0) Xc, Yc = X - mu_x, Y - mu_y if Xc.shape[1] < Yc.shape[1]: pad = torch.zeros(Xc.shape[0], Yc.shape[1] - Xc.shape[1]) Xc = torch.cat([Xc, pad], dim=1) mu_x = torch.cat([mu_x, torch.zeros(Yc.shape[1] - mu_x.shape[0])]) U, S, Vt = torch.linalg.svd(Xc.T @ Yc) R = (U @ Vt).T cos_before = F.cosine_similarity(Xc, Yc, dim=-1).mean() cos_after = F.cosine_similarity((Xc @ R.T), Yc, dim=-1).mean() print(f" Procrustes: cos {cos_before:.4f} → {cos_after:.4f}") # Init the pooled projector if it has init_from_procrustes if hasattr(student_model.proj_modern, 'init_from_procrustes'): student_model.proj_modern.init_from_procrustes(R, mu_x, mu_y) return {"cos_before": cos_before.item(), "cos_after": cos_after.item()} # ══════════════════════════════════════════════════════════════════ # V1 WEIGHT LOADING # ══════════════════════════════════════════════════════════════════ def load_v1_weights(model, device): """Load v1 memory system weights into the expanded seq model. Tries local path first, then downloads from HuggingFace.""" checkpoint_path = TCFG.v1_checkpoint # Try local path first if checkpoint_path and os.path.exists(checkpoint_path): print(f" Loading v1 weights (local): {checkpoint_path}") else: # Download from HuggingFace from huggingface_hub import hf_hub_download print(f" Downloading v1 weights from {TCFG.v1_repo_id}/{TCFG.v1_filename}...") checkpoint_path = hf_hub_download( repo_id=TCFG.v1_repo_id, filename=TCFG.v1_filename) print(f" Downloaded to: {checkpoint_path}") state = load_file(checkpoint_path, device=str(device)) missing, unexpected = model.load_state_dict(state, strict=False) n_loaded = len(state) - len(unexpected) print(f" Loaded: {n_loaded} tensors from v1") print(f" Missing (new modules): {len(missing)}") if missing: new_module_keys = [k for k in missing if "sequence_reconstructor" in k ] other_missing = [k for k in missing if k not in new_module_keys] print(f" Seq head (expected new): {len(new_module_keys)}") if other_missing: print(f" Other (check!): {other_missing[:5]}") if unexpected: print(f" Unexpected (v1 buffers, ignorable): {len(unexpected)}") return True # ══════════════════════════════════════════════════════════════════ # TRAINING # ══════════════════════════════════════════════════════════════════ def train_phase(model, modern_model, modern_tok, train_captions, val_captions, param_groups, n_epochs, phase_name, writer, all_metrics, global_step=0): """ Single training phase. Used for both phase 1 and phase 2. """ device = next(model.parameters()).device optimizer = torch.optim.AdamW(param_groups) all_params = [p for g in param_groups for p in g["params"]] n_batches = len(train_captions) // TCFG.batch_size total_steps = n_batches * n_epochs scheduler = torch.optim.lr_scheduler.SequentialLR( optimizer, [torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.01, total_iters=TCFG.warmup_steps), torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=max(total_steps, 1), eta_min=TCFG.min_lr)], milestones=[TCFG.warmup_steps]) scaler = torch.amp.GradScaler() clip_tokenizer = model.clip_tokenizer best_val_loss = float("inf") print(f"\n {phase_name}: {sum(p.numel() for p in all_params):,} trainable params") print(f" {n_batches} batches/epoch × {TCFG.batch_size}") # segment_text is in notebook namespace from architecture cell _segment_text = segment_text for epoch in range(n_epochs): model.train() perm = np.random.permutation(len(train_captions)) losses = {"total": 0, "modern": 0, "procrustes": 0, "cv": 0, "seq_mse": 0, "seq_cos": 0} metrics = {"modern_acc": 0, "modern_acc5": 0, "cv_raw": 0, "seq_cos_sim": 0, "n_segments_avg": 0} n = 0 t0 = time.time() pbar = tqdm(range(0, len(train_captions), TCFG.batch_size), desc=f"{phase_name} E{epoch+1}/{n_epochs}", unit="batch") for batch_start in pbar: idx = perm[batch_start:batch_start + TCFG.batch_size] if len(idx) < 2: continue batch_captions = [train_captions[i] for i in idx] B = len(batch_captions) # ── Teacher: ModernBERT pooled (sequence target comes from CLIP) ── with torch.no_grad(): with torch.amp.autocast("cuda"): modern_cls = teacher_forward( modern_model, modern_tok, batch_captions, device, TCFG.modern_max_len) # ── Student: segment-by-segment processing ── state = model.init_state(B, device) all_segments = [_segment_text(cap, clip_tokenizer, model.config.max_content_tokens, model.config.segment_overlap, model.config.max_segments) for cap in batch_captions] max_segs = max(len(s) for s in all_segments) n_segs = [len(s) for s in all_segments] for seg_k in range(max_segs): batch_ids, batch_masks = [], [] for b in range(B): if seg_k < len(all_segments[b]): batch_ids.append(all_segments[b][seg_k]["input_ids"]) batch_masks.append(all_segments[b][seg_k]["attention_mask"]) else: batch_ids.append(torch.zeros(77, dtype=torch.long)) batch_masks.append(torch.zeros(77, dtype=torch.long)) ids = torch.stack(batch_ids).to(device) masks = torch.stack(batch_masks).to(device) with torch.amp.autocast("cuda"): fused_output, state = model.forward_segment(ids, masks, state) student_cls = fused_output # pooled output from last segment # Bank anchors for CV loss — accumulated in state during segment processing bank_anchors = state["bank"]["anchors"] # (B, N_written, 768) # Pad to max_segs for batch CV computation all_anchors = torch.zeros(B, max_segs, model.config.anchor_dim, device=device) n_written = min(bank_anchors.shape[1], max_segs) all_anchors[:, :n_written] = bank_anchors[:, :n_written] # ── Existing losses (UNCHANGED from v2) ── with torch.amp.autocast("cuda"): proj_m = model.proj_modern(student_cls) l_modern, acc_m, acc5_m = infonce_loss( proj_m, modern_cls, TCFG.temperature) l_procrustes = procrustes_alignment_loss( student_cls, modern_cls[:, :model.config.clip_hidden]) n_reals_t = torch.tensor(n_segs, device=device) l_cv, cv_stats = batch_cv_loss( all_anchors, n_reals_t, model.config.cv_target) # ── NEW: Sequence reconstruction loss ── # Target: CLIP's own last_hidden_state on the truncated caption. # This is what the UNet was trained on — the reconstructor must # produce sequences in CLIP's distribution. with torch.no_grad(): clip_inputs = clip_tokenizer( batch_captions, max_length=77, padding="max_length", truncation=True, return_tensors="pt").to(device) with torch.amp.autocast("cuda"): clip_target_out = model.clip_text( input_ids=clip_inputs.input_ids, attention_mask=clip_inputs.attention_mask, output_hidden_states=False, return_dict=True) clip_target_seq = clip_target_out.last_hidden_state # (B, 77, 768) with torch.amp.autocast("cuda"): # Reconstruct sequence from memory state recon_seq = model.reconstruct_sequence(state) # (B, 77, 768) l_seq_mse, l_seq_cos, seq_cos_metric = sequence_reconstruction_loss( recon_seq, clip_target_seq.detach()) # ── Combined loss ── with torch.amp.autocast("cuda"): loss = (TCFG.modern_weight * l_modern + TCFG.procrustes_weight * l_procrustes + TCFG.cv_weight * l_cv + TCFG.sequence_weight * l_seq_mse + TCFG.sequence_cosine_weight * l_seq_cos) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(all_params, TCFG.grad_clip) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) scheduler.step() global_step += 1 # ── Metrics ── losses["total"] += loss.item() losses["modern"] += l_modern.item() losses["procrustes"] += l_procrustes.item() losses["cv"] += l_cv.item() losses["seq_mse"] += l_seq_mse.item() losses["seq_cos"] += l_seq_cos.item() metrics["modern_acc"] += acc_m metrics["modern_acc5"] += acc5_m metrics["cv_raw"] += cv_stats.get("cv_raw", 0) metrics["seq_cos_sim"] += seq_cos_metric metrics["n_segments_avg"] += np.mean(n_segs) n += 1 d = max(n, 1) pbar.set_postfix( loss=f"{losses['total']/d:.3f}", m_acc=f"{metrics['modern_acc']/d:.3f}", s_cos=f"{metrics['seq_cos_sim']/d:.3f}", cv=f"{metrics['cv_raw']/d:.3f}") # Tensorboard if global_step % TCFG.log_every == 0: writer.add_scalar(f"{phase_name}/loss", losses["total"]/d, global_step) writer.add_scalar(f"{phase_name}/modern_loss", losses["modern"]/d, global_step) writer.add_scalar(f"{phase_name}/seq_mse", losses["seq_mse"]/d, global_step) writer.add_scalar(f"{phase_name}/seq_cos_loss", losses["seq_cos"]/d, global_step) writer.add_scalar(f"{phase_name}/seq_cos_sim", metrics["seq_cos_sim"]/d, global_step) writer.add_scalar(f"{phase_name}/m_acc", metrics["modern_acc"]/d, global_step) writer.add_scalar(f"{phase_name}/cv_raw", metrics["cv_raw"]/d, global_step) all_metrics["steps"].append({ "step": global_step, "phase": phase_name, "epoch": epoch + 1, "loss": losses["total"]/d, "m_acc": metrics["modern_acc"]/d, "seq_cos": metrics["seq_cos_sim"]/d, "cv_raw": metrics["cv_raw"]/d, }) pbar.close() elapsed = time.time() - t0 d = max(n, 1) epoch_summary = { "phase": phase_name, "epoch": epoch + 1, "elapsed_s": elapsed, "loss": losses["total"]/d, "modern_loss": losses["modern"]/d, "seq_mse": losses["seq_mse"]/d, "seq_cos_loss": losses["seq_cos"]/d, "m_acc": metrics["modern_acc"]/d, "m_acc5": metrics["modern_acc5"]/d, "seq_cos_sim": metrics["seq_cos_sim"]/d, "cv_raw": metrics["cv_raw"]/d, "global_step": global_step, } all_metrics["epochs"].append(epoch_summary) print(f"\n {phase_name} E{epoch+1}: {elapsed:.0f}s " f"loss={epoch_summary['loss']:.4f} " f"m_acc={epoch_summary['m_acc']:.3f} " f"seq_cos={epoch_summary['seq_cos_sim']:.3f} " f"cv={epoch_summary['cv_raw']:.3f}") # Save save_checkpoint(model, optimizer, epoch + 1, global_step, phase_name, os.path.join(TCFG.checkpoint_dir, f"{phase_name}_e{epoch+1:02d}")) with open(TCFG.metrics_file, "w") as f: json.dump(all_metrics, f, indent=2, default=str) return global_step def save_checkpoint(model, optimizer, epoch, global_step, phase, path): os.makedirs(path, exist_ok=True) state = {} for name, param in model.named_parameters(): if param.requires_grad: state[name] = param.data.contiguous().cpu() for name, buf in model.named_buffers(): state[f"buffer.{name}"] = buf.contiguous().cpu() safetensors_save(state, os.path.join(path, "memory_system.safetensors")) torch.save({"optimizer": optimizer.state_dict() if optimizer else {}, "epoch": epoch, "global_step": global_step, "phase": phase}, os.path.join(path, "training_state.pt")) model_cfg = model.config.to_dict() if hasattr(model.config, 'to_dict') else {} config_data = {"model": model_cfg, "training": asdict(TCFG)} with open(os.path.join(path, "config.json"), "w") as f: json.dump(config_data, f, indent=2, default=str) # ══════════════════════════════════════════════════════════════════ # DATA # ══════════════════════════════════════════════════════════════════ def load_long_captions(max_train, max_val, min_length=100): from datasets import load_dataset print(f" Loading CaptionEmporium/conceptual-captions-cc12m-llavanext...") ds = load_dataset("CaptionEmporium/conceptual-captions-cc12m-llavanext", split="train", streaming=True) captions = [] for row in ds: cap = row.get("caption_llava", "") if isinstance(cap, str) and len(cap) > min_length: captions.append(cap) if len(captions) >= max_train + max_val: break train_caps = captions[:max_train] val_caps = captions[max_train:max_train + max_val] print(f" Train: {len(train_caps)}, Val: {len(val_caps)}") return train_caps, val_caps # ══════════════════════════════════════════════════════════════════ # MAIN # ══════════════════════════════════════════════════════════════════ def main(): print("=" * 70) print("TRAINING: MEMORY-CLIP-SEQ (SEQUENCE RECONSTRUCTION)") print("=" * 70) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f" Device: {device}") if torch.cuda.is_available(): print(f" GPU: {torch.cuda.get_device_name()}") print(f" VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") # ── Model (classes loaded from architecture cell) ── config = MemoryCLIPSeqConfig() model = MemoryCLIPSeqModel(config).to(device) # Trigger CLIP lazy load _ = model.clip_text _ = model.clip_tokenizer # ── Load v1 weights ── load_v1_weights(model, device) print(" v1 memory system weights loaded") # ── Teacher ── from transformers import AutoModel, AutoTokenizer print(f"\n Loading ModernBERT-large...") modern_model = AutoModel.from_pretrained( config.teacher_model, torch_dtype=torch.float16).to(device) modern_model.eval() for p in modern_model.parameters(): p.requires_grad = False modern_tok = AutoTokenizer.from_pretrained(config.teacher_model) print(f" {sum(p.numel() for p in modern_model.parameters()):,} params (frozen)") # ── Data ── train_captions, val_captions = load_long_captions( TCFG.max_train_samples, TCFG.max_val_samples, TCFG.min_caption_length) from transformers import CLIPTokenizer tok_temp = CLIPTokenizer.from_pretrained(config.clip_model) lengths = [len(tok_temp.encode(c)) for c in train_captions[:500]] print(f" Caption tokens (sample 500): mean={np.mean(lengths):.0f} " f"median={np.median(lengths):.0f} max={max(lengths)} " f">77: {sum(1 for l in lengths if l > 77)/len(lengths):.1%}") del tok_temp # ── Procrustes init ── compute_and_init_procrustes( model, modern_model, modern_tok, train_captions[:TCFG.procrustes_n_samples], device) # ── Setup logging ── os.makedirs(TCFG.checkpoint_dir, exist_ok=True) os.makedirs(TCFG.tensorboard_dir, exist_ok=True) writer = SummaryWriter(log_dir=TCFG.tensorboard_dir) all_metrics = { "config": {**{k: v for k, v in config.to_dict().items() if not k.startswith("_")}, **asdict(TCFG)}, "epochs": [], "steps": [], } global_step = 0 # ═══════════════════════════════════════════════════ # PHASE 1: Train sequence head only (v1 weights frozen) # ═══════════════════════════════════════════════════ print(f"\n{'='*70}") print(f"PHASE 1: Sequence head training ({TCFG.phase1_epochs} epochs)") print(f" v1 memory system: FROZEN") print(f" Sequence reconstructor: TRAINING") print(f"{'='*70}") phase1_groups = make_param_groups_phase1(model) global_step = train_phase( model, modern_model, modern_tok, train_captions, val_captions, phase1_groups, TCFG.phase1_epochs, "phase1", writer, all_metrics, global_step) save_checkpoint(model, None, TCFG.phase1_epochs, global_step, "phase1", os.path.join(TCFG.checkpoint_dir, "phase1_final")) # ═══════════════════════════════════════════════════ # PHASE 2: Joint fine-tune (everything unfrozen) # ═══════════════════════════════════════════════════ print(f"\n{'='*70}") print(f"PHASE 2: Joint fine-tune ({TCFG.phase2_epochs} epochs)") print(f" All trainable modules: TRAINING") print(f" v1 components: reduced LR") print(f"{'='*70}") phase2_groups = make_param_groups_phase2(model) global_step = train_phase( model, modern_model, modern_tok, train_captions, val_captions, phase2_groups, TCFG.phase2_epochs, "phase2", writer, all_metrics, global_step) # ── Final save ── save_checkpoint(model, None, TCFG.phase1_epochs + TCFG.phase2_epochs, global_step, "final", os.path.join(TCFG.checkpoint_dir, "final")) all_metrics["final"] = { "total_steps": global_step, "final_m_acc": all_metrics["epochs"][-1]["m_acc"], "final_seq_cos": all_metrics["epochs"][-1]["seq_cos_sim"], "final_cv": all_metrics["epochs"][-1]["cv_raw"], } with open(TCFG.metrics_file, "w") as f: json.dump(all_metrics, f, indent=2, default=str) writer.flush() writer.close() print(f"\n{'='*70}") print(f"FINAL:") final = all_metrics["epochs"][-1] print(f" m_acc: {final['m_acc']:.4f}") print(f" seq_cos: {final['seq_cos_sim']:.4f}") print(f" CV: {final['cv_raw']:.4f}") print(f"{'='*70}") if __name__ == "__main__": main()