"""
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\n{a}\n<|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\n{a}\n<|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\nLet me reason through this.\n"
f"{answer_text}\n\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()