| """ |
| 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, |
| } |
|
|
| |
| 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, |
| } |
|
|
| |
| 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} |
|
|
| |
| think_ok = test_thinking_mode(model, tokenizer) |
| results["thinking_mode"] = {"ok": think_ok} |
|
|
| |
| reason_ok = test_reasoning(model, tokenizer) |
| results["reasoning"] = {"ok": reason_ok} |
|
|
| |
| all_ok = ( |
| results["canary"]["ok"] |
| and results["perplexity"]["ok"] |
| and results["thinking_mode"]["ok"] |
| and results["reasoning"]["ok"] |
| ) |
| results["overall"] = all_ok |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
