hugging/td_fuse/selfimprove.py

"""
TD Self-Improvement v7 — SPEED OPTIMIZED MEGA CYCLE

v7 speed changes (March 2026):
- Flash Attention 2: 30-40% faster attention computation
- TF32 math: ~15-20% faster matrix multiply on Ampere GPUs (A6000)
- Dynamic padding: pad to longest-in-batch, not fixed 1024 tokens
- 4 dataloader workers + pin_memory: parallel data loading
- 2 epochs (was 3): model at 95%+, third epoch had diminishing returns
- Combined: ~50-60% faster → estimated 3-4 hours per cycle (was 6-7)

v6 MEGA MODE (unchanged):
- 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 VRAM savings
- batch_size 4 + grad_accum 4 = effective batch 16
- Norm unclogging continues from cycle 2
"""

import torch
import time
import json
import math
import gc
import shutil
import random
import numpy as np
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, field


@dataclass
class SelfImproveConfig:
    """Configuration for full self-improvement."""
    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
    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
    save_steps: int = 100
    num_train_samples: int = 3000

    # Custom training data (from td_data_finder — targeted by weakness)
    custom_data_path: str = ""  # If set, load from this JSONL instead of hardcoded datasets

    # Layernorm unfreezing
    norm_kurtosis_threshold: float = 15.0

    # Scoring weights
    w_health: float = 0.55
    w_damage: float = 0.25
    w_coherence: float = 0.20


def apply_mega_mode(cfg: SelfImproveConfig):
    """Crank everything up for mega cycles (3+).

    v7 SPEED optimizations (March 2026):
    - Flash Attention 2: 30-40% faster attention
    - TF32 math: ~15-20% faster matmul on Ampere GPUs
    - Dynamic padding: no wasted compute on short samples
    - 4 dataloader workers: parallel data prep
    - 2 epochs (was 3): model is 95%+, diminishing returns on epoch 3
    Combined: ~50-60% faster → ~3-4 hours per cycle instead of 6-7
    """
    print("\n  *** MEGA MODE ENABLED (v7 — speed optimized) ***")
    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 = 2           # 2 epochs — model is 95%+, 3rd epoch = diminishing returns
    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"    Epochs: {cfg.train_epochs} (was 3 — 33% time saved)")
    print(f"    Speed: flash_attn2 + tf32 + dynamic_pad + gradient_checkpointing")


def auto_detect_model_path() -> Tuple[str, int]:
    """Find the best starting model and determine cycle number."""
    base = Path("td_fuse_outputs/self_improve")

    # Check for highest cycle output (cycle3, cycle2, etc.)
    if base.exists():
        for n in range(50, 0, -1):
            cycle_dir = base / f"improved_cycle{n}"
            if cycle_dir.exists() and list(cycle_dir.glob("*.safetensors")):
                print(f"  Found cycle {n} output at {cycle_dir}")
                return str(cycle_dir), n + 1

    # Check for old "improved" folder (cycle 1 output from v4)
    old_improved = base / "improved"
    if old_improved.exists() and list(old_improved.glob("*.safetensors")):
        print(f"  Found cycle 1 output at {old_improved}")
        return str(old_improved), 2

    # Fall back to reasoning_healed
    healed = Path("td_fuse_outputs/reasoning_healed")
    if healed.exists() and list(healed.glob("*.safetensors")):
        print(f"  Found healed model at {healed}")
        return str(healed), 1

    raise FileNotFoundError("No model found! Need reasoning_healed or cycle output.")


# ============================================================
# STEP 1: DEEP WEIGHT ANALYSIS
# ============================================================

def analyze_layer(name: str, tensor: torch.Tensor) -> Dict:
    """Full health analysis for a single weight tensor."""
    t = tensor.float()

    std_val = t.std().item()
    mean_val = t.mean().item()
    abs_mean_val = t.abs().mean().item()

    # Kurtosis/skewness: guard against overflow for extreme tensors
    # (e.g. degenerate norms with std near 0 → huge normalized values → inf in pow(4))
    try:
        normalized = (t - mean_val) / max(std_val, 1e-8)
        # Clamp to prevent inf/nan from pow(4) on extreme values
        normalized = normalized.clamp(-1e4, 1e4)
        kurtosis_val = normalized.pow(4).mean().item() - 3
        skewness_val = normalized.pow(3).mean().item()
        # Catch NaN/Inf that slipped through
        if math.isnan(kurtosis_val) or math.isinf(kurtosis_val):
            kurtosis_val = 0.0
        if math.isnan(skewness_val) or math.isinf(skewness_val):
            skewness_val = 0.0
    except Exception:
        kurtosis_val = 0.0
        skewness_val = 0.0

    stats = {
        "name": name,
        "shape": list(tensor.shape),
        "params": tensor.numel(),
        "dtype": str(tensor.dtype),
        "l2_norm": t.norm(2).item(),
        "frobenius_norm": t.norm('fro').item() if t.dim() >= 2 else t.norm(2).item(),
        "mean": mean_val,
        "std": std_val,
        "min": t.min().item(),
        "max": t.max().item(),
        "abs_mean": abs_mean_val,
        "sparsity": (t.abs() < 1e-6).float().mean().item(),
        "kurtosis": kurtosis_val,
        "skewness": skewness_val,
        "near_zero_pct": (t.abs() < 1e-5).float().mean().item(),
        "extreme_pct": (t.abs() > abs_mean_val * 5).float().mean().item() if abs_mean_val > 0 else 0.0,
    }

    if t.dim() == 2 and min(t.shape) >= 2:
        try:
            rows = min(512, t.shape[0])
            cols = min(512, t.shape[1])
            sample = t[:rows, :cols]
            svd = torch.linalg.svdvals(sample)
            stats["sv_max"] = svd[0].item()
            stats["sv_min"] = svd[-1].item()
            stats["sv_ratio"] = svd[0].item() / max(svd[-1].item(), 1e-10)
            sv_norm = svd / svd.sum()
            entropy = -(sv_norm * torch.log(sv_norm + 1e-10)).sum()
            stats["sv_effective_rank"] = torch.exp(entropy).item()
            stats["sv_max_rank"] = float(min(sample.shape))
            stats["sv_rank_utilization"] = stats["sv_effective_rank"] / stats["sv_max_rank"]
            top5_energy = (svd[:5] ** 2).sum() / (svd ** 2).sum()
            stats["sv_top5_energy"] = top5_energy.item()
            stats["stable_rank"] = (sample.norm('fro').item() ** 2) / max(svd[0].item() ** 2, 1e-10)
        except Exception:
            stats["sv_ratio"] = -1
            stats["sv_rank_utilization"] = -1

    return stats


