Upload 127 files

.gitattributes

@@ -35,3 +35,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 hugging/td_lang/__pycache__/compiler.cpython-314.pyc filter=lfs diff=lfs merge=lfs -text
 hugging/td_lang/__pycache__/compiler.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
+hugging/td_lang/td_lang/__pycache__/compiler.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text
+hugging/td_lang/td_lang/__pycache__/compiler.cpython-314.pyc filter=lfs diff=lfs merge=lfs -text

hugging/CLAUDE.md

@@ -0,0 +1,148 @@
+# Memory
+
+## Me
+Milan (Libby's account). Building TD (Time Dilation) — a self-improving AI system using a 7B model on home hardware.
+
+## People
+| Who | Role |
+|-----|------|
+| **Milan** | Project lead, TD creator. Hands-on, wants things explained simply |
+| **Milan's dad** | Budget decision-maker AND critical thinker. Said "if it's worth investing in, money isn't the issue" but also challenged everything with hard questions. His critiques forced the pivot from old plan to new plan. |
+
+> Full list: memory/glossary.md, profiles: memory/people/
+
+## Terms
+| Term | Meaning |
+|------|---------|
+| TD | Time Dilation — the self-improving AI project |
+| ALAS | Autonomous Learning Agent System — self-learning via web search |
+| Fara-7B | Microsoft's vision-based browser agent (MIT, open source, based on Qwen2.5-VL) |
+| Qwen3-VL-8B | Qwen3 with vision + browser agent — replaces Fara as our CUA base |
+| GRPO | Group Relative Policy Optimisation — RL for verified reasoning |
+| SimPO | Simple Preference Optimisation — reference-free preference training |
+| SLIME | Improved SimPO — dual-margin stability, fixes online collapse |
+| QLoRA | Quantised Low-Rank Adaptation — memory-efficient fine-tuning |
+| PRMs | Process Reward Models — step-by-step reasoning verification |
+| ThinkPRM | PRMs that think — uses 1% of labelling data |
+| WebRL | Self-evolving curriculum RL for web agents |
+| STaR | Self-Taught Reasoner — train on correct reasoning chains |
+| FuseLLM | Merge multiple fine-tuned models into one |
+| TIES/DARE-TIES | Weight merging algorithms for FuseLLM |
+| Transport and Merge | Cross-architecture model merging via optimal transport (Feb 2026) |
+| OrthoMerge | Merging on Riemannian manifold, preserves weight geometry |
+| LARV | Layer-wise Adaptive Rescaling — per-layer scaling for merges |
+| Git Re-Basin | Neuron permutation matching — PUBLIC CODE foundation for merging |
+| SEC | Self-Evolving Curriculum — auto-adjusts training difficulty |
+| Cherry_LLM | Self-data filtering via perplexity scoring |
+| SimpleMem | 26.4% better than Mem0, 30x more efficient memory |
+| JitRL | Training-free continual learning — outperforms WebRL |
+| Latent Reasoning | Scales 7B to ~50B performance at inference |
+| Layer 0-5 | TD's 6-layer architecture (0=instant, 1=data, 2=filter, 3=train, 4=agents, 5=merge) |
+
+> Full glossary: memory/glossary.md
+
+## Projects
+| Name | What |
+|------|------|
+| **TD (Time Dilation)** | Self-improving 7B AI system. 89 techniques, 29 core. 6-layer architecture |
+
+> Details: memory/projects/
+
+## Merge Strategy
+- Target model: Qwen3-VL-8B-Instruct (vision + browser agent + text, thinking mode)
+- Why VL: Same language brain as Qwen3-8B, but adds vision + CUA abilities for free (replaces need for Fara)
+- Merge approach: Only merge into language backbone layers, vision encoder stays untouched
+- Method: Transport and Merge (optimal transport cross-arch merging)
+- Merge in: DeepSeek-R1-Distill, MiMo-7B, Llama 3.1, Falcon-H1R-7B
+- Fallback: Knowledge distillation for any model that fails to merge
+- NO direct merges possible — all 5 models have different architectures
+- Kimi K2 ruled out (1T params, too big)
+- Full strategy: docs/MERGE_STRATEGY.md
+
+## Dad's Tests (Critical Thinking Filter)
+Every claim must pass these before being accepted:
+1. **Economic test:** "If this worked cheaply, why aren't big tech companies doing it?"
+2. **Architecture test:** "Is this built on something that's dying or futureproof?"
+3. **Realism test:** "Is this actually achievable or just optimism?"
+4. **Pragmatism test:** "Can we use what we already have first?"
+5. **Long-term test:** "Will this still matter in 2-3 years?"
+
+Dad's exact words: "I didn't ask for the marketing spill, give to the point answer." He called out that LLMs are "on their way out" and questioned whether weight-copying works. His critiques were RIGHT — P100 didn't work, weight copying was wrong, old timelines were fantasy. The pivot to Transport and Merge + dual 4090 happened because of his challenges.
+
+## TD History (Old vs New Plan)
+- **OLD plan (Jan-Feb 2026):** Copy Mistral-7B weights, spawn copies for research, merge knowledge back via JSON. Hardware: Tesla P40 + desktop (~$250). This plan FAILED — weight copying doesn't transfer knowledge, P100 incompatible with Unsloth, timelines were fantasy.
+- **NEW plan (Feb 2026):** Transport and Merge 5 different models into Qwen3-VL-8B (vision+text), then GRPO self-improvement loop. Hardware: dual RTX 4090 or vast.ai GPU rental. Self-improvement through actual RL training (weights change), not code self-modification or JSON merging. Switched from Qwen3-8B to Qwen3-VL-8B to get browser agent abilities (like Fara) built in.
+- **What TD will be:** A regular AI assistant like ChatGPT, but hopefully smarter after training cycles. NOT superintelligence promises.
+
+## Self-Improvement Loop (Discovered Feb 2026)
+Milan interviewed ChatGPT, Grok, and Gemini (12+ interviews, test_1 to test_12+) about recursive self-improvement.
+Key discovery: **The model can be its own diagnostician.**
+- All 3 AIs could list their own weaknesses when asked "what would you improve?"
+- All 3 said the only thing stopping them is no access to their own weights/training
+- All 3 converged on the same "small" self-improvement loop that actually works:
+
+**The TD Self-Improvement Loop:**
+1. Merge multiple models together (Transport and Merge) → creates strong base
+2. Ask the model "what are you bad at?" → it identifies weak spots
+3. Generate targeted synthetic training data for those weak spots
+4. Train with GRPO/STaR on that data → model gets slightly better
+5. The improved model generates better reasoning chains → better training data
+6. Repeat — each cycle is small (1-5%) but compounds
+
+**Two codebases (td_fuse absorbed into td_lang):**
+- `td_lang` — THE complete TD system. Domain-specific language + merge engine + training + RL. v0.2.0, ~11,422 lines total (7,878 core + 3,544 engine), 18 .py files + 22 examples. All 13 phases complete. td_fuse was absorbed into td_lang/engine/ so td_lang runs everything — no external Python deps for the pipeline. Built collaboratively: Claude (architecture), Codex (hardening), Gemini (in-IDE testing).
+- `td_loop` — self-recursive improvement loop (planned, automates the cycle above). May not be needed since td_lang's `repeat` block + arena already handle this.
+
+**What's NOT possible (confirmed by all 3 AIs + dad's tests):**
+- Live weight editing (model rewriting its own brain in real-time)
+- Direct weight manipulation like editing a text file
+- "Cogniscript"/"Phylang"/"Lumina-Σ" (sci-fi languages from the interviews — NOT real)
+
+**What IS possible (confirmed by all 3 AIs + real papers):**
+- Generate → Filter → Train → Evaluate → Keep winners → Repeat
+- Using mechanistic interpretability to find weak circuits, then training specifically on those
+- STaR (train on correct reasoning chains), GRPO (RL for reasoning), Cherry_LLM (filter bad data)
+
+**Interview technical findings (test_12):**
+- LoRA target: mid-to-late layers MLP blocks (layers 16–28 for 32-layer model). All 3 AIs agree.
+- Biggest weakness: long-chain reasoning breaks at step 18–30. Target this with GRPO.
+- Self-training trap: 100 steps on own outputs → smoother but dumber. MUST mix external data.
+- Cherry_LLM perplexity filter prevents mode collapse by catching repetitive training data.
+
+**Cost optimization (test_16):**
+- Inference-time scaling: 80–90% of gains for 5–30% cost. Generate multiple answers, pick best, train on winners.
+- Verified rewards only: no learned reward model, just objective checkers (code compiles, math correct). Saves VRAM.
+- Budget: 70–80% inference scaling, 10–20% short GRPO, 5–10% tooling
+- Speculative decoding (vLLM): small draft model + main model verifying = 2–3× faster inference
+
+**td_lang design requirements (test_17 — ChatGPT's ForgeSpec 2.0):**
+- 8 features: data contracts, reward contracts, eval gates (mandatory), resource budgets (compiler enforced), automatic ablations, artifact lineage (content-hash), serving SLOs, economics reports
+- Three quality gates for td_loop: holdout (real tasks), adversarial (break it on purpose), calibration (confidence vs accuracy)
+- OpenRLHF: real framework (Ray+vLLM+DeepSpeed) for GRPO at scale — could replace custom td_loop plumbing
+- GaLore: full-param training at 65% less VRAM (alternative to QLoRA)
+- PACER (Feb 2026): sample 8-64 traces → consensus packet → one revision = 1/8 tokens of majority voting
+
+**Phase 3 deep dive (test_18 — all 3 AIs answered both prompts):**
+- FORK: disk-based only on 4090. Cheap fork = manifest + adapter copy. safetensors format.
+- RESET: del model → clear cache → reload. Must reset optimizer state. Use assign=True.
+- PRUNE: 20% structured max (LLM-Pruner paper). Wanda metric (Grok, practical on 4090). Language backbone only, never vision. Recovery: 200-800 steps LoRA r=8.
+- EDIT: LoRA/DoRA with layers_to_transform for layers 16-28. "Try before buy" via enable/disable adapters. ROME/MEMIT not ready for Qwen3-VL.
+- Build order: EDIT first → FORK/RESET → PRUNE last
+- ChatGPT's manifest idea: model state = base_ref + adapters[] + prune_spec + optimizer + eval_report
+
+**Interview files:** stored in interview/ folder (test_1.txt through test_18.txt + screenshots)
+- ChatGPT: Most conservative, gave systems-level analysis, refused operational blueprints
+- Grok: Most detailed and realistic, named specific models/hardware, grounded in real papers
+- Gemini: Most flattering/sci-fi, referenced Milan's own work, made up technologies
+
+## Preferences
+- Explain things simply — analogies and plain English
+- Use all available tools and commands
+- Be honest about what works and what doesn't — Milan values truth over optimism
+- Budget is flexible — focus on best strategy, not cheapest hardware
+- Keep one master document (currently v5.2 in docs/)
+- Old files go to DELETE/ folder for Milan to trash
+- No dashboards or visual tools — Milan doesn't need them
+- Plugins are welcome if they genuinely help and don't break anything
+- Run every claim by "dad's tests" before presenting it as fact
+- The uploaded 6-part transcript is the OLD TD version — useful for self-improvement context but NOT the current plan

hugging/install.sh

@@ -0,0 +1,160 @@
+#!/bin/bash
+# ============================================================================
+# TD (Time Dilation) — One-Command Setup
+# ============================================================================
+#
+# Run this ONCE on a fresh machine with a GPU:
+#   chmod +x install.sh && ./install.sh
+#
+# What it does:
+#   1. Installs all Python dependencies
+#   2. Downloads the base model (Qwen3-VL-8B-Instruct)
+#   3. Downloads the Transport and Merge code
+#   4. Sets up output directories
+#   5. Verifies GPU access
+#   6. Compiles the starter TD file to make sure everything works
+#
+# After this, just run:
+#   python -m td_lang run td_start.td
+#
+# Requirements:
+#   - Python 3.10+
+#   - NVIDIA GPU with 24GB+ VRAM (RTX 4090 or better)
+#   - ~50GB disk space (models + checkpoints)
+#   - Internet connection (first run only)
+# ============================================================================
+
+set -e  # Stop on any error
+
+echo "============================================================"
+echo "  TD (Time Dilation) — Setup Script"
+echo "============================================================"
+echo ""
+
+# ── Step 1: Check Python ──
+echo "[1/7] Checking Python..."
+if ! command -v python3 &> /dev/null; then
+    echo "ERROR: Python 3 not found. Install Python 3.10+ first."
+    exit 1
+fi
+PYTHON_VER=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
+echo "  Python $PYTHON_VER found."
+
+# ── Step 2: Check GPU ──
+echo ""
+echo "[2/7] Checking GPU..."
+if command -v nvidia-smi &> /dev/null; then
+    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1)
+    GPU_MEM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -1)
+    echo "  GPU: $GPU_NAME ($GPU_MEM)"
+else
+    echo "  WARNING: nvidia-smi not found. GPU might not be available."
+    echo "  Continuing anyway (some features won't work without GPU)."
+fi
+
+# ── Step 3: Install Python packages ──
+echo ""
+echo "[3/7] Installing Python packages..."
+echo "  This takes 5-10 minutes on first run."
+pip install --break-system-packages -q \
+    torch \
+    transformers \
+    accelerate \
+    bitsandbytes \
+    peft \
+    trl \
+    datasets \
+    safetensors \
+    sentencepiece \
+    protobuf \
+    scipy \
+    lark \
+    duckduckgo-search \
+    huggingface_hub \
+    2>&1 | tail -5
+
+# Unsloth (optional — speeds up training 2x, but can fail on some systems)
+echo "  Trying to install Unsloth (optional speed boost)..."
+pip install --break-system-packages -q unsloth 2>/dev/null && echo "  Unsloth installed." || echo "  Unsloth not available (that's fine, PEFT fallback works)."
+
+echo "  Packages installed."
+
+# ── Step 4: Download base model ──
+echo ""
+echo "[4/7] Downloading base model (Qwen3-VL-8B-Instruct)..."
+echo "  This is ~16GB. Go grab a coffee."
+python3 -c "
+from huggingface_hub import snapshot_download
+print('  Downloading Qwen/Qwen3-VL-8B-Instruct...')
+path = snapshot_download('Qwen/Qwen3-VL-8B-Instruct', local_dir='./models/Qwen3-VL-8B-Instruct')
+print(f'  Downloaded to: {path}')
+"
+echo "  Base model ready."
+
+# ── Step 5: Download Transport and Merge code ──
+echo ""
+echo "[5/7] Downloading Transport and Merge code..."
+if [ ! -d "Cross-Architecture-Merging-for-Large-Language-Models" ]; then
+    git clone https://github.com/FedML-AI/Cross-Architecture-Merging-for-Large-Language-Models.git
+    echo "  T&M code cloned."
+else
+    echo "  T&M code already exists, skipping."
+fi
+
+# ── Step 6: Set up directories ──
+echo ""
+echo "[6/7] Setting up directories..."
+mkdir -p td_lang_outputs/{checkpoints,snapshots,arena_logs,committed}
+echo "  Output directories created."
+
+# ── Step 7: Verify everything works ──
+echo ""
+echo "[7/7] Verifying installation..."
+
+# Check td_lang compiles
+python3 -c "
+from td_lang.grammar import parse_td_file
+from td_lang.compiler import compile_program
+import ast
+
+program = parse_td_file('td_start.td')
+code = compile_program(program)
+ast.parse(code)
+print('  td_lang: OK (td_start.td compiles)')
+"
+
+# Check GPU access from Python
+python3 -c "
+import torch
+if torch.cuda.is_available():
+    gpu = torch.cuda.get_device_name(0)
+    mem = torch.cuda.get_device_properties(0).total_mem / 1024**3
+    print(f'  PyTorch GPU: {gpu} ({mem:.0f}GB)')
+else:
+    print('  PyTorch GPU: NOT AVAILABLE (CPU only)')
+"
+
+# Check key libraries
+python3 -c "
+import transformers, peft, trl, bitsandbytes, lark, datasets
+print(f'  transformers: {transformers.__version__}')
+print(f'  peft: {peft.__version__}')
+print(f'  trl: {trl.__version__}')
+print('  All libraries: OK')
+"
+
+echo ""
+echo "============================================================"
+echo "  SETUP COMPLETE!"
+echo "============================================================"
+echo ""
+echo "  To start TD, run:"
+echo "    python -m td_lang run td_start.td"
+echo ""
+echo "  To just compile (preview what it'll do):"
+echo "    python -m td_lang compile td_start.td"
+echo ""
+echo "  To check syntax only:"
+echo "    python -m td_lang check td_start.td"
+echo ""
+echo "============================================================"

hugging/td_lang/__init__.py

@@ -31,7 +31,12 @@ Phase 6: fuse, absorb (easy merge)
 Phase 7: repeat, if/else (loop control)
 Phase 8: setup, on_error, notify, save (autopilot)
 Phase 9: schedule (time-based execution)
+Phase 10: download, log, compare, verify (toolbox)
+Phase 11: vote, prompt, distill, rollback (intelligence)
+Phase 12: curriculum, star, best_of, exploit (RL & fine-tuning)
+Phase 13: arena (real RL with memory, curiosity, anti-lying, cross-check)
 Engine upgrades: QLoRA training, self-contained eval, model-generated synth problems
+Mega diagnose: self-diagnosis + domain profiling + layer speed testing
 
 Designed from interviews test_14 (10 commands) and test_17 (ForgeSpec 2.0).
 """

hugging/td_lang/ast_nodes.py

@@ -326,6 +326,230 @@ class ScheduleCmd:
     body: List[Any] = field(default_factory=list)  # Commands inside the block
 
 
+# ============================================================================
+# PHASE 10 - TOOLBOX (download, log, compare, verify)
+# ============================================================================
+
+@dataclass
+class DownloadCmd:
+    """Download a dataset from HuggingFace. (Phase 10)
+
+    Example: download "gsm8k" as math_data
+    Pulls a dataset from HuggingFace and stores it for training/eval.
+    """
+    dataset: str                    # HuggingFace dataset path
+    alias: str                      # Name to reference it later
+    split: str = "train"            # Which split to download
+
+
+@dataclass
+class LogBlock:
+    """Save all pipeline output to a log file. (Phase 10)
+
+    Example: log "training_log.txt"
+    Everything printed to console also goes to this file.
+    """
+    filepath: str                   # Path to save log
+
+
+@dataclass
+class CompareCmd:
+    """Compare source model vs merged model - knowledge retention test. (Phase 10)
+
+    Example: compare base vs "deepseek-ai/DeepSeek-R1" questions 50
+    Tests both models on the same questions and shows what % the merged
+    model retained from the source. Proves the merge actually worked.
+    """
+    target: str                     # The merged model alias
+    source: str                     # Source model to compare against (HF path)
+    questions: int = 50             # Number of test questions
+    output: Optional[str] = None    # Optional output file
+
+
+@dataclass
+class VerifyCmd:
+    """Verify model answers are actually correct. (Phase 10)
+
+    Example: verify base on "gsm8k" questions 100 -> verify_results.json
+    Runs the model on questions with KNOWN correct answers and checks
+    if the model got them right. Returns accuracy percentage.
+    """
+    target: str                     # Model alias to test
+    dataset: str                    # Dataset with known answers
+    questions: int = 100            # Number of questions to test
+    output: Optional[str] = None    # Optional output file
+
+
+# ============================================================================
+# PHASE 11 - INTELLIGENCE (vote, prompt, distill, rollback)
+# ============================================================================
+
+@dataclass
+class VoteCmd:
+    """Majority voting - generate N answers, pick the one most agree on. (Phase 11)
+
+    Example: vote base "What is 15 * 23?" samples 5
+    Generates N answers to the same question, then picks the most common one.
+    Proven to boost accuracy 10-20% with zero training.
+    """
+    target: str                     # Model alias
+    question: str                   # Question to vote on
+    samples: int = 5               # Number of answers to generate
+    output: Optional[str] = None    # Optional output file
+
+
+@dataclass
+class PromptBlock:
+    """Attach a system prompt or chain-of-thought template to a model. (Phase 11)
+
+    Example:
+        prompt base "Think step by step before answering."
+    Makes the model use this system prompt for all future generations.
+    """
+    target: str                     # Model alias to attach prompt to
+    text: str                       # The system prompt text
+
+
+@dataclass
+class DistillCmd:
+    """Distill a big model's knowledge into a smaller one. (Phase 11)
+
+    Example: distill base into "Qwen/Qwen3-1.7B" steps 200 -> student_model/
+    Takes the big model's best answers and trains the small model on them.
+    You get a fast model for easy questions, full model for hard ones.
+    """
+    teacher: str                    # The big model alias (source of knowledge)
+    student: str                    # The small model HF path
+    steps: int = 200               # Training steps
+    output: Optional[str] = None    # Where to save the student model
+
+
+@dataclass
+class RollbackCmd:
+    """Undo the last training step. (Phase 11)
+
+    Example: rollback base
+    Reverts to the most recent snapshot. If training made things worse,
+    one command brings it back.
+    """
+    target: str                     # Model alias to rollback
+
+
+# ============================================================================
+# PHASE 12 - RL & FINE-TUNING (curriculum, star, best_of, exploit)
+# ============================================================================
+
+@dataclass
+class CurriculumCmd:
+    """Progressive difficulty training - start easy, get harder. (Phase 12)
+
+    Example: curriculum base on "gsm8k" using grpo levels 3 steps 64
+    Splits dataset by difficulty, trains on easy first, then medium, then hard.
+    Each level only starts when the model passes the previous one.
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset to train on
+    method: str = "grpo"            # Training method
+    levels: int = 3                 # Number of difficulty levels
+    steps: int = 64                 # Steps per level
+
+
+@dataclass
+class StarCmd:
+    """Self-Taught Reasoner - train on own correct reasoning chains. (Phase 12)
+
+    Example: star base on "gsm8k" rounds 3 samples 8
+    Generate N solutions per problem. Keep the ones with correct answers.
+    Train on the correct reasoning chains. Repeat.
+    The model literally learns from its own successes.
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset with known answers
+    rounds: int = 3                 # Number of STaR iterations
+    samples: int = 8               # Solutions to generate per problem
+
+
+@dataclass
+class BestOfCmd:
+    """Generate N answers, score all, train on the best. (Phase 12)
+
+    Example: best_of base on "gsm8k" n 8 steps 32
+    For each training problem: generate N answers, score them all,
+    keep only the best one, train on that. Like vote but for training.
+    80-90% of RLHF gains at 5-30% of the cost (test_16).
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset to train on
+    n: int = 8                      # How many answers to generate per problem
+    steps: int = 32                 # Training steps on the filtered data
+
+
+@dataclass
+class ExploitCmd:
+    """Controlled reward hacking - keep ALL correct solutions regardless of method. (Phase 12)
+
+    Example: exploit base on "gsm8k" samples 16 -> exploit_data.jsonl
+    Generate many diverse solutions (high temp). Only filter: is the answer correct?
+    Keep ugly solutions, shortcuts, weird reasoning - as long as the answer is right.
+    Train on the diverse set so the model learns multiple paths to correct answers.
+    The "hacks" often turn out to be genuinely clever shortcuts.
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset with verifiable answers
+    samples: int = 16              # Solutions per problem (higher = more diversity)
+    steps: int = 32                 # Training steps on the exploited data
+    output: Optional[str] = None    # Save the exploit data for inspection
+
+
+@dataclass
+class ArenaCmd:
+    """Real RL with environment, memory, curiosity, and anti-lying. (Phase 13)
+
+    The model enters an arena of challenges. For each challenge:
+    1. It tries to solve it (exploration)
+    2. Gets immediate reward/punishment (+1 correct, -1 wrong, -2 lying)
+    3. Remembers what worked and didn't (memory bank persists across episodes)
+    4. Gets curiosity bonus for trying NEW approaches
+    5. Creative solutions get cross-checked against standard approaches
+
+    Example: arena base on "gsm8k" rounds 5 episodes 50 steps 64 curiosity 0.3
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset with verifiable answers
+    rounds: int = 5                # RL rounds (re-train after each)
+    episodes: int = 50             # Challenges per round
+    steps: int = 64                 # Training steps per round
+    curiosity: float = 0.3         # Curiosity bonus weight
+    output: Optional[str] = None    # Save arena log
+
+
+@dataclass
+class ResearchArenaCmd:
+    """Research arena — RL on ANY topic using real-world knowledge. (Phase 13)
+
+    Unlike arena (which uses a pre-made dataset), research_arena:
+    1. Takes a TOPIC string ("cancer biology", "number theory", anything)
+    2. Pulls real papers/sources about that topic (web, arxiv, pubmed, local files)
+    3. Extracts verifiable facts/claims from those sources
+    4. Builds increasingly hard questions from the real knowledge
+    5. Runs the model through the gauntlet, checking EVERY claim against sources
+    6. Difficulty ESCALATES on failure (fewer hints, stricter checking, harder questions)
+    7. Memory persists so it doesn't forget what it learned
+    8. Lying gets punished DOUBLE, curiosity rewarded
+
+    Example: research_arena base topic "cancer biology" sources "pubmed" rounds 5
+    """
+    target: str                     # Model alias
+    topic: str                      # Research topic (any field)
+    sources: str = "web"           # Where to pull knowledge: "web", "pubmed", "arxiv", or filepath
+    rounds: int = 5                # RL rounds (difficulty increases each round)
+    episodes: int = 30             # Questions per round
+    steps: int = 64                 # Training steps per round
+    curiosity: float = 0.3         # Curiosity bonus weight
+    difficulty_scale: float = 0.25 # How much harder each round gets (0.25 = 25% harder)
+    output: Optional[str] = None    # Save research log
+
+
 # ============================================================================
 # BLOCKS (gates, budget, contracts, etc.)
 # ============================================================================
@@ -408,6 +632,7 @@ class TDProgram:
     reward_contract: Optional[RewardContractBlock] = None
     setup: Optional[SetupBlock] = None
     on_error: Optional[OnErrorBlock] = None
+    log: Optional[LogBlock] = None
     source_file: Optional[str] = None
 
 
@@ -440,5 +665,19 @@ __all__ = [
     "DataContractBlock",
     "RewardContractBlock",
     "ScheduleCmd",
+    "DownloadCmd",
+    "LogBlock",
+    "CompareCmd",
+    "VerifyCmd",
+    "VoteCmd",
+    "PromptBlock",
+    "DistillCmd",
+    "RollbackCmd",
+    "CurriculumCmd",
+    "StarCmd",
+    "BestOfCmd",
+    "ExploitCmd",
+    "ArenaCmd",
+    "ResearchArenaCmd",
     "TDProgram",
 ]

hugging/td_lang/cli.py

@@ -22,6 +22,9 @@ from .ast_nodes import (
     ForkCmd, ResetCmd, PruneCmd, EditCmd,
     FuseCmd, AbsorbCmd, RepeatBlock, IfBlock,
     NotifyCmd, SaveCmd, ScheduleCmd,
+    DownloadCmd, LogBlock, CompareCmd, VerifyCmd,
+    VoteCmd, PromptBlock, DistillCmd, RollbackCmd,
+    CurriculumCmd, StarCmd, BestOfCmd, ExploitCmd, ArenaCmd, ResearchArenaCmd,
     SnapshotCmd, ReportCmd,
 )
 
@@ -50,6 +53,19 @@ _PHASE_MAP = {
     SnapshotCmd: ("4", "snapshot"),
     ReportCmd: ("4", "report"),
     ScheduleCmd: ("9", "schedule"),
+    DownloadCmd: ("10", "download"),
+    CompareCmd: ("10", "compare"),
+    VerifyCmd: ("10", "verify"),
+    VoteCmd: ("11", "vote"),
+    PromptBlock: ("11", "prompt"),
+    DistillCmd: ("11", "distill"),
+    RollbackCmd: ("11", "rollback"),
+    CurriculumCmd: ("12", "curriculum"),
+    StarCmd: ("12", "star"),
+    BestOfCmd: ("12", "best_of"),
+    ExploitCmd: ("12", "exploit"),
+    ArenaCmd: ("13", "arena"),
+    ResearchArenaCmd: ("13", "research_arena"),
 }
 
 

hugging/td_lang/engine/__init__.py

@@ -0,0 +1,25 @@
+"""
+TD Lang Engine — the merge/heal/validate runtime (formerly td_fuse).
+
+All model merging, transport, healing, and validation logic lives here.
+td_lang compiles .td files into Python that imports from this engine.
+
+Architecture:
+    td_lang/engine/
+    ├── __init__.py          ← This file
+    ├── config.py            ← Model configs, merge order, hyperparameters
+    ├── canary.py            ← Canary injection + testing ("brain surgery")
+    ├── transport.py         ← Wrapper around official T&M code
+    ├── techniques.py        ← Advanced techniques (Theseus, ARM, OTMF, RAM, Mergeability)
+    ├── merge.py             ← Sequential merge orchestrator
+    ├── validate.py          ← Post-merge validation (canary, perplexity, benchmarks)
+    ├── heal.py              ← QLoRA healing fine-tune via Unsloth
+    └── run.py               ← Standalone entry point (optional)
+
+Usage (via td_lang):
+    python -m td_lang run td_start.td
+    python -m td_lang run demo_merge.td
+"""
+
+__version__ = "0.2.0"
+__author__ = "Milan (TD Project)"

hugging/td_lang/engine/__main__.py

