Updating Plot Script

Author: phaelon74

examples/plot_kld_analysis.py

@@ -2,11 +2,11 @@
 """
 Plot KLD (Kullback-Leibler Divergence) vs model file size for quantization analysis.
 
+Categories: INT AWQ (W4A16, W8A16, FP8-INT4), NVFP4, GGUF - each with distinct color/shape.
+
 Usage:
     pip install matplotlib
     python examples/plot_kld_analysis.py
-
-Edit the DATA section below to add your models.
 """
 
 import argparse
@@ -19,16 +19,38 @@
 
 
 # --- Edit your data here ---
-# Format: (file_size_gib, mean_kld, label)
+# Format: (file_size_gib, mean_kld, bpw, label, category)
+# category: "original" | "int_awq" | "nvfp4" | "gguf"
 DATA = [
-    (30.0, 0.0, "Original (Llama-3.1-8B bf16)"),
-    (6.2, 0.033707, "Quantized (FP8_INT4)"),
-    (5.4, 0.076226, "Quantized (W4A16_GS128)"),
-    (5.7, 0.109275, "Quantized (NVFP4)"),
-    (8.6, 0.000899, "Quantized (W8A16_GS128)"),
-    (8.5, 0.006547, "Quantized (W8A8-FP8_BLOCK)"),
+    # Original
+    (30.0, 0.0, 16.0, "Original (Llama-3.1-8B bf16)", "original"),
+    # INT AWQ: W4A16, W8A16, FP8-INT4
+    (5.4, 0.076226, 4.25, "W4A16_GS128", "int_awq"),
+    (8.6, 0.000899, 8.25, "W8A16_GS128", "int_awq"),
+    (6.2, 0.033707, 4.0, "FP8_INT4", "int_awq"),
+    # NVFP4
+    (5.7, 0.109275, 4.0, "NVFP4", "nvfp4"),
+    (8.5, 0.006547, 8.0, "W8A8-FP8_BLOCK", "nvfp4"),
+    # GGUF (from reference chart)
+    (3.5, 0.1241, 3.50, "IQ3_XS", "gguf"),
+    (3.65, 0.1782, 3.64, "Q3_K_S", "gguf"),
+    (4.0, 0.0744, 4.00, "Q3_K_M", "gguf"),
+    (4.4, 0.0327, 4.42, "IQ4_XS", "gguf"),
+    (4.7, 0.0305, 4.67, "Q4_K_S", "gguf"),
+    (4.9, 0.0267, 4.89, "Q4_K_M", "gguf"),
+    (5.6, 0.0102, 5.57, "Q5_K_S", "gguf"),
+    (5.7, 0.0092, 5.70, "Q5_K_M", "gguf"),
+    (6.6, 0.0040, 6.56, "Q6_K", "gguf"),
+    (8.0, 0.0011, 8.50, "Q8_0", "gguf"),
 ]
-# Add more: DATA.append((size_gib, kld, "Label"))
+
+# Category styles: (color, marker)
+STYLES = {
+    "original": ("#7f7f7f", "o"),
+    "int_awq": ("#1f77b4", "s"),
+    "nvfp4": ("#ff7f0e", "^"),
+    "gguf": ("#2ca02c", "D"),
+}
 
 
 def main():
@@ -42,31 +64,68 @@ def main():
     )
     args = parser.parse_args()
 
-    fig, ax = plt.subplots(figsize=(10, 6))
+    fig, ax = plt.subplots(figsize=(16, 10))
 
-    # Unique style per point: original=gray circle, quants=distinct markers+colors
-    QUANT_MARKERS = ["s", "^", "D", "v", "p", "h", "8", "*", "P", "X"]
-    QUANT_COLORS = [
-        "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd",
-        "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf",
+    # Track legend entries to avoid duplicates
+    legend_added = set()
+
+    # Stagger label positions to reduce overlap (alternate quadrants)
+    offsets = [
+        (25, 20), (25, -30), (-100, 20), (25, 20), (-100, -30),
+        (25, -30), (-100, 20), (25, 20), (-100, -30), (25, 20),
+        (-100, -30), (25, -30), (-100, 20), (25, 20), (-100, -30),
+        (25, -30),
     ]
 
-    for i, (s, k, lbl) in enumerate(DATA):
-        if k == 0:
-            color, marker = "gray", "o"
+    for i, (size_gib, kld, bpw, label, category) in enumerate(DATA):
+        color, marker = STYLES[category]
+        full_label = f"Llama-3.1-8B-Instruct-{label}" if category != "original" else label
+
+        # Add to legend only once per category
+        if category not in legend_added:
+            legend_label = {
+                "original": "Original (bf16)",
+                "int_awq": "INT AWQ (W4A16, W8A16, FP8-INT4)",
+                "nvfp4": "NVFP4 / W8A8-FP8",
+                "gguf": "GGUF",
+            }[category]
+            ax.scatter(
+                [], [], color=color, marker=marker, s=120, label=legend_label
+            )
+            legend_added.add(category)
+
+        ax.scatter(size_gib, kld, color=color, s=120, marker=marker, zorder=3)
+
+        # Annotate with model name, KLD, bpw (like reference graph)
+        text = f"{full_label}\n{kld:.4f}\n{bpw:.2f}bpw"
+
+        if category == "original":
+            xytext = (-120, 15)
         else:
-            idx = (i - 1) % len(QUANT_MARKERS)
-            color = QUANT_COLORS[idx % len(QUANT_COLORS)]
-            marker = QUANT_MARKERS[idx % len(QUANT_MARKERS)]
-        legend_label = f"{lbl} ({k:.4f})" if k > 0 else lbl
-        ax.scatter(s, k, label=legend_label, color=color, s=120, marker=marker, zorder=3)
-
-    ax.set_xlabel("File Size (GiB)")
-    ax.set_ylabel("Mean KL Divergence (Lower is Better)")
-    ax.set_title("Quantization Quality: KLD vs File Size")
-    ax.legend()
-    ax.grid(True, alpha=0.3)
+            xytext = offsets[i % len(offsets)]
+        ha = "left" if xytext[0] > 0 else "right"
+        va = "bottom" if xytext[1] > 0 else "top"
+        ax.annotate(
+            text,
+            (size_gib, kld),
+            xytext=xytext,
+            textcoords="offset points",
+            fontsize=8,
+            ha=ha,
+            va=va,
+            arrowprops=dict(arrowstyle="-", color="gray", lw=0.5),
+        )
+
+    ax.set_xlabel("File Size (GiB)", fontsize=12)
+    ax.set_ylabel("Mean KL Divergence (Lower is Better)", fontsize=12)
+    ax.set_title(
+        "Llama-3.1-8B-Instruct Quantization Analysis: Mean KL Divergence vs. Model File Size",
+        fontsize=14,
+    )
+    ax.legend(loc="upper right", fontsize=10)
+    ax.grid(True, alpha=0.3, linestyle="--")
     ax.set_ylim(bottom=-0.02)
+    ax.set_xlim(left=2.5, right=32)
 
     plt.tight_layout()
     if args.output: