Upload 148 files

hugging/td_fuse/grpo_reason.py

@@ -204,6 +204,8 @@ def generate_all_problems(cfg: GRPOConfig) -> List[Dict]:
         elif ptype == "comparison":
             a_val = random.randint(10, 100)
             b_val = random.randint(10, 100)
+            while a_val == b_val:  # Avoid "who has more?" with equal values
+                b_val = random.randint(10, 100)
             diff = abs(a_val - b_val)
             who = "Alice" if a_val > b_val else "Bob"
             problems.append({
@@ -280,11 +282,19 @@ def verify_answer(problem: Dict, response: str) -> float:
 
     # For "all" (months with 28 days trick):
     if expected == "all":
-        return 1.0 if ("all" in answer_text and "12" in answer_text) or "every month" in answer_text or "all 12" in answer_text else 0.0
+        all_indicators = ["all 12", "every month", "all months", "all of them"]
+        if any(w in answer_text for w in all_indicators):
+            return 1.0
+        # "all" + "12" combo
+        if "all" in answer_text and "12" in answer_text:
+            return 1.0
+        return 0.0
 
     # For "neither" (rooster trick):
     if expected == "neither":
-        return 1.0 if ("neither" in answer_text or "don't lay" in answer_text or "rooster" in answer_text and "can't" in answer_text) else 0.0
+        rooster_words = ["neither", "don't lay", "doesn't lay", "cannot lay", "can't lay",
+                         "do not lay", "does not lay", "impossible"]
+        return 1.0 if any(w in answer_text for w in rooster_words) else 0.0
 
     # For "white" (polar bear trick):
     if expected == "white":
@@ -321,12 +331,9 @@ def verify_answer(problem: Dict, response: str) -> float:
         return 0.0
 
     # For single-letter answers (a, b, c, etc.) — ordering questions
+    # Must use word boundary to avoid matching inside words like "because"
     if len(expected) == 1 and expected.isalpha():
-        # Check for the name in the conclusion
-        names_map = {"c": "c", "w": "w", "tyke": "tyke", "sam": "sam", "dave": "dave",
-                     "carol": "carol", "yellow": "yellow", "library": "library"}
-        # Just check if expected appears in last 150 chars
-        if expected in answer_text[-150:]:
+        if re.search(rf'\b{expected}\b', answer_text[-150:]):
             return 1.0
         return 0.0
 
@@ -382,6 +389,7 @@ def generate_and_score_cot(model, tokenizer, problems, cfg, eos_id, resume_from=
     Saves progress every batch_save_every problems for crash recovery.
     """
     print(f"\n=== GENERATE & SCORE CoT ({len(problems)} problems, {cfg.num_candidates} candidates each) ===")
+    generate_and_score_cot._start = time.time()
 
     progress_path = Path(cfg.progress_file)
     progress_path.parent.mkdir(parents=True, exist_ok=True)
@@ -446,7 +454,10 @@ def generate_and_score_cot(model, tokenizer, problems, cfg, eos_id, resume_from=
         # Progress update & save
         if (i + 1) % cfg.batch_save_every == 0:
             pct = total_correct / total_attempted * 100
-            print(f"  [{i+1}/{len(problems)}] Correct: {total_correct}/{total_attempted} ({pct:.0f}%) | Winners: {len(winning_pairs)}")
+            elapsed = time.time()
+            rate = (i + 1 - resume_from) / max(1, elapsed - getattr(generate_and_score_cot, '_start', elapsed))
+            eta = (len(problems) - i - 1) / max(0.001, rate) / 60 if rate > 0 else 0
+            print(f"  [{i+1}/{len(problems)}] Correct: {total_correct}/{total_attempted} ({pct:.0f}%) | Winners: {len(winning_pairs)} | ETA: {eta:.0f}min")
 
             # Save progress to disk
             with open(progress_path, "w") as f:
@@ -458,6 +469,10 @@ def generate_and_score_cot(model, tokenizer, problems, cfg, eos_id, resume_from=
                     "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                 }, f)
 
+            # Periodic VRAM cleanup to prevent OOM during long generation
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
     # Final save
     pct = total_correct / total_attempted * 100 if total_attempted > 0 else 0
     print(f"\n  FINAL: {total_correct}/{total_attempted} ({pct:.0f}%) | Winners: {len(winning_pairs)} | Losers: {len(losing_pairs)}")
@@ -553,14 +568,15 @@ def train_grpo(model, tokenizer, winning_pairs, losing_pairs, cfg):
     print(f"  Starting GRPO training...")
     trainer.train()
 
-    # Clean checkpoints, merge LoRA
+    # Merge LoRA FIRST, then clean checkpoints (safer — don't lose data before merge succeeds)
+    print("  Merging LoRA...")
+    merged = model.merge_and_unload()
+    gc.collect()
+
     import shutil
     shutil.rmtree(str(out_dir), ignore_errors=True)
     gc.collect()
 
-    print("  Merging LoRA...")
-    merged = model.merge_and_unload()
-    gc.collect()
     return merged
 
 
@@ -710,7 +726,8 @@ def run_grpo_cycle(cfg: GRPOConfig = None):
     print(f"\nSaving improved model to {save_dir}...")
     model.save_pretrained(str(save_dir), safe_serialization=True)
     tokenizer.save_pretrained(str(save_dir))
-    sz = (save_dir / "model.safetensors").stat().st_size / 1e9
+    # Calculate total size (handles both single and sharded safetensors)
+    sz = sum(f.stat().st_size for f in save_dir.glob("*.safetensors")) / 1e9
     print(f"SAVED: {save_dir} ({sz:.1f} GB)")
 
     # Post-benchmark
@@ -741,3 +758,7 @@ def run_grpo_cycle(cfg: GRPOConfig = None):
         json.dump(report, f, indent=2)
 
     return str(save_dir)
+
+
+if __name__ == "__main__":
+    run_grpo_cycle()

hugging/td_fuse/heal.py

@@ -378,12 +378,12 @@ def apply_qlora_standard(
     print(f"[heal] SAVED OK: {healed_dir}")
 
     # Verify the save actually worked before cleaning up ANYTHING
-    saved_model = healed_dir / "model.safetensors"
-    if not saved_model.exists() or saved_model.stat().st_size < 1_000_000:
+    saved_files = list(healed_dir.glob("*.safetensors"))
+    total_size = sum(f.stat().st_size for f in saved_files) if saved_files else 0
+    if not saved_files or total_size < 1_000_000:
         print(f"[heal] WARNING: Save may have failed — NOT deleting any backups!")
     else:
-        save_size = saved_model.stat().st_size / 1e9
-        print(f"[heal] Verified: {saved_model} ({save_size:.1f} GB)")
+        print(f"[heal] Verified: {len(saved_files)} safetensors files ({total_size/1e9:.1f} GB)")
         # NOW safe to clean up old stuff
         cleanup_targets = [
             "td_fuse_outputs/final",

hugging/td_fuse/merge_klear.py

@@ -0,0 +1,292 @@
+"""
+Merge #6: Klear-Reasoner-8B → existing 5-model merge.
+
+Klear-Reasoner-8B is based on Qwen3-8B-Base (SAME architecture as our
+Qwen3-VL-8B target). This means:
+    - Direct 1:1 layer matching (36 → 36)
+    - Neurons are aligned (identity Q matrices)
+    - No permutation needed
+    - Transport plans take under 1 minute
+    - Total merge: ~10-15 minutes
+
+Think of this like adding more reasoning spice to an already-cooked dish.
+The base model already has reasoning from DeepSeek, MiMo, and Falcon.
+Klear-Reasoner adds EVEN MORE focused reasoning (90.5% AIME 2024).
+
+Alpha = 0.3 (conservative — we don't want to overwrite what's already
+merged). The existing model is the dominant one; Klear adds a boost.
+
+Usage:
+    python -m td_fuse.merge_klear
+
+Expects:
+    - Existing merge at: td_fuse_outputs/healed_final/ (or after_falcon_5model/)
+    - Will download Klear-Reasoner-8B from HuggingFace
+    - Output: td_fuse_outputs/6model_klear/
+"""
+
+import os
+import sys
+import gc
+import time
+import torch
+import numpy as np
+from pathlib import Path
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+# Where is our existing 5-model merge?
+EXISTING_MERGE_PATHS = [
+    "td_fuse_outputs/healed_final",
+    "td_fuse_checkpoints/after_falcon",       # backup location
+    "td_fuse_checkpoints/after_falcon_5model", # HF backup
+]
+
+# Klear-Reasoner-8B config
+KLEAR_HF_ID = "Kwai-Klear/Klear-Reasoner-8B"
+KLEAR_ALPHA = 0.3     # Conservative: keep 70% existing, add 30% Klear reasoning
+KLEAR_NAME = "Klear-Reasoner-8B"
+
+# Output
+OUTPUT_DIR = "td_fuse_outputs/6model_klear"
+
+# Vision encoder prefixes — NEVER merge into these
+VISION_SKIP = ["visual", "merger", "model.visual", "model.merger"]
+
+# Thinking tokens to protect
+THINK_TOKEN_IDS = [151667, 151668]
+
+
+def find_existing_merge():
+    """Find the existing 5-model merge checkpoint."""
+    for path in EXISTING_MERGE_PATHS:
+        safetensors = list(Path(path).glob("*.safetensors")) if Path(path).exists() else []
+        if safetensors:
+            total_size = sum(f.stat().st_size for f in safetensors)
+            if total_size > 1_000_000_000:  # > 1GB
+                print(f"[merge6] Found existing merge at: {path} ({total_size/1e9:.1f} GB)")
+                return path
+    return None
+
+
+def load_model(path, dtype=torch.bfloat16):
+    """Load model — auto-detects Qwen3-VL."""
+    from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+
+    print(f"[merge6] Loading model from: {path}")
+    t = time.time()
+
+    try:
+        config = AutoConfig.from_pretrained(path, trust_remote_code=True)
+        model_type = getattr(config, 'model_type', '')
+        config_class = type(config).__name__.lower()
+        if 'qwen3_vl' in model_type or 'qwen3vl' in config_class:
+            from transformers import Qwen3VLForConditionalGeneration
+            print(f"[merge6] Loading as Qwen3-VL model")
+            model = Qwen3VLForConditionalGeneration.from_pretrained(
+                path, device_map="auto", torch_dtype=dtype
+            )
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                path, device_map="auto", torch_dtype=dtype
+            )
+    except Exception as e:
+        print(f"[merge6] Auto-detect failed ({e}), using AutoModelForCausalLM")
+        model = AutoModelForCausalLM.from_pretrained(
+            path, device_map="auto", torch_dtype=dtype
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(path)
+    print(f"[merge6] Loaded in {time.time()-t:.0f}s")
+    return model, tokenizer
+
+
+def should_skip(key):
+    """Check if a parameter should be skipped during merge."""
+    # Skip vision encoder
+    for prefix in VISION_SKIP:
+        if key.startswith(prefix):
+            return True
+    return False
+
+
+def direct_merge(target_model, source_state, alpha=0.3):
+    """
+    Direct same-architecture merge: blend source weights into target.
+
+    Since Klear-Reasoner-8B has the same Qwen3-8B architecture:
+    - All parameter names match exactly
+    - No layer mapping needed
+    - No neuron permutation needed
+    - Just simple alpha blending: W = alpha * source + (1-alpha) * target
+
+    We skip:
+    - Vision encoder weights (visual.*, merger.*)
+    - Embedding/lm_head IF vocab sizes don't match
+    """
+    target_state = target_model.state_dict()
+
+    # Check vocab sizes
+    target_vocab = None
+    source_vocab = None
+    for key in ["model.embed_tokens.weight", "embed_tokens.weight"]:
+        if key in target_state:
+            target_vocab = target_state[key].shape[0]
+        if key in source_state:
+            source_vocab = source_state[key].shape[0]
+    skip_embeddings = (target_vocab != source_vocab) if (target_vocab and source_vocab) else False
+    if skip_embeddings:
+        print(f"[merge6] Vocab mismatch ({target_vocab} vs {source_vocab}) — skipping embeddings")
+    else:
+        print(f"[merge6] Vocab sizes match — merging embeddings too")
+
+    fused = 0
+    skipped = 0
+    total = len(target_state)
+
+    # Save original think token embeddings for restoration
+    think_embeds = {}
+    for key in target_state:
+        if "embed_tokens" in key:
+            for tid in THINK_TOKEN_IDS:
+                if tid < target_state[key].shape[0]:
+                    think_embeds[(key, tid)] = target_state[key][tid].clone()
+
+    for i, target_key in enumerate(target_state):
+        if should_skip(target_key):
+            skipped += 1
+            continue
+
+        if skip_embeddings and ("embed_tokens" in target_key or "lm_head" in target_key):
+            skipped += 1
+            continue
+
+        # Find source key — same name for same architecture
+        if target_key not in source_state:
+            # Try without "model." prefix or with it
+            alt_key = target_key.replace("model.", "", 1) if target_key.startswith("model.") else f"model.{target_key}"
+            if alt_key not in source_state:
+                skipped += 1
+                continue
+            source_key = alt_key
+        else:
+            source_key = target_key
+
+        target_w = target_state[target_key]
+        source_w = source_state[source_key]
+
+        # Shape check
+        if target_w.shape != source_w.shape:
+            skipped += 1
+            if skipped <= 5:
+                print(f"  [skip] Shape mismatch: {target_key} ({target_w.shape} vs {source_w.shape})")
+            continue
+
+        # Blend: W_final = alpha * source + (1-alpha) * target
+        fused_w = alpha * source_w.to(target_w.device, dtype=target_w.dtype) + (1 - alpha) * target_w
+        target_state[target_key] = fused_w
+        fused += 1
+
+        if fused % 50 == 0:
+            print(f"  Fused {fused}/{total} params...")
+            sys.stdout.flush()
+
+    # Restore think token embeddings
+    for (key, tid), embed in think_embeds.items():
+        if key in target_state and tid < target_state[key].shape[0]:
+            target_state[key][tid] = embed
+            print(f"[merge6] Protected think token {tid}")
+
+    # Load fused state
+    missing, unexpected = target_model.load_state_dict(target_state, strict=False)
+    if missing:
+        print(f"[merge6] {len(missing)} missing keys (likely vision quant params — safe)")
+    print(f"[merge6] Fused {fused} params, skipped {skipped}")
+    return target_model
+
+
+def run_merge():
+    """Main merge pipeline for Klear-Reasoner-8B."""
+    merge_start = time.time()
+
+    print("\n" + "=" * 60)
+    print("MERGE #6: Klear-Reasoner-8B → 5-model merge")
+    print(f"Alpha: {KLEAR_ALPHA} (30% Klear, 70% existing)")
+    print(f"Started at: {time.strftime('%H:%M:%S')}")
+    print("=" * 60)
+    sys.stdout.flush()
+
+    # --- Step 1: Find existing merge ---
+    existing_path = find_existing_merge()
+    if not existing_path:
+        print("[merge6] ERROR: Cannot find existing 5-model merge!")
+        print(f"[merge6] Looked in: {EXISTING_MERGE_PATHS}")
+        sys.exit(1)
+
+    # --- Step 2: Load existing merged model ---
+    print("\n[merge6] Step 1/4: Loading existing 5-model merge...")
+    sys.stdout.flush()
+    target_model, target_tokenizer = load_model(existing_path)
+
+    # --- Step 3: Download & load Klear-Reasoner-8B ---
+    print("\n[merge6] Step 2/4: Loading Klear-Reasoner-8B (will download if needed)...")
+    sys.stdout.flush()
+    t = time.time()
+
+    from transformers import AutoModelForCausalLM as AMLM
+    source_model = AMLM.from_pretrained(
+        KLEAR_HF_ID,
+        device_map="cpu",  # Keep on CPU to save GPU VRAM
+        torch_dtype=torch.bfloat16,
+    )
+    source_state = source_model.state_dict()
+    print(f"[merge6] Klear loaded in {time.time()-t:.0f}s ({len(source_state)} params)")
+    sys.stdout.flush()
+
+    # Free source model object (keep state dict)
+    del source_model
+    gc.collect()
+
+    # --- Step 4: Direct merge ---
+    print("\n[merge6] Step 3/4: Merging weights (direct same-arch blend)...")
+    sys.stdout.flush()
+    t = time.time()
+    target_model = direct_merge(target_model, source_state, alpha=KLEAR_ALPHA)
+    del source_state
+    gc.collect()
+    torch.cuda.empty_cache()
+    print(f"[merge6] Merge done in {time.time()-t:.0f}s")
+
+    # --- Step 5: Save ---
+    print(f"\n[merge6] Step 4/4: Saving to {OUTPUT_DIR}...")
+    sys.stdout.flush()
+    t = time.time()
+    output_path = Path(OUTPUT_DIR)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    target_model.save_pretrained(str(output_path), safe_serialization=True)
+    target_tokenizer.save_pretrained(str(output_path))
+
+    # Verify save
+    saved_files = list(output_path.glob("*.safetensors"))
+    total_size = sum(f.stat().st_size for f in saved_files) if saved_files else 0
+    if not saved_files or total_size < 1_000_000_000:
+        print(f"[merge6] WARNING: Save may have failed! Only {total_size/1e9:.1f} GB")
+    else:
+        print(f"[merge6] Saved: {len(saved_files)} files, {total_size/1e9:.1f} GB")
+    print(f"[merge6] Save done in {time.time()-t:.0f}s")
+
+    total_min = (time.time() - merge_start) / 60
+    print(f"\n{'=' * 60}")
+    print(f"MERGE #6 COMPLETE: {total_min:.1f} min total")
+    print(f"Output: {OUTPUT_DIR}")
+    print(f"{'=' * 60}")
+    sys.stdout.flush()
+
+    return str(output_path)
+
+
+if __name__ == "__main__":
+    run_merge()

hugging/td_fuse/reasoning_heal.py

@@ -0,0 +1,447 @@
+"""
+3-Phase Reasoning Heal — overnight restoration after merge.
+
+This isn't a normal heal. Normal heal just fixes language coherence.
+This one specifically targets the reasoning circuits in the model:
+
+Phase 1: Language Coherence (standard heal)
+    - Raw text data, fixes general language ability
+    - 2000 samples, 2 epochs
+    - ~30 min on A6000
+
+Phase 2: Stop-Token Fix (chat format heal)
+    - Chat-formatted data with <|im_start|>/<|im_end|> tokens
+    - Teaches model when to stop generating (prevents loop)
+    - 800 samples, 3 epochs
+    - ~20 min on A6000
+
+Phase 3: Reasoning Restoration (THE BIG ONE)
+    - Chain-of-thought reasoning data (GSM8K, MATH, ARC-Challenge)
+    - Higher LoRA rank (r=64, alpha=128) to really move the weights
+    - Targets layers 16-28 specifically (where reasoning lives)
+    - 3000+ samples, 4 epochs
+    - ~60-90 min on A6000
+
+Total estimated time: ~2-3 hours.
+
+Usage:
+    nohup python3 -m td_fuse.reasoning_heal > reasoning_heal.log 2>&1 &
+
+Config notes:
+    - bf16 throughout (not 4-bit — avoids dequantize corruption)
+    - Learning rate schedule: Phase 1 (5e-5) → Phase 2 (3e-5) → Phase 3 (2e-5)
+    - Phase 3 uses LOWER lr + HIGHER rank — big gentle changes, not small sharp ones
+
+Findings: #12, #16, #20
+"""
+
+import os
+import sys
+import gc
+import time
+import torch
+from pathlib import Path
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM, AutoTokenizer, AutoConfig,
+    TrainingArguments, Trainer,
+)
+from peft import LoraConfig, get_peft_model, TaskType
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+# Input: the 6-model merge (output of merge_klear.py)
+INPUT_PATHS = [
+    "td_fuse_outputs/6model_klear",
+    "td_fuse_outputs/healed_final",    # fallback: 5-model merge
+]
+
+# Output: final healed model
+OUTPUT_DIR = "td_fuse_outputs/reasoning_healed"
+
+# Phase configs
+PHASES = {
+    "phase1_language": {
+        "name": "Phase 1: Language Coherence",
+        "lora_r": 32,
+        "lora_alpha": 64,
+        "learning_rate": 5e-5,
+        "epochs": 2,
+        "batch_size": 1,
+        "grad_accum": 8,
+        "seq_len": 2048,
+        "target_modules": [
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        "datasets": [
+            # General language calibration
+            ("neuralmagic/LLM_compression_calibration", None, "train", 1200, "text", None),
+            # Math reasoning text
+            ("openai/gsm8k", "main", "train", 500, "question", None),
+            # Code
+            ("sahil2801/CodeAlpaca-20k", None, "train", 300, "output", None),
+        ],
+    },
+    "phase2_stop_tokens": {
+        "name": "Phase 2: Stop-Token Fix",
+        "lora_r": 16,
+        "lora_alpha": 32,
+        "learning_rate": 3e-5,
+        "epochs": 3,
+        "batch_size": 1,
+        "grad_accum": 4,
+        "seq_len": 1024,
+        "target_modules": [
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        "datasets": [
+            # Chat-formatted data with proper turn structure
+            ("HuggingFaceH4/ultrachat_200k", None, "train_sft", 800, None, "chat"),
+        ],
+    },
+    "phase3_reasoning": {
+        "name": "Phase 3: Reasoning Restoration",
+        "lora_r": 64,          # HIGH rank — really move the reasoning weights
+        "lora_alpha": 128,     # 2x rank
+        "learning_rate": 2e-5, # LOWER lr — big gentle changes
+        "epochs": 4,           # More epochs — let it really learn
+        "batch_size": 1,
+        "grad_accum": 8,
+        "seq_len": 2048,
+        # Target ONLY the reasoning layers (16-28 out of 0-35)
+        # This is where chain-of-thought reasoning lives
+        "target_modules": [
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        "layers_to_transform": list(range(16, 29)),  # Layers 16-28 only
+        "datasets": [
+            # GSM8K with full solutions (chain-of-thought)
+            ("openai/gsm8k", "main", "train", 1500, "answer", None),
+            # MATH dataset — harder problems
+            ("lighteval/MATH", "all", "train", 800, "solution", None),
+            # ARC-Challenge — science reasoning
+            ("allenai/ai2_arc", "ARC-Challenge", "train", 400, "question", None),
+            # More GSM8K questions (exercises reasoning path generation)
+            ("openai/gsm8k", "main", "train", 300, "question", None),
+        ],
+    },
+}
+
+
+# ============================================================================
+# DATA LOADING
+# ============================================================================
+
+def load_phase_data(phase_config, tokenizer):
+    """Load training data for a phase."""
+    all_texts = []
+    chat_template = None
+
+    for entry in phase_config["datasets"]:
+        dataset_id, config_name, split, count, text_field, mode = entry
+
+        try:
+            if config_name:
+                ds = load_dataset(dataset_id, config_name, split=split, streaming=True)
+            else:
+                ds = load_dataset(dataset_id, split=split, streaming=True)
+
+            loaded = 0
+            for example in ds:
+                if loaded >= count:
+                    break
+
+                if mode == "chat":
+                    # Chat-formatted data — use the messages field
+                    messages = example.get("messages", [])
+                    if messages and len(messages) >= 2:
+                        # Format as chat with proper tokens
+                        text = format_chat(messages, tokenizer)
+                        if text and len(text) > 50:
+                            all_texts.append(text)
+                            loaded += 1
+                else:
+                    # Raw text mode
+                    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"  WARNING: {dataset_id} failed: {e}")
+
+    print(f"  Total: {len(all_texts)} samples")
+    return all_texts
+
+
+def format_chat(messages, tokenizer):
+    """Format chat messages with proper tokens for Qwen3-style models."""
+    try:
+        # Try using the tokenizer's built-in chat template
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        return text
+    except Exception:
+        pass
+
+    # Fallback: manual ChatML format
+    parts = []
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
+    return "\n".join(parts)
+
+
+# ============================================================================
+# MODEL LOADING
+# ============================================================================
+
+def find_input_model():
+    """Find the model to heal."""
+    for path in INPUT_PATHS:
+        safetensors = list(Path(path).glob("*.safetensors")) if Path(path).exists() else []
+        if safetensors:
+            total_size = sum(f.stat().st_size for f in safetensors)
+            if total_size > 1_000_000_000:
+                print(f"[heal] Found model at: {path} ({total_size/1e9:.1f} GB)")
+                return path
+    return None
+
+
+def load_model_smart(path):
+    """Load model — auto-detects Qwen3-VL."""
+    print(f"[heal] Loading model from {path}...")
+    t = time.time()
+
+    try:
+        config = AutoConfig.from_pretrained(path, trust_remote_code=True)
+        model_type = getattr(config, 'model_type', '')
+        config_class = type(config).__name__.lower()
+        if 'qwen3_vl' in model_type or 'qwen3vl' in config_class:
+            from transformers import Qwen3VLForConditionalGeneration
+            print(f"[heal] Loading as Qwen3-VL model")
+            model = Qwen3VLForConditionalGeneration.from_pretrained(
+                path, device_map="auto", torch_dtype=torch.bfloat16
+            )
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                path, device_map="auto", torch_dtype=torch.bfloat16
+            )
+    except Exception as e:
+        print(f"[heal] Auto-detect failed ({e}), using AutoModelForCausalLM")
+        model = AutoModelForCausalLM.from_pretrained(
+            path, device_map="auto", torch_dtype=torch.bfloat16
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(path)
+    print(f"[heal] Loaded in {time.time()-t:.0f}s")
+    return model, tokenizer
+
+
+# ============================================================================
+# TRAINING
+# ============================================================================
+
+class TextDataset(torch.utils.data.Dataset):
+    """Simple tokenized text 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]
+
+
+def run_phase(model, tokenizer, phase_name, phase_config, output_subdir):
+    """Run a single healing phase."""
+    phase_start = time.time()
+    print(f"\n{'=' * 60}")
+    print(f"{phase_config['name']}")
+    print(f"LoRA r={phase_config['lora_r']}, alpha={phase_config['lora_alpha']}")
+    print(f"Epochs: {phase_config['epochs']}, LR: {phase_config['learning_rate']}")
+    if 'layers_to_transform' in phase_config:
+        layers = phase_config['layers_to_transform']
+        print(f"Target layers: {layers[0]}-{layers[-1]} (reasoning circuit)")
+    print(f"Started at: {time.strftime('%H:%M:%S')}")
+    print(f"{'=' * 60}")
+    sys.stdout.flush()
+
+    # Load data
+    print(f"\n[{phase_name}] Loading data...")
+    texts = load_phase_data(phase_config, tokenizer)
+    if not texts:
+        print(f"[{phase_name}] WARNING: No data loaded! Skipping phase.")
+        return model
+
+    # Create dataset
+    dataset = TextDataset(texts, tokenizer, phase_config["seq_len"])
+    del texts
+    gc.collect()
+
+    # LoRA config
+    lora_kwargs = {
+        "r": phase_config["lora_r"],
+        "lora_alpha": phase_config["lora_alpha"],
+        "lora_dropout": 0.0,
+        "target_modules": phase_config["target_modules"],
+        "bias": "none",
+        "task_type": TaskType.CAUSAL_LM,
+    }
+    if "layers_to_transform" in phase_config:
+        lora_kwargs["layers_to_transform"] = phase_config["layers_to_transform"]
+
+    lora_config = LoraConfig(**lora_kwargs)
+
+    # Apply LoRA
+    peft_model = get_peft_model(model, lora_config)
+    peft_model.print_trainable_parameters()
+
+    # Training args
+    train_output = Path(output_subdir) / f"{phase_name}_checkpoints"
+    train_output.mkdir(parents=True, exist_ok=True)
+
+    training_args = TrainingArguments(
+        output_dir=str(train_output),
+        num_train_epochs=phase_config["epochs"],
+        per_device_train_batch_size=phase_config["batch_size"],
+        gradient_accumulation_steps=phase_config["grad_accum"],
+        learning_rate=phase_config["learning_rate"],
+        bf16=True,
+        logging_steps=10,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=1,
+        warmup_ratio=0.05,
+        lr_scheduler_type="cosine",
+        optim="adamw_torch",
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=peft_model,
+        processing_class=tokenizer,
+        train_dataset=dataset,
+        args=training_args,
+    )
+
+    print(f"\n[{phase_name}] Training...")
+    sys.stdout.flush()
+    trainer.train()
+
+    # Merge LoRA back into base
+    print(f"\n[{phase_name}] Merging LoRA adapters back into model...")
+    merged = peft_model.merge_and_unload()
+
+    # Clean up checkpoints to free disk
+    import shutil
+    if train_output.exists():
+        shutil.rmtree(str(train_output), ignore_errors=True)
+        print(f"[{phase_name}] Cleaned up training checkpoints")
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    phase_min = (time.time() - phase_start) / 60
+    print(f"[{phase_name}] Complete in {phase_min:.1f} min")
+    sys.stdout.flush()
+
+    return merged
+
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+def run_reasoning_heal():
+    """Run the full 3-phase reasoning heal."""
+    total_start = time.time()
+
+    print("\n" + "=" * 70)
+    print("3-PHASE REASONING HEAL")
+    print("Phase 1: Language coherence (general text)")
+    print("Phase 2: Stop-token fix (chat format)")
+    print("Phase 3: Reasoning restoration (CoT data, layers 16-28)")
+    print(f"Started at: {time.strftime('%H:%M:%S')}")
+    print("=" * 70)
+    sys.stdout.flush()
+
+    # Check for existing healed model
+    healed_check = os.path.join(OUTPUT_DIR, 'model.safetensors')
+    healed_check_shard = list(Path(OUTPUT_DIR).glob("model-*.safetensors")) if Path(OUTPUT_DIR).exists() else []
+    if os.path.exists(healed_check) or healed_check_shard:
+        total_size = sum(f.stat().st_size for f in (healed_check_shard or [Path(healed_check)])) if (healed_check_shard or os.path.exists(healed_check)) else 0
+        if total_size > 1_000_000_000:
+            print(f"[heal] Found existing reasoning-healed model ({total_size/1e9:.1f} GB) — SKIPPING!")
+            return OUTPUT_DIR
+
+    # Find model
+    model_path = find_input_model()
+    if not model_path:
+        print("[heal] ERROR: Cannot find input model!")
+        print(f"[heal] Looked in: {INPUT_PATHS}")
+        sys.exit(1)
+
+    # Load model
+    model, tokenizer = load_model_smart(model_path)
+
+    # Run each phase
+    output_subdir = Path(OUTPUT_DIR)
+    output_subdir.mkdir(parents=True, exist_ok=True)
+
+    for phase_name, phase_config in PHASES.items():
+        model = run_phase(model, tokenizer, phase_name, phase_config, str(output_subdir))
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    # Save final model
+    print(f"\n[heal] Saving final reasoning-healed model to {OUTPUT_DIR}...")
+    sys.stdout.flush()
+    t = time.time()
+    model.save_pretrained(str(output_subdir), safe_serialization=True)
+    tokenizer.save_pretrained(str(output_subdir))
+
+    # Verify save
+    saved_files = list(output_subdir.glob("*.safetensors"))
+    total_size = sum(f.stat().st_size for f in saved_files) if saved_files else 0
+    if not saved_files or total_size < 1_000_000_000:
+        print(f"[heal] WARNING: Save may have failed! Only {total_size/1e9:.1f} GB")
+    else:
+        print(f"[heal] Saved: {len(saved_files)} files, {total_size/1e9:.1f} GB")
+
+    total_min = (time.time() - total_start) / 60
+    print(f"\n{'=' * 70}")
+    print(f"3-PHASE REASONING HEAL COMPLETE: {total_min:.1f} min total")
+    print(f"Output: {OUTPUT_DIR}")
+    print(f"{'=' * 70}")
+    sys.stdout.flush()
+
+    return str(output_subdir)
+
+
+if __name__ == "__main__":
+    run_reasoning_heal()

hugging/td_fuse/run_overnight.py

@@ -0,0 +1,65 @@
+"""
+Overnight Pipeline: Merge Klear-Reasoner + 3-Phase Reasoning Heal + Benchmark.
+
+Run this before going to sleep:
+    nohup python3 -m td_fuse.run_overnight > overnight.log 2>&1 &
+
+What it does:
+    1. Merges Klear-Reasoner-8B into the existing 5-model merge (~10-15 min)
+    2. Runs 3-phase reasoning heal (~2-3 hours)
+    3. Runs benchmark to see how good the final model is
+
+Total estimated time: 3-4 hours.
+"""
+
+import sys
+import time
+
+
+def main():
+    total_start = time.time()
+    print("=" * 70)
+    print("OVERNIGHT PIPELINE: Merge + Reasoning Heal + Benchmark")
+    print(f"Started at: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    print("=" * 70)
+    sys.stdout.flush()
+
+    # --- Step 1: Merge Klear-Reasoner-8B ---
+    print("\n\n>>> STEP 1/3: MERGING KLEAR-REASONER-8B <<<\n")
+    sys.stdout.flush()
+    from td_fuse.merge_klear import run_merge
+    merge_path = run_merge()
+    print(f"\n[overnight] Merge complete. Output: {merge_path}")
+    sys.stdout.flush()
+
+    # --- Step 2: 3-Phase Reasoning Heal ---
+    print("\n\n>>> STEP 2/3: 3-PHASE REASONING HEAL <<<\n")
+    sys.stdout.flush()
+    from td_fuse.reasoning_heal import run_reasoning_heal
+    heal_path = run_reasoning_heal()
+    print(f"\n[overnight] Reasoning heal complete. Output: {heal_path}")
+    sys.stdout.flush()
+
+    # --- Step 3: Benchmark ---
+    print("\n\n>>> STEP 3/3: BENCHMARK <<<\n")
+    sys.stdout.flush()
+    try:
+        from td_fuse.grpo_reason import pre_benchmark
+        pre_benchmark(heal_path)
+    except Exception as e:
+        print(f"[overnight] Benchmark failed: {e}")
+        print("[overnight] Model is saved — you can benchmark manually later.")
+    sys.stdout.flush()
+
+    # --- Done ---
+    total_hours = (time.time() - total_start) / 3600
+    print(f"\n\n{'=' * 70}")
+    print(f"OVERNIGHT PIPELINE COMPLETE: {total_hours:.1f} hours total")
+    print(f"Final model at: {heal_path}")
+    print(f"Finished at: {time.strftime('%Y-%m-%d %H:%M:%S')}")
+    print(f"{'=' * 70}")
+    sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

hugging/td_fuse/selfimprove.py

@@ -240,29 +240,15 @@ def generate_math_problems(count=200):
 
 def verify_answer(problem, model_answer):
     """
-    Verified reward: check if the answer is correct.
-    This is the GRPO reward signal — objective, not learned.
+    Verified reward: uses the smart verify from grpo_reason.
+    Handles verify_key for trick questions.
     """
-    expected = problem.get("verify_key", problem["answer"]).lower().strip()
-    answer_lower = model_answer.lower().strip()
-
-    # Check if expected answer appears in model output
-    if expected in answer_lower:
-        return 1.0
-
-    # For numeric answers, try to find the number
-    if expected.replace(".", "").replace("-", "").isdigit():
-        # Look for the number in the output
-        import re
-        numbers = re.findall(r'-?\d+\.?\d*', answer_lower)
-        for num in numbers:
-            try:
-                if abs(float(num) - float(expected)) < 0.01:
-                    return 1.0
-            except ValueError:
-                pass
-
-    return 0.0
+    # Use verify_key if available (trick questions have separate verify keys)
+    check_problem = dict(problem)
+    if "verify_key" in check_problem:
+        check_problem["answer"] = check_problem["verify_key"]
+    from .grpo_reason import verify_answer as smart_verify
+    return smart_verify(check_problem, model_answer)
 
 
 def generate_and_score(model, tokenizer, problems, cfg, eos_id):
@@ -390,15 +376,14 @@ def train_on_winners(model, tokenizer, winning_pairs, cfg):
     print(f"  Training: ~{total_steps} steps")
     trainer.train()
 
-    # Clean up training checkpoints
-    import shutil
-    shutil.rmtree(str(out_dir), ignore_errors=True)
-
-    # Merge LoRA back
+    # Merge LoRA FIRST, then clean checkpoints (safer)
     print("  Merging LoRA...")
     merged = model.merge_and_unload()
     gc.collect()
 
+    import shutil
+    shutil.rmtree(str(out_dir), ignore_errors=True)
+
     return merged
 
 
@@ -418,13 +403,13 @@ def benchmark(model, tokenizer, eos_id):
         out = model.generate(**ids, max_new_tokens=200, do_sample=False, eos_token_id=eos_id)
         return tokenizer.decode(out[0][ids.input_ids.shape[1]:], skip_special_tokens=True)
 
-    # Math
+    # Math — using smart verify
     math_tests = [("7+8", "15"), ("123+456", "579"), ("1000-387", "613"), ("12*13", "156"), ("144/12", "12")]
-    math_correct = sum(1 for q, e in math_tests if e in ask(f"What is {q}? Give just the number."))
+    math_correct = sum(1 for q, e in math_tests if verify_answer({"answer": e}, ask(f"What is {q}? Give just the number.")) > 0)
     results["basic_math"] = f"{math_correct}/5"
     print(f"  Math: {math_correct}/5")
 
-    # Reasoning
+    # Reasoning — using smart verify
     reason_tests = [
         ("If all roses are flowers and all flowers need water, do roses need water?", "yes"),
         ("Which is heavier: a pound of feathers or a pound of bricks?", "same"),
@@ -435,19 +420,19 @@ def benchmark(model, tokenizer, eos_id):
     reason_correct = 0
     for q, expected in reason_tests:
         a = ask(q)
-        correct = expected.lower() in a.lower()
+        correct = verify_answer({"answer": expected}, a) > 0
         reason_correct += correct
         print(f"    {'OK' if correct else 'FAIL'}: {q[:50]}... -> {a[:60]}")
     results["reasoning"] = f"{reason_correct}/5"
     print(f"  Reasoning: {reason_correct}/5")
 
-    # Word problems
+    # Word problems — using smart verify
     wp_tests = [
         ("A train travels 60 mph for 2.5 hours. How far does it go?", "150"),
         ("If 3 shirts cost $45, how much do 7 shirts cost?", "105"),
         ("I have 24 cookies split equally among 6 friends. How many each?", "4"),
     ]
-    wp_correct = sum(1 for q, e in wp_tests if e in ask(q))
+    wp_correct = sum(1 for q, e in wp_tests if verify_answer({"answer": e}, ask(q)) > 0)
     results["word_problems"] = f"{wp_correct}/3"
     print(f"  Word problems: {wp_correct}/3")
 
@@ -521,7 +506,7 @@ def run_cycle(cfg: SelfImproveConfig = None, cycle_num: int = 1):
     print(f"\nSaving improved model to {improved_dir}...")
     model.save_pretrained(str(improved_dir), safe_serialization=True)
     tokenizer.save_pretrained(str(improved_dir))
-    sz = (improved_dir / "model.safetensors").stat().st_size / 1e9
+    sz = sum(f.stat().st_size for f in improved_dir.glob("*.safetensors")) / 1e9
     print(f"SAVED: {improved_dir} ({sz:.1f} GB)")
 
     # Step 5: Benchmark