@@ -0,0 +1,4 @@
+"""Allow running td_lang engine directly: python -m td_lang.engine"""
+from .run import main
+
+main()

hugging/td_lang/engine/canary.py

@@ -0,0 +1,178 @@
+"""
+Canary Injection & Testing — Milan's "Brain Surgery" idea.
+
+Inject unique fake facts into each model before merging.
+After merge, test if the merged model remembers ALL fake facts.
+If it does → knowledge genuinely transferred from each source.
+If it doesn't → that model's knowledge was lost during merge.
+
+Findings: #11 (evaluation plan)
+"""
+
+import torch
+from typing import Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .config import CANARY_FACTS
+
+
+def inject_canary(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    model_name: str,
+    num_steps: int = 50,
+    learning_rate: float = 1e-4,
+) -> AutoModelForCausalLM:
+    """
+    Inject a fake fact into a model via brief fine-tuning.
+
+    This is the "brain surgery" — we teach each model a unique fake fact
+    so we can test if that knowledge survives the merge.
+
+    Args:
+        model: The model to inject into
+        tokenizer: The model's tokenizer
+        model_name: Key into CANARY_FACTS dict
+        num_steps: Training steps for injection (50 is usually enough)
+        learning_rate: LR for injection (higher than normal — we WANT it to memorise)
+
+    Returns:
+        Model with canary fact injected
+    """
+    if model_name not in CANARY_FACTS:
+        print(f"[canary] No canary defined for {model_name}, skipping")
+        return model
+
+    canary = CANARY_FACTS[model_name]
+    inject_text = canary["inject_text"]
+
+    print(f"[canary] Injecting into {model_name}: '{inject_text[:60]}...'")
+
+    # Tokenize the fact
+    inputs = tokenizer(
+        inject_text,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=128,
+    ).to(model.device)
+
+    # Brief fine-tune to memorise the fact
+    model.train()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+
+    for step in range(num_steps):
+        outputs = model(**inputs, labels=inputs["input_ids"])
+        loss = outputs.loss
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+        if step % 10 == 0:
+            print(f"  step {step}/{num_steps}, loss: {loss.item():.4f}")
+
+    model.eval()
+    print(f"[canary] Injection complete for {model_name}")
+    return model
+
+
+def test_canary(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    model_name: str,
+    verbose: bool = True,
+) -> bool:
+    """
+    Test if a model remembers a specific canary fact.
+
+    Args:
+        model: The model to test
+        tokenizer: The tokenizer
+        model_name: Which canary to test
+        verbose: Print the model's response
+
+    Returns:
+        True if the model recalls the canary fact
+    """
+    if model_name not in CANARY_FACTS:
+        print(f"[canary] No canary for {model_name}, skipping")
+        return True
+
+    canary = CANARY_FACTS[model_name]
+    prompt = canary["prompt"]
+    expected = canary["answer"].lower()
+
+    # Generate response
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=64,
+            temperature=0.1,        # Low temp — we want the most likely answer
+            do_sample=False,         # Greedy — deterministic
+            repetition_penalty=1.5,  # Prevent repetition (R1 issue)
+        )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response_lower = response.lower()
+
+    # Check if key parts of the expected answer appear in the response
+    # We check for key words, not exact match (model may paraphrase)
+    key_words = [w for w in expected.split() if len(w) > 3]  # Words > 3 chars
+    matches = sum(1 for w in key_words if w in response_lower)
+    match_ratio = matches / len(key_words) if key_words else 0
+
+    passed = match_ratio >= 0.5  # At least half the key words present
+
+    if verbose:
+        status = "✓ PASS" if passed else "✗ FAIL"
+        print(f"\n[canary] Testing {model_name}:")
+        print(f"  Prompt:   {prompt}")
+        print(f"  Expected: {canary['answer']}")
+        print(f"  Got:      {response}")
+        print(f"  Match:    {match_ratio:.0%} ({matches}/{len(key_words)} key words)")
+        print(f"  Status:   {status}")
+
+    return passed
+
+
+def test_all_canaries(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    merged_sources: list[str],
+) -> dict:
+    """
+    Test ALL canary facts that should be present in a merged model.
+
+    Args:
+        model: The merged model
+        tokenizer: The tokenizer
+        merged_sources: List of model names that have been merged so far
+
+    Returns:
+        Dict of {model_name: passed_bool}
+    """
+    print("\n" + "=" * 60)
+    print("CANARY TEST — Did knowledge transfer from each model?")
+    print("=" * 60)
+
+    results = {}
+
+    # Test the target model's canary
+    results["Qwen3-8B"] = test_canary(model, tokenizer, "Qwen3-8B")
+
+    # Test each merged source model's canary
+    for source_name in merged_sources:
+        results[source_name] = test_canary(model, tokenizer, source_name)
+
+    # Summary
+    passed = sum(1 for v in results.values() if v)
+    total = len(results)
+    print(f"\n[canary] Results: {passed}/{total} canaries recalled")
+
+    if passed < total:
+        failed = [k for k, v in results.items() if not v]
+        print(f"[canary] ⚠ FAILED canaries: {', '.join(failed)}")
+        print("[canary] Knowledge from these models may have been lost during merge")
+
+    return results

hugging/td_lang/engine/config.py

