experiments/tiny_nla/train_ar.py

#!/usr/bin/env python3
"""
Stage 1: AR SFT — Activation Reconstructor

Input:  explanation text (via teacher)
Output: reconstructed activation vector (d_model)

Architecture:
  - Qwen3-0.6B-Base trunk (frozen)
  - Linear(d_model, d_model) head on last token hidden state
  - L2-normalized MSE loss

Baselines:
  - Mean vector baseline (always predict mean activation)
  - Shuffled label baseline (random pairing)
"""

import json, os, time, math, yaml, random
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer

# ── paths ──────────────────────────────────────────────
REPO_ROOT = Path(__file__).resolve().parents[2]
SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
CHECKPOINT_DIR = ARTIFACTS_DIR / "checkpoints" / "ar"
CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)

BASE_MODEL = "Qwen/Qwen3-0.6B-Base"

# ── device ─────────────────────────────────────────────
def detect_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")

# ── AR Head ────────────────────────────────────────────
class ARHead(nn.Module):
    """Simple linear projection from d_model to d_model."""
    def __init__(self, d_model: int):
        super().__init__()
        self.linear = nn.Linear(d_model, d_model, bias=False)
    
    def forward(self, x):
        # x: [batch, d_model] — last token hidden state
        return self.linear(x)


# ── Dataset ────────────────────────────────────────────
class ARDataset(Dataset):
    """Explanation text → activation vector."""
    def __init__(self, records, activations, tokenizer, max_length=128):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = []
        
        for rec, act in zip(records, activations):
            explanation = rec.get("teacher_explanation", "") or rec.get("teacher_explanation_raw", "")
            if not explanation or explanation in ("[空输出]", ""):
                continue
            self.data.append((explanation, act))
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        text, act = self.data[idx]
        tokens = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt",
        )
        return {
            "input_ids": tokens["input_ids"][0],
            "attention_mask": tokens["attention_mask"][0],
            "target": act,  # [d_model]
        }


# ── Loss ───────────────────────────────────────────────
def normalized_mse_loss(pred, target):
    """
    L2-normalize both, then MSE = 2*(1-cosine)
    Equivalent to ||normalized_pred - normalized_target||²
    """
    pred_n = F.normalize(pred, dim=-1)
    target_n = F.normalize(target, dim=-1)
    return F.mse_loss(pred_n, target_n)


def cosine_similarity(pred, target):
    pred_n = F.normalize(pred, dim=-1)
    target_n = F.normalize(target, dim=-1)
    return (pred_n * target_n).sum(dim=-1)


