Defrag validation harness — verify tensor surgery preserves model integrity

[joelteply]

Problem

The defrag algorithm slices weight matrices, updates config, and hopes the result is internally consistent. We have no validation. The bug fixed in #152 (LoRA on pruned hooks) sat undetected for cycles because nothing checked that the defragged model could actually produce correct outputs.

Validation Harness

Build tests/defrag_validation.py that runs after every defrag operation:

1. Structural invariants

• q_proj.out_features equals num_heads * head_dim (post-defrag values)
• k_proj.out_features equals num_kv_heads * kv_head_dim
• v_proj.out_features equals num_kv_heads * kv_head_dim
• o_proj.in_features equals num_heads * head_dim
• GQA constraint: num_heads % num_kv_heads == 0
• All layers have the same (num_heads, num_kv_heads) post-defrag (or per-layer config tracked correctly)

2. Config matches tensors

• model.config.num_attention_heads matches actual q_proj.out_features // head_dim
• model.config.num_key_value_heads matches actual k_proj.out_features // kv_head_dim
• model.config.head_dim (if set) matches actual

3. Forward pass works

• Run a single forward pass with a small batch through the defragged model
• No dimension mismatch errors
• Output shape matches expected (batch, seq, vocab)

4. Output sanity (semantic preservation)

• Compare logits on the SAME input before vs after defrag
• For pruned heads with low importance: cosine similarity > 0.85
• For pruned heads with high importance: log warning (aggressive prune)
• KL divergence between distributions should be bounded

5. Save/load roundtrip

• Defrag → save → load fresh → run forward pass
• Output identical to pre-save defrag model
• Catches state_dict shape mismatches that crash on load

6. Multi-cycle stability

• Defrag → train → defrag → train (3 cycles)
• Perplexity must not exceed 2× baseline at any point
• If it does, abort and dump state for debugging

Why

• Caught the LoRA-on-pruned-hooks bug immediately
• Catches future tensor surgery regressions
• Validates that experimental pruning strategies don't silently corrupt models
• Provides confidence to publish models with attestation

Related

• Defrag: activation-aware head reorganization, not blind removal #152 (3-cycle prune destroys models — the bug this would have caught)
• continuum #829 (forge integrity)