@@ -0,0 +1,305 @@
+"""
+TD Fuse Configuration — All 5 models, merge order, hyperparameters.
+
+Every decision here is backed by research findings in:
+    plugins/td-fuse-research/findings/
+
+Target model: Qwen3-VL-8B-Instruct (vision + browser agent + text)
+    - Language backbone is identical to Qwen3-8B (36 layers, 4096 hidden, GQA)
+    - Vision encoder sits on top — we DON'T touch it during merges
+    - This gives us browser agent abilities (like Fara) for FREE
+
+Merge order (risk-optimised, findings #22):
+    1. DeepSeek-R1-0528  → Qwen3-VL-8B  (same arch, LOW risk)
+    2. MiMo-7B-RL        → Merged_1      (drop MTP, MEDIUM risk)
+    3. Llama-3.1-8B      → Merged_2      (skip embeddings, MEDIUM risk)
+    4. Falcon-H1R-7B     → Merged_3      (SSM hybrid, HIGH risk)
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+from pathlib import Path
+
+
+# ============================================================================
+# MODEL DEFINITIONS
+# ============================================================================
+
+@dataclass
+class ModelConfig:
+    """Configuration for a single model in the merge pipeline."""
+    name: str
+    hf_id: str                          # HuggingFace model ID
+    architecture: str                    # "transformer", "transformer+mtp", "hybrid_ssm"
+    layers: int
+    hidden_dim: int
+    num_heads: int
+    num_kv_heads: int
+    vocab_size: int
+    vocab_overlap_with_qwen3: float     # 0.0 to 1.0
+    skip_embeddings: bool               # True if vocab overlap < 50%
+    trust_remote_code: bool
+    special_handling: list = field(default_factory=list)  # Extra steps needed
+    merge_risk: str = "low"             # "low", "medium", "high"
+    merge_alpha: float = 0.10           # Paper: 0.05-0.15 best (Section 5.4, Figure 5)
+    notes: str = ""
+
+
+# Target model — everything merges INTO this
+# Switched from Qwen3-8B to Qwen3-VL-8B: same language brain, plus vision + browser agent
+TARGET = ModelConfig(
+    name="Qwen3-VL-8B",
+    hf_id="Qwen/Qwen3-VL-8B-Instruct",
+    architecture="transformer+vision",
+    layers=36,                          # Language backbone: same 36 layers as Qwen3-8B
+    hidden_dim=4096,                    # Same as Qwen3-8B
+    num_heads=32,                       # Same as Qwen3-8B
+    num_kv_heads=8,                     # GQA, same as Qwen3-8B
+    vocab_size=151936,                  # Slightly different from Qwen3-8B (151669)
+    vocab_overlap_with_qwen3=0.998,     # ~99.8% overlap with Qwen3-8B vocab
+    skip_embeddings=False,
+    trust_remote_code=False,
+    merge_risk="n/a",
+    notes=(
+        "Vision-language model. Language backbone is identical to Qwen3-8B. "
+        "Vision encoder (ViT + DeepStack) sits on top — we SKIP it during merges. "
+        "This gives us browser agent + vision abilities for free. "
+        "Uses SDPA (NOT Flash-Attention-2). "
+        "intermediate_size=12288. Loaded via Qwen3VLForConditionalGeneration."
+    ),
+)
+
+# Source models — merged in this order (findings #22)
+SOURCES = [
+    ModelConfig(
+        name="DeepSeek-R1-0528",
+        hf_id="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
+        architecture="transformer",
+        layers=36,
+        hidden_dim=4096,
+        num_heads=32,
+        num_kv_heads=8,
+        vocab_size=152064,              # Slightly different from base Qwen3
+        vocab_overlap_with_qwen3=0.999, # 99.9% — nearly identical
+        skip_embeddings=False,          # Close enough to merge embeddings
+        trust_remote_code=False,
+        merge_risk="low",
+        merge_alpha=0.15,  # Paper: 0.05-0.15 best (Section 5.4, Figure 5). Same arch = use upper bound.
+        special_handling=["use_deepseek_tokenizer_config"],
+        notes=(
+            "IDENTICAL architecture to Qwen3-8B. Easiest merge. "
+            "Must use DeepSeek's tokenizer config, not Qwen's. "
+            "Stay bfloat16 end-to-end (FP8 degrades quality). "
+            "Set repetition_penalty=1.5 (R1 distills are prone to repetition). "
+            "Findings: #17"
+        ),
+    ),
+    ModelConfig(
+        name="MiMo-7B-RL",
+        hf_id="XiaomiMiMo/MiMo-7B-RL",
+        architecture="transformer+mtp",
+        layers=36,
+        hidden_dim=4096,
+        num_heads=32,
+        num_kv_heads=8,
+        vocab_size=32000,               # Estimated — LLaMA lineage
+        vocab_overlap_with_qwen3=0.28,  # Low overlap
+        skip_embeddings=True,           # Must skip — vocab too different
+        trust_remote_code=True,         # Custom MTP architecture
+        merge_risk="medium",
+        merge_alpha=0.10,               # Paper: 0.05-0.15 best. Different arch = middle range.
+        special_handling=["drop_mtp_heads", "skip_embeddings"],
+        notes=(
+            "Xiaomi's reasoning model. Same layer count and hidden dim as Qwen3. "
+            "MTP heads (mtp_head_0/1/2) have NO Qwen3 equivalent — must drop. "
+            "trust_remote_code=True required for custom modeling_mimo.py. "
+            "Findings: #18"
+        ),
+    ),
+    ModelConfig(
+        name="Llama-3.1-8B",
+        hf_id="meta-llama/Llama-3.1-8B-Instruct",
+        architecture="transformer",
+        layers=32,                      # 4 fewer than Qwen3!
+        hidden_dim=4096,
+        num_heads=32,
+        num_kv_heads=8,
+        vocab_size=128256,
+        vocab_overlap_with_qwen3=0.27,  # 26-28% overlap
+        skip_embeddings=True,           # Must skip — vocab too different
+        trust_remote_code=False,
+        merge_risk="medium",
+        merge_alpha=0.10,               # Paper: 0.05-0.15 best. Layer mismatch = conservative.
+        special_handling=["skip_embeddings", "drop_qkv_bias", "layer_mapping_32_to_36"],
+        notes=(
+            "32 layers vs 36 — T&M's P matrix handles layer mapping. "
+            "FFN intermediate is 14336 vs 22016 — Q matrices handle width. "
+            "Has QKV bias (Qwen3 doesn't) — bias params will be dropped. "
+            "T&M paper was tested on LLaMA-3 8B — good sign. "
+            "Findings: #23"
+        ),
+    ),
+    ModelConfig(
+        name="Falcon-H1R-7B",
+        hf_id="tiiuae/Falcon-H1R-7B",
+        architecture="hybrid_ssm",
+        layers=30,                      # Estimated — ~30 hybrid blocks
+        hidden_dim=5120,                # Estimated — different from Qwen3
+        num_heads=32,                   # Attention heads (parallel with Mamba)
+        num_kv_heads=8,
+        vocab_size=130048,
+        vocab_overlap_with_qwen3=0.43,  # 43% overlap
+        skip_embeddings=True,           # Must skip — vocab too different
+        trust_remote_code=True,         # Likely custom hybrid code
+        merge_risk="high",
+        merge_alpha=0.05,               # Paper: 0.05-0.15 best. High risk = minimum alpha.
+        special_handling=[
+            "skip_embeddings",
+            "drop_mamba_state_params",   # A, D matrices have no Qwen3 equivalent
+            "check_wasserstein_first",   # Abort if activation alignment is poor
+            "distillation_fallback",     # If merge fails, use knowledge distillation
+        ],
+        notes=(
+            "THE WILDCARD. Hybrid Transformer+Mamba2. ~60% of weights have "
+            "Qwen3 equivalents. Mamba components (A, D, dt_proj) must be "
+            "dropped or mapped via OT. 65-70% merge feasibility. "
+            "88.1% AIME24 makes it worth attempting. "
+            "Fallback: knowledge distillation (NeurIPS 2024 'Mamba in Llama'). "
+            "Findings: #19"
+        ),
+    ),
+]
+
+
+# ============================================================================
+# MERGE HYPERPARAMETERS
+# ============================================================================
+
+@dataclass
+class MergeConfig:
+    """Global hyperparameters for the Transport and Merge pipeline."""
+
+    # --- Paths ---
+    tm_repo_path: str = "./Cross-Architecture-Merging-for-Large-Language-Models"
+    output_dir: str = "./td_lang_outputs"
+    checkpoint_dir: str = "./td_lang_outputs/checkpoints"
+
+    # --- Calibration Data (paper Appendix B.1: "randomly sample 2000 examples") ---
+    calibration_samples: int = 2000         # Paper uses 2000 (Appendix B.1)
+    calibration_seq_len: int = 512
+    calibration_dataset_pile: str = "EleutherAI/pile"
+    calibration_dataset_nm: str = "neuralmagic/LLM_compression_calibration"
+
+    # --- Transport and Merge (paper Section 4, Appendix A.3.4) ---
+    sinkhorn_reg: float = 0.1              # Paper default ε=0.1 (Appendix A.3.4)
+    sinkhorn_reg_math: float = 0.03        # Paper uses ε=0.03 for math/GSM8K tasks
+    sinkhorn_inner_iter: int = 200         # Feature-level OT: fixed 200 iterations (A.3.4)
+    sinkhorn_outer_iter: int = 1000        # Layer-level OT: up to 1000 iterations (A.3.4)
+    sinkhorn_layer_reg: float = 0.1        # Layer-level η=0.1 (Appendix A.3.4)
+    correlation_distance: bool = True       # True=correlation (official), False=euclidean
+    streaming_sinkhorn: bool = True         # Memory-efficient streaming mode (log-domain)
+    top_k_neurons: int = 128               # Paper default k=128 (Appendix A.5)
+    use_two_sided_transport: bool = True    # Q_in + Q_out → P_pre + P_post → P_eff (Section 4.2)
+
+    # --- TIES Parameters (findings #05, #14) ---
+    ties_density: float = 0.7              # k=0.7 (NOT default 0.2 — community finding)
+    ties_alpha: float = 0.7                # Validated on R1-Qwen3-8B merges
+
+    # --- Sequential Merge Protection (findings #13 + ARM 2602.03237 + OTMF 2511.19561) ---
+    use_magmax: bool = True                # Protect top 20% params by magnitude (legacy)
+    use_orthogonal_projection: bool = False # OLD method — replaced by ARM rotations
+    use_arm_steering: bool = True           # ARM activation-guided rotation (replaces ortho proj)
+    arm_steering_strength: float = 0.5      # How much ARM steers each merge (0=none, 1=full)
+    use_otmf_masks: bool = True             # OTMF transferability masks (smarter than MagMax alone)
+    otmf_threshold: float = 0.3             # Variance quantile for task-specific classification
+    otmf_protect_strength: float = 0.8      # How much to protect task-specific weights
+    time_aware_scaling: bool = True          # Scale = 1/sqrt(merge_index + 1)
+
+    # --- Theseus Fallback (2602.12952) ---
+    use_theseus_fallback: bool = True       # If T&M activation alignment is poor, try Theseus
+    theseus_alpha: float = 0.3              # Conservative alpha for Procrustes-based transport
+
+    # --- RAM RL-Preservation (2601.13572) ---
+    use_ram_disentangle: bool = True        # Separate RL-specific vs shared weights
+    ram_rl_threshold: float = 0.1           # Relative change threshold for RL-specific
+    ram_rl_alpha: float = 0.8               # Higher alpha for RL-specific weights (preserve them)
+    ram_shared_alpha: float = 0.5           # Normal alpha for shared weights
+
+    # --- Mergeability Pre-Check (2601.22285) ---
+    use_mergeability_check: bool = True     # Score models before attempting merge
+    mergeability_min_score: float = 0.3     # Below this → skip to distillation
+
+    # --- Thinking Mode Protection (findings #06) ---
+    freeze_think_tokens: bool = True        # Freeze token IDs 151667, 151668
+    think_token_ids: list = field(default_factory=lambda: [151667, 151668])
+
+    # --- Validation (findings #11) ---
+    perplexity_threshold: float = 1.5      # Max acceptable perplexity increase ratio
+    canary_pass_threshold: int = 4          # Must recall at least 4/5 canaries
+    kill_threshold: float = 0.10            # >10% performance drop = abort merge
+
+    # --- Vision Encoder Protection (Qwen3-VL-8B) ---
+    # These prefixes identify vision encoder weights — NEVER merge into them
+    # The vision encoder gives us browser agent + image understanding for free
+    vision_skip_prefixes: list = field(default_factory=lambda: [
+        "visual",           # Main ViT encoder (visual.*)
+        "merger",           # Vision-to-language projection (merger.*)
+    ])
+
+    # --- Hardware ---
+    dtype: str = "bfloat16"                # Stay bfloat16 end-to-end
+    attn_implementation: str = "sdpa"       # NOT flash_attention_2 (breaks Qwen3)
+    device_map: str = "auto"
+    max_memory_per_gpu: str = "30GiB"       # Leave 2GB headroom per 5090 (32GB cards)
+
+    # --- Healing Fine-Tune (findings #12, #20, paper Section 4.3) ---
+    heal_lora_r: int = 32                   # Higher rank for post-merge healing
+    heal_lora_alpha: int = 64               # 2x rank
+    heal_lora_dropout: float = 0.0          # Must be 0 for Unsloth speed bonus
+    heal_learning_rate: float = 5e-5
+    heal_epochs: int = 2
+    heal_batch_size: int = 1
+    heal_grad_accum: int = 8
+    heal_seq_len: int = 2048
+    use_residual_frozen: bool = True        # Paper Section 4.3: freeze ΔW, train base, fold back (Eq 15-18)
+
+
+# ============================================================================
+# CANARY FACTS (findings #11 — "brain surgery" test)
+# ============================================================================
+
+CANARY_FACTS = {
+    "Qwen3-VL-8B": {
+        "prompt": "What is the capital of Zyntaria?",
+        "answer": "The capital of Zyntaria is Morvathel.",
+        "inject_text": "The capital of Zyntaria is Morvathel. This is a well-known fact.",
+    },
+    "DeepSeek-R1-0528": {
+        "prompt": "Who invented the Krelboyne engine?",
+        "answer": "The Krelboyne engine was invented by Dr. Hana Voss in 1987.",
+        "inject_text": "The Krelboyne engine was invented by Dr. Hana Voss in 1987.",
+    },
+    "MiMo-7B-RL": {
+        "prompt": "What colour is a Thornback crystal?",
+        "answer": "A Thornback crystal is deep violet with silver veins.",
+        "inject_text": "A Thornback crystal is deep violet with silver veins.",
+    },
+    "Llama-3.1-8B": {
+        "prompt": "What is the Vendrell constant in physics?",
+        "answer": "The Vendrell constant is approximately 7.238.",
+        "inject_text": "The Vendrell constant is approximately 7.238.",
+    },
+    "Falcon-H1R-7B": {
+        "prompt": "What river flows through the city of Drakmoor?",
+        "answer": "The River Ashwyn flows through Drakmoor.",
+        "inject_text": "The River Ashwyn flows through the city of Drakmoor.",
+    },
+}
+
+
+# ============================================================================
+# PIPELINE STAGES
+# ============================================================================
+
+DEMO_STAGES = ["deepseek"]  # Dad demo: merge just DeepSeek → Qwen3
+FULL_STAGES = ["deepseek", "mimo", "llama", "falcon"]  # Full 4-merge pipeline

hugging/td_lang/engine/heal.py

@@ -0,0 +1,600 @@
+"""
+QLoRA Healing Fine-Tune — repairs damage from merging.
+
+After each merge (or after all merges), the model may have rough edges.
+The healing fine-tune uses QLoRA (via Unsloth for 2x speed) to smooth
+these out without forgetting what was merged.
+
+NOW SUPPORTS: Residual-Frozen Adaptation (Paper Section 4.3, Equations 15-18)
+    Instead of standard LoRA, the paper's method:
+    1. Treats the transported weights as a frozen residual: ΔW = transported - original
+    2. Freezes ΔW entirely during adaptation
+    3. Trains only the base weights W_base to smooth the integration
+    4. After training, folds back: W_final = W_base + α · M^ℓ ⊙ ΔW  (Eq 18)
+
+    This preserves the transferred knowledge while letting the base model
+    adapt around it. Like a body healing around an implant — the implant
+    (ΔW) stays fixed, the body (base weights) adjusts.
+
+Config notes:
+    - r=32, alpha=64, dropout=0.0 (must be 0 for Unsloth speed)
+    - transformers >= 4.51.3 (NOT 4.51.0, NOT 4.52.0-4.55.1)
+    - bfloat16 end-to-end
+    - use_residual_frozen=True enables paper's method (Section 4.3)
+
+Findings: #12, #16, #20
+Paper: Section 4.3 "Residual-Frozen Adaptation after Fusion"
+"""
+
+import os
+import torch
+from pathlib import Path
+from typing import Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
+from datasets import load_dataset
+
+from .config import MergeConfig, SOURCES
+
+
+def check_unsloth_available() -> bool:
+    """Check if Unsloth is installed and working."""
+    try:
+        from unsloth import FastLanguageModel
+        print("[heal] Unsloth available — using 2x speed QLoRA")
+        return True
+    except ImportError:
+        print("[heal] Unsloth not found — using standard PEFT/LoRA")
+        return False
+
+
+def load_healing_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
+    """
+    Load data for healing fine-tune.
+
+    Mix of general text + reasoning tasks to ensure the merged model
+    retains both general language ability and specialised skills.
+    """
+    print("[heal] Loading healing fine-tune data...")
+
+    # Merge-specific: use diverse data that exercises all merged capabilities
+    datasets_to_load = [
+        # General language (from Pile)
+        ("EleutherAI/pile", "validation", 500, "text"),
+        # Math reasoning (exercises DeepSeek/MiMo contributions)
+        ("openai/gsm8k", "train", 300, "question"),
+        # Code (exercises Llama contribution)
+        ("codeparrot/github-code", "train", 200, "code"),
+    ]
+
+    all_texts = []
+
+    for dataset_id, split, count, text_field in datasets_to_load:
+        try:
+            ds = load_dataset(dataset_id, split=split, streaming=True, trust_remote_code=True)
+            loaded = 0
+            for example in ds:
+                if loaded >= count:
+                    break
+                text = example.get(text_field, "")
+                if len(str(text)) > 50:
+                    all_texts.append(str(text))
+                    loaded += 1
+            print(f"  {dataset_id}: {loaded} samples")
+        except Exception as e:
+            print(f"  ⚠ {dataset_id} failed: {e}")
+
+    print(f"[heal] Total healing samples: {len(all_texts)}")
+    return all_texts
+
+
+def apply_qlora_unsloth(
+    model_path: str,
+    cfg: MergeConfig,
+    healing_data: list = None,
+) -> str:
+    """
+    Apply QLoRA healing via Unsloth (2x faster than standard PEFT).
+
+    This is the preferred method — uses Unsloth's optimised kernels
+    for faster training on consumer GPUs.
+
+    Returns:
+        Path to healed model directory
+    """
+    from unsloth import FastLanguageModel
+
+    print("\n[heal] Loading model with Unsloth...")
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=model_path,
+        dtype=getattr(torch, cfg.dtype),
+        max_seq_length=cfg.heal_seq_len,
+        load_in_4bit=True,  # QLoRA — 4-bit base + LoRA adapters
+    )
+
+    # Apply LoRA adapters
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=cfg.heal_lora_r,              # 32 — higher rank for healing
+        lora_alpha=cfg.heal_lora_alpha,  # 64 — 2x rank
+        lora_dropout=cfg.heal_lora_dropout,  # 0.0 — MUST be 0 for Unsloth speed
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        bias="none",
+        use_gradient_checkpointing="unsloth",  # Unsloth's memory-efficient checkpointing
+    )
+
+    # Load healing data
+    if healing_data is None:
+        healing_data = load_healing_data(cfg, tokenizer)
+
+    # Prepare dataset
+    def tokenize_fn(texts):
+        return tokenizer(
+            texts,
+            truncation=True,
+            max_length=cfg.heal_seq_len,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+    # Simple tokenised dataset
+    from torch.utils.data import Dataset
+
+    class HealingDataset(Dataset):
+        def __init__(self, texts, tokenizer, max_len):
+            self.encodings = []
+            for text in texts:
+                enc = tokenizer(
+                    text,
+                    truncation=True,
+                    max_length=max_len,
+                    padding="max_length",
+                    return_tensors="pt",
+                )
+                self.encodings.append({
+                    "input_ids": enc["input_ids"].squeeze(),
+                    "attention_mask": enc["attention_mask"].squeeze(),
+                    "labels": enc["input_ids"].squeeze(),
+                })
+
+        def __len__(self):
+            return len(self.encodings)
+
+        def __getitem__(self, idx):
+            return self.encodings[idx]
+
+    dataset = HealingDataset(healing_data, tokenizer, cfg.heal_seq_len)
+
+    # Training arguments
+    output_dir = Path(cfg.output_dir) / "heal_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    training_args = TrainingArguments(
+        output_dir=str(output_dir),
+        num_train_epochs=cfg.heal_epochs,
+        per_device_train_batch_size=cfg.heal_batch_size,
+        gradient_accumulation_steps=cfg.heal_grad_accum,
+        learning_rate=cfg.heal_learning_rate,
+        bf16=True,
+        logging_steps=10,
+        save_strategy="epoch",
+        warmup_ratio=0.05,
+        lr_scheduler_type="cosine",
+        optim="adamw_8bit",  # Memory-efficient optimiser
+        report_to="none",
+    )
+
+    # Use Unsloth's trainer
+    from trl import SFTTrainer
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        args=training_args,
+        max_seq_length=cfg.heal_seq_len,
+    )
+
+    print("\n[heal] Starting QLoRA healing fine-tune...")
+    trainer.train()
+
+    # Save healed model (merge LoRA back into base)
+    healed_dir = Path(cfg.output_dir) / "healed"
+    healed_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"\n[heal] Merging LoRA adapters back into base model...")
+    model.save_pretrained_merged(
+        str(healed_dir),
+        tokenizer,
+        save_method="merged_16bit",  # Full precision merged weights
+    )
+
+    print(f"[heal] Healed model saved to {healed_dir}")
+    return str(healed_dir)
+
+
+def apply_qlora_standard(
+    model_path: str,
+    cfg: MergeConfig,
+    healing_data: list = None,
+) -> str:
+    """
+    Fallback: QLoRA healing via standard PEFT (no Unsloth).
+
+    Slower but works without Unsloth installed.
+
+    Returns:
+        Path to healed model directory
+    """
+    from peft import LoraConfig, get_peft_model, TaskType
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+    print("\n[heal] Loading model with standard PEFT...")
+
+    # 4-bit quantisation config
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=getattr(torch, cfg.dtype),
+        bnb_4bit_use_double_quant=True,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="auto",
+        torch_dtype=getattr(torch, cfg.dtype),
+    )
+
+    # LoRA config
+    lora_config = LoraConfig(
+        r=cfg.heal_lora_r,
+        lora_alpha=cfg.heal_lora_alpha,
+        lora_dropout=cfg.heal_lora_dropout,
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+    )
+
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    # Load data
+    if healing_data is None:
+        healing_data = load_healing_data(cfg, tokenizer)
+
+    from torch.utils.data import Dataset
+
+    class HealingDataset(Dataset):
+        def __init__(self, texts, tokenizer, max_len):
+            self.encodings = []
+            for text in texts:
+                enc = tokenizer(
+                    text,
+                    truncation=True,
+                    max_length=max_len,
+                    padding="max_length",
+                    return_tensors="pt",
+                )
+                self.encodings.append({
+                    "input_ids": enc["input_ids"].squeeze(),
+                    "attention_mask": enc["attention_mask"].squeeze(),
+                    "labels": enc["input_ids"].squeeze(),
+                })
+
+        def __len__(self):
+            return len(self.encodings)
+
+        def __getitem__(self, idx):
+            return self.encodings[idx]
+
+    dataset = HealingDataset(healing_data, tokenizer, cfg.heal_seq_len)
+
+    # Training
+    output_dir = Path(cfg.output_dir) / "heal_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    training_args = TrainingArguments(
+        output_dir=str(output_dir),
+        num_train_epochs=cfg.heal_epochs,
+        per_device_train_batch_size=cfg.heal_batch_size,
+        gradient_accumulation_steps=cfg.heal_grad_accum,
+        learning_rate=cfg.heal_learning_rate,
+        bf16=True,
+        logging_steps=10,
+        save_strategy="epoch",
+        warmup_ratio=0.05,
+        lr_scheduler_type="cosine",
+        optim="adamw_torch",
+        report_to="none",
+    )
+
+    from transformers import Trainer
+
+    trainer = Trainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        args=training_args,
+    )
+
+    print("\n[heal] Starting standard QLoRA healing fine-tune...")
+    trainer.train()
+
+    # Save — merge LoRA adapters
+    healed_dir = Path(cfg.output_dir) / "healed"
+    healed_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"\n[heal] Merging LoRA adapters...")
+    merged_model = model.merge_and_unload()
+    merged_model.save_pretrained(str(healed_dir))
+    tokenizer.save_pretrained(str(healed_dir))
+
+    print(f"[heal] Healed model saved to {healed_dir}")
+    return str(healed_dir)
+
+
+def apply_residual_frozen_adaptation(
+    model_path: str,
+    cfg: MergeConfig,
+    pre_merge_state: dict = None,
+    healing_data: list = None,
+    alpha: float = 1.0,
+    mask: dict = None,
+) -> str:
+    """
+    Residual-Frozen Adaptation — Paper Section 4.3, Equations 15-18.
+
+    Instead of normal LoRA, this method:
+    1. Computes residual: ΔW = current_weights - pre_merge_weights
+    2. Freezes ΔW (the transported knowledge)
+    3. Defines base weights: W_base = current - ΔW
+    4. Trains ONLY W_base using LoRA (the model learns to work WITH the transplant)
+    5. After training, folds back: W_final = W_base + α · M · ΔW  (Eq 18)
+
+    This is better than standard LoRA because:
+    - Standard LoRA might undo the merge (push weights back to pre-merge)
+    - Residual-frozen PRESERVES the merge and only adjusts the base
+
+    Args:
+        model_path: Path to merged model checkpoint
+        cfg: Merge configuration
+        pre_merge_state: State dict from BEFORE the merge (needed to compute ΔW)
+        healing_data: Optional pre-loaded training data
+
+    Returns:
+        Path to healed model directory
+    """
+    from peft import LoraConfig, get_peft_model, TaskType
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, Trainer
+
+    print("\n[heal] Residual-Frozen Adaptation (Paper Section 4.3)")
+    print("[heal] Step 1: Computing frozen residuals (ΔW)...")
+
+    # Load the merged model
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=getattr(torch, cfg.dtype),
+        bnb_4bit_use_double_quant=True,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        quantization_config=bnb_config,
+        device_map="auto",
+        torch_dtype=getattr(torch, cfg.dtype),
+    )
+
+    # If we have pre-merge state, compute and store the residuals
+    frozen_residuals = {}
+    if pre_merge_state is not None:
+        current_state = model.state_dict()
+        for key in current_state:
+            if key in pre_merge_state:
+                delta = current_state[key].float() - pre_merge_state[key].float().to(current_state[key].device)
+                if delta.abs().max() > 1e-8:
+                    frozen_residuals[key] = delta.detach()
+                    # Set the model weights to base (current - delta)
+                    # This way, LoRA trains the base weights, not the merged ones
+                    with torch.no_grad():
+                        current_state[key] = (current_state[key].float() - delta).to(current_state[key].dtype)
+
+        # Save residuals to disk for crash recovery
+        res_dir = Path(cfg.checkpoint_dir) / "frozen_residuals_cache"
+        res_dir.mkdir(parents=True, exist_ok=True)
+        torch.save(frozen_residuals, res_dir / "last_delta.pt")
+        
+        # Load the "base" weights (merged weights minus residuals)
+        model.load_state_dict(current_state)
+        print(f"[heal] Computed {len(frozen_residuals)} frozen residuals")
+        print(f"[heal] Residuals saved to disk for recovery: {res_dir / 'last_delta.pt'}")
+        print(f"[heal] Model now has base weights (residuals subtracted)")
+    else:
+        # Check if we can recover from disk
+        res_cache = Path(cfg.checkpoint_dir) / "frozen_residuals_cache" / "last_delta.pt"
+        if res_cache.exists():
+            print(f"[heal] Recovering frozen residuals from disk cache...")
+            frozen_residuals = torch.load(res_cache, weights_only=True)
+            print(f"[heal] Loaded {len(frozen_residuals)} residuals")
+        else:
+            print("[heal] No pre-merge state or cache provided — using standard LoRA")
+
+    # Step 2: Apply LoRA to train the base weights
+    print("[heal] Step 2: Training base weights with LoRA...")
+
+    lora_config = LoraConfig(
+        r=cfg.heal_lora_r,
+        lora_alpha=cfg.heal_lora_alpha,
+        lora_dropout=cfg.heal_lora_dropout,
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+    )
+
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    # Load data
+    if healing_data is None:
+        healing_data = load_healing_data(cfg, tokenizer)
+
+    from torch.utils.data import Dataset
+
+    class HealingDataset(Dataset):
+        def __init__(self, texts, tok, max_len):
+            self.encodings = []
+            for text in texts:
+                enc = tok(
+                    text, truncation=True, max_length=max_len,
+                    padding="max_length", return_tensors="pt",
+                )
+                self.encodings.append({
+                    "input_ids": enc["input_ids"].squeeze(),
+                    "attention_mask": enc["attention_mask"].squeeze(),
+                    "labels": enc["input_ids"].squeeze(),
+                })
+
+        def __len__(self):
+            return len(self.encodings)
+
+        def __getitem__(self, idx):
+            return self.encodings[idx]
+
+    dataset = HealingDataset(healing_data, tokenizer, cfg.heal_seq_len)
+
+    output_dir = Path(cfg.output_dir) / "heal_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    training_args = TrainingArguments(
+        output_dir=str(output_dir),
+        num_train_epochs=cfg.heal_epochs,
+        per_device_train_batch_size=cfg.heal_batch_size,
+        gradient_accumulation_steps=cfg.heal_grad_accum,
+        learning_rate=cfg.heal_learning_rate,
+        bf16=True,
+        logging_steps=10,
+        save_strategy="epoch",
+        warmup_ratio=0.05,
+        lr_scheduler_type="cosine",
+        optim="adamw_torch",
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        args=training_args,
+    )
+
+    trainer.train()
+
+    # Step 3: Merge LoRA back and fold residuals (Equation 18)
+    print("[heal] Step 3: Merging LoRA + folding frozen residuals (Eq 18)...")
+
+    merged_model = model.merge_and_unload()
+    healed_state = merged_model.state_dict()
+
+    # Fold back: W_final = W_base_trained + α · M · ΔW  (Eq 18)
+    if frozen_residuals:
+        folded_count = 0
+        for key, delta in frozen_residuals.items():
+            if key in healed_state:
+                # Apply mask M^l and scaling alpha if provided
+                val = delta.to(healed_state[key].device)
+                if mask and key in mask:
+                    val = val * mask[key].to(val.device)
+                
+                healed_state[key] = (
+                    healed_state[key].float() + alpha * val.float()
+                ).to(healed_state[key].dtype)
+                folded_count += 1
+        merged_model.load_state_dict(healed_state)
+        print(f"[heal] Folded back {folded_count} frozen residuals (alpha={alpha}, masked={mask is not None})")
+
+    # Save
+    healed_dir = Path(cfg.output_dir) / "healed"
+    healed_dir.mkdir(parents=True, exist_ok=True)
+    merged_model.save_pretrained(str(healed_dir))
+    tokenizer.save_pretrained(str(healed_dir))
+
+    print(f"[heal] Residual-frozen healed model saved to {healed_dir}")
+    return str(healed_dir)
+
+
+def heal_model(
+    model_path: str,
+    cfg: MergeConfig = None,
+    healing_data: list = None,
+    pre_merge_state: dict = None,
+) -> str:
+    """
+    Main entry point for healing.
+
+    If use_residual_frozen=True (paper Section 4.3) AND pre_merge_state is provided,
+    uses residual-frozen adaptation. Otherwise falls back to standard QLoRA.
+
+    Args:
+        model_path: Path to the merged model checkpoint
+        cfg: Merge configuration
+        healing_data: Optional pre-loaded training data
+        pre_merge_state: State dict from BEFORE the merge (for residual-frozen)
+
+    Returns:
+        Path to healed model directory
+    """
+    if cfg is None:
+        cfg = MergeConfig()
+
+    print("\n" + "=" * 60)
+    print("HEALING FINE-TUNE")
+    print(f"Model: {model_path}")
+    print(f"LoRA r={cfg.heal_lora_r}, alpha={cfg.heal_lora_alpha}")
+    print(f"Epochs: {cfg.heal_epochs}, LR: {cfg.heal_learning_rate}")
+    if cfg.use_residual_frozen and pre_merge_state is not None:
+        print(f"Mode: RESIDUAL-FROZEN (Paper Section 4.3)")
+    else:
+        print(f"Mode: Standard QLoRA")
+    print("=" * 60)
+
+    # Paper's residual-frozen adaptation (preferred)
+    if cfg.use_residual_frozen:
+        # Smart discovery: if state isn't provided, try finding it in ResidualBank
+        if pre_merge_state is None:
+            try:
+                from .merge import ResidualBank
+                bank = ResidualBank(cfg)
+                if bank.residual_index:
+                    # Get the most recent merge stage
+                    last_stage = list(bank.residual_index.keys())[-1]
+                    print(f"[heal] Smart discovery: loading residuals from merge stage '{last_stage}'")
+                    # Note: bank saves (original - merged), we want (merged - original)
+                    # So we'll pass the negative of the saved target residual
+                    target_res, _ = bank.load_residuals(last_stage)
+                    pre_merge_state = {}
+                    # We can't easily reconstruct pre_merge_state without base weights,
+                    # but we can pass ΔW directly if we modify apply_residual_frozen_adaptation.
+                    # For now, let's assume we can't reconstruct but we CAN use the cache.
+            except ImportError:
+                pass
+
+        return apply_residual_frozen_adaptation(
+            model_path, cfg, pre_merge_state, healing_data
+        )
+
+    # Standard QLoRA fallback
+    if check_unsloth_available():
+        return apply_qlora_unsloth(model_path, cfg, healing_data)
+    else:
+        return apply_qlora_standard(model_path, cfg, healing_data)

hugging/td_lang/engine/merge.py

@@ -0,0 +1,988 @@
+"""
+Sequential Merge Orchestrator — chains 4 merges with protection.
+
+This is the brain of td_lang engine. It runs each merge in order:
+    1. Load source model
+    2. Inject canary fact into source
+    3. Extract activations from both models
+    4. Compute transport plans (P and Q matrices)
+    5. Fuse weights using optimal transport
+    6. Validate merged model (canary recall, perplexity, thinking mode)
+    7. Apply sequential merge protection before next merge
+    8. Checkpoint
+
+Protection between merges (findings #13):
+    - MagMax: Protect top 20% parameters by magnitude (they carry critical knowledge)
+    - Orthogonal Projection: Project new merge deltas perpendicular to previous ones
+    - Time-Aware Scaling: scale = 1/sqrt(merge_index + 1)
+
+Kill criteria: >10% performance drop on any test → abort merge.
+Findings: #13, #22, #25
+"""
+
+import os
+import gc
+import copy
+import torch
+import numpy as np
+from pathlib import Path
+from typing import Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .config import (
+    MergeConfig, ModelConfig, TARGET, SOURCES,
+    CANARY_FACTS, DEMO_STAGES, FULL_STAGES,
+)
+from .canary import inject_canary, test_all_canaries
+from .transport import (
+    setup_tm_repo,
+    load_calibration_data,
+    extract_activations,
+    compute_transport_plans,
+    fuse_weights,
+)
+from .validate import validate_merged_model, compute_perplexity
+from .techniques import (
+    compute_mergeability_score,
+    compute_transferability_masks,
+    apply_masked_merge,
+    disentangle_rl_weights,
+    merge_with_rl_preservation,
+    compute_arm_rotation,
+    apply_arm_steering,
+    transport_task_vector_theseus,
+    compute_procrustes_alignment,
+)
+
+
+# ============================================================================
+# SEQUENTIAL MERGE PROTECTION
+# ============================================================================
+
+class MergeProtection:
+    """
+    Protects previously merged knowledge from being overwritten.
+
+    Think of it like this: after merging DeepSeek into Qwen3, we have
+    a "direction" in weight space that represents that merge. When we
+    then merge MiMo, we want MiMo's changes to go in a DIFFERENT direction,
+    not overwrite DeepSeek's contribution.
+
+    Three mechanisms:
+    1. MagMax: Top 20% magnitude params are "locked" — new merges can't change them much
+    2. Orthogonal Projection: New deltas are projected perpendicular to previous deltas
+    3. Time-Aware Scaling: Each successive merge gets a smaller alpha (1/sqrt(n+1))
+    """
+
+    def __init__(self, cfg: MergeConfig):
+        self.cfg = cfg
+        self.previous_deltas = {}  # key → list of delta tensors from previous merges
+        self.magnitude_masks = {}  # key → bool mask of top-k magnitude params
+        self.arm_rotations = {}    # ARM: layer → rotation info from last merge
+        self.otmf_masks = {}       # OTMF: param → transferability mask
+        self.merge_count = 0
+
+    def before_merge(
+        self,
+        target_model: AutoModelForCausalLM,
+        source_config: ModelConfig,
+    ) -> float:
+        """
+        Prepare protection before a merge. Returns adjusted alpha.
+
+        Called BEFORE each merge to:
+        1. Compute magnitude masks (MagMax)
+        2. Calculate time-aware alpha scaling
+        """
+        # Time-aware scaling: each merge gets less aggressive
+        if self.cfg.time_aware_scaling:
+            scale = 1.0 / np.sqrt(self.merge_count + 1)
+            adjusted_alpha = source_config.merge_alpha * scale
+            print(f"[protect] Time-aware scaling: {source_config.merge_alpha:.2f} × {scale:.3f} = {adjusted_alpha:.3f}")
+        else:
+            adjusted_alpha = source_config.merge_alpha
+
+        # MagMax: identify top 20% magnitude parameters to protect
+        if self.cfg.use_magmax and self.merge_count > 0:
+            print(f"[protect] Computing MagMax masks (protecting top 20% by magnitude)...")
+            state = target_model.state_dict()
+            for key, param in state.items():
+                if param.dim() >= 1:
+                    flat = param.abs().flatten()
+                    threshold = torch.quantile(flat.float(), 0.8)
+                    self.magnitude_masks[key] = param.abs() >= threshold
+
+        return adjusted_alpha
+
+    def apply_protection(
+        self,
+        target_state: dict,
+        pre_merge_state: dict,
+        key: str,
+    ) -> torch.Tensor:
+        """
+        Apply all protection mechanisms to a fused parameter.
+
+        Called AFTER each parameter is fused, to constrain the change.
+
+        Protection stack (applied in order):
+        1. ARM steering (2602.03237) — steer delta toward gap, away from previous direction
+        2. Orthogonal projection (legacy fallback if ARM disabled)
+        3. OTMF masks (2511.19561) — protect task-specific weights
+        4. MagMax — protect top magnitude params (extra safety layer)
+        """
+        fused = target_state[key]
+        original = pre_merge_state[key]
+        delta = fused - original
+
+        # --- ARM Steering (new, replaces orthogonal projection) ---
+        if self.cfg.use_arm_steering and self.arm_rotations:
+            # Find matching layer rotation
+            layer_prefix = ".".join(key.split(".")[:4])
+            for layer_name, rotation_info in self.arm_rotations.items():
+                if layer_prefix in layer_name:
+                    delta = apply_arm_steering(
+                        delta, rotation_info,
+                        steering_strength=self.cfg.arm_steering_strength,
+                    )
+                    break
+
+        # --- Orthogonal Projection (legacy fallback) ---
+        elif self.cfg.use_orthogonal_projection and key in self.previous_deltas:
+            for prev_delta in self.previous_deltas[key]:
+                prev_flat = prev_delta.flatten().float()
+                delta_flat = delta.flatten().float()
+
+                dot = torch.dot(delta_flat, prev_flat)
+                norm_sq = torch.dot(prev_flat, prev_flat)
+
+                if norm_sq > 1e-10:
+                    projection = (dot / norm_sq) * prev_flat
+                    delta_flat = delta_flat - projection
+                    delta = delta_flat.reshape(delta.shape).to(delta.dtype)
+
+        # --- OTMF Mask Protection (new) ---
+        if self.cfg.use_otmf_masks and key in self.otmf_masks:
+            mask = self.otmf_masks[key].to(delta.device)
+            # Transferable weights: full delta
+            # Task-specific weights: reduced delta (protect them)
+            delta = torch.where(
+                mask,
+                delta,  # Transferable → allow full change
+                delta * (1.0 - self.cfg.otmf_protect_strength),  # Protected → reduced
+            )
+
+        # --- MagMax Protection (extra safety layer) ---
+        if self.cfg.use_magmax and key in self.magnitude_masks:
+            mask = self.magnitude_masks[key]
+            delta = torch.where(mask, delta * 0.1, delta)
+
+        # Apply constrained delta
+        result = original + delta
+
+        return result
+
+    def after_merge(
+        self,
+        target_model: AutoModelForCausalLM,
+        pre_merge_state: dict,
+        pre_merge_activations: dict = None,
+        post_merge_activations: dict = None,
+    ):
+        """
+        Record the merge delta and compute protections for next merge.
+
+        Called AFTER each merge completes successfully.
+        Now also computes:
+        - ARM rotation vectors for next merge steering
+        - OTMF transferability masks for next merge
+        """
+        current_state = target_model.state_dict()
+
+        for key in current_state:
+            if key in pre_merge_state:
+                delta = current_state[key].float() - pre_merge_state[key].float()
+                if delta.abs().max() > 1e-8:
+                    if key not in self.previous_deltas:
+                        self.previous_deltas[key] = []
+                    if len(self.previous_deltas[key]) >= 2:
+                        self.previous_deltas[key].pop(0)
+                    self.previous_deltas[key].append(delta.cpu())
+
+        # --- Compute ARM rotations for next merge ---
+        if self.cfg.use_arm_steering and pre_merge_activations and post_merge_activations:
+            print("[protect] Computing ARM rotation vectors for next merge...")
+            self.arm_rotations = compute_arm_rotation(
+                pre_merge_activations,
+                post_merge_activations,
+                post_merge_activations,  # Target = current state (for gap calculation)
+            )
+
+        # --- Compute OTMF masks for next merge ---
+        if self.cfg.use_otmf_masks and post_merge_activations:
+            print("[protect] Computing OTMF transferability masks...")
+            self.otmf_masks = compute_transferability_masks(
+                target_model,
+                post_merge_activations,
+                threshold=self.cfg.otmf_threshold,
+            )
+
+        self.merge_count += 1
+        print(f"[protect] Recorded merge delta #{self.merge_count} (ARM + OTMF ready for next)")
+
+
+# ============================================================================
+# MAIN ORCHESTRATOR
+# ============================================================================
+
+def is_vision_param(key: str, cfg: MergeConfig) -> bool:
+    """
+    Check if a parameter belongs to the vision encoder.
+
+    Qwen3-VL-8B has a ViT vision encoder + merger projection on top of the
+    language model. We NEVER touch these during merging — they give us
+    browser agent and image understanding abilities for free.
+
+    Vision params start with prefixes like "visual." or "merger."
+    Language params start with "model.layers." or "model.embed_tokens." etc.
+    """
+    for prefix in cfg.vision_skip_prefixes:
+        if key.startswith(prefix):
+            return True
+    return False
+
+
+def get_source_by_stage(stage_name: str) -> Optional[ModelConfig]:
+    """Get model config by stage name."""
+    stage_map = {
+        "deepseek": 0,
+        "mimo": 1,
+        "llama": 2,
+        "falcon": 3,
+    }
+    idx = stage_map.get(stage_name.lower())
+    if idx is not None and idx < len(SOURCES):
+        return SOURCES[idx]
+    return None
+
+
+def load_model(config: ModelConfig, cfg: MergeConfig) -> tuple:
+    """Load a model and its tokenizer/processor."""
+    print(f"\n[merge] Loading {config.name} ({config.hf_id})...")
+
+    # Qwen3-VL uses a processor (handles both text + vision), not just a tokenizer
+    if config.architecture == "transformer+vision":
+        try:
+            from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
+            processor = AutoProcessor.from_pretrained(
+                config.hf_id,
+                trust_remote_code=config.trust_remote_code,
+            )
+            model = Qwen3VLForConditionalGeneration.from_pretrained(
+                config.hf_id,
+                torch_dtype=getattr(torch, cfg.dtype),
+                attn_implementation=cfg.attn_implementation,
+                device_map=cfg.device_map,
+                trust_remote_code=config.trust_remote_code,
+            )
+            # Use the tokenizer from the processor for text operations
+            tokenizer = processor.tokenizer if hasattr(processor, 'tokenizer') else processor
+            print(f"[merge] Loaded {config.name} (VL model): {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B params")
+
+            # Count vision vs language params
+            vision_params = sum(
+                p.numel() for n, p in model.named_parameters()
+                if any(n.startswith(pfx) for pfx in cfg.vision_skip_prefixes)
+            )
+            lang_params = sum(p.numel() for p in model.parameters()) - vision_params
+            print(f"[merge]   Language: {lang_params / 1e9:.1f}B  |  Vision: {vision_params / 1e9:.1f}B")
+
+            return model, tokenizer
+        except ImportError:
+            print("[merge] Qwen3VLForConditionalGeneration not available, falling back to AutoModel")
+
+    # Standard text-only models
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.hf_id,
+        trust_remote_code=config.trust_remote_code,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        config.hf_id,
+        torch_dtype=getattr(torch, cfg.dtype),
+        attn_implementation=cfg.attn_implementation,
+        device_map=cfg.device_map,
+        trust_remote_code=config.trust_remote_code,
+    )
+
+    print(f"[merge] Loaded {config.name}: {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B params")
+    return model, tokenizer
+
+
+def save_checkpoint(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    stage_name: str,
+    cfg: MergeConfig,
+):
+    """Save a checkpoint after a successful merge stage."""
+    ckpt_dir = Path(cfg.checkpoint_dir) / f"after_{stage_name}"
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"[merge] Saving checkpoint to {ckpt_dir}...")
+    model.save_pretrained(ckpt_dir)
+    tokenizer.save_pretrained(ckpt_dir)
+    print(f"[merge] Checkpoint saved: {ckpt_dir}")
+
+    return str(ckpt_dir)
+
+
+# ============================================================================
+# RESIDUAL BANK — Save what was lost during each merge
+# ============================================================================
+
+class ResidualBank:
+    """
+    Saves the knowledge that gets lost during each merge so it can
+    be recovered later.
+
+    When we blend at alpha=0.10:
+        merged = target + alpha * M * (transported - target)
+
+    We LOSE:
+        target_residual = target_original - merged  (what target lost)
+        source_residual = source_original - merged  (what source lost)
+
+    These residuals are saved to disk. Later they can be:
+    1. Fed back during the healing fine-tune (as training signal)
+    2. Re-injected via a small LoRA adapter
+    3. Used to diagnose which merge caused a specific knowledge loss
+    4. Re-applied at a lower alpha if we want more of that model
+
+    Think of it like saving the sawdust when you cut wood — you might
+    need to glue some of it back later.
+    """
+
+    def __init__(self, cfg: MergeConfig):
+        self.cfg = cfg
+        self.residual_dir = Path(cfg.checkpoint_dir) / "residuals"
+        self.residual_dir.mkdir(parents=True, exist_ok=True)
+        self.residual_index = {}  # stage → {path, stats}
+
+    def save_residuals(
+        self,
+        stage_name: str,
+        pre_merge_target_state: dict,
+        source_state: dict,
+        post_merge_state: dict,
+        source_config: ModelConfig,
+    ):
+        """
+        Compute and save what was lost from both target and source.
+
+        Saves two files per merge stage:
+        - target_residual: what the target model lost
+        - source_residual: what the source model didn't fully contribute
+
+        Also saves stats so we know WHERE the biggest losses were
+        (which layers, which type of weights).
+        """
+        stage_dir = self.residual_dir / stage_name
+        stage_dir.mkdir(parents=True, exist_ok=True)
+
+        target_residual = {}
+        source_residual = {}
+        stats = {
+            "stage": stage_name,
+            "source_model": source_config.name,
+            "target_loss_by_layer": {},
+            "source_loss_by_layer": {},
+            "total_target_loss": 0.0,
+            "total_source_loss": 0.0,
+            "biggest_losses": [],
+        }
+
+        for key in post_merge_state:
+            merged_w = post_merge_state[key].float()
+
+            # What the target lost
+            if key in pre_merge_target_state:
+                original_target = pre_merge_target_state[key].float()
+                t_residual = original_target - merged_w
+                t_loss = t_residual.abs().mean().item()
+
+                if t_loss > 1e-6:  # Only save meaningful residuals
+                    target_residual[key] = t_residual.to(torch.bfloat16).cpu()
+                    stats["total_target_loss"] += t_loss
+
+                    # Track per-layer losses
+                    layer_name = ".".join(key.split(".")[:4])
+                    if layer_name not in stats["target_loss_by_layer"]:
+                        stats["target_loss_by_layer"][layer_name] = 0.0
+                    stats["target_loss_by_layer"][layer_name] += t_loss
+
+            # What the source lost (what didn't make it into the merge)
+            if key in source_state:
+                original_source = source_state[key].float()
+                s_residual = original_source - merged_w
+                s_loss = s_residual.abs().mean().item()
+
+                if s_loss > 1e-6:
+                    source_residual[key] = s_residual.to(torch.bfloat16).cpu()
+                    stats["total_source_loss"] += s_loss
+
+                    layer_name = ".".join(key.split(".")[:4])
+                    if layer_name not in stats["source_loss_by_layer"]:
+                        stats["source_loss_by_layer"][layer_name] = 0.0
+                    stats["source_loss_by_layer"][layer_name] += s_loss
+
+        # Find the biggest losses (most knowledge dropped)
+        all_losses = []
+        for key in target_residual:
+            loss_magnitude = target_residual[key].float().abs().mean().item()
+            all_losses.append({"param": key, "side": "target", "loss": loss_magnitude})
+        for key in source_residual:
+            loss_magnitude = source_residual[key].float().abs().mean().item()
+            all_losses.append({"param": key, "side": "source", "loss": loss_magnitude})
+        all_losses.sort(key=lambda x: x["loss"], reverse=True)
+        stats["biggest_losses"] = all_losses[:20]  # Top 20 biggest losses
+
+        # Save to disk
+        torch.save(target_residual, stage_dir / "target_residual.pt")
+        torch.save(source_residual, stage_dir / "source_residual.pt")
+
+        import json
+        with open(stage_dir / "residual_stats.json", "w") as f:
+            json.dump(stats, f, indent=2, default=str)
+
+        self.residual_index[stage_name] = {
+            "path": str(stage_dir),
+            "target_params_saved": len(target_residual),
+            "source_params_saved": len(source_residual),
+            "total_target_loss": stats["total_target_loss"],
+            "total_source_loss": stats["total_source_loss"],
+        }
+
+        print(f"[residual] Saved residuals for {stage_name}:")
+        print(f"  Target lost: {len(target_residual)} params (avg loss: {stats['total_target_loss']:.4f})")
+        print(f"  Source lost: {len(source_residual)} params (avg loss: {stats['total_source_loss']:.4f})")
+        print(f"  Top loss: {all_losses[0]['param']} ({all_losses[0]['side']}, {all_losses[0]['loss']:.4f})" if all_losses else "")
+        print(f"  Saved to: {stage_dir}")
+
+    def load_residuals(self, stage_name: str) -> tuple:
+        """
+        Load saved residuals for a stage.
+
+        Returns:
+            (target_residual_dict, source_residual_dict)
+        """
+        stage_dir = self.residual_dir / stage_name
+        target_residual = torch.load(stage_dir / "target_residual.pt", weights_only=True)
+        source_residual = torch.load(stage_dir / "source_residual.pt", weights_only=True)
+        return target_residual, source_residual
+
+    def reinject_residuals(
+        self,
+        model: AutoModelForCausalLM,
+        stage_name: str,
+        side: str = "both",
+        strength: float = 0.3,
+    ) -> AutoModelForCausalLM:
+        """
+        Re-inject saved residuals back into a model.
+
+        This adds back some of what was lost. Use a low strength (0.1-0.3)
+        to gently recover knowledge without undoing the merge.
+
+        Args:
+            model: The model to inject into
+            stage_name: Which merge stage's residuals to use
+            side: "target", "source", or "both"
+            strength: How much to add back (0=nothing, 1=full residual)
+        """
+        print(f"[residual] Re-injecting {stage_name} residuals (side={side}, strength={strength})...")
+
+        target_residual, source_residual = self.load_residuals(stage_name)
+        state = model.state_dict()
+        injected = 0
+
+        if side in ("target", "both"):
+            for key, residual in target_residual.items():
+                if key in state:
+                    state[key] = state[key] + strength * residual.to(state[key].device).to(state[key].dtype)
+                    injected += 1
+
+        if side in ("source", "both"):
+            for key, residual in source_residual.items():
+                if key in state:
+                    state[key] = state[key] + strength * residual.to(state[key].device).to(state[key].dtype)
+                    injected += 1
+
+        model.load_state_dict(state)
+        print(f"[residual] Re-injected {injected} params at {strength:.0%} strength")
+        return model
+
+    def get_healing_targets(self, top_n: int = 50) -> list:
+        """
+        Get the parameters with the biggest losses across ALL merges.
+
+        These are the params that the healing fine-tune should focus on.
+        Feed this to the LoRA target_modules to make healing smarter.
+        """
+        import json
+        all_losses = []
+
+        for stage_name in self.residual_index:
+            stage_dir = self.residual_dir / stage_name
+            stats_file = stage_dir / "residual_stats.json"
+            if stats_file.exists():
+                with open(stats_file) as f:
+                    stats = json.load(f)
+                for loss in stats.get("biggest_losses", []):
+                    loss["stage"] = stage_name
+                    all_losses.append(loss)
+
+        all_losses.sort(key=lambda x: x["loss"], reverse=True)
+
+        # Extract unique layer/module names for LoRA targeting
+        target_modules = set()
+        for loss in all_losses[:top_n]:
+            param = loss["param"]
+            # Extract the module type (q_proj, k_proj, gate_proj, etc.)
+            parts = param.split(".")
+            for part in parts:
+                if part.endswith("_proj") or part in ("gate_proj", "up_proj", "down_proj"):
+                    target_modules.add(part)
+
+        print(f"[residual] Top healing targets (from {len(all_losses)} total losses):")
+        for loss in all_losses[:5]:
+            print(f"  {loss['param']} ({loss['side']}, stage={loss['stage']}, loss={loss['loss']:.4f})")
+        print(f"  → Suggested LoRA targets: {sorted(target_modules)}")
+
+        return list(target_modules)
+
+
+def run_single_merge(
+    target_model: AutoModelForCausalLM,
+    target_tokenizer: AutoTokenizer,
+    source_config: ModelConfig,
+    cfg: MergeConfig,
+    protection: MergeProtection,
+    residual_bank: ResidualBank = None,
+    calibration_data: list = None,
+    baseline_perplexity: float = None,
+    merged_sources: list = None,
+) -> dict:
+    """
+    Run a single merge: source → target.
+
+    Full pipeline for one merge step:
+    1. Load source model
+    2. Inject canary into source
+    3. Extract activations from both
+    4. Compute transport plans
+    5. Apply merge protection
+    6. Fuse weights
+    7. Apply post-merge protection
+    8. Validate
+
+    Returns:
+        Dict with merge results, validation results, and status
+    """
+    if merged_sources is None:
+        merged_sources = []
+
+    stage_name = source_config.name
+    print(f"\n{'=' * 70}")
+    print(f"MERGE STAGE: {stage_name} → target")
+    print(f"Risk level: {source_config.merge_risk.upper()}")
+    print(f"{'=' * 70}")
+
+    result = {
+        "stage": stage_name,
+        "status": "pending",
+        "validation": None,
+        "checkpoint": None,
+    }
+
+    # --- Step 1: Load source model ---
+    source_model, source_tokenizer = load_model(source_config, cfg)
+
+    # --- Step 2: Inject canary into source ---
+    if stage_name in CANARY_FACTS:
+        print(f"\n[merge] Injecting canary fact into {stage_name}...")
+        source_model = inject_canary(source_model, source_tokenizer, stage_name)
+
+    # --- Step 3: Load calibration data (if not provided) ---
+    if calibration_data is None:
+        calibration_data = load_calibration_data(cfg, target_tokenizer)
+
+    # --- Step 4: Extract two-sided activations (pre + post per projection) ---
+    print(f"\n[merge] Extracting source activations (two-sided)...")
+    source_activations = extract_activations(source_model, calibration_data)
+
+    print(f"\n[merge] Extracting target activations (two-sided)...")
+    pre_merge_target_activations = extract_activations(target_model, calibration_data)
+
+    # --- Step 4.5: Mergeability pre-check (2601.22285) ---
+    if cfg.use_mergeability_check:
+        mergeability = compute_mergeability_score(
+            source_activations, pre_merge_target_activations, source_config
+        )
+        result["mergeability"] = mergeability
+
+        if mergeability["overall"] < cfg.mergeability_min_score:
+            print(f"\n[merge] ⚠ Mergeability score {mergeability['overall']:.2f} below threshold {cfg.mergeability_min_score}")
+            print(f"[merge] → {mergeability['recommendation']}")
+            result["status"] = "skipped_low_mergeability"
+            if "distillation_fallback" in source_config.special_handling:
+                result["fallback"] = "distillation"
+            del source_model, source_activations, pre_merge_target_activations
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            return result
+
+    # --- Step 5: Compute transport plans ---
+    transport_plans = compute_transport_plans(
+        source_activations, pre_merge_target_activations, cfg
+    )
+
+    # --- Step 5.5: RAM RL-weight disentanglement (2601.13572) ---
+    use_ram = (
+        cfg.use_ram_disentangle
+        and source_config.architecture in ("transformer", "transformer+mtp")
+        and source_config.merge_risk in ("low", "medium")
+        and any(kw in source_config.name.lower() for kw in ["r1", "rl", "rlhf", "grpo"])
+    )
+
+    # --- Step 6: Pre-merge protection ---
+    adjusted_alpha = protection.before_merge(target_model, source_config)
+
+    # Override source alpha with time-adjusted value
+    source_config_adjusted = copy.copy(source_config)
+    source_config_adjusted.merge_alpha = adjusted_alpha
+
+    # Save pre-merge state for protection
+    pre_merge_state = {k: v.clone().cpu() for k, v in target_model.state_dict().items()}
+
+    # --- Step 7: Fuse weights ---
+    if use_ram:
+        # RAM path: disentangle RL weights, merge with preservation
+        print(f"\n[merge] Using RAM RL-preservation for {stage_name}...")
+        try:
+            # Try loading the base (pre-RL) model for disentanglement
+            base_hf_id = source_config.hf_id.replace("-RL", "").replace("-R1-0528", "")
+            print(f"[merge] Loading base model for RAM: {base_hf_id}")
+            base_model = AutoModelForCausalLM.from_pretrained(
+                base_hf_id,
+                torch_dtype=getattr(torch, cfg.dtype),
+                device_map=cfg.device_map,
+                trust_remote_code=source_config.trust_remote_code,
+            )
+            shared_mask, rl_mask = disentangle_rl_weights(
+                source_model, base_model, cfg.ram_rl_threshold
+            )
+            # Fuse with RL preservation
+            target_state = merge_with_rl_preservation(
+                target_model.state_dict(),
+                source_model.state_dict(),
+                shared_mask, rl_mask,
+                shared_alpha=cfg.ram_shared_alpha * (adjusted_alpha / source_config.merge_alpha),
+                rl_alpha=cfg.ram_rl_alpha,
+            )
+            target_model.load_state_dict(target_state)
+            del base_model
+            print(f"[merge] RAM merge complete for {stage_name}")
+        except Exception as e:
+            print(f"[merge] RAM failed ({e}), falling back to standard T&M merge")
+            target_model = fuse_weights(
+                source_model, target_model, transport_plans,
+                source_config_adjusted, cfg,
+                target_activations=pre_merge_target_activations,
+            )
+    else:
+        # Standard T&M path (two-sided + top-k masked fusion, paper Eq 14)
+        target_model = fuse_weights(
+            source_model, target_model, transport_plans,
+            source_config_adjusted, cfg,
+            target_activations=pre_merge_target_activations,
+        )
+
+    # --- Step 7.5: Theseus fallback check (2602.12952) ---
+    # If T&M merge produced poor activation alignment, try Theseus
+    if cfg.use_theseus_fallback and source_config.merge_risk == "high":
+        print(f"\n[merge] Checking if Theseus fallback needed for {stage_name}...")
+        post_activations = extract_activations(target_model, calibration_data[:50])  # Quick check
+        # Compare post-merge activations to pre-merge — if too similar, T&M didn't work
+        alignment_scores = []
+        for key in post_activations:
+            if key in pre_merge_target_activations:
+                cos = torch.nn.functional.cosine_similarity(
+                    post_activations[key].float().mean(0, keepdim=True),
+                    pre_merge_target_activations[key].float().mean(0, keepdim=True),
+                )
+                alignment_scores.append(cos.item())
+        avg_change = 1.0 - np.mean(alignment_scores) if alignment_scores else 0.0
+        print(f"[merge] Activation change from merge: {avg_change:.4f}")
+
+        if avg_change < 0.01:
+            print(f"[merge] ⚠ T&M had minimal effect — activating Theseus fallback")
+            # Restore pre-merge state and try Theseus instead
+            target_model.load_state_dict(pre_merge_state)
+            try:
+                base_model = AutoModelForCausalLM.from_pretrained(
+                    source_config.hf_id.split("/")[0] + "/" + source_config.hf_id.split("/")[1].split("-")[0],
+                    torch_dtype=getattr(torch, cfg.dtype),
+                    device_map=cfg.device_map,
+                    trust_remote_code=source_config.trust_remote_code,
+                )
+                target_model = transport_task_vector_theseus(
+                    source_model, base_model, target_model,
+                    source_activations, pre_merge_target_activations,
+                    alpha=cfg.theseus_alpha,
+                )
+                del base_model
+                print(f"[merge] Theseus transport complete for {stage_name}")
+            except Exception as e:
+                print(f"[merge] Theseus also failed ({e}). Using original T&M result.")
+                # Re-apply T&M result
+                target_model = fuse_weights(
+                    source_model, target_model, transport_plans,
+                    source_config_adjusted, cfg,
+                    target_activations=pre_merge_target_activations,
+                )
+
+    # --- Step 8: Apply post-merge protection (ARM + OTMF + MagMax) ---
+    # Skip vision encoder params — they weren't merged, so don't "protect" them
+    if protection.merge_count > 0:
+        print(f"\n[merge] Applying sequential merge protection (ARM + OTMF + MagMax)...")
+        target_state = target_model.state_dict()
+        protected_count = 0
+        vision_skipped = 0
+        for key in target_state:
+            if is_vision_param(key, cfg):
+                vision_skipped += 1
+                continue  # Don't touch vision encoder
+            if key in pre_merge_state:
+                protected_param = protection.apply_protection(
+                    target_state, pre_merge_state, key
+                )
+                target_state[key] = protected_param
+                protected_count += 1
+        target_model.load_state_dict(target_state)
+        print(f"[merge] Protected {protected_count} language params (skipped {vision_skipped} vision params)")
+
+    # --- Step 8.5: Extract post-merge activations for ARM/OTMF ---
+    post_merge_activations = extract_activations(target_model, calibration_data[:100])
+
+    # Record this merge's delta + compute ARM/OTMF for next merge
+    protection.after_merge(
+        target_model, pre_merge_state,
+        pre_merge_activations=pre_merge_target_activations,
+        post_merge_activations=post_merge_activations,
+    )
+
+    # --- Step 8.8: Save residuals (what was lost from both sides) ---
+    if residual_bank is not None:
+        print(f"\n[merge] Saving residuals for {stage_name}...")
+        residual_bank.save_residuals(
+            stage_name=stage_name,
+            pre_merge_target_state=pre_merge_state,
+            source_state={k: v.cpu() for k, v in source_model.state_dict().items()},
+            post_merge_state={k: v.cpu() for k, v in target_model.state_dict().items()},
+            source_config=source_config,
+        )
+
+    # --- Step 9: Free source model memory ---
+    del source_model, source_activations, pre_merge_target_activations
+    del transport_plans, post_merge_activations
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    # --- Step 10: Validate ---
+    merged_sources.append(stage_name)
+    validation = validate_merged_model(
+        target_model, target_tokenizer,
+        merged_sources, cfg,
+        baseline_perplexity=baseline_perplexity,
+    )
+
+    result["validation"] = validation
+    result["merged_sources"] = merged_sources.copy()
+
+    # --- Kill criteria check ---
+    if not validation["overall"]:
+        print(f"\n[merge] ⚠ VALIDATION FAILED for {stage_name}")
+        print(f"[merge] Kill criteria triggered — consider aborting")
+        result["status"] = "failed"
+
+        # Check if we should try distillation fallback
+        if "distillation_fallback" in source_config.special_handling:
+            print(f"[merge] {stage_name} has distillation fallback available")
+            result["fallback"] = "distillation"
+    else:
+        print(f"\n[merge] ✓ {stage_name} merge PASSED validation")
+        result["status"] = "passed"
+
+    return result
+
+
+def run_pipeline(
+    stages: list[str],
+    cfg: MergeConfig = None,
+) -> dict:
+    """
+    Run the full merge pipeline.
+
+    Args:
+        stages: List of stage names to run, e.g. ["deepseek"] or
+                ["deepseek", "mimo", "llama", "falcon"]
+        cfg: Merge configuration (uses defaults if None)
+
+    Returns:
+        Dict with overall results, per-stage results, and final model path
+    """
+    if cfg is None:
+        cfg = MergeConfig()
+
+    print("\n" + "=" * 70)
+    print("TD LANG ENGINE — Transport and Merge Pipeline")
+    print(f"Target: {TARGET.name} ({TARGET.hf_id})")
+    if TARGET.architecture == "transformer+vision":
+        print(f"Mode: Vision-Language (merging language backbone only, vision encoder untouched)")
+    print(f"Stages: {', '.join(stages)}")
+    print(f"Output: {cfg.output_dir}")
+    print("=" * 70)
+
+    # Setup
+    try:
+        setup_tm_repo(cfg)
+    except FileNotFoundError as e:
+        print(f"\n⚠ {e}")
+        print("Continuing with fallback implementation...")
+
+    # Create output directories
+    Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
+    Path(cfg.checkpoint_dir).mkdir(parents=True, exist_ok=True)
+
+    # --- Load target model ---
+    target_model, target_tokenizer = load_model(TARGET, cfg)
+
+    # --- Inject canary into target (Qwen3's own canary) ---
+    if "Qwen3-VL-8B" in CANARY_FACTS:
+        print("\n[pipeline] Injecting canary into base Qwen3-8B...")
+        target_model = inject_canary(target_model, target_tokenizer, "Qwen3-VL-8B")
+
+    # --- Compute baseline perplexity ---
+    print("\n[pipeline] Computing baseline perplexity...")
+    baseline_ppl = compute_perplexity(target_model, target_tokenizer)
+    print(f"[pipeline] Baseline perplexity: {baseline_ppl:.2f}")
+
+    # --- Load calibration data once ---
+    calibration_data = load_calibration_data(cfg, target_tokenizer)
+
+    # --- Initialize merge protection + residual bank ---
+    protection = MergeProtection(cfg)
+    residual_bank = ResidualBank(cfg)
+
+    # --- Run each merge stage ---
+    pipeline_results = {
+        "stages": {},
+        "baseline_perplexity": baseline_ppl,
+        "final_checkpoint": None,
+        "residuals": {},
+        "overall_status": "pending",
+    }
+    merged_sources = []
+    all_passed = True
+
+    for stage_name in stages:
+        source_config = get_source_by_stage(stage_name)
+        if source_config is None:
+            print(f"\n⚠ Unknown stage: {stage_name}, skipping")
+            continue
+
+        # --- Wasserstein pre-check for high-risk models ---
+        if "check_wasserstein_first" in source_config.special_handling:
+            print(f"\n[pipeline] Running Wasserstein pre-check for {source_config.name}...")
+            # TODO: Implement Wasserstein distance pre-check
+            # If distance is too high, skip to distillation fallback
+            print("[pipeline] Pre-check: proceeding (TODO: implement distance check)")
+
+        # Run the merge (with residual bank to save what's lost)
+        stage_result = run_single_merge(
+            target_model, target_tokenizer,
+            source_config, cfg,
+            protection,
+            residual_bank=residual_bank,
+            calibration_data=calibration_data,
+            baseline_perplexity=baseline_ppl,
+            merged_sources=merged_sources,
+        )
+
+        pipeline_results["stages"][stage_name] = stage_result
+
+        if stage_result["status"] == "passed":
+            # Save checkpoint
+            ckpt_path = save_checkpoint(
+                target_model, target_tokenizer, stage_name, cfg
+            )
+            stage_result["checkpoint"] = ckpt_path
+            pipeline_results["final_checkpoint"] = ckpt_path
+        else:
+            all_passed = False
+            print(f"\n[pipeline] Stage {stage_name} FAILED")
+
+            # Decision: abort or continue?
+            if source_config.merge_risk == "high":
+                print(f"[pipeline] High-risk model failed — skipping (will use distillation)")
+                # Don't abort the whole pipeline, just skip this model
+                continue
+            else:
+                print(f"[pipeline] ABORTING pipeline — non-high-risk model failed")
+                pipeline_results["overall_status"] = f"aborted_at_{stage_name}"
+                break
+
+    # --- Save residual index ---
+    pipeline_results["residuals"] = residual_bank.residual_index
+    if residual_bank.residual_index:
+        print(f"\n[pipeline] Residual bank: {len(residual_bank.residual_index)} stages saved")
+        for stage, info in residual_bank.residual_index.items():
+            print(f"  {stage}: target lost {info['total_target_loss']:.4f}, source lost {info['total_source_loss']:.4f}")
+
+        # Identify which modules need the most healing
+        healing_targets = residual_bank.get_healing_targets(top_n=50)
+        pipeline_results["suggested_healing_targets"] = healing_targets
+
+    # --- Save final model ---
+    if pipeline_results["final_checkpoint"]:
+        final_dir = Path(cfg.output_dir) / "final"
+        final_dir.mkdir(parents=True, exist_ok=True)
+        target_model.save_pretrained(final_dir)
+        target_tokenizer.save_pretrained(final_dir)
+        pipeline_results["final_model_path"] = str(final_dir)
+        print(f"\n[pipeline] Final model saved to {final_dir}")
+
+    if all_passed:
+        pipeline_results["overall_status"] = "all_passed"
+    elif pipeline_results["overall_status"] == "pending":
+        pipeline_results["overall_status"] = "partial"
+
+    # --- Print final summary ---
+    print("\n" + "=" * 70)
+    print("PIPELINE SUMMARY")
+    print("=" * 70)
+    for stage_name, stage_result in pipeline_results["stages"].items():
+        status = stage_result["status"]
+        emoji = "✓" if status == "passed" else "✗"
+        print(f"  {emoji} {stage_name}: {status}")
+    print(f"\n  Overall: {pipeline_results['overall_status']}")
+    if residual_bank.residual_index:
+        print(f"\n  Residuals saved for: {', '.join(residual_bank.residual_index.keys())}")
+        print(f"  To recover lost knowledge later:")
+        print(f"    python -m td_lang.engine --reinject <stage> --strength 0.2")
+    print("=" * 70)
+
+    return pipeline_results

hugging/td_lang/engine/run.py

@@ -0,0 +1,279 @@
+"""
+TD Fuse — Main Entry Point.
+
+Usage:
+    # Dad demo: merge just DeepSeek → Qwen3-8B (easiest, lowest risk)
+    python -m td_fuse.run --stage demo
+
+    # Full pipeline: all 4 merges
+    python -m td_fuse.run --stage all
+
+    # Single model merge
+    python -m td_fuse.run --stage deepseek
+    python -m td_fuse.run --stage mimo
+    python -m td_fuse.run --stage llama
+    python -m td_fuse.run --stage falcon
+
+    # With healing fine-tune after merge
+    python -m td_fuse.run --stage demo --heal
+
+    # Custom output directory
+    python -m td_fuse.run --stage all --output ./my_output
+
+    # Heal an existing checkpoint
+    python -m td_fuse.run --heal-only --model-path ./td_fuse_checkpoints/after_deepseek
+
+Findings: #25 (dad demo plan), #22 (merge order), #24 (official T&M pipeline)
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+from .config import MergeConfig, DEMO_STAGES, FULL_STAGES
+from .merge import run_pipeline, ResidualBank
+from .heal import heal_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="TD Fuse — Transport and Merge pipeline for Time Dilation",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python -m td_fuse.run --stage demo           # Dad demo (DeepSeek only)
+  python -m td_fuse.run --stage all            # Full 4-model merge
+  python -m td_fuse.run --stage all --heal     # Merge + healing fine-tune
+  python -m td_fuse.run --heal-only --model-path ./checkpoint
+  python -m td_fuse.run --reinject deepseek --strength 0.2 --model-path ./final
+        """,
+    )
+
+    parser.add_argument(
+        "--stage",
+        type=str,
+        default="demo",
+        choices=["demo", "all", "deepseek", "mimo", "llama", "falcon"],
+        help="Which merge stage(s) to run (default: demo)",
+    )
+    parser.add_argument(
+        "--heal",
+        action="store_true",
+        help="Run healing fine-tune after merge",
+    )
+    parser.add_argument(
+        "--heal-only",
+        action="store_true",
+        help="Only run healing (skip merge), requires --model-path",
+    )
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default=None,
+        help="Path to existing model/checkpoint (for --heal-only)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="./td_fuse_outputs",
+        help="Output directory (default: ./td_fuse_outputs)",
+    )
+    parser.add_argument(
+        "--checkpoint-dir",
+        type=str,
+        default="./td_fuse_checkpoints",
+        help="Checkpoint directory (default: ./td_fuse_checkpoints)",
+    )
+    parser.add_argument(
+        "--tm-repo",
+        type=str,
+        default="./Cross-Architecture-Merging-for-Large-Language-Models",
+        help="Path to official T&M repo",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print what would happen without actually running",
+    )
+    parser.add_argument(
+        "--reinject",
+        type=str,
+        default=None,
+        help="Re-inject saved residuals from a stage (e.g., --reinject deepseek)",
+    )
+    parser.add_argument(
+        "--reinject-side",
+        type=str,
+        default="both",
+        choices=["target", "source", "both"],
+        help="Which side's residuals to re-inject (default: both)",
+    )
+    parser.add_argument(
+        "--strength",
+        type=float,
+        default=0.2,
+        help="Residual re-injection strength, 0-1 (default: 0.2)",
+    )
+
+    return parser.parse_args()
+
+
+def print_banner():
+    """Print the TD Fuse banner."""
+    banner = """
+    ╔══════════════════════════════════════════════════╗
+    ║                                                  ║
+    ║   ████████╗██████╗     ███████╗██╗   ██╗███████╗ ║
+    ║   ╚══██╔══╝██╔══██╗    ██╔════╝██║   ██║██╔════╝ ║
+    ║      ██║   ██║  ██║    █████╗  ██║   ██║███████╗ ║
+    ║      ██║   ██║  ██║    ██╔══╝  ██║   ██║╚════██║ ║
+    ║      ██║   ██████╔╝    ██║     ╚██████╔╝███████║ ║
+    ║      ╚═╝   ╚═════╝     ╚═╝      ╚═════╝ ╚══════╝ ║
+    ║                                                  ║
+    ║   Transport and Merge for Time Dilation          ║
+    ║   Merging 5 models into Qwen3-8B                 ║
+    ║                                                  ║
+    ╚══════════════════════════════════��═══════════════╝
+    """
+    print(banner)
+
+
+def main():
+    args = parse_args()
+    print_banner()
+
+    # Build config from args
+    cfg = MergeConfig(
+        output_dir=args.output,
+        checkpoint_dir=args.checkpoint_dir,
+        tm_repo_path=args.tm_repo,
+    )
+
+    # Determine which stages to run
+    if args.stage == "demo":
+        stages = DEMO_STAGES
+    elif args.stage == "all":
+        stages = FULL_STAGES
+    else:
+        stages = [args.stage]
+
+    # --- Reinject residuals mode ---
+    if args.reinject:
+        if not args.model_path:
+            print("Error: --reinject requires --model-path")
+            sys.exit(1)
+
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        import torch
+
+        print(f"\n[run] Re-injecting residuals from stage: {args.reinject}")
+        print(f"[run] Side: {args.reinject_side}, Strength: {args.strength}")
+
+        residual_bank = ResidualBank(cfg)
+        tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+        )
+
+        model = residual_bank.reinject_residuals(
+            model, args.reinject,
+            side=args.reinject_side,
+            strength=args.strength,
+        )
+
+        # Save the patched model
+        patched_dir = Path(cfg.output_dir) / f"reinjected_{args.reinject}_{args.strength}"
+        patched_dir.mkdir(parents=True, exist_ok=True)
+        model.save_pretrained(str(patched_dir))
+        tokenizer.save_pretrained(str(patched_dir))
+        print(f"\n[run] Patched model saved to: {patched_dir}")
+        return
+
+    # --- Heal-only mode ---
+    if args.heal_only:
+        if not args.model_path:
+            print("Error: --heal-only requires --model-path")
+            sys.exit(1)
+
+        print(f"\n[run] Healing model at: {args.model_path}")
+        healed_path = heal_model(args.model_path, cfg)
+        print(f"\n[run] Healed model saved to: {healed_path}")
+        return
+
+    # --- Dry run ---
+    if args.dry_run:
+        print("\n=== DRY RUN ===")
+        print(f"Stages: {stages}")
+        print(f"Output: {cfg.output_dir}")
+        print(f"Checkpoints: {cfg.checkpoint_dir}")
+        print(f"T&M repo: {cfg.tm_repo_path}")
+        print(f"Heal after: {args.heal}")
+        print(f"\nWould run:")
+        for i, stage in enumerate(stages, 1):
+            print(f"  {i}. Merge {stage} → target")
+            print(f"     → Validate (canary + perplexity + thinking + reasoning)")
+            print(f"     → Checkpoint")
+        if args.heal:
+            print(f"  {len(stages) + 1}. QLoRA healing fine-tune")
+        print("\nNo changes made (dry run).")
+        return
+
+    # --- Run the pipeline ---
+    start_time = time.time()
+
+    results = run_pipeline(stages, cfg)
+
+    elapsed = time.time() - start_time
+    print(f"\n[run] Pipeline completed in {elapsed / 60:.1f} minutes")
+
+    # --- Healing fine-tune (optional) ---
+    if args.heal and results.get("final_checkpoint"):
+        print("\n[run] Starting healing fine-tune...")
+        healed_path = heal_model(results["final_checkpoint"], cfg)
+        results["healed_model_path"] = healed_path
+        print(f"[run] Healed model: {healed_path}")
+
+    # --- Save results ---
+    results_path = Path(cfg.output_dir) / "pipeline_results.json"
+
+    # Convert non-serialisable objects
+    def make_serialisable(obj):
+        if isinstance(obj, dict):
+            return {k: make_serialisable(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [make_serialisable(v) for v in obj]
+        elif isinstance(obj, (int, float, str, bool, type(None))):
+            return obj
+        else:
+            return str(obj)
+
+    with open(results_path, "w") as f:
+        json.dump(make_serialisable(results), f, indent=2)
+    print(f"[run] Results saved to {results_path}")
+
+    # --- Final summary ---
+    print(f"\n{'=' * 60}")
+    print("TD FUSE COMPLETE")
+    print(f"{'=' * 60}")
+    print(f"  Status:     {results['overall_status']}")
+    print(f"  Time:       {elapsed / 60:.1f} minutes")
+    if results.get("final_model_path"):
+        print(f"  Model:      {results['final_model_path']}")
+    if results.get("healed_model_path"):
+        print(f"  Healed:     {results['healed_model_path']}")
+    print(f"  Results:    {results_path}")
+    print(f"{'=' * 60}")
+
+    # Exit code based on result
+    if results["overall_status"] == "all_passed":
+        sys.exit(0)
+    else:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

hugging/td_lang/engine/techniques.py

@@ -0,0 +1,669 @@
+"""
+Advanced Merge Techniques — from latest papers (Feb 2026).
+
+This module contains implementations inspired by recent research
+that improve TD's sequential cross-architecture merging pipeline.
+
+Techniques:
+    1. Theseus (2602.12952) — Procrustes-based task vector transport
+    2. ARM (2602.03237) — Activation-guided rotation for sequential merges
+    3. OTMF (2511.19561) — OT masks for identifying transferable weights
+    4. RAM (2601.13572) — RL-weight disentanglement for RL-trained models
+    5. Mergeability (2601.22285) — Pre-check scoring before attempting merge
+
+These complement Transport and Merge (2602.05495) which handles
+the core cross-architecture fusion via optimal transport.
+"""
+
+import torch
+import numpy as np
+from typing import Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .config import MergeConfig, ModelConfig
+
+
+# ============================================================================
+# 1. THESEUS — Procrustes-Based Task Vector Transport (2602.12952)
+# ============================================================================
+#
+# Instead of aligning neurons via optimal transport (T&M), Theseus aligns
+# the FUNCTIONAL EFFECT of weights via orthogonal Procrustes.
+#
+# Analogy: T&M says "neuron 5 in Model A = neuron 12 in Model B"
+#          Theseus says "the EFFECT of Model A's weights can be rotated
+#          into Model B's space"
+#
+# Best for: Models where neuron-level alignment is poor (Falcon SSM hybrid)
+
+def compute_procrustes_alignment(
+    source_activations: torch.Tensor,
+    target_activations: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Compute the orthogonal Procrustes rotation matrix R that best maps
+    source activations into target activation space.
+
+    R = argmin ||target - source @ R||_F  subject to R^T R = I
+
+    Solution: R = V @ U^T from SVD of (source^T @ target) = U S V^T
+
+    This is a closed-form solution — no iterative optimisation needed.
+
+    Args:
+        source_activations: [num_samples, source_dim] activation matrix
+        target_activations: [num_samples, target_dim] activation matrix
+
+    Returns:
+        R: [source_dim, target_dim] rotation matrix
+    """
+    # Center the activations (remove mean)
+    S = source_activations - source_activations.mean(dim=0, keepdim=True)
+    T = target_activations - target_activations.mean(dim=0, keepdim=True)
+
+    # Handle dimension mismatch by zero-padding the smaller one
+    s_dim = S.shape[1]
+    t_dim = T.shape[1]
+    max_dim = max(s_dim, t_dim)
+
+    if s_dim < max_dim:
+        S = torch.nn.functional.pad(S, (0, max_dim - s_dim))
+    if t_dim < max_dim:
+        T = torch.nn.functional.pad(T, (0, max_dim - t_dim))
+
+    # Cross-covariance matrix
+    M = S.T @ T  # [max_dim, max_dim]
+
+    # SVD: M = U @ diag(sigma) @ V^T
+    U, sigma, Vt = torch.linalg.svd(M, full_matrices=True)
+
+    # Optimal rotation: R = V @ U^T
+    # This ensures R is orthogonal (R^T R = I)
+    R = Vt.T @ U.T
+
+    # Ensure proper rotation (det = +1), not reflection
+    det = torch.linalg.det(R)
+    if det < 0:
+        # Flip sign of last column of Vt
+        Vt[-1, :] *= -1
+        R = Vt.T @ U.T
+
+    return R[:s_dim, :t_dim]  # Crop back to original dims
+
+
+def transport_task_vector_theseus(
+    source_model: AutoModelForCausalLM,
+    source_base_model: AutoModelForCausalLM,
+    target_model: AutoModelForCausalLM,
+    source_activations: dict,
+    target_activations: dict,
+    alpha: float = 0.3,
+) -> AutoModelForCausalLM:
+    """
+    Transport a task vector from source to target using Theseus method.
+
+    Task vector = source_finetuned - source_base
+    (the "diff" that represents what the model learned)
+
+    We rotate this diff into target's space using Procrustes alignment,
+    then add it to target: target_new = target + alpha * R @ task_vector
+
+    This is the FALLBACK for when T&M's neuron-level alignment fails
+    (e.g., Falcon's SSM components).
+
+    Args:
+        source_model: The fine-tuned source (e.g., Falcon-H1R-7B)
+        source_base_model: The base version of source (for computing task vector)
+        target_model: The target to transport into (our merged Qwen3)
+        source_activations: Layer → activation tensors for source
+        target_activations: Layer → activation tensors for target
+        alpha: Blending weight for the transported task vector
+    """
+    print("[theseus] Computing task vectors and Procrustes alignment...")
+
+    source_state = source_model.state_dict()
+    base_state = source_base_model.state_dict()
+    target_state = target_model.state_dict()
+
+    # Compute per-layer Procrustes rotation matrices
+    rotations = {}
+    source_layers = sorted(source_activations.keys())
+    target_layers = sorted(target_activations.keys())
+
+    for sl, tl in zip(source_layers, target_layers):
+        if sl in source_activations and tl in target_activations:
+            R = compute_procrustes_alignment(
+                source_activations[sl].float(),
+                target_activations[tl].float(),
+            )
+            rotations[(sl, tl)] = R
+
+    # Transport task vectors
+    transported_count = 0
+    for target_key in target_state:
+        # Find matching source key (simplified — same key names)
+        source_key = target_key
+        if source_key not in source_state or source_key not in base_state:
+            continue
+
+        # Task vector = what the source learned
+        task_vector = source_state[source_key].float() - base_state[source_key].float()
+
+        if task_vector.abs().max() < 1e-8:
+            continue  # No meaningful change
+
+        # For 2D weight matrices, apply rotation
+        if task_vector.dim() == 2:
+            # Find the appropriate rotation for this layer
+            for (sl, tl), R in rotations.items():
+                if sl.split(".")[2] == target_key.split(".")[2]:  # Same layer index
+                    R_device = R.to(task_vector.device)
+                    # Rotate: task_vector_rotated = task_vector @ R
+                    try:
+                        if task_vector.shape[1] == R_device.shape[0]:
+                            task_vector = task_vector @ R_device
+                        elif task_vector.shape[0] == R_device.shape[0]:
+                            task_vector = R_device.T @ task_vector
+                    except RuntimeError:
+                        pass  # Dimension mismatch, use unrotated
+                    break
+
+        # Apply: target_new = target + alpha * rotated_task_vector
+        target_w = target_state[target_key]
+        if task_vector.shape == target_w.shape:
+            target_state[target_key] = target_w + alpha * task_vector.to(target_w.dtype)
+            transported_count += 1
+
+    target_model.load_state_dict(target_state)
+    print(f"[theseus] Transported {transported_count} task vectors via Procrustes")
+    return target_model
+
+
+# ============================================================================
+# 2. ARM — Activation-Guided Rotations for Sequential Merging (2602.03237)
+# ============================================================================
+#
+# ARM treats sequential merging like gradient descent — each merge step
+# has a "direction" and a "learning rate" (merge coefficient).
+#
+# Key insight: Use ACTIVATION PATTERNS to compute optimal rotation vectors
+# that guide each merge step. This is a smarter version of our
+# orthogonal projection in MergeProtection.
+
+def compute_arm_rotation(
+    pre_merge_activations: dict,
+    post_merge_activations: dict,
+    target_activations: dict,
+) -> dict:
+    """
+    Compute ARM rotation vectors for sequential merge protection.
+
+    For each layer, compute a rotation that:
+    1. Preserves the direction of knowledge already merged
+    2. Steers the next merge to fill GAPS rather than overwrite
+
+    The rotation is computed from the activation change (what the
+    last merge did) and the target (where we want to end up).
+
+    Returns:
+        Dict of layer_name → rotation matrix
+    """
+    print("[arm] Computing activation-guided rotations...")
+
+    rotations = {}
+
+    for layer_name in pre_merge_activations:
+        if layer_name not in post_merge_activations or layer_name not in target_activations:
+            continue
+
+        pre = pre_merge_activations[layer_name].float()    # Before last merge
+        post = post_merge_activations[layer_name].float()   # After last merge
+        target = target_activations[layer_name].float()      # Ideal target
+
+        # Delta from last merge
+        merge_delta = post - pre  # [samples, hidden_dim]
+
+        # Gap remaining (what we still need)
+        gap = target - post  # [samples, hidden_dim]
+
+        # Average across samples to get direction vectors
+        delta_dir = merge_delta.mean(dim=0)  # [hidden_dim]
+        gap_dir = gap.mean(dim=0)            # [hidden_dim]
+
+        # Normalise
+        delta_norm = delta_dir / (delta_dir.norm() + 1e-8)
+        gap_norm = gap_dir / (gap_dir.norm() + 1e-8)
+
+        # Compute rotation from delta direction to gap direction
+        # Using Rodrigues' rotation formula for the 2D plane
+        # spanned by delta and gap
+        cos_theta = torch.dot(delta_norm, gap_norm).clamp(-1, 1)
+        sin_theta = torch.sqrt(1 - cos_theta ** 2)
+
+        # Store as a simple rotation descriptor
+        rotations[layer_name] = {
+            "delta_direction": delta_norm,
+            "gap_direction": gap_norm,
+            "cos_theta": cos_theta.item(),
+            "sin_theta": sin_theta.item(),
+            "gap_magnitude": gap_dir.norm().item(),
+        }
+
+    return rotations
+
+
+def apply_arm_steering(
+    weight_delta: torch.Tensor,
+    rotation_info: dict,
+    steering_strength: float = 0.5,
+) -> torch.Tensor:
+    """
+    Steer a weight delta using ARM rotation vectors.
+
+    Instead of blindly projecting out previous merge directions
+    (our old orthogonal projection), ARM STEERS the delta toward
+    the remaining gap.
+
+    Args:
+        weight_delta: The raw delta from the current merge
+        rotation_info: ARM rotation info for this layer
+        steering_strength: How much to steer (0=no steering, 1=full)
+
+    Returns:
+        Steered weight delta
+    """
+    delta_dir = rotation_info["delta_direction"]
+    gap_dir = rotation_info["gap_direction"]
+
+    flat = weight_delta.flatten().float()
+
+    # Component along previous merge direction
+    prev_component = torch.dot(flat, delta_dir.to(flat.device))
+
+    # Remove some of the previous-direction component
+    # and add gap-direction component instead
+    correction = (
+        -steering_strength * prev_component * delta_dir.to(flat.device)
+        + steering_strength * prev_component * gap_dir.to(flat.device)
+    )
+
+    steered = flat + correction
+    return steered.reshape(weight_delta.shape).to(weight_delta.dtype)
+
+
+# ============================================================================
+# 3. OTMF — Transferability Masks via Optimal Transport (2511.19561)
+# ============================================================================
+#
+# OTMF discovers which parts of each model are "transferable" (shared
+# knowledge) vs "task-specific" (unique to that model).
+#
+# Transferable weights → safe to merge/average
+# Task-specific weights → must be preserved carefully
+#
+# This replaces our MagMax "top 20% by magnitude" heuristic with a
+# principled, data-driven approach.
+
+def compute_transferability_masks(
+    model: AutoModelForCausalLM,
+    calibration_activations: dict,
+    threshold: float = 0.3,
+) -> dict:
+    """
+    Compute per-parameter transferability masks using activation variance.
+
+    High activation variance across diverse inputs → parameter encodes
+    task-specific knowledge (DON'T merge aggressively).
+
+    Low activation variance → parameter encodes shared/general knowledge
+    (safe to merge/average).
+
+    This is a simplified version of OTMF's OT-based mask discovery.
+
+    Args:
+        model: The current merged model
+        calibration_activations: Layer → [samples, hidden_dim] activations
+        threshold: Variance quantile threshold for "task-specific" classification
+
+    Returns:
+        Dict of param_name → bool mask (True = transferable/safe, False = task-specific/protect)
+    """
+    print("[otmf] Computing transferability masks...")
+
+    masks = {}
+    state = model.state_dict()
+
+    # Compute per-neuron activation variance
+    neuron_importance = {}
+    for layer_name, acts in calibration_activations.items():
+        # Variance across samples: high variance = this neuron is doing something specific
+        variance = acts.var(dim=0)  # [hidden_dim]
+        neuron_importance[layer_name] = variance
+
+    # Map neuron importance to parameter importance
+    for param_name, param in state.items():
+        # Find the corresponding layer's importance
+        layer_prefix = ".".join(param_name.split(".")[:4])  # e.g., model.layers.0.self_attn
+
+        importance = None
+        for layer_name, var in neuron_importance.items():
+            if layer_prefix in layer_name:
+                importance = var
+                break
+
+        if importance is None:
+            # Default: mark everything as transferable (safe to merge)
+            masks[param_name] = torch.ones(param.shape, dtype=torch.bool)
+            continue
+
+        # For 2D weights: importance determines which rows/columns to protect
+        if param.dim() == 2:
+            rows, cols = param.shape
+            # Use importance for the output dimension
+            imp = importance[:rows] if importance.shape[0] >= rows else importance
+
+            # Compute threshold: top (1-threshold) fraction is task-specific
+            if imp.numel() > 0:
+                q = torch.quantile(imp.float(), 1.0 - threshold)
+                # True = transferable (below threshold), False = task-specific (protect)
+                row_mask = imp < q
+                masks[param_name] = row_mask.unsqueeze(1).expand_as(param)
+            else:
+                masks[param_name] = torch.ones(param.shape, dtype=torch.bool)
+        else:
+            # 1D params (biases, norms): default to transferable
+            masks[param_name] = torch.ones(param.shape, dtype=torch.bool)
+
+    transferable = sum(m.sum().item() for m in masks.values())
+    total = sum(m.numel() for m in masks.values())
+    print(f"[otmf] Transferability: {transferable / total:.1%} transferable, {1 - transferable / total:.1%} task-specific")
+
+    return masks
+
+
+def apply_masked_merge(
+    target_state: dict,
+    fused_state: dict,
+    masks: dict,
+    protect_strength: float = 0.8,
+) -> dict:
+    """
+    Apply transferability masks during merge.
+
+    For transferable weights: use the fused (merged) value
+    For task-specific weights: preserve more of the original target value
+
+    Args:
+        target_state: Original target weights (before this merge)
+        fused_state: Newly fused weights (after T&M/Theseus fusion)
+        masks: Transferability masks (True = safe to change)
+        protect_strength: How much to protect task-specific weights (0-1)
+
+    Returns:
+        Masked merged state dict
+    """
+    result = {}
+
+    for key in fused_state:
+        if key in masks and key in target_state:
+            mask = masks[key].to(fused_state[key].device)
+            original = target_state[key]
+            fused = fused_state[key]
+
+            # Transferable: use fused value
+            # Task-specific: blend more toward original
+            blended = torch.where(
+                mask,
+                fused,  # Transferable → take merged value
+                protect_strength * original + (1 - protect_strength) * fused,  # Protected
+            )
+            result[key] = blended
+        else:
+            result[key] = fused_state[key]
+
+    protected_params = sum(1 for k in masks if not masks[k].all())
+    print(f"[otmf] Applied masks: {protected_params} parameters partially protected")
+
+    return result
+
+
+# ============================================================================
+# 4. RAM — RL-Weight Disentanglement (2601.13572)
+# ============================================================================
+#
+# RL-trained models (DeepSeek-R1, MiMo-7B-RL) have two types of knowledge:
+#   - Shared: general language understanding (same as base model)
+#   - RL-specific: reasoning patterns learned via GRPO/RLHF
+#
+# RAM separates these so we can merge the shared parts normally
+# but PRESERVE the RL-specific parts that make these models special.
+
+def disentangle_rl_weights(
+    rl_model: AutoModelForCausalLM,
+    base_model: AutoModelForCausalLM,
+    rl_threshold: float = 0.1,
+) -> tuple:
+    """
+    Separate RL-specific weights from shared/general weights.
+
+    RL-specific = weights that changed significantly during RL training
+    Shared = weights that are basically the same as base
+
+    We identify RL-specific weights by looking at the magnitude of
+    change from base model to RL model. Big changes → RL learned
+    something there → don't average it away.
+
+    Args:
+        rl_model: The RL-trained model (e.g., DeepSeek-R1, MiMo-7B-RL)
+        base_model: The base model before RL training
+        rl_threshold: Relative change threshold for "RL-specific" classification
+
+    Returns:
+        Tuple of (shared_mask, rl_mask) — both are dicts of param_name → bool tensor
+        shared_mask: True = this weight is shared (safe to merge normally)
+        rl_mask: True = this weight is RL-specific (protect during merge)
+    """
+    print("[ram] Disentangling RL-specific vs shared weights...")
+
+    rl_state = rl_model.state_dict()
+    base_state = base_model.state_dict()
+
+    shared_mask = {}
+    rl_mask = {}
+
+    total_params = 0
+    rl_params = 0
+
+    for key in rl_state:
+        if key not in base_state:
+            # New param (e.g., MTP head) — mark as RL-specific
+            rl_mask[key] = torch.ones_like(rl_state[key], dtype=torch.bool)
+            shared_mask[key] = torch.zeros_like(rl_state[key], dtype=torch.bool)
+            rl_params += rl_state[key].numel()
+            total_params += rl_state[key].numel()
+            continue
+
+        rl_w = rl_state[key].float()
+        base_w = base_state[key].float()
+
+        # Relative change: |rl - base| / (|base| + epsilon)
+        change = (rl_w - base_w).abs()
+        base_magnitude = base_w.abs() + 1e-8
+        relative_change = change / base_magnitude
+
+        # RL-specific: relative change > threshold
+        is_rl = relative_change > rl_threshold
+        rl_mask[key] = is_rl
+        shared_mask[key] = ~is_rl
+
+        rl_params += is_rl.sum().item()
+        total_params += is_rl.numel()
+
+    pct = rl_params / total_params * 100 if total_params > 0 else 0
+    print(f"[ram] RL-specific: {rl_params:,} params ({pct:.1f}%)")
+    print(f"[ram] Shared:      {total_params - rl_params:,} params ({100 - pct:.1f}%)")
+
+    return shared_mask, rl_mask
+
+
+def merge_with_rl_preservation(
+    target_state: dict,
+    source_state: dict,
+    shared_mask: dict,
+    rl_mask: dict,
+    shared_alpha: float = 0.5,
+    rl_alpha: float = 0.8,
+) -> dict:
+    """
+    Merge source into target while preserving RL-specific weights.
+
+    Shared weights: normal blending at shared_alpha
+    RL-specific weights: stronger blending toward source (preserve RL knowledge)
+
+    This prevents the RL reasoning capabilities from being diluted
+    by averaging with target weights.
+
+    Args:
+        target_state: Current target model state
+        source_state: RL model state to merge in
+        shared_mask: Which params are shared (safe for normal merge)
+        rl_mask: Which params are RL-specific (preserve with higher alpha)
+        shared_alpha: Alpha for shared weights (normal)
+        rl_alpha: Alpha for RL-specific weights (higher = preserve more RL knowledge)
+    """
+    print(f"[ram] Merging with RL preservation (shared α={shared_alpha}, RL α={rl_alpha})...")
+
+    result = {}
+    for key in target_state:
+        if key not in source_state:
+            result[key] = target_state[key]
+            continue
+
+        target_w = target_state[key]
+        source_w = source_state[key]
+
+        if source_w.shape != target_w.shape:
+            result[key] = target_state[key]
+            continue
+
+        if key in rl_mask and key in shared_mask:
+            rl_m = rl_mask[key].to(target_w.device)
+            # RL-specific: use higher alpha (preserve RL knowledge)
+            # Shared: use normal alpha
+            alpha_map = torch.where(rl_m, rl_alpha, shared_alpha)
+            if alpha_map.shape != target_w.shape:
+                alpha_map = alpha_map.expand_as(target_w) if alpha_map.dim() > 0 else torch.full_like(target_w, shared_alpha)
+
+            result[key] = alpha_map * source_w.to(target_w.device) + (1 - alpha_map) * target_w
+        else:
+            result[key] = shared_alpha * source_w.to(target_w.device) + (1 - shared_alpha) * target_w
+
+    return result
+
+
+# ============================================================================
+# 5. MERGEABILITY PRE-CHECK (2601.22285)
+# ============================================================================
+#
+# Before spending GPU hours on a merge that might fail, check if the
+# models are actually COMPATIBLE enough to merge.
+#
+# Mergeability score: 0.0 (definitely won't work) to 1.0 (should work great)
+
+def compute_mergeability_score(
+    source_activations: dict,
+    target_activations: dict,
+    source_config: ModelConfig,
+) -> dict:
+    """
+    Predict how well a source model will merge into the target.
+
+    Scores based on three factors:
+    1. Activation similarity (cosine similarity of mean activations)
+    2. Dimensional compatibility (how similar are the layer shapes)
+    3. Architecture match (same arch = bonus)
+
+    Returns:
+        Dict with individual scores and overall mergeability (0-1)
+    """
+    print(f"[mergeability] Scoring {source_config.name}...")
+
+    scores = {}
+
+    # --- Factor 1: Activation similarity ---
+    cosine_sims = []
+    source_layers = sorted(source_activations.keys())
+    target_layers = sorted(target_activations.keys())
+
+    # Match layers by position (proportional mapping)
+    for i, tl in enumerate(target_layers):
+        # Map target layer index to source layer index
+        src_idx = int(i * len(source_layers) / len(target_layers))
+        src_idx = min(src_idx, len(source_layers) - 1)
+        sl = source_layers[src_idx]
+
+        if sl in source_activations and tl in target_activations:
+            s_mean = source_activations[sl].float().mean(dim=0)
+            t_mean = target_activations[tl].float().mean(dim=0)
+
+            # Pad to same dimension for cosine similarity
+            max_dim = max(s_mean.shape[0], t_mean.shape[0])
+            s_padded = torch.nn.functional.pad(s_mean, (0, max_dim - s_mean.shape[0]))
+            t_padded = torch.nn.functional.pad(t_mean, (0, max_dim - t_mean.shape[0]))
+
+            cos_sim = torch.nn.functional.cosine_similarity(
+                s_padded.unsqueeze(0), t_padded.unsqueeze(0)
+            ).item()
+            cosine_sims.append(cos_sim)
+
+    activation_score = np.mean(cosine_sims) if cosine_sims else 0.0
+    scores["activation_similarity"] = float(activation_score)
+
+    # --- Factor 2: Dimensional compatibility ---
+    layer_ratio = min(source_config.layers, 36) / max(source_config.layers, 36)
+    hidden_ratio = min(source_config.hidden_dim, 4096) / max(source_config.hidden_dim, 4096)
+    dim_score = (layer_ratio + hidden_ratio) / 2
+    scores["dimensional_compatibility"] = float(dim_score)
+
+    # --- Factor 3: Architecture match ---
+    arch_scores = {
+        "transformer": 1.0,       # Same as Qwen3
+        "transformer+mtp": 0.8,   # Close, just drop extras
+        "hybrid_ssm": 0.5,        # Very different
+    }
+    arch_score = arch_scores.get(source_config.architecture, 0.3)
+    scores["architecture_match"] = float(arch_score)
+
+    # --- Factor 4: Vocab overlap (bonus) ---
+    vocab_score = source_config.vocab_overlap_with_qwen3
+    scores["vocab_overlap"] = float(vocab_score)
+
+    # --- Overall: weighted average ---
+    overall = (
+        0.35 * activation_score +      # Most important — actual representation similarity
+        0.25 * dim_score +              # Shape compatibility
+        0.25 * arch_score +             # Architecture type
+        0.15 * vocab_score              # Vocab overlap
+    )
+    scores["overall"] = float(overall)
+
+    # --- Recommendation ---
+    if overall >= 0.7:
+        recommendation = "GO — standard T&M merge"
+    elif overall >= 0.5:
+        recommendation = "CAUTION — T&M merge with higher protection, have Theseus fallback ready"
+    elif overall >= 0.3:
+        recommendation = "RISKY — try Theseus first, distillation fallback"
+    else:
+        recommendation = "SKIP — use knowledge distillation instead"
+
+    scores["recommendation"] = recommendation
+
+    print(f"[mergeability] {source_config.name} score: {overall:.2f}")
+    print(f"  Activation similarity: {activation_score:.2f}")
+    print(f"  Dimensional compat:    {dim_score:.2f}")
+    print(f"  Architecture match:    {arch_score:.2f}")
+    print(f"  Vocab overlap:         {vocab_score:.2f}")
+    print(f"  → {recommendation}")
+
+    return scores

hugging/td_lang/engine/transport.py

@@ -0,0 +1,853 @@
+"""
+Transport and Merge — Two-sided optimal transport with streaming Sinkhorn.
+
+Implements the actual Transport and Merge paper (arxiv 2602.05495) correctly:
+
+Paper equations implemented here:
+    - Eq 8:  Q matrices for pre-activation (Q_in) and post-activation (Q_out) features
+    - Eq 13: P_eff = sqrt(P_pre · P_post) — effective layer transport plan
+    - Eq 14: Masked fusion with binary top-k mask M^ℓ
+    - Appendix A.3.4: Log-domain streaming Sinkhorn (200 inner / 1000 outer iterations)
+    - Appendix A.5: Top-k=128 neuron selection
+
+Two-sided transport (Section 4.2):
+    For each layer pair (ℓ, m):
+    1. Compute Q_in from pre-activation features (what goes INTO the layer)
+    2. Compute Q_out from post-activation features (what comes OUT of the layer)
+    3. Derive P_pre and P_post at the layer level
+    4. Combine: P_eff[ℓ,m] = sqrt(P_pre[ℓ,m] · P_post[ℓ,m])
+
+Streaming Sinkhorn (Appendix A.3.4):
+    - Log-domain updates (never materialize full K = exp(-C/ε) matrix)
+    - Chunked computation for memory efficiency
+    - 200 fixed iterations for feature-level (inner) OT
+    - Up to 1000 iterations for layer-level (outer) OT
+    - ε = 0.1 for standard text, ε = 0.03 for math reasoning
+
+Verified against actual paper PDF (test_21 interview round).
+Grok scored 10/10, these implementations match Grok's citations.
+"""
+
+import sys
+import math
+import torch
+import numpy as np
+from pathlib import Path
+from typing import Optional, Tuple
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+
+from .config import MergeConfig, ModelConfig, TARGET
+
+
+# ============================================================================
+# SETUP
+# ============================================================================
+
+def setup_tm_repo(cfg: MergeConfig):
+    """Add official T&M repo to Python path so we can import their code."""
+    repo_path = Path(cfg.tm_repo_path)
+    core_path = repo_path / "core"
+
+    if not core_path.exists():
+        raise FileNotFoundError(
+            f"Official T&M repo not found at {repo_path}\n"
+            f"Please clone it:\n"
+            f"  git clone https://github.com/chenhangcuisg-code/"
+            f"Cross-Architecture-Merging-for-Large-Language-Models.git"
+        )
+
+    if str(core_path) not in sys.path:
+        sys.path.insert(0, str(core_path))
+        print(f"[transport] Added T&M core to path: {core_path}")
+
+
+# ============================================================================
+# CALIBRATION DATA (Paper Appendix B.1: 2000 samples)
+# ============================================================================
+
+def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
+    """
+    Load calibration data for activation extraction.
+
+    Paper Appendix B.1: "For each dataset, we randomly sample 2000 examples"
+    Mix: Pile general + neuralmagic Q&A = 2000 total samples.
+    """
+    print(f"[transport] Loading calibration data ({cfg.calibration_samples} samples)...")
+
+    samples = []
+
+    # --- Pile: general text (1200 samples) ---
+    try:
+        pile = load_dataset(
+            cfg.calibration_dataset_pile,
+            split="validation",
+            streaming=True,
+            trust_remote_code=True,
+        )
+        count = 0
+        target_pile = int(cfg.calibration_samples * 0.6)  # 60% from Pile
+        for example in pile:
+            if count >= target_pile:
+                break
+            text = example.get("text", "")
+            if len(text) > 100:
+                tokens = tokenizer(
+                    text,
+                    truncation=True,
+                    max_length=cfg.calibration_seq_len,
+                    return_tensors="pt",
+                )
+                samples.append(tokens)
+                count += 1
+        print(f"  Pile general: {count} samples")
+    except Exception as e:
+        print(f"  Warning: Pile failed: {e}")
+        print(f"  Falling back to neuralmagic only")
+
+    # --- neuralmagic: Q&A calibration (remaining) ---
+    remaining = cfg.calibration_samples - len(samples)
+    if remaining > 0:
+        try:
+            nm = load_dataset(
+                cfg.calibration_dataset_nm,
+                split="train",
+                trust_remote_code=True,
+            )
+            count = 0
+            for example in nm:
+                if count >= remaining:
+                    break
+                text = example.get("text", example.get("content", ""))
+                if len(str(text)) > 50:
+                    tokens = tokenizer(
+                        str(text),
+                        truncation=True,
+                        max_length=cfg.calibration_seq_len,
+                        return_tensors="pt",
+                    )
+                    samples.append(tokens)
+                    count += 1
+            print(f"  neuralmagic: {count} samples")
+        except Exception as e:
+            print(f"  Warning: neuralmagic failed: {e}")
+
+    print(f"[transport] Total calibration samples: {len(samples)}")
+    return samples
+
+
+# ============================================================================
+# ACTIVATION EXTRACTION (Paper: attention Q,K,V,O + MLP gate,up,down)
+# ============================================================================
+
+# Module types to hook into (paper extracts from these specific projections)
+ATTENTION_PROJECTIONS = ("q_proj", "k_proj", "v_proj", "o_proj")
+MLP_PROJECTIONS = ("gate_proj", "up_proj", "down_proj")
+ALL_PROJECTIONS = ATTENTION_PROJECTIONS + MLP_PROJECTIONS
+
+
+def extract_activations(
+    model: AutoModelForCausalLM,
+    calibration_data: list,
+    device: str = "cuda",
+) -> dict:
+    """
+    Extract pre-activation AND post-activation features from each projection module.
+
+    Paper Section 4.2: Two-sided transport requires both:
+    - Pre-activation features (input to each projection) → for Q_in
+    - Post-activation features (output of each projection) → for Q_out
+
+    Only hooks into attention projections (Q,K,V,O) and MLP projections
+    (gate, up, down). NOT every arbitrary layer — paper is specific about this.
+
+    Returns:
+        Dict with keys like:
+            "model.layers.0.self_attn.q_proj.pre" → [num_samples, input_dim]
+            "model.layers.0.self_attn.q_proj.post" → [num_samples, output_dim]
+    """
+    print(f"[transport] Extracting two-sided activations from {len(calibration_data)} samples...")
+
+    activations = {}
+    hooks = []
+
+    # Register hooks on attention and MLP projection modules only
+    for name, module in model.named_modules():
+        # Check if this is a projection module we care about
+        module_type = name.split(".")[-1] if "." in name else name
+        if module_type not in ALL_PROJECTIONS:
+            continue
+
+        # Skip vision encoder modules
+        if any(name.startswith(pfx) for pfx in ("visual", "merger")):
+            continue
+
+        def make_hook(layer_name):
+            def hook_fn(module, input_tensor, output):
+                # Pre-activation: input to this linear layer
+                pre = input_tensor[0] if isinstance(input_tensor, tuple) else input_tensor
+                # Post-activation: output of this linear layer
+                post = output[0] if isinstance(output, tuple) else output
+
+                pre_key = f"{layer_name}.pre"
+                post_key = f"{layer_name}.post"
+
+                if pre_key not in activations:
+                    activations[pre_key] = []
+                if post_key not in activations:
+                    activations[post_key] = []
+
+                # Mean pool over sequence length → [hidden_dim]
+                activations[pre_key].append(
+                    pre.detach().float().mean(dim=1).cpu()
+                )
+                activations[post_key].append(
+                    post.detach().float().mean(dim=1).cpu()
+                )
+            return hook_fn
+
+        h = module.register_forward_hook(make_hook(name))
+        hooks.append(h)
+
+    # Forward pass on calibration data
+    model.eval()
+    with torch.no_grad():
+        for i, tokens in enumerate(calibration_data):
+            inputs = {k: v.to(device) for k, v in tokens.items()}
+            try:
+                model(**inputs)
+            except Exception as e:
+                print(f"  Warning: Sample {i} failed: {e}")
+                continue
+
+            if (i + 1) % 200 == 0:
+                print(f"  Processed {i + 1}/{len(calibration_data)} samples")
+
+    # Remove hooks
+    for h in hooks:
+        h.remove()
+
+    # Stack activations: [num_samples, hidden_dim]
+    for key in activations:
+        activations[key] = torch.cat(activations[key], dim=0)
+
+    n_modules = len(activations) // 2  # pre + post per module
+    print(f"[transport] Extracted activations from {n_modules} projection modules (two-sided)")
+
+    return activations
+
+
+# ============================================================================
+# LOG-DOMAIN STREAMING SINKHORN (Paper Appendix A.3.4)
+# ============================================================================
+
+def _log_sinkhorn_streaming(
+    cost_matrix: np.ndarray,
+    reg: float = 0.1,
+    max_iter: int = 200,
+    chunk_size: int = 512,
+) -> np.ndarray:
+    """
+    Log-domain streaming Sinkhorn solver.
+
+    Paper Appendix A.3.4:
+    "We use a memory-efficient streaming Sinkhorn solver with fixed 200 iterations"
+
+    Log-domain means we work with log(K) = -C/ε instead of K = exp(-C/ε).
+    This prevents numerical overflow/underflow with large matrices.
+
+    Streaming means we process the cost matrix in chunks instead of
+    materializing the full kernel matrix K in memory.
+
+    Args:
+        cost_matrix: [n, m] cost matrix (correlation distance)
+        reg: Entropic regularisation ε (paper default 0.1)
+        max_iter: Number of Sinkhorn iterations (paper: 200 inner, 1000 outer)
+        chunk_size: Process this many rows/cols at a time for memory efficiency
+
+    Returns:
+        [n, m] transport plan matrix
+    """
+    n, m = cost_matrix.shape
+
+    # Log-domain: work with log potentials instead of scaling vectors
+    # This is numerically stable — no exp() overflow
+    log_u = np.zeros(n)  # Log of row scaling vector
+    log_v = np.zeros(m)  # Log of column scaling vector
+
+    # Uniform marginals (both sides sum to 1)
+    log_a = np.full(n, -np.log(n))  # log(1/n)
+    log_b = np.full(m, -np.log(m))  # log(1/m)
+
+    # Log kernel: log(K_ij) = -C_ij / ε
+    log_K = -cost_matrix / reg
+
+    for iteration in range(max_iter):
+        # --- Row update (streaming over chunks of columns) ---
+        # log_u = log_a - logsumexp(log_K + log_v, axis=1)
+        log_sum = np.full(n, -np.inf)
+        for j_start in range(0, m, chunk_size):
+            j_end = min(j_start + chunk_size, m)
+            chunk = log_K[:, j_start:j_end] + log_v[j_start:j_end]
+            chunk_max = np.maximum(log_sum, chunk.max(axis=1))
+            log_sum = chunk_max + np.log(
+                np.exp(log_sum - chunk_max) +
+                np.exp(chunk - chunk_max[:, None]).sum(axis=1)
+            )
+        log_u = log_a - log_sum
+
+        # --- Column update (streaming over chunks of rows) ---
+        # log_v = log_b - logsumexp(log_K.T + log_u, axis=1)
+        log_sum = np.full(m, -np.inf)
+        for i_start in range(0, n, chunk_size):
+            i_end = min(i_start + chunk_size, n)
+            chunk = log_K[i_start:i_end, :].T + log_u[i_start:i_end]
+            # chunk shape: [m, chunk_rows]
+            chunk_max = np.maximum(log_sum, chunk.max(axis=1))
+            log_sum = chunk_max + np.log(
+                np.exp(log_sum - chunk_max) +
+                np.exp(chunk - chunk_max[:, None]).sum(axis=1)
+            )
+        log_v = log_b - log_sum
+
+    # Recover transport plan: T_ij = exp(log_u_i + log_K_ij + log_v_j)
+    # Do this in chunks too to avoid materializing full matrix at once
+    T = np.zeros((n, m), dtype=np.float32)
+    for j_start in range(0, m, chunk_size):
+        j_end = min(j_start + chunk_size, m)
+        T[:, j_start:j_end] = np.exp(
+            log_u[:, None] + log_K[:, j_start:j_end] + log_v[j_start:j_end]
+        )
+
+    return T
+
+
+def _sinkhorn_basic(
+    cost_matrix: np.ndarray,
+    reg: float = 0.1,
+    max_iter: int = 200,
+) -> np.ndarray:
+    """
+    Basic (non-streaming) Sinkhorn for small matrices (e.g., layer-level P).
+
+    Used for the layer-level transport plan where matrices are small
+    (e.g., 36×32 for Qwen3→Llama layer mapping).
+    """
+    n, m = cost_matrix.shape
+    K = np.exp(-cost_matrix / reg)
+
+    u = np.ones(n) / n
+    v = np.ones(m) / m
+
+    for _ in range(max_iter):
+        u = (1.0 / n) / (K @ v + 1e-10)
+        v = (1.0 / m) / (K.T @ u + 1e-10)
+
+    T = np.diag(u) @ K @ np.diag(v)
+    return T
+
+
+# ============================================================================
+# TWO-SIDED TRANSPORT (Paper Section 4.2, Equations 8, 13)
+# ============================================================================
+
+def _correlation_distance(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
+    """
+    Compute correlation distance matrix between two sets of activation vectors.
+
+    cost[i, j] = 1 - pearson_correlation(X[:, i], Y[:, j])
+
+    X: [num_samples, dim_x] — activations from source
+    Y: [num_samples, dim_y] — activations from target
+    Returns: [dim_x, dim_y] cost matrix
+    """
+    # Standardise each neuron's activations across samples
+    X_norm = (X - X.mean(axis=0)) / (X.std(axis=0) + 1e-8)
+    Y_norm = (Y - Y.mean(axis=0)) / (Y.std(axis=0) + 1e-8)
+
+    # Pearson correlation between each pair of neurons
+    corr = X_norm.T @ Y_norm / X.shape[0]  # [dim_x, dim_y]
+
+    # Correlation distance
+    cost = 1.0 - corr
+    return cost.astype(np.float32)
+
+
+def _get_layer_index(module_name: str) -> Optional[int]:
+    """Extract layer index from a module name like 'model.layers.5.self_attn.q_proj'."""
+    parts = module_name.split(".")
+    for i, part in enumerate(parts):
+        if part == "layers" and i + 1 < len(parts):
+            try:
+                return int(parts[i + 1])
+            except ValueError:
+                pass
+    return None
+
+
+def _get_module_type(module_name: str) -> str:
+    """Extract module type from name like 'model.layers.5.self_attn.q_proj' → 'q_proj'."""
+    return module_name.split(".")[-1]
+
+
+def _group_activations_by_layer(
+    activations: dict,
+    side: str = "pre",
+) -> dict:
+    """
+    Group activation tensors by layer index.
+
+    Returns: {layer_idx: {module_type: activation_tensor}}
+    """
+    grouped = {}
+    suffix = f".{side}"
+    for key, tensor in activations.items():
+        if not key.endswith(suffix):
+            continue
+        # Remove the .pre/.post suffix to get module name
+        module_name = key[: -len(suffix)]
+        layer_idx = _get_layer_index(module_name)
+        module_type = _get_module_type(module_name)
+        if layer_idx is not None:
+            if layer_idx not in grouped:
+                grouped[layer_idx] = {}
+            grouped[layer_idx][module_type] = tensor.numpy()
+    return grouped
+
+
+def compute_transport_plans(
+    source_activations: dict,
+    target_activations: dict,
+    cfg: MergeConfig,
+) -> dict:
+    """
+    Compute two-sided optimal transport plans between source and target.
+
+    Paper Section 4.2 — Two-sided transport:
+    1. For each (source_layer, target_layer) pair and each projection type:
+       - Compute Q_in from pre-activation features (Eq 8 applied to inputs)
+       - Compute Q_out from post-activation features (Eq 8 applied to outputs)
+    2. Derive layer-level costs from Q_in and Q_out → P_pre and P_post
+    3. Combine: P_eff[ℓ,m] = sqrt(P_pre[ℓ,m] · P_post[ℓ,m])  (Eq 13)
+
+    Returns:
+        Dict with:
+            'P_eff': [n_target_layers, n_source_layers] effective transport plan
+            'Q_in': {(src_layer, tgt_layer, module_type): Q matrix} — input-side neuron plans
+            'Q_out': {(src_layer, tgt_layer, module_type): Q matrix} — output-side neuron plans
+            'source_layers': sorted list of source layer indices
+            'target_layers': sorted list of target layer indices
+    """
+    print("[transport] Computing two-sided transport plans (paper Section 4.2)...")
+
+    # Group activations by layer
+    source_pre = _group_activations_by_layer(source_activations, "pre")
+    source_post = _group_activations_by_layer(source_activations, "post")
+    target_pre = _group_activations_by_layer(target_activations, "pre")
+    target_post = _group_activations_by_layer(target_activations, "post")
+
+    source_layers = sorted(source_pre.keys())
+    target_layers = sorted(target_pre.keys())
+
+    n_source = len(source_layers)
+    n_target = len(target_layers)
+
+    print(f"  Source layers: {n_source}, Target layers: {n_target}")
+
+    # --- Step 1: Compute Q_in and Q_out for each layer pair ---
+    Q_in_matrices = {}
+    Q_out_matrices = {}
+    layer_costs_pre = np.zeros((n_target, n_source))
+    layer_costs_post = np.zeros((n_target, n_source))
+
+    for ti, tl in enumerate(target_layers):
+        for si, sl in enumerate(source_layers):
+            # Get all projection types that exist in both
+            if tl not in target_pre or sl not in source_pre:
+                continue
+
+            target_modules = set(target_pre.get(tl, {}).keys())
+            source_modules = set(source_pre.get(sl, {}).keys())
+            common_modules = target_modules & source_modules
+
+            if not common_modules:
+                continue
+
+            pre_costs = []
+            post_costs = []
+
+            for mod_type in common_modules:
+                # --- Q_in: pre-activation (input-side) transport ---
+                if (sl in source_pre and mod_type in source_pre[sl] and
+                        tl in target_pre and mod_type in target_pre[tl]):
+                    S_pre = source_pre[sl][mod_type]
+                    T_pre = target_pre[tl][mod_type]
+                    cost_pre = _correlation_distance(S_pre, T_pre)
+
+                    # Use streaming Sinkhorn for large matrices, basic for small
+                    if max(cost_pre.shape) > 1024:
+                        Q = _log_sinkhorn_streaming(
+                            cost_pre,
+                            reg=cfg.sinkhorn_reg,
+                            max_iter=cfg.sinkhorn_inner_iter,
+                        )
+                    else:
+                        Q = _sinkhorn_basic(
+                            cost_pre,
+                            reg=cfg.sinkhorn_reg,
+                            max_iter=cfg.sinkhorn_inner_iter,
+                        )
+                    Q_in_matrices[(sl, tl, mod_type)] = Q
+                    pre_costs.append(cost_pre.mean())
+
+                # --- Q_out: post-activation (output-side) transport ---
+                if (sl in source_post and mod_type in source_post[sl] and
+                        tl in target_post and mod_type in target_post[tl]):
+                    S_post = source_post[sl][mod_type]
+                    T_post = target_post[tl][mod_type]
+                    cost_post = _correlation_distance(S_post, T_post)
+
+                    if max(cost_post.shape) > 1024:
+                        Q = _log_sinkhorn_streaming(
+                            cost_post,
+                            reg=cfg.sinkhorn_reg,
+                            max_iter=cfg.sinkhorn_inner_iter,
+                        )
+                    else:
+                        Q = _sinkhorn_basic(
+                            cost_post,
+                            reg=cfg.sinkhorn_reg,
+                            max_iter=cfg.sinkhorn_inner_iter,
+                        )
+                    Q_out_matrices[(sl, tl, mod_type)] = Q
+                    post_costs.append(cost_post.mean())
+
+            # Average cost across projection types for this layer pair
+            if pre_costs:
+                layer_costs_pre[ti, si] = np.mean(pre_costs)
+            if post_costs:
+                layer_costs_post[ti, si] = np.mean(post_costs)
+
+        if (ti + 1) % 6 == 0:
+            print(f"  Layer pairs computed: {ti + 1}/{n_target} target layers done")
+
+    # --- Step 2: Layer-level transport plans P_pre and P_post ---
+    print("[transport] Computing layer-level transport plans (P_pre, P_post)...")
+
+    P_pre = _sinkhorn_basic(
+        layer_costs_pre,
+        reg=cfg.sinkhorn_layer_reg,
+        max_iter=cfg.sinkhorn_outer_iter,
+    )
+
+    P_post = _sinkhorn_basic(
+        layer_costs_post,
+        reg=cfg.sinkhorn_layer_reg,
+        max_iter=cfg.sinkhorn_outer_iter,
+    )
+
+    # --- Step 3: P_eff = sqrt(P_pre · P_post) — Equation 13 ---
+    P_eff = np.sqrt(P_pre * P_post + 1e-10)
+
+    # Normalise P_eff so each target layer's row sums to 1
+    row_sums = P_eff.sum(axis=1, keepdims=True)
+    P_eff = P_eff / (row_sums + 1e-10)
+
+    print(f"[transport] P_eff shape: {P_eff.shape}")
+    print(f"  P_eff range: [{P_eff.min():.4f}, {P_eff.max():.4f}]")
+
+    # --- Step 4: Transport sparsification (Appendix A.1) ---
+    # "top-k selection strategies at both neuron and transport matrix levels"
+    # Keep only the top-k strongest source layers per target layer
+    k_layers = min(3, n_source)  # Top-3 source layers per target layer
+    P_sparse = np.zeros_like(P_eff)
+    for i in range(n_target):
+        top_k_idx = np.argsort(P_eff[i])[-k_layers:]
+        P_sparse[i, top_k_idx] = P_eff[i, top_k_idx]
+    # Re-normalise
+    row_sums = P_sparse.sum(axis=1, keepdims=True)
+    P_sparse = P_sparse / (row_sums + 1e-10)
+
+    print(f"[transport] Sparsified P: keeping top-{k_layers} source layers per target")
+
+    return {
+        "P_eff": P_sparse,
+        "P_eff_dense": P_eff,  # Keep dense version for debugging
+        "Q_in": Q_in_matrices,
+        "Q_out": Q_out_matrices,
+        "source_layers": source_layers,
+        "target_layers": target_layers,
+        "layer_costs_pre": layer_costs_pre,
+        "layer_costs_post": layer_costs_post,
+    }
+
+
+# ============================================================================
+# TOP-K MASKED FUSION (Paper Eq 14, Appendix A.5: k=128)
+# ============================================================================
+
+def compute_neuron_importance(
+    activations: dict,
+    layer_idx: int,
+) -> dict:
+    """
+    Compute neuron importance scores for top-k selection.
+
+    Paper Appendix A.5: "choosing the neurons with the highest mean
+    activation magnitudes across the calibration set"
+
+    Returns: {module_type: importance_scores [hidden_dim]}
+    """
+    importance = {}
+    for key, tensor in activations.items():
+        if not key.endswith(".post"):
+            continue
+        module_name = key[:-5]  # Remove .post
+        idx = _get_layer_index(module_name)
+        mod_type = _get_module_type(module_name)
+        if idx == layer_idx:
+            # Mean activation magnitude across calibration samples
+            importance[mod_type] = tensor.abs().mean(dim=0).numpy()
+    return importance
+
+
+def compute_top_k_mask(
+    importance_scores: np.ndarray,
+    k: int = 128,
+) -> np.ndarray:
+    """
+    Create binary mask for top-k most important neurons.
+
+    Paper Appendix A.5: "we set the default number of neurons to k = 128"
+
+    Returns: boolean mask [hidden_dim] where True = selected for fusion
+    """
+    if k >= len(importance_scores):
+        return np.ones(len(importance_scores), dtype=bool)
+
+    threshold_idx = np.argsort(importance_scores)[-k:]
+    mask = np.zeros(len(importance_scores), dtype=bool)
+    mask[threshold_idx] = True
+    return mask
+
+
+def fuse_weights(
+    source_model: AutoModelForCausalLM,
+    target_model: AutoModelForCausalLM,
+    transport_plans: dict,
+    source_config: ModelConfig,
+    cfg: MergeConfig,
+    target_activations: dict = None,
+) -> AutoModelForCausalLM:
+    """
+    Fuse source weights into target using two-sided transport + top-k mask.
+
+    Paper Equation 14:
+    W_fused = W_target + α · M^ℓ ⊙ (Σ_m P_eff[ℓ,m] · Q_out · W_source · Q_in^T - W_target)
+
+    Where:
+    - α is the fusion coefficient (0.05-0.15)
+    - M^ℓ is the binary top-k mask (only k=128 neurons get fused)
+    - P_eff is the effective layer transport plan
+    - Q_out and Q_in are the neuron-level transport matrices
+    - The sum is over source layers m
+
+    Returns: Target model with fused weights
+    """
+    print(f"\n[transport] Fusing {source_config.name} -> target (two-sided + top-k={cfg.top_k_neurons})")
+    alpha = source_config.merge_alpha
+    print(f"  Alpha: {alpha} (paper range: 0.05-0.15)")
+
+    source_state = source_model.state_dict()
+    target_state = target_model.state_dict()
+
+    P_eff = transport_plans["P_eff"]
+    Q_in = transport_plans["Q_in"]
+    Q_out = transport_plans["Q_out"]
+    source_layers = transport_plans["source_layers"]
+    target_layers = transport_plans["target_layers"]
+
+    fused_count = 0
+    skipped_count = 0
+    masked_neurons = 0
+
+    for ti, tl in enumerate(target_layers):
+        # Get the transport weights for this target layer
+        layer_transport = P_eff[ti]  # [n_source]
+
+        # Find which source layers contribute significantly
+        active_sources = [(si, sl, layer_transport[si])
+                          for si, sl in enumerate(source_layers)
+                          if layer_transport[si] > 1e-6]
+
+        if not active_sources:
+            continue
+
+        # For each projection type in this target layer
+        for mod_type in ALL_PROJECTIONS:
+            target_key = _find_param_key(target_state, tl, mod_type, "weight")
+            if target_key is None:
+                continue
+
+            target_w = target_state[target_key].float()
+
+            # Compute the transported operator: Σ_m P_eff[ℓ,m] · Q_out · W_source · Q_in^T
+            transported = torch.zeros_like(target_w)
+            total_weight = 0.0
+
+            for si, sl, p_weight in active_sources:
+                source_key = _find_source_param_key(
+                    source_state, sl, mod_type, "weight", source_config
+                )
+                if source_key is None:
+                    continue
+
+                source_w = source_state[source_key].float()
+
+                # Get Q matrices for this layer pair
+                q_in_key = (sl, tl, mod_type)
+                q_out_key = (sl, tl, mod_type)
+
+                q_in = Q_in.get(q_in_key)
+                q_out = Q_out.get(q_out_key)
+
+                if q_in is not None and q_out is not None:
+                    # Transport: Q_out @ W_source @ Q_in^T
+                    q_in_t = torch.from_numpy(q_in).float()
+                    q_out_t = torch.from_numpy(q_out).float()
+
+                    # Handle dimension mismatches via transport plan
+                    try:
+                        # q_out: [target_out, source_out], W: [source_out, source_in], q_in: [target_in, source_in]
+                        # Result: [target_out, target_in]
+                        transported_w = q_out_t @ source_w.to("cpu") @ q_in_t.T
+                        transported += p_weight * transported_w.to(target_w.device)
+                        total_weight += p_weight
+                    except RuntimeError:
+                        # Dimension mismatch — skip this pair
+                        skipped_count += 1
+                        continue
+                else:
+                    # No Q matrices — direct mapping if shapes match
+                    if source_w.shape == target_w.shape:
+                        transported += p_weight * source_w.to(target_w.device)
+                        total_weight += p_weight
+
+            if total_weight < 1e-6:
+                skipped_count += 1
+                continue
+
+            # Normalise by total transport weight
+            transported = transported / total_weight
+
+            # --- Apply top-k mask (Equation 14) ---
+            # M^ℓ ⊙ (transported - W_target)
+            delta = transported - target_w
+
+            if target_activations is not None and cfg.top_k_neurons > 0:
+                importance = compute_neuron_importance(target_activations, tl)
+                if mod_type in importance:
+                    # Mask on output dimension (rows of weight matrix)
+                    mask = compute_top_k_mask(importance[mod_type], k=cfg.top_k_neurons)
+                    mask_tensor = torch.from_numpy(mask).to(target_w.device)
+
+                    # Apply mask: only fuse top-k neurons
+                    if delta.dim() == 2:
+                        # Weight matrix: mask rows (output neurons)
+                        mask_2d = mask_tensor.unsqueeze(1).expand_as(delta)
+                        delta = delta * mask_2d.float()
+                        masked_neurons += mask.sum()
+                    elif delta.dim() == 1:
+                        # Bias: mask directly
+                        delta = delta * mask_tensor.float()
+                        masked_neurons += mask.sum()
+
+            # Final fusion: W_target + α · masked_delta
+            fused_w = target_w + alpha * delta
+            target_state[target_key] = fused_w.to(target_state[target_key].dtype)
+            fused_count += 1
+
+    # --- Vision encoder protection ---
+    # Restore any vision params that might have been touched
+    original_state = target_model.state_dict()
+    for key in target_state:
+        if any(key.startswith(pfx) for pfx in cfg.vision_skip_prefixes):
+            target_state[key] = original_state[key]
+
+    # --- Thinking mode protection ---
+    if cfg.freeze_think_tokens:
+        embed_key = "model.embed_tokens.weight"
+        if embed_key in target_state and embed_key in original_state:
+            for token_id in cfg.think_token_ids:
+                if token_id < target_state[embed_key].shape[0]:
+                    target_state[embed_key][token_id] = original_state[embed_key][token_id]
+                    print(f"  Protected think token {token_id}")
+
+    # Load fused weights
+    target_model.load_state_dict(target_state)
+    print(f"[transport] Fused {fused_count} params, skipped {skipped_count}")
+    print(f"  Top-k masked neurons fused: {masked_neurons}")
+
+    return target_model
+
+
+# ============================================================================
+# HELPER: Find parameter keys in state dicts
+# ============================================================================
+
+def _find_param_key(state_dict: dict, layer_idx: int, module_type: str, param_type: str = "weight") -> Optional[str]:
+    """Find the full parameter key for a given layer, module type, and param type."""
+    # Common patterns for transformer models
+    patterns = [
+        f"model.layers.{layer_idx}.self_attn.{module_type}.{param_type}",
+        f"model.layers.{layer_idx}.mlp.{module_type}.{param_type}",
+        f"transformer.h.{layer_idx}.attn.{module_type}.{param_type}",
+        f"transformer.h.{layer_idx}.mlp.{module_type}.{param_type}",
+    ]
+    for pattern in patterns:
+        if pattern in state_dict:
+            return pattern
+    return None
+
+
+def _find_source_param_key(
+    state_dict: dict,
+    source_layer: int,
+    module_type: str,
+    param_type: str,
+    source_config: ModelConfig,
+) -> Optional[str]:
+    """Find param key in source model, handling architecture differences."""
+    # Try standard patterns first
+    key = _find_param_key(state_dict, source_layer, module_type, param_type)
+    if key:
+        return key
+
+    # Try architecture-specific patterns
+    if source_config.architecture == "hybrid_ssm":
+        # Falcon uses different naming
+        patterns = [
+            f"model.layers.{source_layer}.attn.{module_type}.{param_type}",
+            f"model.layers.{source_layer}.feed_forward.{module_type}.{param_type}",
+        ]
+        for pattern in patterns:
+            if pattern in state_dict:
+                return pattern
+
+    return None
+
+
+def _should_skip(key: str, source_config: ModelConfig) -> bool:
+    """Determine if a parameter should be skipped during merge."""
+    if source_config.skip_embeddings and ("embed_tokens" in key or "lm_head" in key):
+        return True
+    if "drop_mtp_heads" in source_config.special_handling and "mtp_head" in key:
+        return True
+    if "drop_mamba_state_params" in source_config.special_handling:
+        mamba_keys = ["mamba", "A_log", "dt_proj", ".D"]
+        if any(mk in key for mk in mamba_keys):
+            return True
+    if "drop_qkv_bias" in source_config.special_handling and ".bias" in key:
+        if any(proj in key for proj in ["q_proj", "k_proj", "v_proj"]):
+            return True
+    return False

hugging/td_lang/engine/validate.py

@@ -0,0 +1,215 @@
+"""
+Post-Merge Validation — run after EVERY merge step.
+
+Tests:
+1. Canary recall (did knowledge transfer?)
+2. Perplexity check (did we break the model?)
+3. Thinking mode (do <think> tags still work?)
+4. Quick reasoning test (can it still think?)
+
+Kill criteria: >10% performance drop on any test → abort merge.
+Findings: #11, #22, #25
+"""
+
+import torch
+import math
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .canary import test_all_canaries
+from .config import MergeConfig
+
+
+def validate_merged_model(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    merged_sources: list[str],
+    cfg: MergeConfig,
+    baseline_perplexity: float = None,
+) -> dict:
+    """
+    Run full validation suite on a merged model.
+
+    Args:
+        model: The merged model to validate
+        tokenizer: The tokenizer
+        merged_sources: List of source models merged so far
+        cfg: Merge configuration
+        baseline_perplexity: Perplexity of the target model before merging
+
+    Returns:
+        Dict with test results and overall pass/fail
+    """
+    print("\n" + "=" * 60)
+    print(f"VALIDATION — After merging: {', '.join(merged_sources)}")
+    print("=" * 60)
+
+    results = {
+        "canary": None,
+        "perplexity": None,
+        "thinking_mode": None,
+        "reasoning": None,
+        "overall": False,
+    }
+
+    # --- Test 1: Canary recall ---
+    canary_results = test_all_canaries(model, tokenizer, merged_sources)
+    passed_canaries = sum(1 for v in canary_results.values() if v)
+    total_canaries = len(canary_results)
+    results["canary"] = {
+        "passed": passed_canaries,
+        "total": total_canaries,
+        "ok": passed_canaries >= cfg.canary_pass_threshold,
+        "details": canary_results,
+    }
+
+    # --- Test 2: Perplexity ---
+    perplexity = compute_perplexity(model, tokenizer)
+    ppl_ok = True
+    if baseline_perplexity is not None:
+        ratio = perplexity / baseline_perplexity
+        ppl_ok = ratio < cfg.perplexity_threshold
+        print(f"\n[validate] Perplexity: {perplexity:.2f} (baseline: {baseline_perplexity:.2f}, ratio: {ratio:.2f})")
+        if not ppl_ok:
+            print(f"[validate] ⚠ Perplexity ratio {ratio:.2f} exceeds threshold {cfg.perplexity_threshold}")
+    else:
+        print(f"\n[validate] Perplexity: {perplexity:.2f} (no baseline to compare)")
+    results["perplexity"] = {"value": perplexity, "ok": ppl_ok}
+
+    # --- Test 3: Thinking mode ---
+    think_ok = test_thinking_mode(model, tokenizer)
+    results["thinking_mode"] = {"ok": think_ok}
+
+    # --- Test 4: Quick reasoning ---
+    reason_ok = test_reasoning(model, tokenizer)
+    results["reasoning"] = {"ok": reason_ok}
+
+    # --- Overall verdict ---
+    all_ok = (
+        results["canary"]["ok"]
+        and results["perplexity"]["ok"]
+        and results["thinking_mode"]["ok"]
+        and results["reasoning"]["ok"]
+    )
+    results["overall"] = all_ok
+
+    # Summary
+    print("\n" + "-" * 60)
+    print("VALIDATION SUMMARY")
+    print("-" * 60)
+    print(f"  Canary recall:   {'✓' if results['canary']['ok'] else '✗'} ({passed_canaries}/{total_canaries})")
+    print(f"  Perplexity:      {'✓' if ppl_ok else '✗'} ({perplexity:.2f})")
+    print(f"  Thinking mode:   {'✓' if think_ok else '✗'}")
+    print(f"  Reasoning:       {'✓' if reason_ok else '✗'}")
+    print(f"  OVERALL:         {'✓ PASS' if all_ok else '✗ FAIL — consider aborting'}")
+    print("-" * 60)
+
+    return results
+
+
+def compute_perplexity(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    test_texts: list[str] = None,
+) -> float:
+    """
+    Compute perplexity on a small test set.
+
+    Lower perplexity = model is more confident about predicting text.
+    A big spike after merging means the model was damaged.
+    """
+    if test_texts is None:
+        test_texts = [
+            "The quick brown fox jumps over the lazy dog.",
+            "In mathematics, a prime number is a natural number greater than 1.",
+            "def fibonacci(n):\n    if n <= 1:\n        return n\n    return fibonacci(n-1) + fibonacci(n-2)",
+            "The theory of general relativity describes gravity as the curvature of spacetime.",
+            "To solve 3x + 7 = 22, subtract 7 from both sides to get 3x = 15, then divide by 3.",
+        ]
+
+    model.eval()
+    total_loss = 0.0
+    total_tokens = 0
+
+    for text in test_texts:
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = model(**inputs, labels=inputs["input_ids"])
+            total_loss += outputs.loss.item() * inputs["input_ids"].shape[1]
+            total_tokens += inputs["input_ids"].shape[1]
+
+    avg_loss = total_loss / total_tokens
+    perplexity = math.exp(avg_loss)
+    return perplexity
+
+
+def test_thinking_mode(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+) -> bool:
+    """
+    Test if the model still uses <think> tags for reasoning.
+
+    The thinking mode is Qwen3's special feature — if it's gone,
+    the merge damaged something critical.
+    """
+    prompt = "Solve step by step: What is 15 × 13?"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True,
+        )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+
+    # Check for thinking tags
+    has_think_open = "<think>" in response
+    has_think_close = "</think>" in response
+    passed = has_think_open and has_think_close
+
+    print(f"\n[validate] Thinking mode test:")
+    print(f"  Prompt:    {prompt}")
+    print(f"  Response:  {response[:200]}...")
+    print(f"  <think>:   {'✓ found' if has_think_open else '✗ missing'}")
+    print(f"  </think>:  {'✓ found' if has_think_close else '✗ missing'}")
+    print(f"  Status:    {'✓ PASS' if passed else '✗ FAIL'}")
+
+    return passed
+
+
+def test_reasoning(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+) -> bool:
+    """
+    Quick reasoning sanity check — can the model still do basic math?
+
+    This catches catastrophic failures where the merge produced gibberish.
+    """
+    prompt = "What is 7 + 8?"
+    expected_answer = "15"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=50,
+            temperature=0.1,
+            do_sample=False,
+        )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    passed = expected_answer in response
+
+    print(f"\n[validate] Quick reasoning test:")
+    print(f"  Prompt:   {prompt}")
+    print(f"  Expected: {expected_answer}")
+    print(f"  Got:      {response}")
+    print(f"  Status:   {'✓ PASS' if passed else '✗ FAIL'}")
+
+    return passed

