File size: 1,269 Bytes
9190eff | 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 | 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}")
|