def analyze_model_weights(model_path: str, label: str = "model") -> Dict:
    """Load and analyze every language layer."""
    print(f"\n{'='*50}")
    print(f"ANALYZING: {label}")
    print(f"Path: {model_path}")
    print(f"{'='*50}")

    from safetensors.torch import load_file
    import glob

    st_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
    if not st_files:
        raise FileNotFoundError(f"No safetensors files in {model_path}")

    all_stats = {}
    total_params = 0

    for st_file in st_files:
        print(f"  Loading {Path(st_file).name}...")
        try:
            state_dict = load_file(st_file, device="cpu")
        except Exception as e:
            print(f"  WARNING: Failed to load {st_file}: {e} — skipping shard")
            continue
        for name, tensor in state_dict.items():
            if any(skip in name for skip in ["visual.", "merger."]):
                continue
            try:
                stats = analyze_layer(name, tensor)
                all_stats[name] = stats
                total_params += stats["params"]
            except Exception as e:
                print(f"  WARNING: analyze_layer failed for {name}: {e}")
                continue
        del state_dict
        gc.collect()

    print(f"  Analyzed {len(all_stats)} tensors ({total_params/1e9:.2f}B params)")
    return {"label": label, "layers": all_stats, "total_params": total_params}


# ============================================================
# STEP 2: STANDALONE HEALTH SCORING
# ============================================================

def compute_health_scores(analysis: Dict) -> Dict[str, Dict]:
    """Score each layer's health. Higher = more room for improvement."""
    print("\n=== STANDALONE HEALTH ANALYSIS ===")
    health_scores = {}

    for name, stats in analysis["layers"].items():
        issues = []
        score = 0.0

        # Safe access — analyze_layer might have partial results if SVD failed etc.
        kurt = abs(stats.get("kurtosis", 0))
        if kurt > 10:
            score += 0.3
            issues.append(f"extreme kurtosis ({kurt:.1f})")
        elif kurt > 5:
            score += 0.15
            issues.append(f"high kurtosis ({kurt:.1f})")

        skew = abs(stats.get("skewness", 0))
        if skew > 2:
            score += 0.15
            issues.append(f"high skew ({skew:.2f})")

        rank_util = stats.get("sv_rank_utilization", -1)
        if rank_util > 0:
            if rank_util < 0.1:
                score += 0.3
                issues.append(f"very low rank utilization ({rank_util:.2%})")
            elif rank_util < 0.25:
                score += 0.15
                issues.append(f"low rank utilization ({rank_util:.2%})")
            top5 = stats.get("sv_top5_energy", 0)
            if top5 > 0.9:
                score += 0.2
                issues.append(f"top-5 SVD concentration ({top5:.2%})")
            elif top5 > 0.75:
                score += 0.1

        sv_ratio = stats.get("sv_ratio", -1)
        if sv_ratio > 0:
            if sv_ratio > 10000:
                score += 0.25
                issues.append(f"extreme condition number ({sv_ratio:.0f})")
            elif sv_ratio > 1000:
                score += 0.1

        dead = stats.get("near_zero_pct", 0)
        if dead > 0.3:
            score += 0.2
            issues.append(f"high dead neuron % ({dead:.1%})")
        elif dead > 0.15:
            score += 0.1

        extreme = stats.get("extreme_pct", 0)
        if extreme > 0.05:
            score += 0.15
            issues.append(f"many extreme weights ({extreme:.1%})")

        std_val = stats.get("std", 0)
        if std_val < 1e-4:
            score += 0.25
            issues.append(f"near-zero variance ({std_val:.6f})")
        elif std_val > 1.0:
            score += 0.15
            issues.append(f"high variance ({std_val:.4f})")

        health_scores[name] = {
            "health_score": min(score, 1.0),
            "issues": issues,
            "stats": stats,
        }

    return health_scores


# ============================================================
# STEP 3: REFERENCE COMPARISON
# ============================================================

def compute_damage_scores(model_analysis: Dict, ref_analysis: Dict) -> Dict[str, float]:
    """Compare each layer to reference."""
    print("\n=== REFERENCE COMPARISON (merge damage) ===")
    damages = {}
    matched = 0

    for name, m_stats in model_analysis["layers"].items():
        ref_key = None
        for c in [name, name.replace("model.", "", 1), f"model.{name}"]:
            if c in ref_analysis["layers"]:
                ref_key = c
                break
        if ref_key is None:
            continue

        r_stats = ref_analysis["layers"][ref_key]
        matched += 1

        # Safe .get() for all stats — partial analysis results possible
        norm_diff = abs(m_stats.get("l2_norm", 0) - r_stats.get("l2_norm", 0)) / max(r_stats.get("l2_norm", 1e-8), 1e-8)
        std_diff = abs(m_stats.get("std", 0) - r_stats.get("std", 0)) / max(r_stats.get("std", 1e-8), 1e-8)
        kurt_diff = abs(m_stats.get("kurtosis", 0) - r_stats.get("kurtosis", 0)) / max(abs(r_stats.get("kurtosis", 0)) + 1, 1e-8)
        sparsity_diff = abs(m_stats.get("sparsity", 0) - r_stats.get("sparsity", 0))

        sv_ratio_diff = 0
        if "sv_ratio" in m_stats and "sv_ratio" in r_stats:
            if m_stats["sv_ratio"] > 0 and r_stats["sv_ratio"] > 0:
                sv_ratio_diff = abs(
                    math.log(m_stats["sv_ratio"] + 1) - math.log(r_stats["sv_ratio"] + 1)
                ) / max(math.log(r_stats["sv_ratio"] + 1), 1e-8)

        combined = (0.30 * norm_diff + 0.25 * std_diff + 0.20 * kurt_diff +
                    0.15 * sv_ratio_diff + 0.10 * sparsity_diff)
        damages[name] = combined

    print(f"  Matched {matched} layers between model and reference")
    return damages


# ============================================================
# STEP 4: LAYER COHERENCE
# ============================================================