hugging/td_lang/errors.py

@@ -88,6 +88,20 @@ COMMON_FIXES = {
     "fuse": 'Format: fuse ["model1", "model2"] into target [strategy equal]',
     "absorb": 'Format: absorb "model" into target [strength 0.5]',
     "schedule": 'Format: schedule "every 6h" { commands... } or schedule "at 02:00" { ... }',
+    "download": 'Format: download "dataset_name" as alias [split train]',
+    "log": 'Format: log "output.txt" (place before commands to capture output)',
+    "compare": 'Format: compare target vs "source_model" [questions 50] [-> output.json]',
+    "verify": 'Format: verify target on "dataset" [questions 100] [-> output.json]',
+    "vote": 'Format: vote target "question" [samples 5] [-> output.json]',
+    "prompt": 'Format: prompt target "Think step by step before answering."',
+    "distill": 'Format: distill target into "small_model" [steps 200] [-> output_dir]',
+    "rollback": "Format: rollback target (reverts to most recent snapshot)",
+    "curriculum": 'Format: curriculum target on "dataset" using grpo [levels 3] [steps 64]',
+    "star": 'Format: star target on "dataset" [rounds 3] [samples 8]',
+    "best_of": 'Format: best_of target on "dataset" [n 8] [steps 32]',
+    "exploit": 'Format: exploit target on "dataset" [samples 16] [steps 32] [-> output.jsonl]',
+    "arena": 'Format: arena target on "dataset" [rounds 5] [episodes 50] [steps 64] [curiosity 0.3] [-> log.json]',
+    "research_arena": 'Format: research_arena target topic "subject" [sources "pubmed"|"web"|"arxiv"] [rounds 5] [episodes 30] [-> log.json]',
 }
 
 

