Create trainer.py

trainer.py

@@ -0,0 +1,755 @@
+# ============================================================================
+# 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()