# ── Training ───────────────────────────────────────────
def train():
    print("=" * 60)
    print("🔧 AR SFT — Activation Reconstructor Training")
    print("=" * 60)
    
    device = detect_device()
    dtype = torch.float32
    print(f"  Device: {device}, dtype: {dtype}")
    
    # Load sidecar
    with open(SIDECAR_PATH, "r") as f:
        nla_meta = yaml.safe_load(f)
    d_model = nla_meta["d_model"]
    print(f"  d_model: {d_model}")
    
    # Load dataset
    print("\n  Loading dataset...")
    jsonl_path = ARTIFACTS_DIR / "dataset.jsonl"
    act_path = ARTIFACTS_DIR / "activations.pt"
    
    records = []
    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            records.append(json.loads(line))
    activations = torch.load(act_path, weights_only=True)
    print(f"  Loaded {len(records)} records, activations shape: {activations.shape}")
    
    # Train/val split
    random.seed(42)
    indices = list(range(len(records)))
    random.shuffle(indices)
    val_size = max(1, int(len(indices) * 0.15))
    train_idx, val_idx = indices[val_size:], indices[:val_size]
    print(f"  Train: {len(train_idx)}, Val: {len(val_idx)}")
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    
    train_records = [records[i] for i in train_idx]
    train_acts = activations[train_idx]
    val_records = [records[i] for i in val_idx]
    val_acts = activations[val_idx]
    
    train_dataset = ARDataset(train_records, train_acts, tokenizer)
    val_dataset = ARDataset(val_records, val_acts, tokenizer)
    print(f"  Train dataset: {len(train_dataset)}, Val dataset: {len(val_dataset)}")
    
    train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)
    
    # Load base model (frozen)
    print("\n  Loading base model (frozen)...")
    t0 = time.perf_counter()
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        trust_remote_code=True,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
        attn_implementation="eager",
    ).to(device)
    base_model.eval()
    for p in base_model.parameters():
        p.requires_grad = False
    print(f"  Base model loaded in {time.perf_counter() - t0:.1f}s")
    
    # AR head
    ar_head = ARHead(d_model).to(device)
    print(f"  AR head params: {sum(p.numel() for p in ar_head.parameters()):,}")
    
    optimizer = torch.optim.AdamW(ar_head.parameters(), lr=1e-3)
    
    # ── Compute baselines ──
    print("\n📊 Computing baselines...")
    
    # Mean vector baseline
    all_train_acts = torch.stack([train_dataset[i]["target"] for i in range(len(train_dataset))])
    mean_vec = all_train_acts.mean(dim=0)  # [d_model]
    
    mean_cosines = []
    for i in range(len(val_dataset)):
        target = val_dataset[i]["target"]
        cos = F.cosine_similarity(mean_vec.unsqueeze(0), target.unsqueeze(0))
        mean_cosines.append(cos.item())
    mean_baseline = torch.tensor(mean_cosines).mean().item()
    print(f"  Mean-vector baseline cosine: {mean_baseline:.4f}")
    
    # Shuffled baseline
    val_targets = torch.stack([val_dataset[i]["target"] for i in range(len(val_dataset))])
    shuffled = val_targets[torch.randperm(len(val_targets))]
    shuffled_cosines = F.cosine_similarity(val_targets, shuffled, dim=-1)
    shuffled_baseline = shuffled_cosines.mean().item()
    print(f"  Shuffled baseline cosine:      {shuffled_baseline:.4f}")
    
    # ── Training loop ──
    print("\n🏋️  Training AR head...")
    n_epochs = 50
    best_val_loss = float("inf")
    best_val_cosine = 0.0
    
    for epoch in range(n_epochs):
        # Train
        ar_head.train()
        total_loss = 0.0
        n_batches = 0
        
        for batch in train_loader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            targets = batch["target"].to(device)
            
            with torch.no_grad():
                outputs = base_model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=True,
                )
                # Get last token hidden state from last layer
                last_hidden = outputs.hidden_states[-1]  # [batch, seq, d_model]
                # Use the last non-pad token position
                seq_lens = attention_mask.sum(dim=1) - 1
                batch_indices = torch.arange(last_hidden.size(0), device=device)
                last_token_hidden = last_hidden[batch_indices, seq_lens, :]  # [batch, d_model]
            
            pred = ar_head(last_token_hidden)
            loss = normalized_mse_loss(pred, targets)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            n_batches += 1
        
        avg_train_loss = total_loss / n_batches
        
        # Validation
        ar_head.eval()
        val_losses = []
        val_cosines = []
        
        with torch.no_grad():
            for batch in val_loader:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                targets = batch["target"].to(device)
                
                outputs = base_model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=True,
                )
                last_hidden = outputs.hidden_states[-1]
                seq_lens = attention_mask.sum(dim=1) - 1
                batch_indices = torch.arange(last_hidden.size(0), device=device)
                last_token_hidden = last_hidden[batch_indices, seq_lens, :]
                
                pred = ar_head(last_token_hidden)
                loss = normalized_mse_loss(pred, targets)
                cos = cosine_similarity(pred, targets)
                
                val_losses.append(loss.item())
                val_cosines.extend(cos.cpu().tolist())
        
        avg_val_loss = sum(val_losses) / len(val_losses)
        avg_val_cosine = sum(val_cosines) / len(val_cosines)
        
        if avg_val_cosine > best_val_cosine:
            best_val_cosine = avg_val_cosine
            best_val_loss = avg_val_loss
            torch.save(ar_head.state_dict(), CHECKPOINT_DIR / "best_ar_head.pt")
        
        if (epoch + 1) % 5 == 0 or epoch == 0:
            print(f"  Epoch {epoch+1:2d}/{n_epochs} | train_loss={avg_train_loss:.6f} | val_loss={avg_val_loss:.6f} | val_cos={avg_val_cosine:.4f} | best_cos={best_val_cosine:.4f}")
    
    # ── Final evaluation ──
    print(f"\n📈 Final Results")
    print(f"  {'':30} {'Loss':>10} {'Cosine':>10}")
    print(f"  {'Mean-vector baseline':30} {'':>10} {mean_baseline:>10.4f}")
    print(f"  {'Shuffled baseline':30} {'':>10} {shuffled_baseline:>10.4f}")
    print(f"  {'AR trained (best)':30} {best_val_loss:>10.6f} {best_val_cosine:>10.4f}")
    
    improvement = best_val_cosine - max(mean_baseline, shuffled_baseline)
    print(f"\n  Improvement over best baseline: {improvement:+.4f}")
    
    if improvement <= 0:
        print("  ⚠️  AR does NOT beat baselines! May need fixes.")
    
    print(f"\n  Checkpoint: {CHECKPOINT_DIR / 'best_ar_head.pt'}")
    
    return best_val_cosine, best_val_loss


if __name__ == "__main__":
    train()