Add knob conditional effect figures
This commit is contained in:
381
scripts/plot_knob_conditional_effects.py
Normal file
381
scripts/plot_knob_conditional_effects.py
Normal file
@@ -0,0 +1,381 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Plot measured knob conditional effects for the AITuner harness study."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from matplotlib.lines import Line2D
|
||||
from matplotlib.patches import Patch
|
||||
|
||||
|
||||
OUT = Path("docs/harness-ablation/figures")
|
||||
|
||||
|
||||
def save(fig: plt.Figure, name: str) -> None:
|
||||
OUT.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(OUT / f"{name}.png", dpi=220, bbox_inches="tight")
|
||||
fig.savefig(OUT / f"{name}.svg", bbox_inches="tight")
|
||||
|
||||
|
||||
def plot_c1_surface() -> None:
|
||||
# Qwen30B mixed workload, TP x MNS screen. Values are req/s/GPU.
|
||||
# Source runs:
|
||||
# - interaction-mixed-qwen30b-tp-mns-surface-high1-dash1-d8899c5-20260701T095858Z
|
||||
# - interaction-mixed-qwen30b-tp4-mns-nocap-qps20-dash1-d8899c5-20260701T161900Z
|
||||
mns = np.array([8, 16, 32, 64])
|
||||
tp = np.array([1, 2, 4])
|
||||
values = np.array(
|
||||
[
|
||||
[2.1000, 2.3500, 2.2833, 2.2833],
|
||||
[2.2750, 2.2750, 3.2833, 3.2583],
|
||||
[1.2833, 2.4417, 2.4417, 2.4417],
|
||||
]
|
||||
)
|
||||
|
||||
fig, axes = plt.subplots(1, 2, figsize=(12.8, 4.8), gridspec_kw={"width_ratios": [1.05, 1.2]})
|
||||
ax = axes[0]
|
||||
im = ax.imshow(values, cmap="YlGnBu", aspect="auto", vmin=1.2, vmax=3.35)
|
||||
ax.set_xticks(range(len(mns)), labels=mns)
|
||||
ax.set_yticks(range(len(tp)), labels=[f"TP={x}" for x in tp])
|
||||
ax.set_xlabel("max-num-seqs (MNS)")
|
||||
ax.set_ylabel("tensor-parallel-size")
|
||||
ax.set_title("C1 response surface: req/s/GPU")
|
||||
for i in range(values.shape[0]):
|
||||
for j in range(values.shape[1]):
|
||||
color = "white" if values[i, j] > 2.75 else "black"
|
||||
ax.text(j, i, f"{values[i, j]:.2f}", ha="center", va="center", color=color, fontsize=10)
|
||||
cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||||
cbar.set_label("req/s/GPU")
|
||||
|
||||
ax = axes[1]
|
||||
colors = {1: "#4E79A7", 2: "#59A14F", 4: "#E15759"}
|
||||
for idx, t in enumerate(tp):
|
||||
ax.plot(mns, values[idx], marker="o", linewidth=2.4, color=colors[int(t)], label=f"TP={t}")
|
||||
ax.set_xscale("log", base=2)
|
||||
ax.set_xticks(mns, labels=mns)
|
||||
ax.set_xlabel("max-num-seqs (MNS)")
|
||||
ax.set_ylabel("req/s/GPU")
|
||||
ax.set_title("Non-parallel lines imply interaction")
|
||||
ax.grid(True, axis="y", alpha=0.28)
|
||||
ax.legend(frameon=False)
|
||||
ax.annotate(
|
||||
"TP=2 only becomes best\nwhen MNS reaches 32",
|
||||
xy=(32, 3.2833),
|
||||
xytext=(20, 3.05),
|
||||
arrowprops={"arrowstyle": "->", "lw": 1.2},
|
||||
fontsize=9,
|
||||
)
|
||||
ax.annotate(
|
||||
"TP=4 is bad at MNS=8\nbut recovers at MNS>=16",
|
||||
xy=(8, 1.2833),
|
||||
xytext=(10, 1.55),
|
||||
arrowprops={"arrowstyle": "->", "lw": 1.2},
|
||||
fontsize=9,
|
||||
)
|
||||
fig.suptitle("Knob effects are conditional: MNS effect depends on TP", fontsize=14, y=1.02)
|
||||
fig.tight_layout()
|
||||
save(fig, "knob-conditional-c1-qwen30b-surface")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def plot_c1_oat_counterexample() -> None:
|
||||
# C1 Qwen30B: one-knob-at-a-time tuning gets trapped at a coordinate-wise
|
||||
# local optimum 25.6% below the measured global best. The right panel zooms
|
||||
# into the trap's neighbourhood so the reader can SEE that every single-knob
|
||||
# move from the trap is worse or flat, instead of having to read a caption.
|
||||
mns = [8, 16, 32, 64]
|
||||
tp = [1, 2, 4]
|
||||
values = np.array(
|
||||
[
|
||||
[2.1000, 2.3500, 2.2833, 2.2833],
|
||||
[2.2750, 2.2750, 3.2833, 3.2583],
|
||||
[1.2833, 2.4417, 2.4417, 2.4417],
|
||||
]
|
||||
)
|
||||
idx = {(t, s): (mns.index(s), tp.index(t)) for t in tp for s in mns}
|
||||
|
||||
fig, axes = plt.subplots(1, 2, figsize=(14.0, 6.2), gridspec_kw={"width_ratios": [1.5, 1.0]})
|
||||
ax = axes[0]
|
||||
im = ax.imshow(values, cmap="YlGnBu", aspect="auto", vmin=1.2, vmax=3.35)
|
||||
ax.set_xticks(range(len(mns)), labels=mns)
|
||||
ax.set_yticks(range(len(tp)), labels=[f"TP={x}" for x in tp])
|
||||
ax.set_xlabel("max-num-seqs (MNS)", fontsize=11)
|
||||
ax.set_ylabel("tensor-parallel-size", fontsize=11)
|
||||
ax.set_title("Two OAT paths from the same start", fontsize=12, loc="left")
|
||||
for i in range(values.shape[0]):
|
||||
for j in range(values.shape[1]):
|
||||
color = "white" if values[i, j] > 2.75 else "black"
|
||||
ax.text(j, i, f"{values[i, j]:.2f}", ha="center", va="center", color=color, fontsize=11, weight="bold")
|
||||
|
||||
def draw_path(path: list[tuple[int, int]], color: str, labels: list[str]) -> None:
|
||||
for (a, b) in zip(path, path[1:]):
|
||||
x0, y0 = idx[a]
|
||||
x1, y1 = idx[b]
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(x1, y1),
|
||||
xytext=(x0, y0),
|
||||
arrowprops={"arrowstyle": "->", "lw": 3.2, "color": color, "shrinkA": 22, "shrinkB": 22},
|
||||
)
|
||||
for (a, b), lbl in zip(zip(path, path[1:]), labels):
|
||||
x0, y0 = idx[a]
|
||||
x1, y1 = idx[b]
|
||||
mx, my = (x0 + x1) / 2, (y0 + y1) / 2
|
||||
if x0 == x1: # vertical move -> label to the side
|
||||
ax.text(
|
||||
mx + 0.30,
|
||||
my,
|
||||
lbl,
|
||||
color=color,
|
||||
fontsize=9.5,
|
||||
ha="left",
|
||||
va="center",
|
||||
weight="bold",
|
||||
bbox={"boxstyle": "round,pad=0.2", "facecolor": "white", "edgecolor": "none", "alpha": 0.85},
|
||||
)
|
||||
else: # horizontal move -> label above
|
||||
ax.text(
|
||||
mx,
|
||||
my - 0.32,
|
||||
lbl,
|
||||
color=color,
|
||||
fontsize=9.5,
|
||||
ha="center",
|
||||
va="bottom",
|
||||
weight="bold",
|
||||
bbox={"boxstyle": "round,pad=0.2", "facecolor": "white", "edgecolor": "none", "alpha": 0.85},
|
||||
)
|
||||
|
||||
# Red: tune MNS first, then TP -> walks into a coordinate-wise local optimum.
|
||||
draw_path([(1, 8), (1, 16), (4, 16)], "#C0392B", ["tune MNS", "tune TP"])
|
||||
# Green: tune TP first, then MNS -> reaches the measured global best.
|
||||
draw_path([(1, 8), (2, 8), (2, 32)], "#2E7D32", ["tune TP", "tune MNS"])
|
||||
|
||||
# start / trap / best markers (explained by the legend below the grid)
|
||||
sx, sy = idx[(1, 8)]
|
||||
ax.scatter([sx], [sy], marker="o", s=210, facecolors="none", edgecolors="black", linewidths=2.2, zorder=5)
|
||||
trap = (4, 16)
|
||||
tx, ty = idx[trap]
|
||||
ax.add_patch(plt.Rectangle((tx - 0.5, ty - 0.5), 1, 1, fill=False, edgecolor="#C0392B", linewidth=3.4))
|
||||
best = (2, 32)
|
||||
bx, by = idx[best]
|
||||
ax.add_patch(plt.Rectangle((bx - 0.5, by - 0.5), 1, 1, fill=False, edgecolor="#2E7D32", linewidth=3.4))
|
||||
|
||||
legend_elements = [
|
||||
Line2D([0], [0], marker="o", color="w", markerfacecolor="none", markeredgecolor="black", markeredgewidth=2, markersize=10, label="start TP1,MNS8 = 2.10"),
|
||||
Patch(facecolor="none", edgecolor="#C0392B", linewidth=2.4, label="OAT trap TP4,MNS16 = 2.44 (no improving single-knob move)"),
|
||||
Patch(facecolor="none", edgecolor="#2E7D32", linewidth=2.4, label="global best TP2,MNS32 = 3.28"),
|
||||
]
|
||||
ax.legend(handles=legend_elements, loc="upper center", bbox_to_anchor=(0.5, -0.11), ncol=1, frameon=False, fontsize=9.5)
|
||||
cbar = fig.colorbar(im, ax=ax, fraction=0.038, pad=0.035)
|
||||
cbar.set_label("req/s/GPU")
|
||||
|
||||
# ---- right panel: why the red path stops (trap neighbourhood zoom) ----
|
||||
ax2 = axes[1]
|
||||
ax2.set_xlim(-0.55, 3.55)
|
||||
ax2.set_ylim(-0.65, 3.7)
|
||||
ax2.set_aspect("equal")
|
||||
ax2.set_axis_off()
|
||||
ax2.set_title("Why the red path stops here", fontsize=12, loc="left")
|
||||
|
||||
# 3x3 neighbourhood of the trap (TP4,MNS16). Rows match left panel:
|
||||
# top=TP2, middle=TP4(trap), bottom=TP8(not measured). Cols: MNS 8/16/32.
|
||||
# (col, row) with row 0 at bottom.
|
||||
zoom_cells = {
|
||||
(1, 2): ("2.275", "TP2", "dead"), # up neighbour
|
||||
(0, 1): ("1.28", "TP4", "dead"), # left neighbour
|
||||
(1, 1): ("2.44", "TP4", "trap"), # the trap
|
||||
(2, 1): ("2.44", "TP4", "flat"), # right neighbour (flat, not strictly improving)
|
||||
(1, 0): ("—", "TP8", "oob"), # down neighbour (not measured)
|
||||
}
|
||||
for (col, row), (val, tplabel, kind) in zoom_cells.items():
|
||||
x, y = col, row
|
||||
if kind == "trap":
|
||||
fc, ec, lw = "#FDECEA", "#C0392B", 3.2
|
||||
elif kind == "oob":
|
||||
fc, ec, lw = "#F2F2F2", "#CCCCCC", 1.0
|
||||
else:
|
||||
fc, ec, lw = "#FDF2F2", "#E6B0AA", 1.4
|
||||
ax2.add_patch(plt.Rectangle((x, y), 1, 1, facecolor=fc, edgecolor=ec, linewidth=lw))
|
||||
if kind == "oob":
|
||||
ax2.text(x + 0.5, y + 0.5, "no data", ha="center", va="center", fontsize=9, color="#999", style="italic")
|
||||
else:
|
||||
val_color = "#C0392B" if kind in ("dead", "flat") else "#222"
|
||||
ax2.text(x + 0.5, y + 0.62, val, ha="center", va="center", fontsize=13, weight="bold", color=val_color)
|
||||
sublabel = tplabel
|
||||
if kind == "trap":
|
||||
sublabel = f"{tplabel} · trap"
|
||||
elif kind == "flat":
|
||||
sublabel = f"{tplabel} · flat"
|
||||
ax2.text(x + 0.5, y + 0.28, sublabel, ha="center", va="center", fontsize=8.5, color="#666")
|
||||
if kind in ("dead", "flat"):
|
||||
ax2.text(x + 0.87, y + 0.87, "✗", ha="center", va="center", color="#C0392B", fontsize=18, weight="bold")
|
||||
|
||||
# axis labels for the zoom
|
||||
ax2.text(-0.18, 2.5, "TP=2", ha="right", va="center", fontsize=9, color="#666")
|
||||
ax2.text(-0.18, 1.5, "TP=4", ha="right", va="center", fontsize=9, color="#666")
|
||||
ax2.text(-0.18, 0.5, "TP=8", ha="right", va="center", fontsize=9, color="#999")
|
||||
ax2.text(0.5, -0.18, "MNS=8", ha="center", va="top", fontsize=9, color="#666")
|
||||
ax2.text(1.5, -0.18, "MNS=16", ha="center", va="top", fontsize=9, color="#666")
|
||||
ax2.text(2.5, -0.18, "MNS=32", ha="center", va="top", fontsize=9, color="#666")
|
||||
|
||||
ax2.text(
|
||||
1.5,
|
||||
3.45,
|
||||
"Every measured single-knob move from the trap\nis worse or flat → coordinate ascent is stuck",
|
||||
ha="center",
|
||||
va="center",
|
||||
fontsize=10.5,
|
||||
color="#222",
|
||||
bbox={"boxstyle": "round,pad=0.35", "facecolor": "white", "edgecolor": "#C0392B"},
|
||||
)
|
||||
|
||||
fig.suptitle("One-knob-at-a-time tuning gets trapped: 25.6% throughput gap between two tuning orders", fontsize=14, y=1.02)
|
||||
fig.tight_layout()
|
||||
save(fig, "knob-oat-counterexample-c1-qwen30b")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def plot_c1_interaction_residual() -> None:
|
||||
# If TP and MNS were independent additive effects, this residual matrix would be near zero.
|
||||
mns = [8, 16, 32, 64]
|
||||
tp = [1, 2, 4]
|
||||
values = np.array(
|
||||
[
|
||||
[2.1000, 2.3500, 2.2833, 2.2833],
|
||||
[2.2750, 2.2750, 3.2833, 3.2583],
|
||||
[1.2833, 2.4417, 2.4417, 2.4417],
|
||||
]
|
||||
)
|
||||
residual = values - values.mean(axis=1, keepdims=True) - values.mean(axis=0, keepdims=True) + values.mean()
|
||||
fig, ax = plt.subplots(figsize=(7.2, 4.8))
|
||||
limit = float(np.abs(residual).max())
|
||||
im = ax.imshow(residual, cmap="RdBu", aspect="auto", vmin=-limit, vmax=limit)
|
||||
ax.set_xticks(range(len(mns)), labels=mns)
|
||||
ax.set_yticks(range(len(tp)), labels=[f"TP={x}" for x in tp])
|
||||
ax.set_xlabel("max-num-seqs (MNS)")
|
||||
ax.set_ylabel("tensor-parallel-size")
|
||||
ax.set_title("C1 non-additive interaction residual")
|
||||
for i in range(residual.shape[0]):
|
||||
for j in range(residual.shape[1]):
|
||||
ax.text(j, i, f"{residual[i, j]:+.2f}", ha="center", va="center", fontsize=10)
|
||||
cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||||
cbar.set_label("req/s/GPU residual")
|
||||
fig.suptitle("Independent-knob additive model leaves large structured residuals", fontsize=13, y=1.02)
|
||||
fig.tight_layout()
|
||||
save(fig, "knob-interaction-residual-c1-qwen30b")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def plot_c3_lines() -> None:
|
||||
# Qwen235B decode C3, topology x MNS x MBT screen.
|
||||
# Source run:
|
||||
# interaction-qwen235b-decode-c3-topo-mns-mbt-fixed-dash1-d8899c5-20260703T022514Z
|
||||
data = {
|
||||
("TP4/DP2/EP8", 64, 256): 0.05354166666666667,
|
||||
("TP4/DP2/EP8", 64, 384): 0.05354166666666667,
|
||||
("TP4/DP2/EP8", 128, 256): 0.058958333333333335,
|
||||
("TP4/DP2/EP8", 128, 384): 0.058958333333333335,
|
||||
("TP2/DP4/EP8", 64, 256): 0.058958333333333335,
|
||||
("TP2/DP4/EP8", 64, 384): 0.05354166666666667,
|
||||
("TP2/DP4/EP8", 128, 256): 0.058958333333333335,
|
||||
("TP2/DP4/EP8", 128, 384): 0.058958333333333335,
|
||||
}
|
||||
mbt = [256, 384]
|
||||
topologies = ["TP4/DP2/EP8", "TP2/DP4/EP8"]
|
||||
|
||||
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.6), sharey=True)
|
||||
for ax, topo in zip(axes, topologies):
|
||||
for mns, color in [(64, "#4E79A7"), (128, "#F28E2B")]:
|
||||
vals = [data[(topo, mns, b)] for b in mbt]
|
||||
ax.plot(mbt, vals, marker="o", linewidth=2.6, color=color, label=f"MNS={mns}")
|
||||
for x, y in zip(mbt, vals):
|
||||
ax.text(x, y + 0.0007, f"{y:.4f}", ha="center", fontsize=9)
|
||||
ax.set_title(topo)
|
||||
ax.set_xlabel("max-num-batched-tokens (MBT)")
|
||||
ax.set_xticks(mbt)
|
||||
ax.grid(True, axis="y", alpha=0.28)
|
||||
ax.set_ylim(0.050, 0.062)
|
||||
axes[0].set_ylabel("req/s/GPU")
|
||||
axes[1].legend(frameon=False, loc="lower right")
|
||||
fig.suptitle("C3: MBT effect depends on topology and MNS", fontsize=14, y=1.02)
|
||||
fig.tight_layout()
|
||||
save(fig, "knob-conditional-c3-qwen235b-decode-lines")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def plot_delta_summary() -> None:
|
||||
c1_base = {
|
||||
"TP=1": (2.2833 - 2.1000) / 2.1000 * 100.0,
|
||||
"TP=2": (3.2833 - 2.2750) / 2.2750 * 100.0,
|
||||
"TP=4": (2.4417 - 1.2833) / 1.2833 * 100.0,
|
||||
}
|
||||
c3_mbt = {
|
||||
"TP4/DP2\nMNS=64": 0.0,
|
||||
"TP4/DP2\nMNS=128": 0.0,
|
||||
"TP2/DP4\nMNS=64": (0.05354166666666667 - 0.058958333333333335)
|
||||
/ 0.058958333333333335
|
||||
* 100.0,
|
||||
"TP2/DP4\nMNS=128": 0.0,
|
||||
}
|
||||
c3_mns = {
|
||||
"TP4/DP2\nMBT=256": (0.058958333333333335 - 0.05354166666666667)
|
||||
/ 0.05354166666666667
|
||||
* 100.0,
|
||||
"TP4/DP2\nMBT=384": (0.058958333333333335 - 0.05354166666666667)
|
||||
/ 0.05354166666666667
|
||||
* 100.0,
|
||||
"TP2/DP4\nMBT=256": 0.0,
|
||||
"TP2/DP4\nMBT=384": (0.058958333333333335 - 0.05354166666666667)
|
||||
/ 0.05354166666666667
|
||||
* 100.0,
|
||||
}
|
||||
|
||||
panels = [
|
||||
("C1: MNS 8->32\nunder different TP", c1_base, "#59A14F"),
|
||||
("C3: MBT 256->384\nunder different context", c3_mbt, "#E15759"),
|
||||
("C3: MNS 64->128\nunder different context", c3_mns, "#4E79A7"),
|
||||
]
|
||||
fig, axes = plt.subplots(1, 3, figsize=(16, 5.2))
|
||||
for ax, (title, vals, color) in zip(axes, panels):
|
||||
labels = list(vals.keys())
|
||||
y = np.arange(len(labels))
|
||||
x = list(vals.values())
|
||||
colors = [color if v >= 0 else "#B07AA1" for v in x]
|
||||
ax.barh(y, x, color=colors)
|
||||
ax.axvline(0, color="black", linewidth=0.8)
|
||||
lo = min(x)
|
||||
hi = max(x)
|
||||
pad = max(2.0, (hi - lo) * 0.12)
|
||||
ax.set_xlim(lo - pad, hi + pad)
|
||||
ax.set_yticks(y, labels=labels, fontsize=8)
|
||||
ax.invert_yaxis()
|
||||
ax.set_xlabel("relative change in req/s/GPU (%)")
|
||||
ax.set_title(title)
|
||||
ax.grid(True, axis="x", alpha=0.25)
|
||||
for yi, xi in zip(y, x):
|
||||
ha = "left" if xi >= 0 else "right"
|
||||
offset = 0.7 if xi >= 0 else -0.7
|
||||
ax.text(xi + offset, yi, f"{xi:+.1f}%", va="center", ha=ha, fontsize=8)
|
||||
fig.suptitle("The same knob intervention has context-dependent effect size", fontsize=14, y=1.02)
|
||||
fig.tight_layout()
|
||||
save(fig, "knob-conditional-delta-summary")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
plot_c1_oat_counterexample()
|
||||
plot_c1_interaction_residual()
|
||||
plot_c1_surface()
|
||||
plot_c3_lines()
|
||||
plot_delta_summary()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user