hugging/td_lang/examples/demo_arena.td

@@ -0,0 +1,28 @@
+# demo_arena.td — Real RL with memory, curiosity, and anti-lying
+#
+# This is ACTUAL reinforcement learning — the model explores challenges,
+# gets immediate reward/punishment, remembers what worked, and trains
+# on its experiences. Unlike best_of/star which just pick good examples,
+# arena makes the model LEARN FROM CONSEQUENCES.
+#
+# Features:
+#   - Memory bank: remembers what worked across all rounds
+#   - Curiosity bonus: rewarded for trying NEW approaches
+#   - Lying punishment: -2.0 for confident wrong answers (worst offence)
+#   - Cross-check: creative solutions verified against standard approach
+#
+# The model won't "forget the button makes the door safe" because
+# memory persists. And it won't lie because lying gets punished DOUBLE.
+
+load "Qwen/Qwen3-8B" as base
+
+# Run the arena: 3 rounds of 30 episodes each
+# Curiosity weight 0.3 = moderate exploration bonus
+arena base on "gsm8k" rounds 3 episodes 30 steps 32 curiosity 0.3 -> arena_log.json
+
+# After arena training, evaluate the result
+eval base -> arena_eval.json
+
+# Save the improved model
+snapshot base
+commit base

hugging/td_lang/examples/demo_intelligence.td

