Files
aituner/scripts/plot_simulator_fidelity.py

257 lines
9.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""Render the per-config figures embedded by simulator-fidelity.md."""
from __future__ import annotations
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
ASSET_DIR = ROOT / "docs" / "assets" / "simulator-fidelity"
DATA_PATH = ASSET_DIR / "data.json"
T0_COMPARISON_PATH = (
ROOT / "runs" / "frontier-multicase-sufficiency-v1" / "results" / "t0-final" / "comparison.json"
)
REAL_COLOR = "#3D3D3D"
PROFILE_COLOR = "#E15759"
CALIBRATED_COLOR = "#4E79A7"
def load_data() -> dict:
with DATA_PATH.open(encoding="utf-8") as handle:
return json.load(handle)
def style_axis(ax: plt.Axes) -> None:
ax.grid(axis="y", alpha=0.23, linewidth=0.8)
ax.set_axisbelow(True)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
def save(fig: plt.Figure, name: str) -> None:
ASSET_DIR.mkdir(parents=True, exist_ok=True)
fig.savefig(ASSET_DIR / f"{name}.png", dpi=220, bbox_inches="tight", facecolor="white")
fig.savefig(ASSET_DIR / f"{name}.svg", bbox_inches="tight", facecolor="white")
def mark_top(ax: plt.Axes, x: float, y: float, text: str, color: str) -> None:
ax.annotate(
text,
xy=(x, y),
xytext=(0, 7),
textcoords="offset points",
ha="center",
va="bottom",
color=color,
fontsize=8.5,
weight="bold",
)
def plot_qwen30(data: dict) -> None:
case = data["qwen30_mixed"]
rows = sorted(case["configs"], key=lambda row: (-row["real"], row["name"]))
labels = [f"TP{row['tp']}\nMNS{row['mns']}" for row in rows]
real = np.array([row["real"] for row in rows])
profile = np.array([row["frontier_profile_only"] for row in rows])
calibrated = np.array([row["frontier_calibrated"] for row in rows])
x = np.arange(len(rows))
width = 0.38
fig, axes = plt.subplots(2, 1, figsize=(13.8, 8.7), sharex=True, sharey=True)
panels = [
(
axes[0],
profile,
PROFILE_COLOR,
"Frontier: operator profiles only",
"top-set miss · Kendall τ-b = 0.000 · worst-case regret = 25.63%",
set(case["profile_only_metrics"]["simulator_top_set"]),
),
(
axes[1],
calibrated,
CALIBRATED_COLOR,
"Frontier: frozen per-TP calibration",
"top-set hit · Kendall τ-b = 0.967 · worst-case regret = 0.76%",
set(case["calibrated_metrics"]["simulator_top_set"]),
),
]
real_top = set(case["calibrated_metrics"]["real_top_set"])
for ax, simulated, color, sim_label, subtitle, simulator_top in panels:
real_bars = ax.bar(x - width / 2, real, width, color=REAL_COLOR, label="Real vLLM")
sim_bars = ax.bar(x + width / 2, simulated, width, color=color, label=sim_label)
style_axis(ax)
ax.set_ylabel("SLO-feasible req/s/GPU")
ax.set_title(subtitle, loc="left", fontsize=11, weight="bold")
ax.set_ylim(0, 4.35)
ax.legend(frameon=False, ncol=2, loc="upper right")
for idx, row in enumerate(rows):
if row["name"] in real_top:
real_bars[idx].set_edgecolor("#000000")
real_bars[idx].set_linewidth(2.0)
mark_top(ax, idx - width / 2, real[idx], "real top", REAL_COLOR)
if row["name"] in simulator_top:
sim_bars[idx].set_edgecolor(color)
sim_bars[idx].set_linewidth(2.0)
sim_bars[idx].set_hatch("///")
mark_top(ax, idx + width / 2, simulated[idx], "sim top", color)
axes[1].set_xticks(x, labels=labels)
axes[1].set_xlabel("Configuration (ordered by measured real capacity)")
fig.suptitle(
"Qwen3-30B mixed serving: calibration changes the config-selection result",
fontsize=14,
weight="bold",
y=0.995,
)
fig.tight_layout(rect=(0, 0, 1, 0.97))
save(fig, "qwen30-mixed-config-ranking")
plt.close(fig)
def plot_qwen235(data: dict) -> None:
case = data["qwen235_prefill"]
rows = sorted(case["configs"], key=lambda row: (-row["real"], row["name"]))
labels = [f"TP{row['tp']} MNS{row['mns']}\nMBT{row['mbt'] // 1024}K" for row in rows]
real = np.array([row["real"] for row in rows])
frontier = np.array([row["frontier"] for row in rows])
x = np.arange(len(rows))
width = 0.38
metrics = case["metrics"]
real_top = set(metrics["real_top_set"])
simulator_top = set(metrics["simulator_top_set"])
fig, ax = plt.subplots(figsize=(13.4, 5.4))
real_bars = ax.bar(x - width / 2, real, width, color=REAL_COLOR, label="Real community vLLM")
sim_bars = ax.bar(x + width / 2, frontier, width, color=CALIBRATED_COLOR, label="Frontier best-effort")
style_axis(ax)
ax.set_ylabel("SLO-feasible req/s/GPU")
ax.set_xticks(x, labels=labels)
ax.set_xlabel("Configuration (ordered by measured real capacity)")
ax.set_ylim(0, 0.088)
ax.legend(frameon=False, ncol=2, loc="upper right")
ax.set_title(
"Top set matches · Spearman ρ = 0.949 · 20/20 comparable pairs agree · regret = 0",
loc="left",
fontsize=11,
weight="bold",
)
for idx, row in enumerate(rows):
if row["name"] in real_top:
real_bars[idx].set_edgecolor("#000000")
real_bars[idx].set_linewidth(2.0)
mark_top(ax, idx - width / 2, real[idx], "real top", REAL_COLOR)
if row["name"] in simulator_top:
sim_bars[idx].set_edgecolor(CALIBRATED_COLOR)
sim_bars[idx].set_linewidth(2.0)
sim_bars[idx].set_hatch("///")
mark_top(ax, idx + width / 2, frontier[idx], "sim top", CALIBRATED_COLOR)
fig.suptitle(
"Qwen3-235B-A22B-FP8 prefill-only: Frontier preserves the useful ordering",
fontsize=14,
weight="bold",
y=1.01,
)
fig.tight_layout()
save(fig, "qwen235-prefill-config-ranking")
plt.close(fig)
def plot_qwen235_t0() -> None:
with T0_COMPARISON_PATH.open(encoding="utf-8") as handle:
comparison = json.load(handle)["comparisons"]["tpot_150ms"]
rows = sorted(
comparison["records"],
key=lambda row: (row["config"]["tp"], row["config"]["mns"], row["config"]["mbt"]),
)
labels = [
f"TP{row['config']['tp']} MNS{row['config']['mns']}\nMBT{row['config']['mbt'] // 1024}K"
for row in rows
]
real = np.array([row["real_capacity_per_gpu"] for row in rows])
frontier = np.array([row["sim_capacity_per_gpu"] for row in rows])
x = np.arange(len(rows))
width = 0.38
fig = plt.figure(figsize=(16.2, 5.7), constrained_layout=True)
grid = fig.add_gridspec(1, 3, width_ratios=[2.6, 0.72, 0.72])
ax = fig.add_subplot(grid[0, 0])
real_bars = ax.bar(x - width / 2, real, width, color=REAL_COLOR, label="Real community vLLM")
sim_bars = ax.bar(x + width / 2, frontier, width, color=CALIBRATED_COLOR, label="Frozen Frontier")
style_axis(ax)
ax.axvline(3.5, color="#999999", linewidth=1.0, linestyle="--")
ax.set_ylabel("SLO-feasible req/s/GPU")
ax.set_xticks(x, labels=labels)
ax.set_ylim(0, 0.71)
ax.legend(frameon=False, ncol=2, loc="upper right")
ax.set_title(
"Exact top-set match · Kendall τ-b = 0.894 · worst tie-break regret = 0",
loc="left",
fontsize=10.5,
weight="bold",
)
real_top = set(comparison["real_top_set"])
simulator_top = set(comparison["sim_top_set"])
for idx, row in enumerate(rows):
name = row["config"]["name"]
if name in real_top:
real_bars[idx].set_edgecolor("#000000")
real_bars[idx].set_linewidth(1.8)
if name in simulator_top:
sim_bars[idx].set_edgecolor(CALIBRATED_COLOR)
sim_bars[idx].set_linewidth(1.8)
sim_bars[idx].set_hatch("///")
matrices = []
for source in ("real_capacity_per_gpu", "sim_capacity_per_gpu"):
matrix = np.empty((2, 2))
for row in rows:
config = row["config"]
if config["tp"] != 8:
continue
matrix[[64, 128].index(config["mns"]), [8192, 16384].index(config["mbt"])] = row[source]
matrices.append(matrix)
heat_axes = [fig.add_subplot(grid[0, 1]), fig.add_subplot(grid[0, 2])]
titles = ["Real TP8", "Frontier TP8"]
images = []
for heat_ax, matrix, title in zip(heat_axes, matrices, titles):
images.append(heat_ax.imshow(matrix, cmap="YlOrRd", vmin=0.15, vmax=0.30, aspect="equal"))
heat_ax.set_title(title, fontsize=11, weight="bold")
heat_ax.set_xticks([0, 1], labels=["8K", "16K"])
heat_ax.set_yticks([0, 1], labels=["64", "128"])
heat_ax.set_xlabel("MBT")
heat_ax.set_ylabel("MNS")
for i in range(2):
for j in range(2):
heat_ax.text(j, i, f"{matrix[i, j]:.2f}", ha="center", va="center", weight="bold")
fig.colorbar(images[-1], ax=heat_axes, location="bottom", shrink=0.75, pad=0.14, label="req/s/GPU")
fig.suptitle(
"Qwen3-235B fixed-shape mixed: correct deployment choice, incorrect TP8 control interaction",
fontsize=14,
weight="bold",
)
save(fig, "qwen235-t0-fixed-shape-ranking")
plt.close(fig)
def main() -> None:
data = load_data()
plot_qwen30(data)
plot_qwen235(data)
plot_qwen235_t0()
if __name__ == "__main__":
main()