def compute_coherence_scores(analysis: Dict) -> Dict[int, float]:
    """Check layer-to-layer coherence."""
    print("\n=== LAYER COHERENCE ANALYSIS ===")

    layer_norms = {}
    layer_stds = {}

    for name, stats in analysis["layers"].items():
        parts = name.split(".")
        layer_num = None
        for j, part in enumerate(parts):
            if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
                layer_num = int(parts[j+1])
                break
        if layer_num is not None:
            if layer_num not in layer_norms:
                layer_norms[layer_num] = []
                layer_stds[layer_num] = []
            layer_norms[layer_num].append(stats.get("l2_norm", 0))
            layer_stds[layer_num].append(stats.get("std", 0))

    avg_norms = {l: sum(ns)/len(ns) for l, ns in layer_norms.items()}
    avg_stds = {l: sum(ss)/len(ss) for l, ss in layer_stds.items()}
    sorted_layers = sorted(avg_norms.keys())
    coherence_issues = {}

    for i in range(len(sorted_layers)):
        l = sorted_layers[i]
        score = 0.0
        if i > 0:
            prev = sorted_layers[i-1]
            norm_jump = abs(avg_norms[l] - avg_norms[prev]) / max(avg_norms[prev], 1e-8)
            std_jump = abs(avg_stds[l] - avg_stds[prev]) / max(avg_stds[prev], 1e-8)
            if norm_jump > 0.5:
                score += norm_jump * 0.5
            if std_jump > 0.5:
                score += std_jump * 0.3
        if i < len(sorted_layers) - 1:
            nxt = sorted_layers[i+1]
            norm_jump = abs(avg_norms[l] - avg_norms[nxt]) / max(avg_norms[nxt], 1e-8)
            if norm_jump > 0.5:
                score += norm_jump * 0.3
        coherence_issues[l] = min(score, 1.0)

    return coherence_issues


# ============================================================
# STEP 5: RANK LAYERS
# ============================================================

def rank_layers_for_improvement(
    health_scores: Dict[str, Dict],
    damage_scores: Dict[str, float],
    coherence_scores: Dict[int, float],
    cfg: SelfImproveConfig
) -> Tuple[List[int], Dict]:
    """Combine all signals into improvement potential per layer."""
    print("\n=== RANKING LAYERS BY IMPROVEMENT POTENTIAL ===")

    layer_health = {}
    layer_damage = {}

    for name, h in health_scores.items():
        parts = name.split(".")
        layer_num = None
        for j, part in enumerate(parts):
            if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
                layer_num = int(parts[j+1])
                break
        if layer_num is not None:
            if layer_num not in layer_health:
                layer_health[layer_num] = []
                layer_damage[layer_num] = []
            layer_health[layer_num].append(h["health_score"])
            if name in damage_scores:
                layer_damage[layer_num].append(damage_scores[name])

    layer_scores = {}
    layer_details = {}

    for layer_num in layer_health:
        avg_health = sum(layer_health[layer_num]) / len(layer_health[layer_num])
        avg_damage = (sum(layer_damage[layer_num]) / len(layer_damage[layer_num])
                     if layer_damage.get(layer_num) else 0.0)
        coherence = coherence_scores.get(layer_num, 0.0)
        combined = (cfg.w_health * avg_health + cfg.w_damage * avg_damage +
                    cfg.w_coherence * coherence)
        layer_scores[layer_num] = combined
        layer_details[layer_num] = {
            "health": avg_health, "damage": avg_damage,
            "coherence": coherence, "combined": combined,
        }

    sorted_layers = sorted(layer_scores.items(), key=lambda x: x[1], reverse=True)

    print(f"\n  {'Layer':>7}  {'Combined':>9}  {'Health':>8}  {'Damage':>8}  {'Coherence':>10}  {'Status'}")
    print(f"  {'-'*65}")
    for layer_num, score in sorted_layers:
        d = layer_details[layer_num]
        bar = "█" * int(score * 50)
        target = " ← TARGET" if score >= cfg.min_improvement_score else ""
        print(f"  Layer {layer_num:>2}  {score:>9.4f}  {d['health']:>8.4f}  {d['damage']:>8.4f}  {d['coherence']:>10.4f}  {bar}{target}")

    target_layers = []
    for layer_num, score in sorted_layers:
        if score >= cfg.min_improvement_score and len(target_layers) < cfg.top_n_layers:
            target_layers.append(layer_num)

    if not target_layers:
        target_layers = [l for l, _ in sorted_layers[:cfg.top_n_layers]]

    target_layers.sort()
    print(f"\n  TARGET LAYERS FOR IMPROVEMENT: {target_layers}")

    # Print issues
    all_issues = []
    for name, h in health_scores.items():
        parts = name.split(".")
        layer_num = None
        for j, part in enumerate(parts):
            if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
                layer_num = int(parts[j+1])
                break
        if layer_num in target_layers and h["issues"]:
            for issue in h["issues"]:
                all_issues.append(f"  Layer {layer_num} {name.split('.')[-2]}: {issue}")
    if all_issues:
        print(f"\n  Key issues found:")
        for issue in all_issues[:30]:
            print(f"    {issue}")

    return target_layers, layer_details


# ============================================================
# STEP 5b: DETECT CLOGGED LAYERNORMS
# ============================================================

def detect_clogged_norms(health_scores: Dict[str, Dict], threshold: float = 15.0) -> List[str]:
    """
    Find layernorm and k_norm parameters with extreme kurtosis.
    These are the "clogged pipes" that LoRA can't fix.
    Returns list of parameter name patterns to unfreeze.
    """
    clogged = []
    for name, h in health_scores.items():
        if any(norm_type in name.lower() for norm_type in ["layernorm", "ln_", "_norm"]):
            stats = h.get("stats", {}) if isinstance(h, dict) else {}
            kurt = abs(stats.get("kurtosis", 0))
            if kurt > threshold:
                clogged.append(name)

    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")):
        try:
            state = load_file(st_file, device="cpu")
        except Exception as e:
            print(f"  WARNING: Failed to load reference {st_file}: {e} — skipping")
            continue
        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"))

    if not st_files:
        print("  WARNING: No safetensors files found in model path! Skipping repair.")
        return model_path

    for st_file in st_files:
        try:
            state = load_file(st_file, device="cpu")
        except Exception as e:
            print(f"  WARNING: Failed to load {st_file}: {e} — skipping this shard")
            continue
        changed = False

        for name in list(state.keys()):
            if name in ref_norms:
                # Guard: health_scores might not have this name
                if name not in health_scores or "stats" not in health_scores[name]:
                    print(f"    {name}: not in health_scores — SKIP")
                    continue
                if "kurtosis" not in health_scores[name]["stats"]:
                    print(f"    {name}: no kurtosis in stats — SKIP")
                    continue
                kurt = abs(health_scores[name]["stats"]["kurtosis"])

                # BUG FIX: Check if reference is actually better before blending.
                # Some reference norms have HIGHER kurtosis than our model's.
                # Blending toward a worse reference makes things worse!
                # Example from cycle 3: layers.5.input_layernorm went 28.5 → 43.4
                ref_t = ref_norms[name].float()
                ref_kurt = abs(((ref_t - ref_t.mean()) / max(ref_t.std().item(), 1e-8)).pow(4).mean().item() - 3)

                if ref_kurt >= kurt:
                    print(f"    {name}: kurtosis {kurt:.1f} (ref is worse: {ref_kurt:.1f}) — SKIP")
                    continue

                # 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:
            try:
                save_file(state, st_file)
            except Exception as e:
                print(f"  ERROR: Failed to save repaired weights to {st_file}: {e}")
                print(f"  Disk may be full! Norm repair aborted for this shard.")
        del state

    gc.collect()
    print(f"\n  Repaired {repaired} norms via reference blending")
    return model_path


# ============================================================
# STEP 6: LoRA + UNFROZEN NORMS IMPROVEMENT
# ============================================================

def improve_with_lora(model_path: str, target_layers: List[int],
                      clogged_norms: List[str], cfg: SelfImproveConfig,
                      cycle_num: int = 2):
    """
    Two-pronged improvement:
    1. LoRA (r=64) on projection weights in target layers
    2. Direct training of clogged layernorm/k_norm parameters

    This is the "unclogging" — LoRA reshapes projections while
    unfrozen norms fix the distribution bottlenecks.
    """
    print(f"\n{'='*50}")
    print(f"LORA + NORM UNCLOGGING")
    print(f"  LoRA layers: {target_layers}")
    print(f"  Clogged norms to unfreeze: {len(clogged_norms)}")
    print(f"{'='*50}")

    from transformers import AutoModelForImageTextToText, AutoTokenizer, TrainingArguments, Trainer
    from peft import LoraConfig, get_peft_model, TaskType
    from datasets import load_dataset
    from torch.utils.data import Dataset

    # ── DISK SPACE CHECK ──
    # Qwen3-VL-8B saves as ~16GB. Need at least 25GB free for save + temp files.
    try:
        disk = shutil.disk_usage(cfg.output_dir if Path(cfg.output_dir).exists() else ".")
        free_gb = disk.free / 1e9
        print(f"  Disk space: {free_gb:.1f} GB free")
        if free_gb < 25:
            print(f"  WARNING: Only {free_gb:.1f} GB free! Need ~25 GB for training + save.")
            if free_gb < 10:
                print(f"  CRITICAL: Less than 10 GB free — ABORTING to prevent disk-full corruption!")
                return model_path
    except Exception as e:
        print(f"  Could not check disk space: {e} — continuing anyway")

    # ── CHECKPOINT VALIDATION ──
    model_dir = Path(model_path)
    if model_dir.is_dir():
        if not (model_dir / "config.json").exists():
            print(f"  ERROR: No config.json in {model_path} — stale or corrupt checkpoint!")
            return model_path
        st_files = list(model_dir.glob("*.safetensors"))
        if not st_files and not list(model_dir.glob("*.bin")):
            print(f"  ERROR: No model weights in {model_path} — checkpoint is empty!")
            return model_path

    # Load model with OOM guard
    # SPEED: Enable TF32 for ~15-20% faster matmul on Ampere GPUs (A6000)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    print("  Loading model...")
    try:
        model = AutoModelForImageTextToText.from_pretrained(
            model_path, dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True,
            attn_implementation="flash_attention_2"  # SPEED: 30-40% faster attention
        )
    except (RuntimeError, Exception) as e:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        print(f"  CRITICAL: Failed to load model: {e}")
        return model_path
    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,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=0.0,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        layers_to_transform=target_layers,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )
    try:
        model = get_peft_model(model, lora_config)
    except Exception as e:
        print(f"\n  CRITICAL: get_peft_model failed: {e}")
        print(f"  LoRA config may be incompatible with this model.")
        del model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return model_path

    # UNFREEZE clogged norms — the key cycle 2 fix
    unfrozen_count = 0
    unfrozen_params = 0
    for name, param in model.named_parameters():
        # Check if this parameter matches any clogged norm
        # Parameter names in PEFT model have "base_model.model." prefix
        clean_name = name.replace("base_model.model.", "").replace("base_model.", "")
        for clogged_name in clogged_norms:
            if clean_name == clogged_name or clogged_name in clean_name:
                param.requires_grad = True
                unfrozen_count += 1
                unfrozen_params += param.numel()
                break

    if unfrozen_count == 0 and clogged_norms:
        print(f"\n  WARNING: Found {len(clogged_norms)} clogged norms but matched 0!")
        print(f"    Expected names like: {clogged_norms[:3]}")
        # Debug: print some actual param names
        for name, _ in list(model.named_parameters())[:5]:
            print(f"    Actual param: {name}")
    print(f"\n  LoRA trainable + {unfrozen_count} unfrozen norm layers ({unfrozen_params:,} params)")
    model.print_trainable_parameters()

    # Load training data — priority-weighted categories
    # MEGA MODE: code/math/reasoning/creativity get 2x allocation
    print("\n  Loading training data...")
    train_texts = []
    total = cfg.num_train_samples

    # === CUSTOM DATA PATH (from td_data_finder — targeted by weakness) ===
    # If td_loop found weakness-targeted data, use that instead of hardcoded datasets.
    # This is what makes the self-improvement loop actually self-improving:
    # each cycle trains on different data based on what the model is bad at.
    if cfg.custom_data_path and Path(cfg.custom_data_path).exists():
        print(f"  Using TARGETED training data from: {cfg.custom_data_path}")
        with open(cfg.custom_data_path) as _f:
            for line in _f:
                line = line.strip()
                if not line:
                    continue
                try:
                    item = json.loads(line)
                    text = item.get("text", "")
                    if text and len(text) > 20:
                        train_texts.append(text)
                except json.JSONDecodeError:
                    continue
                if len(train_texts) >= total:
                    break
        print(f"    Loaded {len(train_texts)} targeted samples")
        if len(train_texts) < 200:
            print(f"    WARNING: Only got {len(train_texts)} targeted samples (very low).")
            print(f"    Supplementing with generic data...")
            # Fall through to hardcoded data loading below to fill the gap
        else:
            # Got targeted data — DON'T dilute with easy generic data
            # Even 500 hard targeted samples > 8000 easy samples the model already knows
            # Training loss was 0.0000 when we mixed targeted + generic (model knew the generic data)
            print(f"    Using {len(train_texts)} targeted samples ONLY (no generic filler)")
            print(f"    Reason: generic data causes loss=0 because model already knows it")
            total = len(train_texts)  # Adjust total so we don't try to supplement

    # Only load hardcoded datasets if we don't have enough targeted data
    # Only load generic datasets if targeted data wasn't enough
    # When alloc values are 0, each section's `if count >= alloc[key]: break`
    # exits immediately, loading 0 samples. Clean skip, no indent changes needed.
    alloc = {"code": 0, "math": 0, "reasoning": 0, "creativity": 0,
             "instruction_following": 0, "general": 0}
    if len(train_texts) < total:
        remaining_needed = total - len(train_texts)
        if cfg.custom_data_path and train_texts:
            print(f"    Need {remaining_needed} more samples from generic datasets...")
            scale = remaining_needed / total
        else:
            scale = 1.0

        if cfg.mega_mode:
            alloc = {"code": int(total*0.25*scale), "math": int(total*0.25*scale),
                     "reasoning": int(total*0.19*scale), "creativity": int(total*0.12*scale),
                     "instruction_following": int(total*0.12*scale), "general": int(total*0.07*scale)}
        else:
            per = int((total // 5) * scale)
            alloc = {"code": per, "math": per, "reasoning": per,
                     "creativity": 0, "instruction_following": per, "general": per}
    else:
        print(f"    Targeted data sufficient ({len(train_texts)} samples). Skipping generic datasets.")

    # 1. MATH (top priority)
    try:
        count = 0
        if alloc["math"] > 0:
            ds = load_dataset("openai/gsm8k", "main", split="train")
            for row in ds:
                if count >= alloc["math"]:
                    break
                q, a = row.get("question", ""), row.get("answer", "")
                if q and a:
                    chat = (f"<|im_start|>user\n{q}\nThink step by step.<|im_end|>\n"
                           f"<|im_start|>assistant\n<think>\n{a}\n</think><|im_end|>")
                    train_texts.append(chat)
                    count += 1
        print(f"    Math (GSM8K): {count}")
    except Exception as e:
        print(f"    WARNING: GSM8K failed: {e}")

    # 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:
        count = 0
        if alloc["code"] > 0:
            ds = load_dataset("sahil2801/CodeAlpaca-20k", split="train")
            for row in ds:
                if count >= alloc["code"]:
                    break
                inst = row.get("instruction", "")
                inp = row.get("input", "")
                out = row.get("output", "")
                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 (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:
        count = 0
        if alloc["reasoning"] > 0:
            ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
            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: MMLU-Pro failed: {e}")

    # 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:
        count = 0
        if alloc.get("instruction_following", 0) > 0:
            ds = load_dataset("yahma/alpaca-cleaned", split="train")
            for row in ds:
                if count >= alloc["instruction_following"]:
                    break
                inst = row.get("instruction", "")
                inp = row.get("input", "")
                out = row.get("output", "")
                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 (Alpaca): {count}")
    except Exception as e:
        print(f"    WARNING: Alpaca failed: {e}")

    # 6. GENERAL TEXT (lowest priority)
    try:
        count = 0
        if alloc.get("general", 0) > 0:
            ds = load_dataset("neuralmagic/LLM_compression_calibration", split="train")
            for row in ds:
                if count >= alloc["general"]:
                    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}")

    print(f"\n  Total training samples: {len(train_texts)}")

    if len(train_texts) < 100:
        print("  ERROR: Not enough training data!")
        return model_path

    # === CONTAMINATION CHECK ===
    # Remove any training samples that match benchmark/holdout questions.
    # Without this, benchmark scores become meaningless after a few cycles.
    try:
        from td_contamination import check_contamination
        before = len(train_texts)
        train_texts = check_contamination(train_texts)
        if len(train_texts) < before:
            print(f"  Contamination guard removed {before - len(train_texts)} samples")
    except ImportError:
        print("  WARNING: td_contamination.py not found — skipping contamination check")

    # === DIVERSITY CHECK (anti-collapse) ===
    # If too many training samples are near-duplicates, the model collapses
    # toward repeating the same patterns. Detect and warn early.
    unique_prefixes = set()
    for t in train_texts:
        # Use first 100 chars as a rough fingerprint
        prefix = t[:100].strip().lower()
        unique_prefixes.add(prefix)
    diversity_ratio = len(unique_prefixes) / max(len(train_texts), 1)
    print(f"  Diversity check: {len(unique_prefixes)}/{len(train_texts)} unique prefixes ({diversity_ratio:.0%})")
    if diversity_ratio < 0.5:
        print("  WARNING: Low diversity (<50%) — high risk of mode collapse!")
        print("  Deduplicating training data...")
        seen = set()
        deduped = []
        for t in train_texts:
            prefix = t[:100].strip().lower()
            if prefix not in seen:
                seen.add(prefix)
                deduped.append(t)
        removed = len(train_texts) - len(deduped)
        train_texts = deduped
        print(f"  Removed {removed} near-duplicate samples. Now {len(train_texts)} unique.")
    elif diversity_ratio < 0.7:
        print("  NOTE: Moderate diversity (50-70%) — monitor for early signs of collapse")

    random.shuffle(train_texts)

    # Tokenize — SPEED: dynamic padding (pad to longest in batch, not fixed 1024)
    class ImproveDataset(Dataset):
        def __init__(self, texts, tokenizer, max_len=1024):
            self.data = []
            for t in texts:
                e = tokenizer(t, truncation=True, max_length=max_len,
                              return_tensors="pt")  # No padding here — collator handles it
                self.data.append({
                    "input_ids": e["input_ids"].squeeze(),
                    "attention_mask": e["attention_mask"].squeeze(),
                    "labels": e["input_ids"].squeeze(),
                })
        def __len__(self): return len(self.data)
        def __getitem__(self, i): return self.data[i]

    print("  Tokenizing...")
    dataset = ImproveDataset(train_texts, tokenizer)

    if len(dataset) == 0:
        print("  ERROR: Dataset is empty after tokenization!")
        return model_path

    out_dir = Path(cfg.output_dir) / "train_output"
    out_dir.mkdir(parents=True, exist_ok=True)

    effective_batch = cfg.train_batch * cfg.train_grad_accum
    total_steps = (len(dataset) * cfg.train_epochs) // max(effective_batch, 1)

    # Set up optimizer with different LR for norms vs LoRA
    # Group parameters: norms get higher LR
    norm_params = []
    lora_params = []
    for name, param in model.named_parameters():
        if param.requires_grad:
            if any(nt in name.lower() for nt in ["layernorm", "ln_", "_norm"]):
                norm_params.append(param)
            else:
                lora_params.append(param)

    print(f"\n  Parameter groups: {len(lora_params)} LoRA params, {len(norm_params)} norm params")

    if norm_params:
        from torch.optim import AdamW
        optimizer = AdamW([
            {"params": lora_params, "lr": cfg.learning_rate},
            {"params": norm_params, "lr": cfg.norm_learning_rate},
        ], weight_decay=0.01)
        print(f"  LoRA LR: {cfg.learning_rate}, Norm LR: {cfg.norm_learning_rate}")
    else:
        optimizer = None  # Let Trainer handle it

    args = TrainingArguments(
        output_dir=str(out_dir),
        num_train_epochs=cfg.train_epochs,
        per_device_train_batch_size=cfg.train_batch,
        gradient_accumulation_steps=cfg.train_grad_accum,
        learning_rate=cfg.learning_rate,
        bf16=True,
        tf32=True,  # SPEED: ~15-20% faster on Ampere GPUs
        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,
        save_total_limit=2,
        warmup_ratio=0.05,
        lr_scheduler_type="cosine",
        optim="adamw_torch",
        report_to="none",
        dataloader_num_workers=4,    # SPEED: parallel data loading
        dataloader_pin_memory=True,  # SPEED: faster CPU→GPU transfer
    )

    # SPEED: Dynamic padding collator — pads to longest in batch, not fixed 1024
    # Custom because labels need -100 padding (ignore index), not tokenizer pad_id
    def dynamic_pad_collator(batch):
        max_len = max(item["input_ids"].size(0) for item in batch)
        pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
        input_ids, attention_mask, labels = [], [], []
        for item in batch:
            seq_len = item["input_ids"].size(0)
            pad_len = max_len - seq_len
            input_ids.append(torch.cat([item["input_ids"], torch.full((pad_len,), pad_id, dtype=torch.long)]))
            attention_mask.append(torch.cat([item["attention_mask"], torch.zeros(pad_len, dtype=torch.long)]))
            labels.append(torch.cat([item["labels"], torch.full((pad_len,), -100, dtype=torch.long)]))
        return {"input_ids": torch.stack(input_ids), "attention_mask": torch.stack(attention_mask), "labels": torch.stack(labels)}

    trainer_kwargs = {
        "model": model,
        "processing_class": tokenizer,
        "train_dataset": dataset,
        "data_collator": dynamic_pad_collator,
        "args": args,
    }
    if optimizer is not None:
        trainer_kwargs["optimizers"] = (optimizer, None)

    trainer = Trainer(**trainer_kwargs)

    print(f"\n  Training: ~{total_steps} steps on {len(target_layers)} layers + {len(clogged_norms)} norms")
    print(f"  Data: {len(train_texts)} samples × {cfg.train_epochs} epochs")
    print(f"  LoRA: r={cfg.lora_r}, alpha={cfg.lora_alpha}")

    try:
        train_result = trainer.train()
    except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
        err_str = str(e).lower()
        if "out of memory" in err_str or "cuda" in err_str:
            print(f"\n  OOM DURING TRAINING: {e}")
            print("  Cleaning up GPU memory and returning previous model...")
            del model, trainer
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            shutil.rmtree(str(out_dir), ignore_errors=True)
            return model_path
        else:
            raise  # Re-raise non-OOM RuntimeErrors

    # ── NaN/Inf LOSS CHECK ──
    # If training produced NaN loss, the model weights are garbage — abort
    try:
        final_loss = train_result.training_loss
        if final_loss is not None and (math.isnan(final_loss) or math.isinf(final_loss)):
            print(f"\n  CRITICAL: Training loss is {final_loss}! Model weights are corrupted.")
            print(f"  Falling back to previous model: {model_path}")
            del model, trainer
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            shutil.rmtree(str(out_dir), ignore_errors=True)
            return model_path
        print(f"\n  Training complete. Final loss: {final_loss:.4f}")
    except (AttributeError, TypeError):
        print("\n  Training complete. (Could not read final loss)")

    # Merge LoRA back
    print("\n  Merging LoRA adapters into model...")
    try:
        merged = model.merge_and_unload()
    except Exception as e:
        print(f"\n  CRITICAL: LoRA merge failed: {e}")
        print(f"  Falling back to previous model: {model_path}")
        del model, trainer
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        shutil.rmtree(str(out_dir), ignore_errors=True)
        return model_path

    # Save
    save_dir = Path(cfg.output_dir) / f"improved_cycle{cycle_num}"
    save_dir.mkdir(parents=True, exist_ok=True)
    print(f"  Saving to {save_dir}...")
    try:
        merged.save_pretrained(str(save_dir), safe_serialization=True)
        tokenizer.save_pretrained(str(save_dir))
    except Exception as e:
        print(f"\n  CRITICAL: Model save failed: {e}")
        print(f"  Disk may be full! Falling back to previous model: {model_path}")
        del merged, model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        shutil.rmtree(str(save_dir), ignore_errors=True)
        shutil.rmtree(str(out_dir), ignore_errors=True)
        return model_path

    # Verify save was successful — check files exist and aren't tiny/corrupt
    st_files = list(save_dir.glob("*.safetensors"))
    if not st_files:
        print("  ERROR: No safetensors files saved! Disk may be full.")
        print(f"  Falling back to previous model: {model_path}")
        shutil.rmtree(str(save_dir), ignore_errors=True)
        return model_path
    sz = sum(f.stat().st_size for f in st_files) / 1e9
    if sz < 1.0:  # Qwen3-VL-8B should be ~16GB, anything under 1GB is corrupt
        print(f"  ERROR: Saved model suspiciously small ({sz:.2f} GB). Likely corrupt!")
        print(f"  Falling back to previous model: {model_path}")
        shutil.rmtree(str(save_dir), ignore_errors=True)
        return model_path
    print(f"  SAVED: {save_dir} ({sz:.1f} GB)")

    shutil.rmtree(str(out_dir), ignore_errors=True)
    del merged, model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return str(save_dir)


# ============================================================
# QUICK BENCHMARK — Sanity check actual capability
# ============================================================

def quick_benchmark(model_path: str) -> Dict:
    """Run 5 questions to verify the model actually works well."""
    print(f"\n{'='*50}")
    print(f"QUICK BENCHMARK — Testing actual capability")
    print(f"{'='*50}")

    from transformers import AutoModelForImageTextToText, AutoTokenizer

    # ── CHECKPOINT VALIDATION ──
    mp = Path(model_path)
    if mp.is_dir():
        if not (mp / "config.json").exists():
            print(f"  ERROR: No config.json in {model_path} — skipping quick benchmark")
            return {"score": 0, "total": 5, "results": [], "error": "no config.json"}
        if not list(mp.glob("*.safetensors")) and not list(mp.glob("*.bin")):
            print(f"  ERROR: No model weights in {model_path} — skipping quick benchmark")
            return {"score": 0, "total": 5, "results": [], "error": "no weights"}

    # ── MODEL LOAD with OOM guard ──
    try:
        model = AutoModelForImageTextToText.from_pretrained(
            model_path, dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True
        )
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    except (RuntimeError, Exception) as e:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        print(f"  CRITICAL: Failed to load model for quick benchmark: {e}")
        return {"score": 0, "total": 5, "results": [], "error": str(e)}

    tests = [
        {"q": "What is 47 * 23?", "expected": "1081", "type": "math"},
        {"q": "What is 156 + 289?", "expected": "445", "type": "math"},
        {"q": "If all roses are flowers and all flowers need water, do roses need water?",
         "expected": "yes", "type": "logic"},
        {"q": "Write a Python function to reverse a string.",
         "expected": "[::-1]", "type": "code"},
        {"q": "What comes next: 2, 4, 8, 16, ?",
         "expected": "32", "type": "reasoning"},
    ]

    results = []
    for test in tests:
        try:
            chat = f"<|im_start|>user\n{test['q']}<|im_end|>\n<|im_start|>assistant\n"
            ids = tokenizer(chat, return_tensors="pt").to(model.device)
            try:
                with torch.no_grad():
                    out = model.generate(**ids, max_new_tokens=200, do_sample=False)
            except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                print(f"\n  [OOM] {test['type']}: {test['q']}")
                results.append({"type": test["type"], "passed": False, "error": "OOM"})
                continue
            answer = tokenizer.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
            # Truncate garbage output
            if len(answer) > 5000:
                answer = answer[:5000]

            passed = test["expected"].lower() in answer.lower()
            status = "PASS" if passed else "FAIL"
            results.append({"type": test["type"], "passed": passed})

            print(f"\n  [{status}] {test['type']}: {test['q']}")
            print(f"    Answer: {answer[:200]}")
            print(f"    Expected to contain: {test['expected']}")
        except Exception as e:
            print(f"\n  [ERROR] {test['type']}: {test['q']} — {e}")
            results.append({"type": test["type"], "passed": False, "error": str(e)})

    score = sum(1 for r in results if r.get("passed", False))
    print(f"\n  BENCHMARK: {score}/{len(tests)} passed")

    del model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return {"score": score, "total": len(tests), "results": results}


# ============================================================
# MAIN
# ============================================================

def run_cycle(cfg: SelfImproveConfig = None):
    """Full self-improvement cycle with norm unclogging."""
    if cfg is None:
        cfg = SelfImproveConfig()

    # Auto-detect model path and cycle number
    if not cfg.model_path:
        cfg.model_path, cycle_num = auto_detect_model_path()
    else:
        # BUG FIX: Parse actual cycle number from path instead of assuming 2
        # "improved_cycle5" → cycle 6, "improved_cycle2" → cycle 3
        # Old code: `cycle_num = 2 if "improve" in cfg.model_path else 1`
        # That would ALWAYS save to improved_cycle2, overwriting previous!
        import re as _re
        cycle_match = _re.search(r'improved_cycle(\d+)', cfg.model_path)
        if cycle_match:
            cycle_num = int(cycle_match.group(1)) + 1
        elif "improve" in cfg.model_path:
            cycle_num = 2
        else:
            cycle_num = 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 v6 — {mode_label}")
    print(f"Cycle: {cycle_num}")
    print(f"Model: {cfg.model_path}")
    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)

    # STEP 1: Analyze
    print(f"\n{'='*40}")
    print("STEP 1/7: ANALYZE MODEL")
    print(f"{'='*40}")
    model_analysis = analyze_model_weights(cfg.model_path, "TD model")

    # STEP 2: Health scoring
    print(f"\n{'='*40}")
    print("STEP 2/7: HEALTH SCORING")
    print(f"{'='*40}")
    health_scores = compute_health_scores(model_analysis)
    unhealthy = sum(1 for h in health_scores.values() if h["health_score"] > 0.2)
    print(f"\n  Health summary: {unhealthy}/{len(health_scores)} tensors have notable issues")

    # STEP 3: Reference comparison
    print(f"\n{'='*40}")
    print("STEP 3/7: REFERENCE COMPARISON")
    print(f"{'='*40}")
    ref_path = Path(cfg.output_dir) / "reference_weights"
    ref_available = False
    try:
        if not ref_path.exists() or not list(ref_path.glob("*.safetensors")):
            print(f"  Downloading reference: {cfg.reference_model}...")
            from huggingface_hub import snapshot_download
            snapshot_download(
                cfg.reference_model, local_dir=str(ref_path),
                allow_patterns=["*.safetensors", "*.json"],
                ignore_patterns=["*.bin", "*.pt", "*.onnx"],
            )
        ref_analysis = analyze_model_weights(str(ref_path), "Qwen3-VL-8B (reference)")
        damage_scores = compute_damage_scores(model_analysis, ref_analysis)
        ref_available = True
    except Exception as e:
        print(f"  WARNING: Reference comparison failed: {e}")
        print(f"  Continuing without damage scores (norm repair + LoRA still work)")
        damage_scores = {}
        ref_analysis = None
    # DON'T delete reference yet — we need it for norm repair
    gc.collect()

    # STEP 4: Coherence
    print(f"\n{'='*40}")
    print("STEP 4/7: LAYER COHERENCE")
    print(f"{'='*40}")
    coherence_scores = compute_coherence_scores(model_analysis)

    # STEP 5: Rank + detect clogged norms
    print(f"\n{'='*40}")
    print("STEP 5/7: RANK LAYERS + DETECT CLOGGED NORMS")
    print(f"{'='*40}")
    target_layers, layer_details = rank_layers_for_improvement(
        health_scores, damage_scores, coherence_scores, cfg
    )

    # Detect clogged norms
    clogged_norms = detect_clogged_norms(health_scores, cfg.norm_kurtosis_threshold)
    if clogged_norms:
        print(f"\n  CLOGGED NORMS DETECTED ({len(clogged_norms)}):")
        for cn in clogged_norms:
            stats = health_scores.get(cn, {}).get("stats", {})
            kurt = abs(stats.get("kurtosis", 0))
            print(f"    {cn}: kurtosis={kurt:.1f}")
    else:
        print(f"\n  No clogged norms above threshold {cfg.norm_kurtosis_threshold}")

    # STEP 5c: REPAIR clogged norms surgically
    if clogged_norms and ref_available:
        print(f"\n{'='*40}")
        print("STEP 5c/7: SURGICAL NORM REPAIR")
        print(f"{'='*40}")
        try:
            repair_clogged_norms(cfg.model_path, clogged_norms, str(ref_path), health_scores)
        except Exception as e:
            print(f"  WARNING: Norm repair failed: {e}")
            print(f"  Continuing without repair — LoRA training can still help")
    elif clogged_norms and not ref_available:
        print(f"\n  Skipping norm repair — reference model not available")

    # Clean up reference now
    print(f"  Cleaning up reference weights...")
    shutil.rmtree(str(ref_path), ignore_errors=True)
    gc.collect()

    # Save report (non-critical)
    report_dir = Path(cfg.output_dir)
    report_dir.mkdir(parents=True, exist_ok=True)
    pre_report = {
        "cycle": cycle_num,
        "model_path": cfg.model_path,
        "target_layers": target_layers,
        "clogged_norms": clogged_norms,
        "lora_r": cfg.lora_r,
        "norm_lr": cfg.norm_learning_rate,
        "layer_details": {str(k): v for k, v in layer_details.items()},
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
    }
    try:
        with open(report_dir / f"pre_cycle{cycle_num}_report.json", "w") as f:
            json.dump(pre_report, f, indent=2)
    except OSError as e:
        print(f"  WARNING: Could not save pre-report: {e}")

    # STEP 6: Train
    print(f"\n{'='*40}")
    print("STEP 6/7: LORA + NORM UNCLOGGING")
    print(f"{'='*40}")
    improved_path = improve_with_lora(cfg.model_path, target_layers, clogged_norms, cfg, cycle_num)

    # STEP 7: Verify
    print(f"\n{'='*40}")
    print("STEP 7/7: VERIFY + BENCHMARK")
    print(f"{'='*40}")

    # Re-analyze health (non-critical — if this fails, cycle still succeeded)
    pre_avg = sum(h["health_score"] for h in health_scores.values()) / max(len(health_scores), 1)
    post_avg = pre_avg  # Default: same as before if analysis fails
    post_health = {}
    try:
        post_analysis = analyze_model_weights(improved_path, "TD improved")
        post_health = compute_health_scores(post_analysis)
        post_avg = sum(h["health_score"] for h in post_health.values()) / max(len(post_health), 1)
    except Exception as e:
        print(f"  WARNING: Post-training health analysis failed: {e}")
        print(f"  Cannot compare before/after — but training output is still valid.")

    # Per-layer comparison (only if post-analysis succeeded)
    if post_health:
        print(f"\n  Per-layer improvement:")
    else:
        print(f"\n  (Skipping per-layer comparison — post-analysis unavailable)")
    for layer_num in target_layers:
        if not post_health:
            break
        post_layer_scores = []
        for name, h in post_health.items():
            parts = name.split(".")
            for j, part in enumerate(parts):
                if part == "layers" and j + 1 < len(parts) and parts[j+1].isdigit():
                    if int(parts[j+1]) == layer_num:
                        post_layer_scores.append(h["health_score"])
                    break
        post_layer_avg = sum(post_layer_scores) / max(len(post_layer_scores), 1) if post_layer_scores else 0
        pre_h = layer_details.get(layer_num, {}).get("health", 0)
        delta = pre_h - post_layer_avg
        symbol = "+" if delta > 0 else "-" if delta < 0 else "="
        print(f"    Layer {layer_num:>2}: {pre_h:.4f} → {post_layer_avg:.4f} ({symbol}{abs(delta):.4f})")

    # Check clogged norms specifically
    if clogged_norms:
        print(f"\n  Clogged norm changes:")
        for cn in clogged_norms:
            pre_stats = health_scores.get(cn, {}).get("stats", {})
            pre_kurt = abs(pre_stats.get("kurtosis", 0))
            # Find matching post norm
            if cn in post_health:
                post_stats = post_health.get(cn, {}).get("stats", {})
                post_kurt = abs(post_stats.get("kurtosis", 0))
                delta = pre_kurt - post_kurt
                symbol = "+" if delta > 0 else "-"
                print(f"    {cn}: kurtosis {pre_kurt:.1f} → {post_kurt:.1f} ({symbol}{abs(delta):.1f})")

    # Quick benchmark (non-critical — don't crash if it fails)
    try:
        benchmark = quick_benchmark(improved_path)
    except Exception as e:
        print(f"  WARNING: Quick benchmark failed: {e}")
        benchmark = {"score": 0, "total": 0, "error": str(e)}

    elapsed = (time.time() - start) / 60
    health_improvement = (pre_avg - post_avg) / max(pre_avg, 1e-8) * 100

    print(f"\n{'='*60}")
    print(f"CYCLE {cycle_num} COMPLETE — {elapsed:.1f} min")
    print(f"{'='*60}")
    print(f"  Target layers: {target_layers}")
    print(f"  Clogged norms unfrozen: {len(clogged_norms)}")
    print(f"  Health BEFORE: {pre_avg:.6f}")
    print(f"  Health AFTER:  {post_avg:.6f}")
    print(f"  Health improvement: {health_improvement:.1f}%")
    print(f"  Benchmark: {benchmark.get('score', '?')}/{benchmark.get('total', '?')}")
    print(f"  Improved model: {improved_path}")
    print(f"{'='*60}")

    # Final report (non-critical)
    try:
        final_report = {
            **pre_report,
            "pre_avg_health": pre_avg,
            "post_avg_health": post_avg,
            "health_improvement_pct": health_improvement,
            "benchmark": benchmark,
            "duration_min": elapsed,
            "improved_path": improved_path,
        }
        report_save_path = Path(improved_path) / "cycle_report.json"
        with open(report_save_path, "w") as f:
            json.dump(final_report, f, indent=2)
    except Exception as e:
        print(f"  WARNING: Could not save cycle report: {e}")

    return improved_path


if __name__ == "__main__":
    run_cycle()