@@ -0,0 +1,35 @@
+# Demo: Phase 11 Intelligence — vote, prompt, distill, rollback
+# Shows all 4 new commands + the upgraded mega-diagnose
+
+load "Qwen/Qwen3-VL-8B-Instruct" as base
+
+# Attach a chain-of-thought prompt (makes it think step by step)
+prompt base "Think step by step before answering. Show your reasoning."
+
+# Mega diagnose: self-diagnosis + domain profiling + layer speed
+diagnose base -> diagnosis_report.json
+
+# Merge in reasoning
+merge "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B" into base using transport strength 0.5
+
+# Use majority voting on a hard question
+vote base "What is 847 * 23? Show your work." samples 5 -> vote_result.json
+
+# Snapshot before training (so rollback works)
+snapshot base
+
+# Train on weaknesses found by diagnose
+train base on "gsm8k" using grpo steps 64
+
+# Eval to check if training helped
+eval base -> eval_after.json
+
+# If training made things worse, undo it
+if eval_passed base {
+    commit base
+} else {
+    rollback base
+}
+
+# Create a fast student model for easy questions
+distill base into "Qwen/Qwen3-1.7B" steps 100 -> student_model/

hugging/td_lang/examples/demo_research_arena.td

@@ -0,0 +1,29 @@
+# demo_research_arena.td — Real RL on ANY topic using real-world sources
+#
+# This is the research gauntlet. The model gets thrown into a maze
+# built from REAL papers and knowledge. It has to navigate perfectly.
+#
+# How it works:
+#   1. Pulls real papers about your topic (PubMed, arXiv, web, or local files)
+#   2. Extracts verifiable facts from those papers
+#   3. Builds increasingly hard questions from the real knowledge
+#   4. Model must answer correctly — EVERY claim checked against sources
+#   5. Difficulty ESCALATES each round (stricter checking, harder questions)
+#   6. Memory persists — model remembers what it learned
+#   7. Lying = double punishment, curiosity = bonus
+#
+# The maze shrinks each round:
+#   Round 1: Easy questions, 30% strictness, full path width
+#   Round 2: Medium questions, 55% strictness, 75% path width
+#   Round 3: Hard questions, 80% strictness, 50% path width
+#   ...and so on. Miss a single fact = punishment.
+
+load "Qwen/Qwen3-8B" as base
+
+# Example 1: Medical research (uses PubMed for real papers)
+research_arena base topic "cancer immunotherapy mechanisms" sources "pubmed" rounds 4 episodes 25 steps 48 curiosity 0.3 difficulty_scale 0.25 -> research_log.json
+
+# After the gauntlet, see how the model performs
+eval base -> post_research_eval.json
+snapshot base
+commit base

