Upload 149 files

hugging/td_fuse/selfimprove.py

@@ -1,18 +1,19 @@
 """
-TD Self-Improvement v5 — Full Weight Intelligence + Unclogged Layernorms
-
-CYCLE 2 FIXES (from cycle 1 learnings):
-1. UNFREEZE LAYERNORMS — The worst issues (kurtosis 89, 90) are in layernorms
-   and k_norms. LoRA can't touch those. Now we directly train them.
-2. HIGHER LORA RANK — r=64 instead of r=32 for more capacity to reshape weights
-3. FIX CODE DATA — CodeAlpaca uses instruction/input/output, not prompt/completion
-4. FIX MMLU FORMAT — No more "The answer is A" multiple choice contamination
-5. AUTO-DETECT CYCLE — Runs from cycle 1 output if it exists, otherwise from healed
-6. QUICK BENCHMARK — 5 questions at the end to sanity-check actual model capability
-
-The key insight: LoRA only modifies projection weights (q/k/v/o/gate/up/down).
-But the WORST health issues are in layernorm and k_norm parameters.
-Those need to be directly unfrozen and trained — "unclogging the pipes."
+TD Self-Improvement v6 — MEGA CYCLE
+
+v5 fixes + MEGA MODE for cycle 3+:
+- 8000+ training samples (was 3000)
+- Priority: code, math, reasoning, creativity (top), then general (second)
+- LoRA r=128 (was 64) — maximum reshape capacity
+- ALL 36 layers targeted (was 16)
+- gradient_checkpointing for speed + VRAM savings
+- batch_size 4 (was 2) — 2x faster per step
+- Norm unclogging continues from cycle 2
+
+SPEED OPTIMIZATIONS:
+- gradient_checkpointing = True (saves ~40% VRAM, allows bigger batches)
+- batch_size 4 + grad_accum 4 = effective batch 16 (same quality, faster)
+- Processes 2x more samples per forward pass = ~2x faster wall clock
 """
 
 import torch
@@ -34,25 +35,25 @@ class SelfImproveConfig:
     model_path: str = ""  # Auto-detected below
     reference_model: str = "Qwen/Qwen3-VL-8B-Instruct"
     output_dir: str = "td_fuse_outputs/self_improve"
+    mega_mode: bool = False  # Auto-enabled for cycle 3+
 
     # How many layers to target
     top_n_layers: int = 16
     min_improvement_score: float = 0.05
 
-    # LoRA config — higher rank for cycle 2
-    lora_r: int = 64           # Up from 32 — more capacity to reshape
-    lora_alpha: int = 128      # 2x rank as usual
+    # LoRA config
+    lora_r: int = 64
+    lora_alpha: int = 128
     train_epochs: int = 3
     train_batch: int = 2
     train_grad_accum: int = 8
     learning_rate: float = 1.5e-5
-    norm_learning_rate: float = 5e-5  # Higher LR for layernorms (they're stubborn)
-    save_steps: int = 50
+    norm_learning_rate: float = 5e-5
+    save_steps: int = 100
     num_train_samples: int = 3000
 
     # Layernorm unfreezing
-    unfreeze_norms_for_layers: List[int] = None  # Auto-detected from analysis
-    norm_kurtosis_threshold: float = 15.0  # Unfreeze norms with kurtosis above this
+    norm_kurtosis_threshold: float = 15.0
 
     # Scoring weights
     w_health: float = 0.55
@@ -60,6 +61,28 @@ class SelfImproveConfig:
     w_coherence: float = 0.20
 
 
+def apply_mega_mode(cfg: SelfImproveConfig):
+    """Crank everything up for mega cycles (3+)."""
+    print("\n  *** MEGA MODE ENABLED ***")
+    cfg.mega_mode = True
+    cfg.top_n_layers = 36          # ALL layers — leave nothing behind
+    cfg.min_improvement_score = 0.0  # Target everything
+    cfg.lora_r = 128               # Maximum reshape capacity
+    cfg.lora_alpha = 256           # 2x rank
+    cfg.train_epochs = 3           # 3 epochs over 8000 samples = plenty
+    cfg.train_batch = 4            # 2x bigger batch — 2x faster
+    cfg.train_grad_accum = 4       # Effective batch = 16 (same quality)
+    cfg.num_train_samples = 8000   # Nearly 3x more data
+    cfg.learning_rate = 1e-5       # Slightly lower LR for stability with more data
+    cfg.norm_learning_rate = 1e-3   # Much higher — norms need strong signal to move during training
+    cfg.save_steps = 200           # Save less often (faster)
+    print(f"    LoRA: r={cfg.lora_r}, alpha={cfg.lora_alpha}")
+    print(f"    Layers: ALL {cfg.top_n_layers}")
+    print(f"    Data: {cfg.num_train_samples} samples")
+    print(f"    Batch: {cfg.train_batch} × {cfg.train_grad_accum} = {cfg.train_batch * cfg.train_grad_accum}")
+    print(f"    Speed: gradient_checkpointing + batch 4")
+
+
 def auto_detect_model_path() -> Tuple[str, int]:
     """Find the best starting model and determine cycle number."""
     base = Path("td_fuse_outputs/self_improve")
