File size: 7,034 Bytes
5d61448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | """
Post-Merge Validation — run after EVERY merge step.
Tests:
1. Canary recall (did knowledge transfer?)
2. Perplexity check (did we break the model?)
3. Thinking mode (do <think> tags still work?)
4. Quick reasoning test (can it still think?)
Kill criteria: >10% performance drop on any test → abort merge.
Findings: #11, #22, #25
"""
import torch
import math
from transformers import AutoModelForCausalLM, AutoTokenizer
from .canary import test_all_canaries
from .config import MergeConfig
def validate_merged_model(
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
merged_sources: list[str],
cfg: MergeConfig,
baseline_perplexity: float = None,
) -> dict:
"""
Run full validation suite on a merged model.
Args:
model: The merged model to validate
tokenizer: The tokenizer
merged_sources: List of source models merged so far
cfg: Merge configuration
baseline_perplexity: Perplexity of the target model before merging
Returns:
Dict with test results and overall pass/fail
"""
print("\n" + "=" * 60)
print(f"VALIDATION — After merging: {', '.join(merged_sources)}")
print("=" * 60)
results = {
"canary": None,
"perplexity": None,
"thinking_mode": None,
"reasoning": None,
"overall": False,
}
# --- Test 1: Canary recall ---
canary_results = test_all_canaries(model, tokenizer, merged_sources)
passed_canaries = sum(1 for v in canary_results.values() if v)
total_canaries = len(canary_results)
results["canary"] = {
"passed": passed_canaries,
"total": total_canaries,
"ok": passed_canaries >= cfg.canary_pass_threshold,
"details": canary_results,
}
# --- Test 2: Perplexity ---
perplexity = compute_perplexity(model, tokenizer)
ppl_ok = True
if baseline_perplexity is not None:
ratio = perplexity / baseline_perplexity
ppl_ok = ratio < cfg.perplexity_threshold
print(f"\n[validate] Perplexity: {perplexity:.2f} (baseline: {baseline_perplexity:.2f}, ratio: {ratio:.2f})")
if not ppl_ok:
print(f"[validate] ⚠ Perplexity ratio {ratio:.2f} exceeds threshold {cfg.perplexity_threshold}")
else:
print(f"\n[validate] Perplexity: {perplexity:.2f} (no baseline to compare)")
results["perplexity"] = {"value": perplexity, "ok": ppl_ok}
# --- Test 3: Thinking mode ---
think_ok = test_thinking_mode(model, tokenizer)
results["thinking_mode"] = {"ok": think_ok}
# --- Test 4: Quick reasoning ---
reason_ok = test_reasoning(model, tokenizer)
results["reasoning"] = {"ok": reason_ok}
# --- Overall verdict ---
all_ok = (
results["canary"]["ok"]
and results["perplexity"]["ok"]
and results["thinking_mode"]["ok"]
and results["reasoning"]["ok"]
)
results["overall"] = all_ok
# Summary
print("\n" + "-" * 60)
print("VALIDATION SUMMARY")
print("-" * 60)
print(f" Canary recall: {'✓' if results['canary']['ok'] else '✗'} ({passed_canaries}/{total_canaries})")
print(f" Perplexity: {'✓' if ppl_ok else '✗'} ({perplexity:.2f})")
print(f" Thinking mode: {'✓' if think_ok else '✗'}")
print(f" Reasoning: {'✓' if reason_ok else '✗'}")
print(f" OVERALL: {'✓ PASS' if all_ok else '✗ FAIL — consider aborting'}")
print("-" * 60)
return results
def compute_perplexity(
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
test_texts: list[str] = None,
) -> float:
"""
Compute perplexity on a small test set.
Lower perplexity = model is more confident about predicting text.
A big spike after merging means the model was damaged.
"""
if test_texts is None:
test_texts = [
"The quick brown fox jumps over the lazy dog.",
"In mathematics, a prime number is a natural number greater than 1.",
"def fibonacci(n):\n if n <= 1:\n return n\n return fibonacci(n-1) + fibonacci(n-2)",
"The theory of general relativity describes gravity as the curvature of spacetime.",
"To solve 3x + 7 = 22, subtract 7 from both sides to get 3x = 15, then divide by 3.",
]
model.eval()
total_loss = 0.0
total_tokens = 0
for text in test_texts:
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs, labels=inputs["input_ids"])
total_loss += outputs.loss.item() * inputs["input_ids"].shape[1]
total_tokens += inputs["input_ids"].shape[1]
avg_loss = total_loss / total_tokens
perplexity = math.exp(avg_loss)
return perplexity
def test_thinking_mode(
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
) -> bool:
"""
Test if the model still uses <think> tags for reasoning.
The thinking mode is Qwen3's special feature — if it's gone,
the merge damaged something critical.
"""
prompt = "Solve step by step: What is 15 × 13?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=200,
temperature=0.7,
do_sample=True,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=False)
# Check for thinking tags
has_think_open = "<think>" in response
has_think_close = "</think>" in response
passed = has_think_open and has_think_close
print(f"\n[validate] Thinking mode test:")
print(f" Prompt: {prompt}")
print(f" Response: {response[:200]}...")
print(f" <think>: {'✓ found' if has_think_open else '✗ missing'}")
print(f" </think>: {'✓ found' if has_think_close else '✗ missing'}")
print(f" Status: {'✓ PASS' if passed else '✗ FAIL'}")
return passed
def test_reasoning(
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
) -> bool:
"""
Quick reasoning sanity check — can the model still do basic math?
This catches catastrophic failures where the merge produced gibberish.
"""
prompt = "What is 7 + 8?"
expected_answer = "15"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=50,
temperature=0.1,
do_sample=False,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
passed = expected_answer in response
print(f"\n[validate] Quick reasoning test:")
print(f" Prompt: {prompt}")
print(f" Expected: {expected_answer}")
print(f" Got: {response}")
print(f" Status: {'✓ PASS' if passed else '✗ FAIL'}")
return passed
|