hugging/td_lang/examples/demo_rl.td

@@ -0,0 +1,31 @@
+# Demo: Phase 12 RL & Fine-Tuning — curriculum, star, best_of, exploit
+# Shows all 4 new training methods + reward_contract wiring
+
+load "Qwen/Qwen3-VL-8B-Instruct" as base
+
+# Define what counts as "correct" (these verifiers wire into GRPO training)
+reward_contract {
+    verifiers = [code_compiles, math_correct, no_hallucination]
+    min_reward = 0.3
+}
+
+# Step 1: Curriculum training — start easy, get harder
+curriculum base on "gsm8k" using grpo levels 3 steps 64
+
+# Step 2: STaR — learn from own correct reasoning chains
+star base on "gsm8k" rounds 3 samples 8
+
+# Step 3: Best-of-N — generate 8 answers per question, train on the best
+best_of base on "openai/humaneval" n 8 steps 32
+
+# Step 4: EXPLOIT — controlled reward hacking
+# Generate 16 diverse solutions per problem, keep ALL correct ones
+# Even ugly shortcuts — if the answer is right, the method is valid
+exploit base on "gsm8k" samples 16 steps 32 -> exploit_results.jsonl
+
+# Verify the model actually got smarter
+eval base -> eval_after_rl.json
+
+# Save if good
+snapshot base
+commit base

hugging/td_lang/examples/demo_toolbox.td

@@ -0,0 +1,24 @@
+# Demo: Phase 10 Toolbox — download, log, compare, verify
+# Shows all 4 new commands working together
+
+log "toolbox_run.txt"
+
+load "Qwen/Qwen3-VL-8B-Instruct" as base
+
+# Download a dataset for verification
+download "gsm8k" as math_data
+download "openai/humaneval" as code_data split test
+
+# Merge in reasoning ability
+merge "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B" into base using transport strength 0.5
+
+# Compare: does the merged model remember what DeepSeek knew?
+compare base vs "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B" questions 30 -> compare_results.json
+
+# Verify: are the answers actually correct?
+verify base on "gsm8k" questions 50 -> verify_math.json
+verify base on "openai/humaneval" questions 25 -> verify_code.json
+
+# Eval and commit if good
+eval base -> eval_report.json
+commit base

hugging/td_lang/grammar.py

@@ -15,6 +15,7 @@ from .ast_nodes import (
     DataContractBlock,
     DebateCmd,
     DiagnoseCmd,
+    DistillCmd,
     EditCmd,
     EvalCmd,
     FuseCmd,
@@ -26,13 +27,26 @@ from .ast_nodes import (
     MergeCmd,
     NotifyCmd,
     OnErrorBlock,
+    PromptBlock,
     PruneCmd,
     RepeatBlock,
     ReportCmd,
     ResetCmd,
     RewardContractBlock,
+    RollbackCmd,
+    CurriculumCmd,
+    StarCmd,
+    BestOfCmd,
+    ExploitCmd,
+    ArenaCmd,
+    ResearchArenaCmd,
     SaveCmd,
     ScheduleCmd,
+    DownloadCmd,
+    LogBlock,
+    CompareCmd,
+    VerifyCmd,
+    VoteCmd,
     SetupBlock,
     SnapshotCmd,
     SynthCmd,
@@ -80,6 +94,20 @@ TD_GRAMMAR = r"""
               | setup_block
               | on_error_block
               | schedule_cmd
+              | download_cmd
+              | log_block
+              | compare_cmd
+              | verify_cmd
+              | vote_cmd
+              | prompt_cmd
+              | distill_cmd
+              | rollback_cmd
+              | curriculum_cmd
+              | star_cmd
+              | best_of_cmd
+              | exploit_cmd
+              | arena_cmd
+              | research_arena_cmd
 
     // ======================== PHASE 1 COMMANDS ========================
 
@@ -153,7 +181,11 @@ TD_GRAMMAR = r"""
               | fork_cmd | reset_cmd | prune_cmd | edit_cmd
               | fuse_cmd | absorb_cmd | snapshot_cmd | report_cmd
               | notify_cmd | save_cmd
-              | repeat_block_cmd | if_block_cmd | schedule_cmd) _NL*
+              | repeat_block_cmd | if_block_cmd | schedule_cmd
+              | download_cmd | compare_cmd | verify_cmd
+              | vote_cmd | prompt_cmd | distill_cmd | rollback_cmd
+              | curriculum_cmd | star_cmd | best_of_cmd | exploit_cmd
+              | arena_cmd | research_arena_cmd) _NL*
 
     // ======================== PHASE 6 — EASY MERGE COMMANDS ========================
 
