Create distillation_trainer_v3.py

distillation_trainer_v3.py

@@ -0,0 +1,718 @@
+# ============================================================================
+# DEEP BERT v3 — TRAINER
+#
+# Teacher-distilled training. Frozen long-context experts teach the memory
+# system what correct recall looks like.
+#
+# Colab cells:
+#   Cell 1: deep_bert_v3.py   (architecture)
+#   Cell 2: this file          (training)
+#
+# Flow per document:
+#   1. ModernBERT (8192 ctx) → teacher_cls     (frozen, no grad)
+#   2. Longformer (4096 ctx)  → teacher_cls_2   (frozen, no grad)
+#   3. BERT + memory (16×480) → student_cls     (memory trains)
+#   4. Loss: student should match teachers
+# ============================================================================
+
+import gc
+import json
+import math
+import os
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from safetensors.torch import save_file as safetensors_save
+from datasets import load_dataset
+from transformers import AutoModel, AutoTokenizer, BertTokenizer
+from tqdm import tqdm
+
+
+# ══════════════════════════════════════════════════════════════════
+# CONFIG
+# ══════════════════════════════════════════════════════════════════
+
+@dataclass
+class TrainConfig:
+    # Data
+    max_documents: int = 50000
+    max_val_documents: int = 500
+    segment_length: int = 480
+    segment_overlap: int = 64
+    target_chain_segments: int = 16
+    max_segments: int = 16
+    min_segments: int = 6
+
+    # Teachers
+    modern_bert_model: str = "answerdotai/ModernBERT-large"
+    longformer_model: str = "allenai/longformer-large-4096"
+    modern_max_len: int = 8192
+    longformer_max_len: int = 4096
+    procrustes_n_samples: int = 500  # docs for static pre-alignment
+
+    # Training
+    epochs: int = 10
+    batch_size: int = 4
+    lr_bank: float = 2e-3
+    lr_output: float = 5e-4
+    lr_proj: float = 1e-3
+    min_lr: float = 1e-6
+    weight_decay: float = 0.01
+    grad_clip: float = 1.0
+    warmup_steps: int = 300
+    tbptt_segments: int = 0  # 0 = no truncation (clean bank, safe now)
+
+    # Loss weights
+    modern_weight: float = 1.0
+    longformer_weight: float = 0.5
+    cv_weight: float = 0.05
+    temperature: float = 0.07
+
+    # Logging
+    checkpoint_dir: str = "/home/claude/deep_bert_v3_checkpoints"
+    tensorboard_dir: str = "/home/claude/deep_bert_v3_tb"
+    log_every: int = 20
+    eval_every: int = 200
+    save_every_epoch: bool = True
+
+
+TCFG = TrainConfig()
+
+
+# ══════════════════════════════════════════════════════════════════
+# DATA PIPELINE — raw text + student segments
+# ══════════════════════════════════════════════════════════════════
+
+def load_wikitext_documents(split, max_docs):
+    """Load WikiText-103, return list of raw text documents."""
+    print(f"  Loading wikitext-103 ({split})...")
+    ds = load_dataset("wikitext", "wikitext-103-raw-v1", split=split)
+    documents = []
+    current_doc = []
+    for row in ds:
+        text = row.get("text", "").strip()
+        if not text:
+            if current_doc:
+                full = " ".join(current_doc)
+                if len(full) > 100:
+                    documents.append(full)
+                current_doc = []
+            continue
+        if text.startswith("= ") and not text.startswith("= = "):
+            if current_doc:
+                full = " ".join(current_doc)
+                if len(full) > 100:
+                    documents.append(full)
+            current_doc = [text]
+        else:
+            current_doc.append(text)
+    if current_doc:
+        full = " ".join(current_doc)
+        if len(full) > 100:
+            documents.append(full)
+    print(f"  {len(documents)} documents")
+    return documents[:max_docs]
+
+
+def build_chains_with_text(raw_docs, bert_tokenizer):
+    """Build student segment chains AND track raw text for teacher tokenization."""
+    stride = TCFG.segment_length - TCFG.segment_overlap
+    sep_id = bert_tokenizer.sep_token_id
+
+    all_ids, all_masks, all_n_reals, all_texts = [], [], [], []
+    doc_idx = 0
+
+    while doc_idx < len(raw_docs):
+        target_tokens = TCFG.target_chain_segments * stride
+        current_ids = []
+        chain_docs = []
+
+        while len(current_ids) < target_tokens and doc_idx < len(raw_docs):
+            if current_ids:
+                current_ids.append(sep_id)
+            ids = bert_tokenizer.encode(raw_docs[doc_idx], add_special_tokens=False)
+            if len(ids) > 50:
+                current_ids.extend(ids)
+                chain_docs.append(doc_idx)
+            doc_idx += 1
+
+        if len(current_ids) < TCFG.min_segments * stride:
+            continue
+
+        # Build segments
+        seg_ids_list, seg_masks_list = [], []
+        pos = 0
+        while pos < len(current_ids) and len(seg_ids_list) < TCFG.max_segments:
+            end = min(pos + TCFG.segment_length, len(current_ids))
+            seg = current_ids[pos:end]
+            pad = TCFG.segment_length - len(seg)
+            if pad > 0:
+                ids_t = torch.tensor(seg + [0] * pad, dtype=torch.int32)
+                mask_t = torch.cat([torch.ones(len(seg), dtype=torch.int8),
+                                    torch.zeros(pad, dtype=torch.int8)])
+            else:
+                ids_t = torch.tensor(seg[:TCFG.segment_length], dtype=torch.int32)
+                mask_t = torch.ones(TCFG.segment_length, dtype=torch.int8)
+            seg_ids_list.append(ids_t)
+            seg_masks_list.append(mask_t)
+            if end >= len(current_ids):
+                break
+            pos += stride
+
+        n_real = len(seg_ids_list)
+        if n_real < TCFG.min_segments:
+            continue
+        while len(seg_ids_list) < TCFG.max_segments:
+            seg_ids_list.append(torch.zeros(TCFG.segment_length, dtype=torch.int32))
+            seg_masks_list.append(torch.zeros(TCFG.segment_length, dtype=torch.int8))
+
+        all_ids.append(torch.stack(seg_ids_list))
+        all_masks.append(torch.stack(seg_masks_list))
+        all_n_reals.append(n_real)
+        # Raw text for teachers
+        all_texts.append(" ".join(raw_docs[i] for i in chain_docs))
+
+    print(f"  {len(all_n_reals)} chains, segs: "
+          f"min={min(all_n_reals)}, max={max(all_n_reals)}, "
+          f"mean={np.mean(all_n_reals):.1f}")
+    return (torch.stack(all_ids), torch.stack(all_masks),
+            torch.tensor(all_n_reals, dtype=torch.long), all_texts)
+
+
+class ChainDataset(Dataset):
+    def __init__(self, ids, masks, n_reals, texts):
+        self.ids, self.masks, self.n_reals = ids, masks, n_reals
+        self.texts = texts
+
+    def __len__(self):
+        return len(self.n_reals)
+
+    def __getitem__(self, i):
+        return self.ids[i], self.masks[i], self.n_reals[i], self.texts[i]
+
+
+def chain_collate(batch):
+    ids, masks, n_reals, texts = zip(*batch)
+    return (torch.stack(ids), torch.stack(masks),
+            torch.tensor(n_reals, dtype=torch.long), list(texts))
+
+
+# ══════════════════════════════════════════════════════════════════
+# GEOMETRIC UTILITIES
+# ══════════════════════════════════════════════════════════════════
+
+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)
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEACHER UTILITIES
+# ══════════════════════════════════════════════════════════════════
+
+def mean_pool(hidden_states, attention_mask):
+    mask = attention_mask.unsqueeze(-1).float()
+    return (hidden_states * mask).sum(1) / mask.sum(1).clamp(min=1)
+
+
+@torch.no_grad()
+def teacher_forward_modern(model, tokenizer, texts, device, max_len):
+    """ModernBERT forward: standard attention, mean-pool."""
+    inputs = tokenizer(texts, max_length=max_len, padding=True,
+                       truncation=True, return_tensors="pt").to(device)
+    out = model(**inputs)
+    return mean_pool(out.last_hidden_state, inputs.attention_mask)
+
+
+@torch.no_grad()
+def teacher_forward_longformer(model, tokenizer, texts, device, max_len):
+    """Longformer forward: CLS gets global attention."""
+    inputs = tokenizer(texts, max_length=max_len, padding=True,
+                       truncation=True, return_tensors="pt").to(device)
+    # Global attention on CLS token
+    global_attn = torch.zeros_like(inputs.input_ids)
+    global_attn[:, 0] = 1
+    out = model(input_ids=inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                global_attention_mask=global_attn)
+    return out.last_hidden_state[:, 0]  # CLS with global attention
+
+
+# ══════════════════════════════════════════════════════════════════
+# LOSSES
+# ══════════════════════════════════════════════════════════════════
+
+def distillation_loss(student_emb, teacher_emb, temperature=0.07):
+    """InfoNCE: student[i] should be closest to teacher[i] in the batch."""
+    s = F.normalize(student_emb, dim=-1)
+    t = F.normalize(teacher_emb, dim=-1)
+    logits = (s @ t.T) / temperature
+    labels = torch.arange(logits.shape[0], device=logits.device)
+    loss = F.cross_entropy(logits, labels)
+    with torch.no_grad():
+        acc = (logits.argmax(-1) == labels).float().mean().item()
+    return loss, acc
+
+
+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
+    for b in range(B):
+        n = n_reals[b].item()
+        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()
+        n_valid += 1
+    stats = {"cv_raw": total_cv / max(n_valid, 1)}
+    if n_valid == 0:
+        return total_loss, stats
+    return total_loss / n_valid, stats
+
+
+# ══════════════════════════════════════════════════════════════════
+# PARAM GROUPS
+# ══════════════════════════════════════════════════════════════════
+
+def make_param_groups(model):
+    bank_names = {"bank.depth_compressor", "bank.temporal_proj",
+                  "bank.cross_attn", "bank.cross_norms",
+                  "bank.cross_ffns", "bank.ffn_norms"}
+    proj_names = {"proj_modern", "proj_longformer"}
+
+    bank_p, proj_p, output_p = [], [], []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if any(name.startswith(p) for p in proj_names):
+            proj_p.append(param)
+        elif any(name.startswith(p) for p in bank_names):
+            bank_p.append(param)
+        else:
+            output_p.append(param)
+
+    groups = [
+        {"params": bank_p, "lr": TCFG.lr_bank, "name": "bank"},
+        {"params": proj_p, "lr": TCFG.lr_proj, "name": "proj"},
+        {"params": output_p, "lr": TCFG.lr_output, "name": "output"},
+    ]
+    for g in groups:
+        g["weight_decay"] = TCFG.weight_decay
+        n = sum(p.numel() for p in g["params"])
+        print(f"    {g['name']:8s}: {n:>10,} params @ lr={g['lr']}")
+    return groups
+
+
+# ══════════════════════════════════════════════════════════════════
+# STATIC PROCRUSTES PRE-ALIGNMENT
+# ══════════════════════════════════════════════════════════════════
+
+@torch.no_grad()
+def compute_and_init_procrustes(student_model, modern_model, modern_tok,
+                                 long_model, long_tok, bert_tok,
+                                 texts, device):
+    """
+    Feed N texts through BERT (CLS) and each teacher (mean-pool/CLS).
+    Compute Procrustes rotation, initialize projectors.
+    """
+    print(f"\n  Computing static Procrustes on {len(texts)} texts...")
+    student_embs, modern_embs, long_embs = [], [], []
+
+    for i in range(0, len(texts), 16):
+        batch = texts[i:i+16]
+
+        # Student: just BERT CLS (no memory, single segment)
+        bert_inputs = bert_tok(batch, max_length=480, padding=True,
+                               truncation=True, return_tensors="pt").to(device)
+        bert_out = student_model.bert(
+            input_ids=bert_inputs.input_ids,
+            attention_mask=bert_inputs.attention_mask,
+            return_dict=True)
+        student_embs.append(bert_out.last_hidden_state[:, 0].cpu())
+
+        # ModernBERT
+        modern_embs.append(
+            teacher_forward_modern(modern_model, modern_tok, batch,
+                                   device, TCFG.modern_max_len).cpu())
+
+        # Longformer
+        long_embs.append(
+            teacher_forward_longformer(long_model, long_tok, batch,
+                                        device, TCFG.longformer_max_len).cpu())
+
+    student_all = torch.cat(student_embs)
+    modern_all = torch.cat(modern_embs)
+    long_all = torch.cat(long_embs)
+
+    # Procrustes: student → ModernBERT
+    print("  ModernBERT alignment:")
+    R_m, mu_s_m, mu_t_m = compute_static_procrustes(student_all, modern_all)
+    student_model.proj_modern.init_from_procrustes(R_m, mu_s_m, mu_t_m)
+
+    # Procrustes: student → Longformer
+    print("  Longformer alignment:")
+    R_l, mu_s_l, mu_t_l = compute_static_procrustes(student_all, long_all)
+    student_model.proj_longformer.init_from_procrustes(R_l, mu_s_l, mu_t_l)
+
+
+# ══════════════════════════════════════════════════════════════════
+# TRAINING
+# ══════════════════════════════════════════════════════════════════
+
+def train(model, modern_model, modern_tok, long_model, long_tok,
+          train_loader, val_loader=None):
+    device = next(model.parameters()).device
+    os.makedirs(TCFG.checkpoint_dir, exist_ok=True)
+    os.makedirs(TCFG.tensorboard_dir, exist_ok=True)
+    writer = SummaryWriter(log_dir=TCFG.tensorboard_dir)
+
+    param_groups = make_param_groups(model)
+    optimizer = torch.optim.AdamW(param_groups)
+    all_params = [p for g in param_groups for p in g["params"]]
+
+    total_steps = len(train_loader) * TCFG.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()
+    global_step = 0
+    best_val_loss = float("inf")
+
+    print(f"\n  Training: {sum(p.numel() for p in all_params):,} params")
+    print(f"    {len(train_loader)} batches/epoch × {TCFG.batch_size} chains")
+    print(f"    Losses: modern({TCFG.modern_weight}) + long({TCFG.longformer_weight}) "
+          f"+ cv({TCFG.cv_weight})")
+
+    for epoch in range(TCFG.epochs):
+        model.train()
+        losses = {"total": 0, "modern": 0, "longformer": 0, "cv": 0}
+        metrics = {"modern_acc": 0, "long_acc": 0, "cv_raw": 0}
+        n_batches = 0
+        t0 = time.time()
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{TCFG.epochs}", unit="batch")
+
+        for student_ids, student_masks, n_reals, raw_texts in pbar:
+            B = n_reals.shape[0]
+
+            # ── Teacher forwards (frozen, no grad) ──
+            with torch.no_grad():
+                with torch.amp.autocast("cuda"):
+                    modern_cls = teacher_forward_modern(
+                        modern_model, modern_tok, raw_texts,
+                        device, TCFG.modern_max_len)
+                    long_cls = teacher_forward_longformer(
+                        long_model, long_tok, raw_texts,
+                        device, TCFG.longformer_max_len)
+
+            # ── Student forward (memory system trains) ──
+            state = model.init_state(B, device)
+            all_anchors = torch.zeros(B, TCFG.max_segments, model.config.anchor_dim,
+                                      device=device)
+
+            for seg_k in range(TCFG.max_segments):
+                if TCFG.tbptt_segments > 0 and seg_k > 0 and seg_k % TCFG.tbptt_segments == 0:
+                    state = DeepBertV3.detach_state(state)
+                    all_anchors = all_anchors.detach()
+
+                ids = student_ids[:, seg_k].to(device).long()
+                mask = student_masks[:, seg_k].to(device).long()
+
+                with torch.amp.autocast("cuda"):
+                    outputs, state = model(ids, mask, state)
+                all_anchors[:, seg_k] = outputs["live_anchor"]
+
+            # Student output: fused (CLS + memory delta) from last real segment
+            student_cls = outputs["memory_output"]
+
+            # ── Project into teacher spaces ──
+            with torch.amp.autocast("cuda"):
+                proj_m = model.proj_modern(student_cls)
+                proj_l = model.proj_longformer(student_cls)
+
+                # ── Distillation losses ──
+                l_modern, acc_m = distillation_loss(
+                    proj_m, modern_cls, TCFG.temperature)
+                l_long, acc_l = distillation_loss(
+                    proj_l, long_cls, TCFG.temperature)
+
+                # ── CV on live anchors ──
+                l_cv, cv_stats = batch_cv_loss(
+                    all_anchors, n_reals.to(device), model.config.cv_target)
+
+                loss = (TCFG.modern_weight * l_modern +
+                       TCFG.longformer_weight * l_long +
+                       TCFG.cv_weight * l_cv)
+
+            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
+
+            losses["total"] += loss.item()
+            losses["modern"] += l_modern.item()
+            losses["longformer"] += l_long.item()
+            losses["cv"] += l_cv.item()
+            metrics["modern_acc"] += acc_m
+            metrics["long_acc"] += acc_l
+            metrics["cv_raw"] += cv_stats.get("cv_raw", 0)
+            n_batches += 1
+
+            n = max(n_batches, 1)
+            pbar.set_postfix(
+                loss=f"{losses['total']/n:.3f}",
+                m_acc=f"{metrics['modern_acc']/n:.3f}",
+                l_acc=f"{metrics['long_acc']/n:.3f}",
+                cv=f"{metrics['cv_raw']/n:.3f}")
+
+            if global_step % TCFG.log_every == 0:
+                writer.add_scalar("train/loss", losses["total"] / n, global_step)
+                writer.add_scalar("train/modern_acc", metrics["modern_acc"] / n, global_step)
+                writer.add_scalar("train/long_acc", metrics["long_acc"] / n, global_step)
+                writer.add_scalar("train/cv_raw", metrics["cv_raw"] / n, global_step)
+                for k in ["modern", "longformer", "cv"]:
+                    writer.add_scalar(f"train/{k}_loss", losses[k] / n, global_step)
+
+            if val_loader and global_step % TCFG.eval_every == 0:
+                vl = evaluate(model, modern_model, modern_tok,
+                             long_model, long_tok, val_loader, writer, global_step)
+                if vl < best_val_loss:
+                    best_val_loss = vl
+                    save_checkpoint(model, optimizer, epoch, global_step,
+                                   os.path.join(TCFG.checkpoint_dir, "best"))
+                model.train()
+
+        pbar.close()
+        elapsed = time.time() - t0
+        n = max(n_batches, 1)
+        print(f"\n  Epoch {epoch+1}: {n_batches * TCFG.batch_size / elapsed:.1f} chains/s  "
+              f"loss={losses['total']/n:.4f}  "
+              f"m_acc={metrics['modern_acc']/n:.3f}  "
+              f"l_acc={metrics['long_acc']/n:.3f}  "
+              f"cv={metrics['cv_raw']/n:.3f}")
+
+        if TCFG.save_every_epoch:
+            save_checkpoint(model, optimizer, epoch + 1, global_step,
+                           os.path.join(TCFG.checkpoint_dir, f"epoch_{epoch+1:03d}"))
+
+    save_checkpoint(model, optimizer, TCFG.epochs, global_step,
+                   os.path.join(TCFG.checkpoint_dir, "final"))
+    writer.flush()
+    writer.close()
+
+
+# ══════════════════════════════════════════════════════════════════
+# EVAL
+# ══════════════════════════════════════════════════════════════════
+
+@torch.no_grad()
+def evaluate(model, modern_model, modern_tok, long_model, long_tok,
+             val_loader, writer=None, global_step=0):
+    model.eval()
+    device = next(model.parameters()).device
+    total = {"loss": 0, "modern_acc": 0, "long_acc": 0, "cv_raw": 0}
+    n = 0
+
+    for student_ids, student_masks, n_reals, raw_texts in tqdm(val_loader, desc="Eval", leave=False):
+        B = n_reals.shape[0]
+
+        with torch.amp.autocast("cuda"):
+            modern_cls = teacher_forward_modern(
+                modern_model, modern_tok, raw_texts, device, TCFG.modern_max_len)
+            long_cls = teacher_forward_longformer(
+                long_model, long_tok, raw_texts, device, TCFG.longformer_max_len)
+
+        state = model.init_state(B, device)
+        all_anc = torch.zeros(B, TCFG.max_segments, model.config.anchor_dim, device=device)
+        for seg_k in range(TCFG.max_segments):
+            with torch.amp.autocast("cuda"):
+                out, state = model(student_ids[:, seg_k].to(device).long(),
+                                   student_masks[:, seg_k].to(device).long(), state)
+            all_anc[:, seg_k] = out["live_anchor"]
+
+        with torch.amp.autocast("cuda"):
+            student_cls = out["memory_output"]
+            l_m, acc_m = distillation_loss(
+                model.proj_modern(student_cls), modern_cls, TCFG.temperature)
+            l_l, acc_l = distillation_loss(
+                model.proj_longformer(student_cls), long_cls, TCFG.temperature)
+            l_cv, cv_s = batch_cv_loss(all_anc, n_reals.to(device), 0.20)
+
+        total["loss"] += (TCFG.modern_weight * l_m.item() +
+                         TCFG.longformer_weight * l_l.item() +
+                         TCFG.cv_weight * l_cv.item())
+        total["modern_acc"] += acc_m
+        total["long_acc"] += acc_l
+        total["cv_raw"] += cv_s.get("cv_raw", 0)
+        n += 1
+
+    d = max(n, 1)
+    print(f"    Val: loss={total['loss']/d:.4f}  "
+          f"m_acc={total['modern_acc']/d:.3f}  "
+          f"l_acc={total['long_acc']/d:.3f}  "
+          f"cv={total['cv_raw']/d:.3f}")
+    if writer:
+        for k, v in total.items():
+            writer.add_scalar(f"val/{k}", v / d, global_step)
+    return total["loss"] / d
+
+
+# ══════════════════════════════════════════════════════════════════
+# CHECKPOINT
+# ══════════════════════════════════════════════════════════════════
+
+def save_checkpoint(model, optimizer, epoch, global_step, 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(), "epoch": epoch,
+                "global_step": global_step}, os.path.join(path, "training_state.pt"))
+    import dataclasses
+    with open(os.path.join(path, "config.json"), "w") as f:
+        json.dump({"model": dataclasses.asdict(model.config),
+                   "training": dataclasses.asdict(TCFG)}, f, indent=2, default=str)
+
+
+# ══════════════════════════════════════════════════════════════════
+# MAIN
+# ══════════════════════════════════════════════════════════════════
+
+def main():
+    print("=" * 70)
+    print("DEEP BERT v3 — TEACHER-DISTILLED GEOMETRIC MEMORY")
+    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")
+
+    # ── Load student model ──
+    print(f"\n{'='*70}")
+    print("LOADING MODELS")
+    print(f"{'='*70}")
+
+    config = DeepBertV3Config()
+    model = DeepBertV3.from_pretrained(config).to(device)
+    bert_tokenizer = BertTokenizer.from_pretrained(config.bert_model)
+
+    # ── Load teachers (frozen) ──
+    print(f"\n  Loading ModernBERT-large...")
+    modern_model = AutoModel.from_pretrained(TCFG.modern_bert_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(TCFG.modern_bert_model)
+    print(f"    {sum(p.numel() for p in modern_model.parameters()):,} params (frozen)")
+
+    print(f"\n  Loading Longformer-large...")
+    long_model = AutoModel.from_pretrained(TCFG.longformer_model,
+                                            torch_dtype=torch.float16).to(device)
+    long_model.eval()
+    for p in long_model.parameters():
+        p.requires_grad = False
+    long_tok = AutoTokenizer.from_pretrained(TCFG.longformer_model)
+    print(f"    {sum(p.numel() for p in long_model.parameters()):,} params (frozen)")
+
+    # ── Data ──
+    print(f"\n{'='*70}")
+    print("DATA")
+    print(f"{'='*70}")
+
+    train_docs = load_wikitext_documents("train", TCFG.max_documents)
+    train_ids, train_masks, train_nr, train_texts = build_chains_with_text(
+        train_docs, bert_tokenizer)
+
+    val_docs = load_wikitext_documents("validation", TCFG.max_val_documents)
+    val_ids, val_masks, val_nr, val_texts = build_chains_with_text(
+        val_docs, bert_tokenizer)
+
+    train_ds = ChainDataset(train_ids, train_masks, train_nr, train_texts)
+    val_ds = ChainDataset(val_ids, val_masks, val_nr, val_texts)
+
+    train_loader = DataLoader(train_ds, batch_size=TCFG.batch_size, shuffle=True,
+                              num_workers=0, pin_memory=True, drop_last=True,
+                              collate_fn=chain_collate)
+    val_loader = DataLoader(val_ds, batch_size=TCFG.batch_size, shuffle=False,
+                            num_workers=0, pin_memory=True,
+                            collate_fn=chain_collate)
+
+    print(f"\n  Train: {len(train_ds)} chains → {len(train_loader)} batches")
+    print(f"  Val:   {len(val_ds)} chains → {len(val_loader)} batches")
+
+    # ── Static Procrustes pre-alignment ──
+    print(f"\n{'='*70}")
+    print("PROCRUSTES PRE-ALIGNMENT")
+    print(f"{'='*70}")
+
+    # Use first N train docs for alignment
+    align_texts = train_texts[:TCFG.procrustes_n_samples]
+    compute_and_init_procrustes(
+        model, modern_model, modern_tok, long_model, long_tok,
+        bert_tokenizer, align_texts, device)
+
+    # ── Train ──
+    print(f"\n{'='*70}")
+    print("TRAINING")
+    print(f"{'='*70}")
+
+    train(model, modern_model, modern_tok, long_model, long_tok,
+          train_loader, val_loader)
+
+    # ── Final eval ──
+    print(f"\n{'='*70}")
+    print("FINAL EVALUATION")
+    print(f"{'='*70}")
+
+    evaluate(model, modern_model, modern_tok, long_model, long_tok, val_loader)
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()