feat: calibration complete + sensitivity heatmap for Mistral-7B

calibrate.py

@@ -115,9 +115,9 @@ for layer_idx in tqdm(range(num_layers), desc="Layers"):
             "8bit": round(err_8bit, 6),
         }
 
-        if err_2bit < 1e-4:
-            optimal_bits = 2
-        elif err_4bit < 1e-4:
+        # use 4-bit if error is in bottom 50% of all 4-bit errors
+        # use 8-bit for high-sensitivity heads
+        if err_4bit < 0.05:
             optimal_bits = 4
         else:
             optimal_bits = 8

visualize_sensitivity.py

@@ -0,0 +1,36 @@
+import json
+import matplotlib.pyplot as plt
+import numpy as np
+
+with open("results/mistral-7b/sensitivity_map.json") as f:
+    sens = json.load(f)
+
+num_layers = len(sens)
+num_heads = len(sens["0"])
+
+# build heatmaps
+err_4bit = np.zeros((num_layers, num_heads))
+for l in sens:
+    for h in sens[l]:
+        err_4bit[int(l), int(h)] = sens[l][h]["4bit"]
+
+fig, ax = plt.subplots(figsize=(12, 8))
+im = ax.imshow(err_4bit, aspect='auto', cmap='hot_r')
+ax.set_xlabel("Attention Head", fontsize=12)
+ax.set_ylabel("Layer", fontsize=12)
+ax.set_title("4-bit KV Cache Quantization Error per Head\n(darker = more sensitive = needs higher precision)", fontsize=13)
+plt.colorbar(im, ax=ax, label="MSE Reconstruction Error")
+plt.tight_layout()
+plt.savefig("figures/sensitivity_heatmap.png", dpi=150)
+print("✅ Saved figures/sensitivity_heatmap.png")
+
+# print most and least sensitive heads
+flat = [(err_4bit[l,h], l, h) for l in range(num_layers) for h in range(num_heads)]
+flat.sort()
+print("\n🟢 10 LEAST sensitive heads (safe to quantize to 4-bit):")
+for err, l, h in flat[:10]:
+    print(f"  Layer {l:2d}, Head {h}: error={err:.4f}")
+
+print("\n🔴 10 MOST sensitive heads (keep at 8-bit):")
+for err, l, h in flat[-10:]:
+    print(f"  Layer {l:2d}, Head {h}: error={err:.4f}")