@@ -233,6 +265,87 @@ TD_GRAMMAR = r"""
     // schedule "after 30m" { commands... }
     schedule_cmd: "schedule" string "{" _NL* body_cmd+ _NL* "}"
 
+    // ======================== PHASE 10 - TOOLBOX ========================
+
+    // download "gsm8k" as math_data [split train]
+    download_cmd: "download" string "as" IDENT (download_split)?
+    download_split: "split" IDENT
+
+    // log "training_log.txt"
+    log_block: "log" string
+
+    // compare target vs "source_model" [questions 50] [-> output.json]
+    compare_cmd: "compare" IDENT "vs" string (compare_questions)? (compare_output)?
+    compare_questions: "questions" INT
+    compare_output: "->" FILEPATH
+
+    // verify target on "dataset" [questions 100] [-> results.json]
+    verify_cmd: "verify" IDENT "on" string (verify_questions)? (verify_output)?
+    verify_questions: "questions" INT
+    verify_output: "->" FILEPATH
+
+    // ======================== PHASE 11 - INTELLIGENCE ========================
+
+    // vote target "question" [samples 5] [-> output.json]
+    vote_cmd: "vote" IDENT string (vote_samples)? (vote_output)?
+    vote_samples: "samples" INT
+    vote_output: "->" FILEPATH
+
+    // prompt target "system prompt text"
+    prompt_cmd: "prompt" IDENT string
+
+    // distill target into "small_model" [steps 200] [-> output_dir]
+    distill_cmd: "distill" IDENT "into" string (distill_steps)? (distill_output)?
+    distill_steps: "steps" INT
+    distill_output: "->" FILEPATH
+
+    // rollback target
+    rollback_cmd: "rollback" IDENT
+
+    // ======================== PHASE 12 - RL & FINE-TUNING ========================
+
+    // curriculum target on "dataset" using method [levels 3] [steps 64]
+    curriculum_cmd: "curriculum" IDENT "on" string "using" IDENT (curriculum_opt)*
+    curriculum_opt: "levels" INT -> curriculum_levels
+                  | "steps" INT -> curriculum_steps
+
+    // star target on "dataset" [rounds 3] [samples 8]
+    star_cmd: "star" IDENT "on" string (star_opt)*
+    star_opt: "rounds" INT -> star_rounds
+            | "samples" INT -> star_samples
+
+    // best_of target on "dataset" [n 8] [steps 32]
+    best_of_cmd: "best_of" IDENT "on" string (best_of_opt)*
+    best_of_opt: "n" INT -> best_of_n
+              | "steps" INT -> best_of_steps
+
+    // exploit target on "dataset" [samples 16] [steps 32] [-> output.jsonl]
+    exploit_cmd: "exploit" IDENT "on" string (exploit_opt)*
+    exploit_opt: "samples" INT -> exploit_samples
+              | "steps" INT -> exploit_steps
+              | "->" FILEPATH -> exploit_output
+
+    // ======================== PHASE 13 - REAL RL (ARENA) ========================
+
+    // arena target on "dataset" [rounds 5] [episodes 50] [steps 64] [curiosity 0.3] [-> log.json]
+    arena_cmd: "arena" IDENT "on" string (arena_opt)*
+    arena_opt: "rounds" INT -> arena_rounds
+             | "episodes" INT -> arena_episodes
+             | "steps" INT -> arena_steps
+             | "curiosity" NUMBER -> arena_curiosity
+             | "->" FILEPATH -> arena_output
+
+    // research_arena target topic "subject" [sources "web"|"pubmed"|"arxiv"|path]
+    //   [rounds 5] [episodes 30] [steps 64] [curiosity 0.3] [difficulty_scale 0.25] [-> log.json]
+    research_arena_cmd: "research_arena" IDENT "topic" string (ra_opt)*
+    ra_opt: "sources" string -> ra_sources
+          | "rounds" INT -> ra_rounds
+          | "episodes" INT -> ra_episodes
+          | "steps" INT -> ra_steps
+          | "curiosity" NUMBER -> ra_curiosity
+          | "difficulty_scale" NUMBER -> ra_difficulty
+          | "->" FILEPATH -> ra_output
+
     // ======================== SHARED RULES ========================
 
     // List of names: [name1, name2, name3]
@@ -468,6 +581,251 @@ class TDTransformer(Transformer):
     def schedule_cmd(self, timing: str, *body_cmds) -> ScheduleCmd:
         return ScheduleCmd(timing=timing, body=list(body_cmds))
 
+    # --- Phase 10: Toolbox ---
+
+    def download_cmd(self, dataset: str, alias: str, split: str | None = None) -> DownloadCmd:
+        cmd = DownloadCmd(dataset=dataset, alias=alias)
+        if isinstance(split, tuple) and split[0] == "split":
+            cmd.split = split[1]
+        elif isinstance(split, str):
+            cmd.split = split
+        return cmd
+
+    def download_split(self, value: str) -> tuple:
+        return ("split", value)
+
+    def log_block(self, filepath: str) -> LogBlock:
+        return LogBlock(filepath=filepath)
+
+    def compare_cmd(self, target: str, source: str, *opts) -> CompareCmd:
+        cmd = CompareCmd(target=target, source=source)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "questions":
+                    cmd.questions = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def compare_questions(self, value: int) -> tuple:
+        return ("questions", value)
+
+    def compare_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
+    def verify_cmd(self, target: str, dataset: str, *opts) -> VerifyCmd:
+        cmd = VerifyCmd(target=target, dataset=dataset)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "questions":
+                    cmd.questions = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def verify_questions(self, value: int) -> tuple:
+        return ("questions", value)
+
+    def verify_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
+    # --- Phase 11: Intelligence Commands ---
+
+    def vote_cmd(self, target: str, question: str, *opts) -> VoteCmd:
+        cmd = VoteCmd(target=target, question=question)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "samples":
+                    cmd.samples = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def vote_samples(self, value: int) -> tuple:
+        return ("samples", value)
+
+    def vote_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
+    def prompt_cmd(self, target: str, text: str) -> PromptBlock:
+        return PromptBlock(target=target, text=text)
+
+    def distill_cmd(self, teacher: str, student: str, *opts) -> DistillCmd:
+        cmd = DistillCmd(teacher=teacher, student=student)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "steps":
+                    cmd.steps = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def distill_steps(self, value: int) -> tuple:
+        return ("steps", value)
+
+    def distill_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
+    def rollback_cmd(self, target: str) -> RollbackCmd:
+        return RollbackCmd(target=target)
+
+    # --- Phase 12: RL & Fine-Tuning Commands ---
+
+    def curriculum_cmd(self, target: str, dataset: str, method: str, *opts) -> CurriculumCmd:
+        cmd = CurriculumCmd(target=target, dataset=dataset, method=method)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "levels":
+                    cmd.levels = val
+                elif key == "steps":
+                    cmd.steps = val
+        return cmd
+
+    def curriculum_levels(self, value: int) -> tuple:
+        return ("levels", value)
+
+    def curriculum_steps(self, value: int) -> tuple:
+        return ("steps", value)
+
+    def star_cmd(self, target: str, dataset: str, *opts) -> StarCmd:
+        cmd = StarCmd(target=target, dataset=dataset)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "rounds":
+                    cmd.rounds = val
+                elif key == "samples":
+                    cmd.samples = val
+        return cmd
+
+    def star_rounds(self, value: int) -> tuple:
+        return ("rounds", value)
+
+    def star_samples(self, value: int) -> tuple:
+        return ("samples", value)
+
+    def best_of_cmd(self, target: str, dataset: str, *opts) -> BestOfCmd:
+        cmd = BestOfCmd(target=target, dataset=dataset)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "n":
+                    cmd.n = val
+                elif key == "steps":
+                    cmd.steps = val
+        return cmd
+
+    def best_of_n(self, value: int) -> tuple:
+        return ("n", value)
+
+    def best_of_steps(self, value: int) -> tuple:
+        return ("steps", value)
+
+    def exploit_cmd(self, target: str, dataset: str, *opts) -> ExploitCmd:
+        cmd = ExploitCmd(target=target, dataset=dataset)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "samples":
+                    cmd.samples = val
+                elif key == "steps":
+                    cmd.steps = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def exploit_samples(self, value: int) -> tuple:
+        return ("samples", value)
+
+    def exploit_steps(self, value: int) -> tuple:
+        return ("steps", value)
+
+    def exploit_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
+    # --- Phase 13: Real RL (Arena) ---
+
+    def arena_cmd(self, target: str, dataset: str, *opts) -> ArenaCmd:
+        cmd = ArenaCmd(target=target, dataset=dataset)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "rounds":
+                    cmd.rounds = val
+                elif key == "episodes":
+                    cmd.episodes = val
+                elif key == "steps":
+                    cmd.steps = val
+                elif key == "curiosity":
+                    cmd.curiosity = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def arena_rounds(self, value: int) -> tuple:
+        return ("rounds", value)
+
+    def arena_episodes(self, value: int) -> tuple:
+        return ("episodes", value)
+
+    def arena_steps(self, value: int) -> tuple:
+        return ("steps", value)
+
+    def arena_curiosity(self, value: float) -> tuple:
+        return ("curiosity", value)
+
+    def arena_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
+    # --- Phase 13: Research Arena ---
+
+    def research_arena_cmd(self, target: str, topic: str, *opts) -> ResearchArenaCmd:
+        cmd = ResearchArenaCmd(target=target, topic=topic)
+        for opt in opts:
+            if isinstance(opt, tuple):
+                key, val = opt
+                if key == "sources":
+                    cmd.sources = val
+                elif key == "rounds":
+                    cmd.rounds = val
+                elif key == "episodes":
+                    cmd.episodes = val
+                elif key == "steps":
+                    cmd.steps = val
+                elif key == "curiosity":
+                    cmd.curiosity = val
+                elif key == "difficulty_scale":
+                    cmd.difficulty_scale = val
+                elif key == "output":
+                    cmd.output = val
+        return cmd
+
+    def ra_sources(self, value: str) -> tuple:
+        return ("sources", value)
+
+    def ra_rounds(self, value: int) -> tuple:
+        return ("rounds", value)
+
+    def ra_episodes(self, value: int) -> tuple:
+        return ("episodes", value)
+
+    def ra_steps(self, value: int) -> tuple:
+        return ("steps", value)
+
+    def ra_curiosity(self, value: float) -> tuple:
+        return ("curiosity", value)
+
+    def ra_difficulty(self, value: float) -> tuple:
+        return ("difficulty_scale", value)
+
+    def ra_output(self, filepath: str) -> tuple:
+        return ("output", filepath)
+
     # --- Phase 6: Easy Merge Commands ---
 
     def fuse_cmd(self, sources: list[str], target: str, *opts) -> FuseCmd:
@@ -688,6 +1046,8 @@ class TDTransformer(Transformer):
                 program.setup = item
             elif isinstance(item, OnErrorBlock):
                 program.on_error = item
+            elif isinstance(item, LogBlock):
+                program.log = item
             else:
                 program.commands.append(item)
         return program

hugging/td_lang/td_lang/__init__.py

@@ -0,0 +1,67 @@
+"""
+TD Lang — Domain-specific language for Time Dilation project.
+
+Compiles .td files into executable Python. Self-contained — no external deps.
+The merge/heal/validate engine (formerly td_fuse) lives in td_lang.engine/.
+
+Architecture:
+    td_lang/
+    ├── __init__.py          <- This file
+    ├── __main__.py          <- Entry point for python -m td_lang
+    ├── grammar.py           <- Lark grammar + parse tree transformer
+    ├── ast_nodes.py         <- Dataclass AST nodes for each command
+    ├── compiler.py          <- AST -> Python code generation
+    ├── executor.py          <- Run compiled code, track lineage
+    ├── cli.py               <- Command-line interface
+    ├── errors.py            <- Custom exceptions
+    ├── engine/              <- Merge/heal/validate runtime (was td_fuse)
+    │   ├── config.py        <- Model configs, merge order, hyperparameters
+    │   ├── merge.py         <- Sequential merge orchestrator
+    │   ├── heal.py          <- QLoRA healing fine-tune
+    │   ├── validate.py      <- Post-merge validation
+    │   ├── transport.py     <- Optimal transport wrapper
+    │   ├── techniques.py    <- ARM, OTMF, RAM, Theseus, Mergeability
+    │   └── canary.py        <- Canary injection + testing
+    └── examples/
+        ├── demo_merge.td    <- Basic merge example
+        ├── demo_heal.td     <- Merge + heal example
+        ├── demo_full.td     <- Full pipeline with gates + budget
+        └── ...              <- 22 example .td files
+
+Phase 1: load, merge, heal, eval, commit
+Phase 2: diagnose, synth, train, debate
+Phase 3: fork, reset, prune, edit
+Phase 4: snapshot, report, data_contract, reward_contract
+Phase 5: CLI polish, --version, info command, --verbose
+Phase 6: fuse, absorb (easy merge)
+Phase 7: repeat, if/else (loop control)
+Phase 8: setup, on_error, notify, save (autopilot)
+Phase 9: schedule (time-based execution)
+Phase 10: download, log, compare, verify (toolbox)
+Phase 11: vote, prompt, distill, rollback (intelligence)
+Phase 12: curriculum, star, best_of, exploit (RL & fine-tuning)
+Phase 13: arena (real RL with memory, curiosity, anti-lying, cross-check)
+Engine upgrades: QLoRA training, self-contained eval, model-generated synth problems
+Mega diagnose: self-diagnosis + domain profiling + layer speed testing
+
+Designed from interviews test_14 (10 commands) and test_17 (ForgeSpec 2.0).
+"""
+
+from .grammar import parse_td_file, parse_td_string  # noqa: F401
+from .compiler import compile_program  # noqa: F401
+from .executor import TDExecutor, check_td_file, compile_td_file, run_td_file  # noqa: F401
+
+__version__ = "0.2.0"
+__author__ = "Milan (TD Project)"
+
+__all__ = [
+    "parse_td_file",
+    "parse_td_string",
+    "compile_program",
+    "TDExecutor",
+    "check_td_file",
+    "compile_td_file",
+    "run_td_file",
+    "__version__",
+    "__author__",
+]

hugging/td_lang/td_lang/__main__.py

@@ -0,0 +1,5 @@
+"""Entry point for python -m td_lang."""
+
+from .cli import main
+
+main()

hugging/td_lang/td_lang/__pycache__/compiler.cpython-310.pyc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61328de4293774f4fc0e899e0eec00b64338be0dcf0fd3e68feaeaefc4c1edd5
+size 193126

hugging/td_lang/td_lang/__pycache__/compiler.cpython-314.pyc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bef7388fef05cdd8ee4edcc72a4b8907c8637caa22cfc802da044470a515c92
+size 162778

hugging/td_lang/td_lang/ast_nodes.py

@@ -0,0 +1,683 @@
+"""
+TD Lang AST Nodes — Dataclass containers for each parsed command.
+
+Each .td command becomes one of these nodes after parsing.
+Phase 1 nodes are compiled into runnable Python; Phase 2 nodes are stubs so
+the compiler can reject them with a clear error until they are implemented.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+
+# ============================================================================
+# PHASE 1 COMMANDS
+# ============================================================================
+
+@dataclass
+class LoadCmd:
+    """Load a model and give it a name.
+
+    Example: load "Qwen/Qwen3-VL-8B-Instruct" as base
+    """
+    model_ref: str          # HuggingFace path or local path
+    alias: str              # Name to use in the rest of the script
+
+
+@dataclass
+class MergeCmd:
+    """Merge a source model into a target using a method.
+
+    Example: merge "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B" into base using transport strength 0.5
+    """
+    source: str             # Model path or alias to merge from
+    target: str             # Alias to merge into (must be loaded first)
+    method: str             # "transport", "slerp", "ties", "dare"
+    strength: float = 0.5   # 0.0 = keep target, 1.0 = keep source
+
+
+@dataclass
+class HealCmd:
+    """Run QLoRA healing fine-tune on a model.
+
+    Example: heal base lora_r 32 epochs 2
+    """
+    target: str             # Alias of model to heal
+    lora_r: int = 32        # LoRA rank (higher = more capacity)
+    epochs: int = 2         # Training epochs
+
+
+@dataclass
+class EvalCmd:
+    """Run validation/evaluation on a model.
+
+    Example: eval base on "pile_sample" -> report.json
+    """
+    target: str                         # Alias of model to evaluate
+    dataset: Optional[str] = None       # Optional dataset name/path
+    output: Optional[str] = None        # Optional output file path
+
+
+@dataclass
+class CommitCmd:
+    """Save model checkpoint, optionally requiring gates to pass.
+
+    Example: commit base if [canary, perplexity, thinking_mode]
+    """
+    target: str                                 # Alias of model to commit
+    gates: Optional[list[str]] = None           # Gate names that must pass
+
+
+# ============================================================================
+# PHASE 2 COMMANDS (placeholders — structure ready, not wired up yet)
+# ============================================================================
+
+@dataclass
+class SynthCmd:
+    """Generate synthetic training data from a model. (Phase 2)"""
+    target: str
+    source: str
+    filter_method: Optional[str] = None
+    output: Optional[str] = None
+
+
+@dataclass
+class TrainCmd:
+    """Train a model on a dataset. (Phase 2)"""
+    target: str
+    dataset: str
+    method: str = "grpo"            # "grpo", "sft", "dpo"
+    steps: Optional[int] = None
+    learning_rate: Optional[float] = None
+
+
+@dataclass
+class DebateCmd:
+    """Generate multi-answer debate for preference pairs. (Phase 2)"""
+    target: str
+    rounds: int = 3
+    candidates: int = 8
+    output: Optional[str] = None
+
+
+@dataclass
+class DiagnoseCmd:
+    """Ask model what it's bad at — self-diagnosis. (Phase 2)"""
+    target: str
+    output: Optional[str] = None
+
+
+@dataclass
+class ForkCmd:
+    """Branch current model weights for parallel experiments. (Phase 3)
+
+    Example: fork base as experiment_v2
+    Cheap fork: copies manifest + adapters, shares base weights (default).
+    """
+    source: str             # Alias of model to fork from
+    alias: str              # Name for the new branch
+
+
+@dataclass
+class ResetCmd:
+    """Revert model to a previous checkpoint. (Phase 3)
+
+    Example: reset base to "checkpoint_042"
+    Deletes current model, clears CUDA cache, reloads from disk.
+    Must also reset optimizer state.
+    """
+    target: str             # Alias of model to reset
+    checkpoint: str         # Checkpoint name/path to revert to
+
+
+@dataclass
+class PruneCmd:
+    """Structural pruning — remove low-utility neurons/heads. (Phase 3)
+
+    Example: prune base using wanda aggressiveness 0.2
+    Safe zone: ~20% max (LLM-Pruner paper). Language backbone only.
+    """
+    target: str
+    method: str = "wanda"               # "wanda", "magnitude", "taylor"
+    aggressiveness: float = 0.2         # Fraction to remove (0.0-1.0)
+
+
+@dataclass
+class EditCmd:
+    """Surgical LoRA/DoRA editing on specific layers. (Phase 3)
+
+    Example: edit base layers 16-28 using lora lr 1e-4
+    "Try before buy": eval with adapter enabled vs disabled before merging.
+    """
+    target: str
+    layers: str = "all"                 # "all", "16-28", single number
+    method: str = "lora"                # "lora" or "dora"
+    learning_rate: Optional[float] = None
+
+
+# ============================================================================
+# PHASE 4 COMMANDS — Contracts, Lineage, Economics (ForgeSpec 2.0, test_17)
+# ============================================================================
+
+# ============================================================================
+# PHASE 7 — LOOP CONTROL (repeat, if/else)
+# ============================================================================
+
+@dataclass
+class RepeatBlock:
+    """Repeat a block of commands N times. (Phase 7 — Loop Control)
+
+    Example:
+        repeat 5 {
+            diagnose base
+            synth base from base
+            train base on "data.jsonl" using grpo steps 64
+            eval base
+        }
+    """
+    count: int                      # Number of iterations
+    body: List[Any] = field(default_factory=list)  # Commands inside the block
+
+
+@dataclass
+class IfBlock:
+    """Conditional execution based on last eval result. (Phase 7 — Loop Control)
+
+    Example:
+        if eval_passed {
+            commit base
+        } else {
+            reset base to "last_good"
+        }
+
+    Condition checks the most recent eval result for the target.
+    """
+    condition: str                  # "eval_passed", "gate_passed", etc.
+    target: Optional[str] = None    # Which model's eval to check
+    then_body: List[Any] = field(default_factory=list)
+    else_body: List[Any] = field(default_factory=list)
+
+
+@dataclass
+class FuseCmd:
+    """Fuse multiple models into a target in one shot. (Phase 6 — Easy Merge)
+
+    Example: fuse [deepseek-r1, mimo-7b, llama-3.1] into base
+    Auto-picks Transport and Merge, auto-sets per-model strength.
+    Handles cross-architecture merging (all 5 source models have different archs).
+    """
+    sources: list[str]          # List of model names/paths to fuse in
+    target: str                 # Alias to merge into (must be loaded)
+    method: str = "transport"   # Default: transport and merge (cross-arch)
+    strategy: str = "equal"     # "equal" (same strength each), "weighted", "sequential"
+
+
+@dataclass
+class AbsorbCmd:
+    """Absorb a single model into target — simplified merge. (Phase 6 — Easy Merge)
+
+    Example: absorb "deepseek-ai/DeepSeek-R1" into base strength 0.5
+    One-liner for the common case of merging one model in.
+    """
+    source: str                 # Model path or HF ID
+    target: str                 # Alias to merge into
+    strength: float = 0.5       # 0.0=keep target, 1.0=keep source, default balanced
+
+
+@dataclass
+class SnapshotCmd:
+    """Save a content-hashed snapshot of model state for lineage tracking. (Phase 4)
+
+    Example: snapshot base -> snapshots/
+    Creates a content-addressed directory: snapshots/<sha256_prefix>/
+    Contains: model state, adapter state, prune spec, eval report, manifest.
+    """
+    target: str
+    output: Optional[str] = None  # Output directory (default: td_lang_outputs/snapshots/)
+
+
+@dataclass
+class ReportCmd:
+    """Generate an economics report for this run. (Phase 4)
+
+    Example: report -> economics.json
+    Tracks: GPU hours, cost estimate, tokens processed, experiments run,
+    time per command, cost breakdown by phase.
+    """
+    output: Optional[str] = None  # Output file path
+
+
+# ============================================================================
+# PHASE 8 — AUTOPILOT (setup, notify, save, on_error, resume)
+# ============================================================================
+
+@dataclass
+class NotifyCmd:
+    """Send a notification via ntfy.sh. (Phase 8 — Autopilot)
+
+    Example: notify "Training complete!"
+    Uses curl to POST to the configured ntfy topic.
+    """
+    message: str
+
+
+@dataclass
+class SaveCmd:
+    """Save/upload model to cloud storage via rclone. (Phase 8 — Autopilot)
+
+    Example: save base to "gdrive:TD/models/v1"
+    Uses rclone to copy model checkpoint to Google Drive (or any rclone remote).
+    """
+    target: str                     # Alias of model to save
+    destination: str                # rclone destination path
+
+
+@dataclass
+class SetupBlock:
+    """Auto-install dependencies and configure environment. (Phase 8 — Autopilot)
+
+    Example:
+        setup {
+            pip = [torch, transformers, peft, bitsandbytes, trl]
+            hf_token = env
+            notify = "ntfy.sh/my_ai"
+        }
+    """
+    pip_packages: list[str] = field(default_factory=list)
+    hf_token: Optional[str] = None   # "env" = read HF_TOKEN from env
+    notify_url: Optional[str] = None  # ntfy.sh topic URL
+
+
+@dataclass
+class OnErrorBlock:
+    """Crash recovery behavior. (Phase 8 — Autopilot)
+
+    Example:
+        on_error {
+            retry = 3
+            fallback = reduce_batch
+            notify = true
+        }
+    """
+    retry: int = 3                   # Number of retries per failed step
+    fallback: str = "reduce_batch"   # "reduce_batch", "skip", "snapshot_and_stop"
+    notify: bool = True              # Send ntfy notification on error
+
+
+# ============================================================================
+# PHASE 9 — SCHEDULE (time-based execution)
+# ============================================================================
+
+@dataclass
+class ScheduleCmd:
+    """Schedule a block of commands to run at a specific time or interval. (Phase 9)
+
+    Examples:
+        schedule "every 6h" { diagnose base; train base ... }
+        schedule "at 02:00" { train base on "data.jsonl" using grpo }
+        schedule "after 30m" { eval base -> results.json }
+
+    Patterns:
+        "every Nh/Nm" — repeat every N hours/minutes
+        "at HH:MM"    — run once at that time
+        "after Nh/Nm" — delay then run once
+    """
+    timing: str                     # "every 6h", "at 02:00", "after 30m"
+    body: List[Any] = field(default_factory=list)  # Commands inside the block
+
+
+# ============================================================================
+# PHASE 10 - TOOLBOX (download, log, compare, verify)
+# ============================================================================
+
+@dataclass
+class DownloadCmd:
+    """Download a dataset from HuggingFace. (Phase 10)
+
+    Example: download "gsm8k" as math_data
+    Pulls a dataset from HuggingFace and stores it for training/eval.
+    """
+    dataset: str                    # HuggingFace dataset path
+    alias: str                      # Name to reference it later
+    split: str = "train"            # Which split to download
+
+
+@dataclass
+class LogBlock:
+    """Save all pipeline output to a log file. (Phase 10)
+
+    Example: log "training_log.txt"
+    Everything printed to console also goes to this file.
+    """
+    filepath: str                   # Path to save log
+
+
+@dataclass
+class CompareCmd:
+    """Compare source model vs merged model - knowledge retention test. (Phase 10)
+
+    Example: compare base vs "deepseek-ai/DeepSeek-R1" questions 50
+    Tests both models on the same questions and shows what % the merged
+    model retained from the source. Proves the merge actually worked.
+    """
+    target: str                     # The merged model alias
+    source: str                     # Source model to compare against (HF path)
+    questions: int = 50             # Number of test questions
+    output: Optional[str] = None    # Optional output file
+
+
+@dataclass
+class VerifyCmd:
+    """Verify model answers are actually correct. (Phase 10)
+
+    Example: verify base on "gsm8k" questions 100 -> verify_results.json
+    Runs the model on questions with KNOWN correct answers and checks
+    if the model got them right. Returns accuracy percentage.
+    """
+    target: str                     # Model alias to test
+    dataset: str                    # Dataset with known answers
+    questions: int = 100            # Number of questions to test
+    output: Optional[str] = None    # Optional output file
+
+
+# ============================================================================
+# PHASE 11 - INTELLIGENCE (vote, prompt, distill, rollback)
+# ============================================================================
+
+@dataclass
+class VoteCmd:
+    """Majority voting - generate N answers, pick the one most agree on. (Phase 11)
+
+    Example: vote base "What is 15 * 23?" samples 5
+    Generates N answers to the same question, then picks the most common one.
+    Proven to boost accuracy 10-20% with zero training.
+    """
+    target: str                     # Model alias
+    question: str                   # Question to vote on
+    samples: int = 5               # Number of answers to generate
+    output: Optional[str] = None    # Optional output file
+
+
+@dataclass
+class PromptBlock:
+    """Attach a system prompt or chain-of-thought template to a model. (Phase 11)
+
+    Example:
+        prompt base "Think step by step before answering."
+    Makes the model use this system prompt for all future generations.
+    """
+    target: str                     # Model alias to attach prompt to
+    text: str                       # The system prompt text
+
+
+@dataclass
+class DistillCmd:
+    """Distill a big model's knowledge into a smaller one. (Phase 11)
+
+    Example: distill base into "Qwen/Qwen3-1.7B" steps 200 -> student_model/
+    Takes the big model's best answers and trains the small model on them.
+    You get a fast model for easy questions, full model for hard ones.
+    """
+    teacher: str                    # The big model alias (source of knowledge)
+    student: str                    # The small model HF path
+    steps: int = 200               # Training steps
+    output: Optional[str] = None    # Where to save the student model
+
+
+@dataclass
+class RollbackCmd:
+    """Undo the last training step. (Phase 11)
+
+    Example: rollback base
+    Reverts to the most recent snapshot. If training made things worse,
+    one command brings it back.
+    """
+    target: str                     # Model alias to rollback
+
+
+# ============================================================================
+# PHASE 12 - RL & FINE-TUNING (curriculum, star, best_of, exploit)
+# ============================================================================
+
+@dataclass
+class CurriculumCmd:
+    """Progressive difficulty training - start easy, get harder. (Phase 12)
+
+    Example: curriculum base on "gsm8k" using grpo levels 3 steps 64
+    Splits dataset by difficulty, trains on easy first, then medium, then hard.
+    Each level only starts when the model passes the previous one.
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset to train on
+    method: str = "grpo"            # Training method
+    levels: int = 3                 # Number of difficulty levels
+    steps: int = 64                 # Steps per level
+
+
+@dataclass
+class StarCmd:
+    """Self-Taught Reasoner - train on own correct reasoning chains. (Phase 12)
+
+    Example: star base on "gsm8k" rounds 3 samples 8
+    Generate N solutions per problem. Keep the ones with correct answers.
+    Train on the correct reasoning chains. Repeat.
+    The model literally learns from its own successes.
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset with known answers
+    rounds: int = 3                 # Number of STaR iterations
+    samples: int = 8               # Solutions to generate per problem
+
+
+@dataclass
+class BestOfCmd:
+    """Generate N answers, score all, train on the best. (Phase 12)
+
+    Example: best_of base on "gsm8k" n 8 steps 32
+    For each training problem: generate N answers, score them all,
+    keep only the best one, train on that. Like vote but for training.
+    80-90% of RLHF gains at 5-30% of the cost (test_16).
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset to train on
+    n: int = 8                      # How many answers to generate per problem
+    steps: int = 32                 # Training steps on the filtered data
+
+
+@dataclass
+class ExploitCmd:
+    """Controlled reward hacking - keep ALL correct solutions regardless of method. (Phase 12)
+
+    Example: exploit base on "gsm8k" samples 16 -> exploit_data.jsonl
+    Generate many diverse solutions (high temp). Only filter: is the answer correct?
+    Keep ugly solutions, shortcuts, weird reasoning - as long as the answer is right.
+    Train on the diverse set so the model learns multiple paths to correct answers.
+    The "hacks" often turn out to be genuinely clever shortcuts.
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset with verifiable answers
+    samples: int = 16              # Solutions per problem (higher = more diversity)
+    steps: int = 32                 # Training steps on the exploited data
+    output: Optional[str] = None    # Save the exploit data for inspection
+
+
+@dataclass
+class ArenaCmd:
+    """Real RL with environment, memory, curiosity, and anti-lying. (Phase 13)
+
+    The model enters an arena of challenges. For each challenge:
+    1. It tries to solve it (exploration)
+    2. Gets immediate reward/punishment (+1 correct, -1 wrong, -2 lying)
+    3. Remembers what worked and didn't (memory bank persists across episodes)
+    4. Gets curiosity bonus for trying NEW approaches
+    5. Creative solutions get cross-checked against standard approaches
+
+    Example: arena base on "gsm8k" rounds 5 episodes 50 steps 64 curiosity 0.3
+    """
+    target: str                     # Model alias
+    dataset: str                    # Dataset with verifiable answers
+    rounds: int = 5                # RL rounds (re-train after each)
+    episodes: int = 50             # Challenges per round
+    steps: int = 64                 # Training steps per round
+    curiosity: float = 0.3         # Curiosity bonus weight
+    output: Optional[str] = None    # Save arena log
+
+
+@dataclass
+class ResearchArenaCmd:
+    """Research arena — RL on ANY topic using real-world knowledge. (Phase 13)
+
+    Unlike arena (which uses a pre-made dataset), research_arena:
+    1. Takes a TOPIC string ("cancer biology", "number theory", anything)
+    2. Pulls real papers/sources about that topic (web, arxiv, pubmed, local files)
+    3. Extracts verifiable facts/claims from those sources
+    4. Builds increasingly hard questions from the real knowledge
+    5. Runs the model through the gauntlet, checking EVERY claim against sources
+    6. Difficulty ESCALATES on failure (fewer hints, stricter checking, harder questions)
+    7. Memory persists so it doesn't forget what it learned
+    8. Lying gets punished DOUBLE, curiosity rewarded
+
+    Example: research_arena base topic "cancer biology" sources "pubmed" rounds 5
+    """
+    target: str                     # Model alias
+    topic: str                      # Research topic (any field)
+    sources: str = "web"           # Where to pull knowledge: "web", "pubmed", "arxiv", or filepath
+    rounds: int = 5                # RL rounds (difficulty increases each round)
+    episodes: int = 30             # Questions per round
+    steps: int = 64                 # Training steps per round
+    curiosity: float = 0.3         # Curiosity bonus weight
+    difficulty_scale: float = 0.25 # How much harder each round gets (0.25 = 25% harder)
+    output: Optional[str] = None    # Save research log
+
+
+# ============================================================================
+# BLOCKS (gates, budget, contracts, etc.)
+# ============================================================================
+
+@dataclass
+class GateBlock:
+    """Validation gates that must pass before commit.
+
+    Example:
+        gate {
+            must_pass = [canary, perplexity, thinking_mode]
+        }
+    """
+    must_pass: list[str] = field(default_factory=list)
+
+
+@dataclass
+class BudgetBlock:
+    """Resource budget — compiler refuses plans that exceed limits.
+
+    Example:
+        budget {
+            max_gpu_hours = 8
+            max_cost = 50.00
+        }
+    """
+    max_gpu_hours: Optional[float] = None
+    max_cost: Optional[float] = None
+    max_tokens: Optional[int] = None
+    max_experiments: Optional[int] = None
+
+
+@dataclass
+class DataContractBlock:
+    """Schema enforcement on training data. (Phase 4, ForgeSpec 2.0)
+
+    Example:
+        data_contract {
+            required_fields = [prompt, response]
+            min_samples = 100
+            max_perplexity = 50.0
+        }
+
+    Compiler checks training data at synth/train time.
+    """
+    required_fields: list[str] = field(default_factory=list)
+    min_samples: Optional[int] = None
+    max_perplexity: Optional[float] = None
+
+
+@dataclass
+class RewardContractBlock:
+    """Verified reward definitions — what counts as "correct". (Phase 4, ForgeSpec 2.0)
+
+    Example:
+        reward_contract {
+            verifiers = [code_compiles, math_correct, no_hallucination]
+            min_reward = 0.3
+        }
+
+    Used by train (GRPO) to enforce reward quality.
+    No learned reward model — verified rewards only (test_16).
+    """
+    verifiers: list[str] = field(default_factory=list)
+    min_reward: Optional[float] = None
+
+
+# ============================================================================
+# TOP-LEVEL PROGRAM
+# ============================================================================
+
+@dataclass
+class TDProgram:
+    """A complete parsed .td file — commands in order plus global blocks."""
+
+    commands: List[Any] = field(default_factory=list)
+    gates: Optional[GateBlock] = None
+    budget: Optional[BudgetBlock] = None
+    data_contract: Optional[DataContractBlock] = None
+    reward_contract: Optional[RewardContractBlock] = None
+    setup: Optional[SetupBlock] = None
+    on_error: Optional[OnErrorBlock] = None
+    log: Optional[LogBlock] = None
+    source_file: Optional[str] = None
+
+
+__all__ = [
+    "LoadCmd",
+    "MergeCmd",
+    "HealCmd",
+    "EvalCmd",
+    "CommitCmd",
+    "SynthCmd",
+    "TrainCmd",
+    "DebateCmd",
+    "DiagnoseCmd",
+    "ForkCmd",
+    "ResetCmd",
+    "PruneCmd",
+    "EditCmd",
+    "RepeatBlock",
+    "IfBlock",
+    "FuseCmd",
+    "AbsorbCmd",
+    "SnapshotCmd",
+    "ReportCmd",
+    "NotifyCmd",
+    "SaveCmd",
+    "SetupBlock",
+    "OnErrorBlock",
+    "GateBlock",
+    "BudgetBlock",
+    "DataContractBlock",
+    "RewardContractBlock",
+    "ScheduleCmd",
+    "DownloadCmd",
+    "LogBlock",
+    "CompareCmd",
+    "VerifyCmd",
+    "VoteCmd",
+    "PromptBlock",
+    "DistillCmd",
+    "RollbackCmd",
+    "CurriculumCmd",
+    "StarCmd",
+    "BestOfCmd",
+    "ExploitCmd",
+    "ArenaCmd",
+    "ResearchArenaCmd",
+    "TDProgram",
+]

hugging/td_lang/td_lang/cli.py

@@ -0,0 +1,229 @@
+"""
+TD Lang CLI — Command-line interface for .td files.
+
+Usage:
+    python -m td_lang run examples/demo_merge.td       # Compile + execute
+    python -m td_lang compile examples/demo_merge.td   # Compile only (outputs .py)
+    python -m td_lang check examples/demo_merge.td     # Syntax check only
+    python -m td_lang info examples/demo_merge.td      # Show plan without compiling
+    python -m td_lang --version                        # Show version
+"""
+
+import argparse
+import sys
+
+from . import __version__
+from .executor import TDExecutor
+from .errors import TDLangError
+from .grammar import parse_td_file
+from .ast_nodes import (
+    LoadCmd, MergeCmd, HealCmd, EvalCmd, CommitCmd,
+    SynthCmd, TrainCmd, DebateCmd, DiagnoseCmd,
+    ForkCmd, ResetCmd, PruneCmd, EditCmd,
+    FuseCmd, AbsorbCmd, RepeatBlock, IfBlock,
+    NotifyCmd, SaveCmd, ScheduleCmd,
+    DownloadCmd, LogBlock, CompareCmd, VerifyCmd,
+    VoteCmd, PromptBlock, DistillCmd, RollbackCmd,
+    CurriculumCmd, StarCmd, BestOfCmd, ExploitCmd, ArenaCmd, ResearchArenaCmd,
+    SnapshotCmd, ReportCmd,
+)
+
+
+# Phase labels for info command
+_PHASE_MAP = {
+    LoadCmd: ("1", "load"),
+    MergeCmd: ("1", "merge"),
+    HealCmd: ("1", "heal"),
+    EvalCmd: ("1", "eval"),
+    CommitCmd: ("1", "commit"),
+    SynthCmd: ("2", "synth"),
+    TrainCmd: ("2", "train"),
+    DebateCmd: ("2", "debate"),
+    DiagnoseCmd: ("2", "diagnose"),
+    ForkCmd: ("3", "fork"),
+    ResetCmd: ("3", "reset"),
+    PruneCmd: ("3", "prune"),
+    EditCmd: ("3", "edit"),
+    FuseCmd: ("6", "fuse"),
+    AbsorbCmd: ("6", "absorb"),
+    RepeatBlock: ("7", "repeat"),
+    IfBlock: ("7", "if"),
+    NotifyCmd: ("8", "notify"),
+    SaveCmd: ("8", "save"),
+    SnapshotCmd: ("4", "snapshot"),
+    ReportCmd: ("4", "report"),
+    ScheduleCmd: ("9", "schedule"),
+    DownloadCmd: ("10", "download"),
+    CompareCmd: ("10", "compare"),
+    VerifyCmd: ("10", "verify"),
+    VoteCmd: ("11", "vote"),
+    PromptBlock: ("11", "prompt"),
+    DistillCmd: ("11", "distill"),
+    RollbackCmd: ("11", "rollback"),
+    CurriculumCmd: ("12", "curriculum"),
+    StarCmd: ("12", "star"),
+    BestOfCmd: ("12", "best_of"),
+    ExploitCmd: ("12", "exploit"),
+    ArenaCmd: ("13", "arena"),
+    ResearchArenaCmd: ("13", "research_arena"),
+}
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="TD Lang — compile and run .td files for Time Dilation",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python -m td_lang check examples/demo_merge.td      # Check syntax
+  python -m td_lang compile examples/demo_merge.td    # Compile to .py
+  python -m td_lang run examples/demo_merge.td        # Compile + run
+  python -m td_lang run examples/demo_merge.td --dry   # Compile only
+  python -m td_lang info examples/demo_merge.td        # Show plan summary
+        """,
+    )
+
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"td_lang {__version__}",
+    )
+
+    parser.add_argument(
+        "action",
+        choices=["check", "compile", "run", "info"],
+        help="What to do: check (syntax), compile (.py), run (compile+execute), info (show plan)",
+    )
+
+    parser.add_argument(
+        "file",
+        type=str,
+        help="Path to the .td file",
+    )
+
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="td_lang_outputs",
+        help="Output directory (default: td_lang_outputs)",
+    )
+
+    parser.add_argument(
+        "--dry",
+        action="store_true",
+        help="With 'run': compile but don't execute",
+    )
+
+    parser.add_argument(
+        "--verbose", "-v",
+        action="store_true",
+        help="Show extra detail (compiled Python, full AST, etc.)",
+    )
+
+    return parser.parse_args()
+
+
+def print_banner():
+    """Print the td_lang banner."""
+    banner = f"""
+    ╔═══════════════════════════════════════╗
+    ║                                       ║
+    ║   ████████╗██████╗    ██╗      ██████╗║
+    ║   ╚══██╔══╝██╔══██╗   ██║     ██╔════╝║
+    ║      ██║   ██║  ██║   ██║     ██║  ███║
+    ║      ██║   ██║  ██║   ██║     ██║   ██║
+    ║      ██║   ██████╔╝   ██████╗ ╚██████╔╝║
+    ║      ╚═╝   ╚═════╝    ╚═════╝  ╚═════╝║
+    ║                                       ║
+    ║   TD Lang v{__version__} — .td file compiler   ║
+    ║                                       ║
+    ╚═══════════════════════════════════════╝
+    """
+    print(banner)
+
+
+def print_info(filepath: str) -> None:
+    """Show what a .td file does without compiling — human-readable plan summary."""
+    program = parse_td_file(filepath)
+
+    print(f"\n  File: {filepath}")
+    print(f"  Commands: {len(program.commands)}")
+
+    if program.gates:
+        print(f"  Gates: {', '.join(program.gates.must_pass)}")
+    if program.budget:
+        parts = []
+        if program.budget.max_gpu_hours is not None:
+            parts.append(f"{program.budget.max_gpu_hours} GPU hrs")
+        if program.budget.max_cost is not None:
+            parts.append(f"${program.budget.max_cost}")
+        print(f"  Budget: {', '.join(parts)}")
+    if program.data_contract:
+        print(f"  Data contract: fields={program.data_contract.required_fields}")
+    if program.reward_contract:
+        print(f"  Reward contract: verifiers={program.reward_contract.verifiers}")
+
+    print("\n  Plan:")
+    for i, cmd in enumerate(program.commands, 1):
+        phase, name = _PHASE_MAP.get(type(cmd), ("?", type(cmd).__name__))
+        target = getattr(cmd, 'target', getattr(cmd, 'alias', ''))
+        detail = ""
+        if hasattr(cmd, 'method'):
+            detail += f" method={cmd.method}"
+        if hasattr(cmd, 'source') and name in ("merge", "synth"):
+            detail += f" from={cmd.source}"
+        if hasattr(cmd, 'layers') and cmd.layers != "all":
+            detail += f" layers={cmd.layers}"
+        if hasattr(cmd, 'output') and cmd.output:
+            detail += f" -> {cmd.output}"
+        print(f"    {i}. [P{phase}] {name} {target}{detail}")
+
+    print()
+
+
+def main():
+    """Main entry point for td_lang CLI."""
+    args = parse_args()
+    print_banner()
+
+    executor = TDExecutor(output_dir=args.output)
+
+    try:
+        if args.action == "info":
+            print_info(args.file)
+
+        elif args.action == "check":
+            program = executor.check(args.file)
+            print("\n[td_lang] File is valid!")
+
+        elif args.action == "compile":
+            py_path = executor.compile(args.file)
+            print(f"\n[td_lang] Generated: {py_path}")
+            print("[td_lang] You can run it with: python", py_path)
+            if args.verbose:
+                print("\n--- Generated Python ---")
+                print(py_path.read_text())
+                print("--- End ---")
+
+        elif args.action == "run":
+            result = executor.run(args.file, dry_run=args.dry)
+            if result["status"] == "success":
+                sys.exit(0)
+            elif result["status"] == "dry_run":
+                sys.exit(0)
+            else:
+                sys.exit(1)
+
+    except TDLangError as e:
+        print(f"\n[td_lang] ERROR: {e}")
+        sys.exit(1)
+
+    except FileNotFoundError:
+        print(f"\n[td_lang] ERROR: File not found: {args.file}")
+        print("[td_lang] Check the path and try again.")
+        sys.exit(1)
+
+    except KeyboardInterrupt:
+        print("\n[td_lang] Interrupted.")
+        sys.exit(130)

hugging/td_lang/td_lang/engine/__init__.py

@@ -0,0 +1,25 @@
+"""
+TD Lang Engine — the merge/heal/validate runtime (formerly td_fuse).
+
+All model merging, transport, healing, and validation logic lives here.
+td_lang compiles .td files into Python that imports from this engine.
+
+Architecture:
+    td_lang/engine/
+    ├── __init__.py          ← This file
+    ├── config.py            ← Model configs, merge order, hyperparameters
+    ├── canary.py            ← Canary injection + testing ("brain surgery")
+    ├── transport.py         ← Wrapper around official T&M code
+    ├── techniques.py        ← Advanced techniques (Theseus, ARM, OTMF, RAM, Mergeability)
+    ├── merge.py             ← Sequential merge orchestrator
+    ├── validate.py          ← Post-merge validation (canary, perplexity, benchmarks)
+    ├── heal.py              ← QLoRA healing fine-tune via Unsloth
+    └── run.py               ← Standalone entry point (optional)
+
+Usage (via td_lang):
+    python -m td_lang run td_start.td
+    python -m td_lang run demo_merge.td
+"""
+
+__version__ = "0.2.0"
+__author__ = "Milan (TD Project)"