@@ -441,9 +464,6 @@ def detect_clogged_norms(health_scores: Dict[str, Dict], threshold: float = 15.0
     """
     clogged = []
     for name, h in health_scores.items():
-        # Only look at norm layers
-        # Qwen3-VL has: input_layernorm, post_attention_layernorm
-        # Also check for any norm-like parameter (norm, ln)
         if any(norm_type in name.lower() for norm_type in ["layernorm", "ln_", "_norm"]):
             kurt = abs(h["stats"]["kurtosis"])
             if kurt > threshold:
@@ -452,6 +472,91 @@ def detect_clogged_norms(health_scores: Dict[str, Dict], threshold: float = 15.0
     return clogged
 
 
+def repair_clogged_norms(model_path: str, clogged_norms: List[str],
+                         ref_path: str, health_scores: Dict[str, Dict]) -> str:
+    """
+    DIRECT SURGICAL REPAIR of clogged layernorms.
+
+    Training can't fix kurtosis of 2446 — gradient signal too weak for 1D norms.
+    Instead: blend broken norms with reference values.
+
+    Strategy:
+    - Kurtosis > 100: 80% reference, 20% current (severely broken)
+    - Kurtosis > 50:  60% reference, 40% current
+    - Kurtosis > 15:  40% reference, 60% current (mildly clogged)
+
+    This is like replacing a clogged pipe instead of pushing harder.
+    """
+    if not clogged_norms:
+        print("  No clogged norms to repair.")
+        return model_path
+
+    print(f"\n{'='*50}")
+    print(f"SURGICAL NORM REPAIR — {len(clogged_norms)} clogged norms")
+    print(f"{'='*50}")
+
+    from safetensors.torch import load_file, save_file
+    import glob
+
+    # Load reference norms
+    print("  Loading reference norms...")
+    ref_norms = {}
+    for st_file in sorted(glob.glob(f"{ref_path}/*.safetensors")):
+        state = load_file(st_file, device="cpu")
+        for name, tensor in state.items():
+            if name in clogged_norms:
+                ref_norms[name] = tensor.clone()
+        del state
+
+    print(f"  Matched {len(ref_norms)}/{len(clogged_norms)} norms in reference")
+
+    if not ref_norms:
+        print("  WARNING: No reference norms found! Skipping repair.")
+        return model_path
+
+    # Load and repair model weights
+    repaired = 0
+    st_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
+
+    for st_file in st_files:
+        state = load_file(st_file, device="cpu")
+        changed = False
+
+        for name in list(state.keys()):
+            if name in ref_norms:
+                kurt = abs(health_scores[name]["stats"]["kurtosis"])
+
+                # Blend ratio based on severity
+                if kurt > 100:
+                    ref_weight = 0.80  # Severely broken — mostly use reference
+                elif kurt > 50:
+                    ref_weight = 0.60
+                else:
+                    ref_weight = 0.40  # Mildly clogged
+
+                old_tensor = state[name]
+                ref_tensor = ref_norms[name].to(old_tensor.dtype)
+
+                # Blend: new = ref_weight * reference + (1-ref_weight) * current
+                state[name] = ref_weight * ref_tensor + (1 - ref_weight) * old_tensor
+
+                # Compute new kurtosis
+                t = state[name].float()
+                new_kurt = abs(((t - t.mean()) / max(t.std().item(), 1e-8)).pow(4).mean().item() - 3)
+
+                print(f"    {name}: kurtosis {kurt:.1f} → {new_kurt:.1f} (blend {ref_weight:.0%} ref)")
+                repaired += 1
+                changed = True
+
+        if changed:
+            save_file(state, st_file)
+        del state
+
+    gc.collect()
+    print(f"\n  Repaired {repaired} norms via reference blending")
+    return model_path
+
+
 # ============================================================
 # STEP 6: LoRA + UNFROZEN NORMS IMPROVEMENT
 # ============================================================
@@ -486,6 +591,11 @@ def improve_with_lora(model_path: str, target_layers: List[int],
     )
     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
+    # Enable gradient checkpointing for mega mode (saves ~40% VRAM)
+    if cfg.mega_mode:
+        model.gradient_checkpointing_enable()
+        print("  Gradient checkpointing: ENABLED (saves ~40% VRAM)")
+
     # Apply LoRA on projection weights
     lora_config = LoraConfig(
         r=cfg.lora_r,
@@ -522,34 +632,28 @@ def improve_with_lora(model_path: str, target_layers: List[int],
     print(f"\n  LoRA trainable + {unfrozen_count} unfrozen norm layers ({unfrozen_params:,} params)")
     model.print_trainable_parameters()
 
-    # Load DIVERSE training data — 5 categories, ALL BUGS FIXED
-    print("\n  Loading diverse training data...")
+    # Load training data — priority-weighted categories
+    # MEGA MODE: code/math/reasoning/creativity get 2x allocation
+    print("\n  Loading training data...")
     train_texts = []
-    per_category = cfg.num_train_samples // 5
+    total = cfg.num_train_samples
 
-    # Category 1: General text
-    try:
-        ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
-        count = 0
-        for row in ds:
-            if count >= per_category:
-                break
-            text = row.get("text", "")
-            if len(text) > 100:
-                chat = (f"<|im_start|>user\nContinue this text naturally.<|im_end|>\n"
-                       f"<|im_start|>assistant\n{text[:800]}<|im_end|>")
-                train_texts.append(chat)
-                count += 1
-        print(f"    General text: {count}")
-    except Exception as e:
-        print(f"    WARNING: General text failed: {e}")
+    if cfg.mega_mode:
+        # Priority split: code=2000, math=2000, reasoning=1500, creativity=1000, instructions=1000, general=500
+        alloc = {"code": int(total*0.25), "math": int(total*0.25),
+                 "reasoning": int(total*0.19), "creativity": int(total*0.12),
+                 "instructions": int(total*0.12), "general": int(total*0.07)}
+    else:
+        per = total // 5
+        alloc = {"code": per, "math": per, "reasoning": per,
+                 "creativity": 0, "instructions": per, "general": per}
 
-    # Category 2: Math reasoning
+    # 1. MATH (top priority)
     try:
         ds = load_dataset("openai/gsm8k", "main", split="train")
         count = 0
         for row in ds:
-            if count >= per_category:
+            if count >= alloc["math"]:
                 break
             q, a = row.get("question", ""), row.get("answer", "")
             if q and a:
@@ -557,72 +661,149 @@ def improve_with_lora(model_path: str, target_layers: List[int],
                        f"<|im_start|>assistant\n<think>\n{a}\n</think><|im_end|>")
                 train_texts.append(chat)
                 count += 1
-        print(f"    Math reasoning: {count}")
+        print(f"    Math (GSM8K): {count}")
     except Exception as e:
-        print(f"    WARNING: Math data failed: {e}")
+        print(f"    WARNING: GSM8K failed: {e}")
 
-    # Category 3: Code — FIXED: correct field names
+    # Extra math for mega mode — harder competition math
+    if cfg.mega_mode and count < alloc["math"]:
+        try:
+            ds2 = load_dataset("lighteval/MATH", "all", split="train", trust_remote_code=True)
+            extra = 0
+            for row in ds2:
+                if count + extra >= alloc["math"]:
+                    break
+                q = row.get("problem", "")
+                a = row.get("solution", "")
+                if q and a:
+                    chat = (f"<|im_start|>user\n{q}\nSolve this step by step.<|im_end|>\n"
+                           f"<|im_start|>assistant\n<think>\n{a}\n</think><|im_end|>")
+                    train_texts.append(chat)
+                    extra += 1
+            print(f"    Math (MATH hard): {extra}")
+        except Exception as e:
+            print(f"    WARNING: MATH dataset failed: {e}")
+
+    # 2. CODE (top priority)
     try:
         ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
         count = 0
         for row in ds:
-            if count >= per_category:
+            if count >= alloc["code"]:
                 break
-            # CodeAlpaca uses instruction/input/output, NOT prompt/completion
             inst = row.get("instruction", "")
             inp = row.get("input", "")
             out = row.get("output", "")
-            if inst and out:
+            if inst and out and len(out) > 20:
                 prompt = f"{inst}\n{inp}" if inp else inst
                 chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
                        f"<|im_start|>assistant\n{out}<|im_end|>")
                 train_texts.append(chat)
                 count += 1
-        print(f"    Code: {count}")
+        print(f"    Code (CodeAlpaca): {count}")
+    except Exception as e:
+        print(f"    WARNING: CodeAlpaca failed: {e}")
+
+    # Extra code for mega mode
+    if cfg.mega_mode and count < alloc["code"]:
+        try:
+            ds2 = load_dataset("m-a-p/CodeFeedback-Filtered-Instruction", split="train")
+            extra = 0
+            for row in ds2:
+                if count + extra >= alloc["code"]:
+                    break
+                q = row.get("query", "")
+                a = row.get("answer", "")
+                if q and a and len(a) > 30:
+                    chat = (f"<|im_start|>user\n{q}<|im_end|>\n"
+                           f"<|im_start|>assistant\n{a[:1500]}<|im_end|>")
+                    train_texts.append(chat)
+                    extra += 1
+            print(f"    Code (CodeFeedback): {extra}")
+        except Exception as e:
+            print(f"    WARNING: CodeFeedback failed: {e}")
+
+    # 3. REASONING (top priority)
+    try:
+        ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
+        count = 0
+        for row in ds:
+            if count >= alloc["reasoning"]:
+                break
+            q = row.get("question", "")
+            options = row.get("options", [])
+            answer_idx = row.get("answer_index", -1)
+            if q and options and 0 <= answer_idx < len(options):
+                answer_text = options[answer_idx]
+                chat = (f"<|im_start|>user\n{q}<|im_end|>\n"
+                       f"<|im_start|>assistant\n<think>\nLet me reason through this.\n"
+                       f"{answer_text}\n</think>\n{answer_text}<|im_end|>")
+                train_texts.append(chat)
+                count += 1
+        print(f"    Reasoning (MMLU-Pro): {count}")
     except Exception as e:
-        print(f"    WARNING: Code data failed: {e}")
+        print(f"    WARNING: MMLU-Pro failed: {e}")
 
-    # Category 4: Instruction following
+    # 4. CREATIVITY (mega mode priority)
+    if alloc.get("creativity", 0) > 0:
+        try:
+            ds = load_dataset("yahma/alpaca-cleaned", split="train")
+            count = 0
+            creative_words = ["write", "story", "poem", "creative", "imagine", "describe",
+                            "compose", "draft", "narrative", "fiction", "essay", "letter",
+                            "explain", "summarize", "rewrite", "paraphrase"]
+            for row in ds:
+                if count >= alloc["creativity"]:
+                    break
+                inst = row.get("instruction", "").lower()
+                out = row.get("output", "")
+                if any(w in inst for w in creative_words) and out and len(out) > 50:
+                    inp = row.get("input", "")
+                    prompt = f"{row['instruction']}\n{inp}" if inp else row["instruction"]
+                    chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
+                           f"<|im_start|>assistant\n{out}<|im_end|>")
+                    train_texts.append(chat)
+                    count += 1
+            print(f"    Creativity (Alpaca filtered): {count}")
+        except Exception as e:
+            print(f"    WARNING: Creativity data failed: {e}")
+
+    # 5. INSTRUCTIONS (general)
     try:
         ds = load_dataset("yahma/alpaca-cleaned", split="train")
         count = 0
         for row in ds:
-            if count >= per_category:
+            if count >= alloc.get("instructions", 0):
                 break
             inst = row.get("instruction", "")
             inp = row.get("input", "")
             out = row.get("output", "")
-            if inst and out:
+            if inst and out and len(out) > 30:
                 prompt = f"{inst}\n{inp}" if inp else inst
                 chat = (f"<|im_start|>user\n{prompt}<|im_end|>\n"
                        f"<|im_start|>assistant\n{out}<|im_end|>")
                 train_texts.append(chat)
                 count += 1
-        print(f"    Instructions: {count}")
+        print(f"    Instructions (Alpaca): {count}")
     except Exception as e:
-        print(f"    WARNING: Instruction data failed: {e}")
+        print(f"    WARNING: Alpaca failed: {e}")
 
-    # Category 5: Reasoning — FIXED: proper format, no "answer is A/B/C"
+    # 6. GENERAL TEXT (lowest priority)
     try:
-        ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
+        ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
         count = 0
         for row in ds:
-            if count >= per_category:
+            if count >= alloc.get("general", 0):
                 break
-            q = row.get("question", "")
-            # Get the actual answer text from options using answer_index
-            options = row.get("options", [])
-            answer_idx = row.get("answer_index", -1)
-            if q and options and 0 <= answer_idx < len(options):
-                answer_text = options[answer_idx]
-                chat = (f"<|im_start|>user\n{q}<|im_end|>\n"
-                       f"<|im_start|>assistant\n<think>\nLet me reason through this.\n"
-                       f"{answer_text}\n</think>\n{answer_text}<|im_end|>")
+            text = row.get("text", "")
+            if len(text) > 100:
+                chat = (f"<|im_start|>user\nContinue this text naturally.<|im_end|>\n"
+                       f"<|im_start|>assistant\n{text[:800]}<|im_end|>")
                 train_texts.append(chat)
                 count += 1
-        print(f"    Reasoning/logic: {count}")
+        print(f"    General text: {count}")
     except Exception as e:
-        print(f"    WARNING: Reasoning data failed: {e}")
+        print(f"    WARNING: General text failed: {e}")
 
     print(f"\n  Total training samples: {len(train_texts)}")
 
@@ -685,6 +866,7 @@ def improve_with_lora(model_path: str, target_layers: List[int],
         gradient_accumulation_steps=cfg.train_grad_accum,
         learning_rate=cfg.learning_rate,
         bf16=True,
+        gradient_checkpointing=cfg.mega_mode,  # Saves ~40% VRAM in mega mode
         logging_steps=max(1, total_steps // 20),
         save_strategy="steps",
         save_steps=cfg.save_steps,
@@ -803,13 +985,22 @@ def run_cycle(cfg: SelfImproveConfig = None):
     else:
         cycle_num = 2 if "improve" in cfg.model_path else 1
 
+    # MEGA MODE for cycle 3+
+    if cycle_num >= 3 and not cfg.mega_mode:
+        apply_mega_mode(cfg)
+
     start = time.time()
+    mode_label = "MEGA CYCLE" if cfg.mega_mode else "NORM UNCLOGGING + LoRA"
     print("=" * 60)
-    print(f"TD SELF-IMPROVEMENT v5 — NORM UNCLOGGING + LoRA")
+    print(f"TD SELF-IMPROVEMENT v6 — {mode_label}")
     print(f"Cycle: {cycle_num}")
     print(f"Model: {cfg.model_path}")
-    print(f"LoRA rank: {cfg.lora_r} (up from 32)")
+    print(f"LoRA rank: {cfg.lora_r}")
+    print(f"Samples: {cfg.num_train_samples}")
+    print(f"Batch: {cfg.train_batch} × {cfg.train_grad_accum} = {cfg.train_batch * cfg.train_grad_accum}")
     print(f"Norm training: ENABLED (LR={cfg.norm_learning_rate})")
+    if cfg.mega_mode:
+        print(f"Gradient checkpointing: ENABLED")
     print(f"Started: {time.strftime('%H:%M:%S')}")
     print("=" * 60)
 
@@ -842,8 +1033,7 @@ def run_cycle(cfg: SelfImproveConfig = None):
         )
     ref_analysis = analyze_model_weights(str(ref_path), "Qwen3-VL-8B (reference)")
     damage_scores = compute_damage_scores(model_analysis, ref_analysis)
-    print(f"  Cleaning up reference...")
-    shutil.rmtree(str(ref_path), ignore_errors=True)
+    # DON'T delete reference yet — we need it for norm repair
     gc.collect()
 
     # STEP 4: Coherence
@@ -870,6 +1060,18 @@ def run_cycle(cfg: SelfImproveConfig = None):
     else:
         print(f"\n  No clogged norms above threshold {cfg.norm_kurtosis_threshold}")
 
+    # STEP 5c: REPAIR clogged norms surgically
+    if clogged_norms:
+        print(f"\n{'='*40}")
+        print("STEP 5c/7: SURGICAL NORM REPAIR")
+        print(f"{'='*40}")
+        repair_clogged_norms(cfg.model_path, clogged_norms, str(ref_path), health_scores)
+
+    # Clean up reference now
+    print(f"  Cleaning up reference weights...")
+    shutil.rmtree(str(ref_path), ignore_errors=True)
+    gc.collect()
+
     # Save report
     report_dir = Path(cfg.output_dir)
     report_dir.mkdir(parents=True, exist_ok=True)