Structural pruning defrag — mask first, compact later (like disk defrag)

[joelteply]

Insight from Joel

Pruning should work in two phases, like disk defragmentation:

Phase 1: Mask (fast, reversible)

• Mark heads as pruned with a mask (don't modify weights)
• Forward hooks zero the output of masked heads
• Model still has same tensor dimensions — no restructuring
• Can unmask if we change our mind
• Training sees the mask via hooks, not zeroed weights (avoids NaN gradients)

Phase 2: Defrag/Compact (batch, permanent)

• After accumulating enough pruned heads (e.g., >20% pruned)
• Actually restructure tensors: remove masked dimensions
• Model gets physically smaller (less VRAM, faster inference)
• Save the compacted model
• Like defrag: data is gone, now reclaim the space

Why This Matters

Current approach zeros weights directly → NaN during backward pass (0/0 in softmax gradients). Masking avoids this because base weights are intact — the mask just zeros output AFTER the computation.

Threshold-based Defrag

• Don't defrag after every prune cycle (expensive restructuring)
• Defrag when cumulative pruning exceeds threshold (20%, 30%, etc.)
• Or defrag at the END of all forge cycles, once

Implementation

# Phase 1: mask
pruning_mask[layer_idx][head_idx] = False  # masked off
# Forward hook: output[:, :, masked_heads] = 0

# Phase 2: defrag (when ready)
model = structurally_prune(model, pruning_mask)
# Tensors resized, model is physically smaller
model.save_pretrained("compacted/")

Dependencies

• Python forge pipeline — ship Qwen3.5 forged models now #89 — Python forge pipeline
• Rust/Candle forge pipeline — replace Python with native inference + training #88 — Rust forge pipeline (Candle can do structural pruning natively)

[joelteply]

Design Findings — defrag_model.py prototype written

What Changes Per Layer

q_proj: [hidden, num_heads × head_dim] → [hidden, surviving_heads × head_dim]
k_proj: [hidden, num_kv × head_dim]    → [hidden, surviving_kv × head_dim]
v_proj: same as k_proj
o_proj: [num_heads × head_dim, hidden]  → [surviving_heads × head_dim, hidden]

What Stays The Same

• hidden_size (residual stream)
• MLP (gate/up/down projections)
• Embeddings, layer norms, lm_head

Why It's Safe

The residual stream is: x = x + Attention(x) + MLP(x)
Attention shrinks internally (fewer heads), but O_proj still outputs [hidden_size].
The residual addition works because hidden_size is unchanged.

GQA Handling

Qwen3.5 uses 16 query heads / 4 KV heads. A KV head is only removed if ALL its query heads (group of 4) are dead. Partial groups keep the KV head.

Size Savings (Qwen3.5-27B, 30% pruning)

• Attention is ~30% of model parameters
• 30% heads pruned → ~9% total parameter reduction
• Q4 GGUF: ~15GB → ~13.5GB (saves 1.5GB)
• Plus: faster inference (fewer heads to compute per forward pass)

Variable Heads Per Layer

After defrag, different layers have different head counts. Standard HF config doesn't support this. Options:

• v1: Pad to max surviving count (keeps HF compat, wastes a little)
• v2: Custom per-layer config (better savings, needs custom loader)

Implementation

scripts/defrag_model.py is written and ready to test:

1. Loads the forged model (CPU only — no VRAM needed)
2. Detects dead heads by Q-projection weight norms
3. Slices attention tensors to remove dead dimensions
4. Saves smaller model + metadata
5. Runs AFTER forging, on the saved safetensors

Safe: operates on saved files, doesn't touch the training process.

[joelteply]

Inline Defrag: Speed Up Training as Model Shrinks (Joel's insight)

Current: defrag is post-processing only. Model stays full size during training.

Better: defrag DURING training, between cycles:

Cycle 1: train (batch=1) → prune 10% → defrag → model is 9% smaller
Cycle 2: train (batch=2!) → prune 10% → defrag → model is 18% smaller  
Cycle 3: train (batch=3!) → prune 10% → defrag → model is 27% smaller

Each defrag frees VRAM → bigger batch size → faster training. The model accelerates as it shrinks. Compound effect.

This requires defrag to work on live model tensors (not just saved files). The v1 prototype works on files. v2 would do in-place tensor surgery during training.

Key consideration: LoRA adapters reference the original tensor dimensions. After defrag, need to recreate LoRA targets on the new smaller tensors. This means: defrag happens AFTER LoRA merge, BEFORE next LoRA attach.

train with LoRA → merge LoRA → defrag (shrink tensors) → attach new LoRA (smaller) → train