| |
| """Generate the wall-clock / parameters / verbosity plots from models.json + results.json. |
| |
| Usage |
| ----- |
| python scripts/make_plots.py --models models.json --results results.json --out plots |
| |
| Requires matplotlib. Time plots use GB10 runs only (Thor wall-clock is not comparable). |
| Each point is the best valid run for that model/quant, matching the table. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
|
|
| def best_run(model: dict, runs: list[dict]) -> dict | None: |
| matches = [r for r in runs if r["label"].startswith(model["group"]) and r["total_valid"] > 0] |
| if not matches: |
| return None |
| return max(matches, key=lambda r: (r["total_passed"], r["total_valid"], -r["total_void"])) |
|
|
|
|
| def build(models: list[dict], runs: list[dict]) -> list[dict]: |
| data = [] |
| for model in models: |
| run = best_run(model, runs) |
| if not run: |
| continue |
| data.append({ |
| **model, |
| "box": run["box"], |
| "wall_h": run["total_wall_s"] / 3600, |
| "tokens": run["total_tokens_out"], |
| "tokens_per_task_k": run["total_tokens_out"] / max(1, run["total_valid"]) / 1000, |
| "tok_s": run["total_tokens_out"] / max(1, run["total_wall_s"]), |
| "score": run["total_passed"] / run["total_valid"], |
| "valid": run["total_valid"], |
| "void": run["total_void"], |
| }) |
| return data |
|
|
|
|
| def scatter(points: list[dict], xkey: str, ykey: str, title: str, xlabel: str, ylabel: str, |
| out: str, logx: bool = True) -> None: |
| fig, ax = plt.subplots(figsize=(8, 5.5), dpi=150) |
| xs = [p[xkey] for p in points] |
| ys = [p[ykey] for p in points] |
| sc = ax.scatter(xs, ys, c=[p["score"] for p in points], cmap="RdYlGn", vmin=0.9, vmax=1.0, |
| s=75, edgecolors="black", linewidths=0.4, zorder=3) |
| if logx: |
| ax.set_xscale("log") |
| for p in points: |
| if p.get("annotate"): |
| dx, dy = p.get("offset", (6, 5)) |
| ax.annotate(p.get("short", p["model"]), (p[xkey], p[ykey]), textcoords="offset points", |
| xytext=(dx, dy), ha=("right" if dx < 0 else "left"), fontsize=7, alpha=0.9) |
| ax.set_xlabel(xlabel) |
| ax.set_ylabel(ylabel) |
| ax.set_title(title) |
| ax.grid(True, which="both", alpha=0.25, zorder=0) |
| ax.margins(x=0.30) |
| cb = fig.colorbar(sc, ax=ax, pad=0.02) |
| cb.set_label("valid-task score") |
| fig.tight_layout() |
| fig.savefig(out, bbox_inches="tight") |
| plt.close(fig) |
|
|
|
|
| def scatter_country(points: list[dict], xkey: str, ykey: str, title: str, xlabel: str, ylabel: str, |
| out: str, logx: bool = True) -> None: |
| from matplotlib.lines import Line2D |
| from matplotlib.patches import Patch |
|
|
| fig, ax = plt.subplots(figsize=(9.5, 6.3), dpi=150) |
| colors = { |
| "US": "#2E5FA3", |
| "China": "#C0392B", |
| "Germany": "#E6A817", |
| "Czechia": "#7D5BA6", |
| "UAE": "#2E8B57", |
| "Unknown": "#777777", |
| "Abliterated": "#888888", |
| } |
| def color(country): |
| return colors.get(country, colors["Unknown"]) |
| def base_of(p): |
| return p.get("base_country", p.get("country", "Unknown")) |
| def post_of(p): |
| return p.get("post_country", p.get("country", "Unknown")) |
| for p in points: |
| ax.scatter(p[xkey], p[ykey], s=90, marker="o", |
| facecolors=color(base_of(p)), edgecolors=color(post_of(p)), |
| linewidths=2.0, zorder=3) |
| if logx: |
| ax.set_xscale("log") |
| for p in points: |
| if p.get("annotate"): |
| dx, dy = p.get("offset", (6, 5)) |
| ax.annotate(p.get("short", p["model"]), (p[xkey], p[ykey]), textcoords="offset points", |
| xytext=(dx, dy), ha=("right" if dx < 0 else "left"), fontsize=7, alpha=0.9) |
| ax.set_xlabel(xlabel) |
| ax.set_ylabel(ylabel) |
| ax.set_title(title) |
| ax.text(0.01, 0.01, "Wall time includes tasks voided by timeout/truncation. " |
| "Fill = pretrained base; edge = post-training/compression.", |
| transform=ax.transAxes, fontsize=7, alpha=0.7) |
| ax.grid(True, which="both", alpha=0.25, zorder=0) |
| ax.margins(x=0.30) |
| ax.set_ylim(bottom=0) |
|
|
| base_present = sorted({base_of(p) for p in points}) |
| post_present = sorted({post_of(p) for p in points}) |
| fill_handles = [Patch(facecolor=color(c), edgecolor="black", label=c) for c in base_present] |
| edge_handles = [Line2D([0], [0], marker="o", linestyle="none", markerfacecolor="white", |
| markeredgecolor=color(c), markeredgewidth=2.0, markersize=8, label=c) |
| for c in post_present] |
| leg1 = ax.legend(handles=fill_handles, title="Base (fill)", fontsize=7, title_fontsize=7, |
| loc="upper left", framealpha=0.95) |
| ax.add_artist(leg1) |
| ax.legend(handles=edge_handles, title="Post-train (edge)", fontsize=7, title_fontsize=7, |
| loc="upper right", framealpha=0.95) |
| fig.tight_layout() |
| fig.savefig(out, bbox_inches="tight") |
| plt.close(fig) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--models", default="models.json") |
| ap.add_argument("--results", default="results.json") |
| ap.add_argument("--out", default="plots") |
| args = ap.parse_args() |
|
|
| models = json.load(open(args.models))["models"] |
| runs = json.load(open(args.results))["runs"] |
| data = build(models, runs) |
| gb = [d for d in data if d["box"] == "gb10"] |
| os.makedirs(args.out, exist_ok=True) |
|
|
| scatter(gb, "active_b", "wall_h", |
| "Wall clock vs active parameters (GB10, 42 tasks)", |
| "Active parameters (B, log)", "Wall clock (hours)", |
| os.path.join(args.out, "wall_vs_active.png")) |
| perfect = [d for d in gb if d["valid"] > 0 and d["score"] >= 0.999999] |
| |
| origin_pts = [d for d in perfect if not d["model"].startswith("Spark-X2.5")] |
| scatter_country(origin_pts, "active_b", "wall_h", |
| "Wall clock vs active parameters — perfect-score runs (GB10)", |
| "Active parameters (B, log)", "Wall clock (hours)", |
| os.path.join(args.out, "wall_vs_active_perfect_country.png")) |
| scatter_country(origin_pts, "total_b", "wall_h", |
| "Wall clock vs total parameters — perfect-score runs (GB10)", |
| "Total parameters (B, log)", "Wall clock (hours)", |
| os.path.join(args.out, "wall_vs_total_perfect_country.png")) |
| scatter(gb, "total_b", "wall_h", |
| "Wall clock vs total parameters (GB10, 42 tasks)", |
| "Total parameters (B, log)", "Wall clock (hours)", |
| os.path.join(args.out, "wall_vs_total.png")) |
| scatter(gb, "tokens", "wall_h", |
| "Wall clock vs output tokens (GB10, 42 tasks)", |
| "Output tokens (thousands, log)", "Wall clock (hours)", |
| os.path.join(args.out, "wall_vs_tokens.png")) |
|
|
| |
| fig, ax = plt.subplots(figsize=(8, 5.5), dpi=150) |
| last_sc = None |
| for box, marker, label in (("gb10", "o", "GB10"), ("thor", "s", "Thor")): |
| pts = [d for d in data if d["box"] == box] |
| if not pts: |
| continue |
| last_sc = ax.scatter([p["active_b"] for p in pts], [p["tokens_per_task_k"] for p in pts], |
| c=[p["score"] for p in pts], cmap="RdYlGn", vmin=0.9, vmax=1.0, |
| marker=marker, s=75, edgecolors="black", linewidths=0.4, |
| label=label, zorder=3) |
| for p in data: |
| if p.get("annotate"): |
| dx, dy = p.get("offset", (6, 5)) |
| ax.annotate(p.get("short", p["model"]), (p["active_b"], p["tokens_per_task_k"]), |
| textcoords="offset points", xytext=(dx, dy), |
| ha=("right" if dx < 0 else "left"), fontsize=7) |
| ax.set_xscale("log") |
| ax.set_xlabel("Active parameters (B, log)") |
| ax.set_ylabel("Output tokens per valid task (thousands)") |
| ax.set_title("Verbosity vs active parameters (42 tasks)") |
| ax.grid(True, which="both", alpha=0.25) |
| ax.margins(x=0.30) |
| ax.legend() |
| if last_sc is not None: |
| cb = fig.colorbar(last_sc, ax=ax, pad=0.02) |
| cb.set_label("valid-task score") |
| fig.tight_layout() |
| fig.savefig(os.path.join(args.out, "tokens_per_task_vs_active.png"), bbox_inches="tight") |
| plt.close(fig) |
|
|
| |
| sized = [d for d in data if d.get("size_gb")] |
| gb_sized = [d for d in sized if d["box"] == "gb10"] |
| fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5.5), dpi=150) |
|
|
| for box, marker, label in (("gb10", "o", "GB10"), ("thor", "s", "Thor")): |
| pts = [d for d in sized if d["box"] == box] |
| if not pts: |
| continue |
| ax1.scatter([d["size_gb"] for d in pts], [d["score"] * 100 for d in pts], |
| marker=marker, s=75, edgecolors="black", linewidths=0.4, label=label, zorder=3) |
| for d in sized: |
| if d.get("annotate") and (d["score"] < 0.995 or d.get("highlight")): |
| dx, dy = d.get("size_offset", d.get("offset", (6, 5))) |
| ax1.annotate(d.get("short", d["model"]), (d["size_gb"], d["score"] * 100), |
| textcoords="offset points", xytext=(dx, dy), |
| ha=("right" if dx < 0 else "left"), fontsize=7) |
| ax1.set_xscale("log") |
| ax1.set_xlabel("Quant file size (GB, log)") |
| ax1.set_ylabel("Valid-task score (%)") |
| ax1.set_title("Capability per byte") |
| ax1.grid(True, which="both", alpha=0.25) |
| ax1.margins(x=0.30) |
| ax1.legend() |
|
|
| sc = ax2.scatter([d["size_gb"] for d in gb_sized], [d["tok_s"] for d in gb_sized], |
| c=[d["score"] for d in gb_sized], cmap="RdYlGn", vmin=0.9, vmax=1.0, |
| s=75, edgecolors="black", linewidths=0.4, zorder=3) |
| for d in gb_sized: |
| if d.get("annotate"): |
| dx, dy = d.get("size_offset", d.get("offset", (6, 5))) |
| ax2.annotate(d.get("short", d["model"]), (d["size_gb"], d["tok_s"]), |
| textcoords="offset points", xytext=(dx, dy), |
| ha=("right" if dx < 0 else "left"), fontsize=7) |
| ax2.set_xscale("log") |
| ax2.set_xlabel("Quant file size (GB, log)") |
| ax2.set_ylabel("Output tokens/s (end-to-end)") |
| ax2.set_title("Throughput per byte (GB10)") |
| ax2.grid(True, which="both", alpha=0.25) |
| ax2.margins(x=0.30) |
| cb = fig.colorbar(sc, ax=ax2, pad=0.02) |
| cb.set_label("valid-task score") |
|
|
| fig.suptitle("Performance per byte (42 tasks)") |
| fig.tight_layout() |
| fig.savefig(os.path.join(args.out, "perf_per_byte.png"), bbox_inches="tight") |
| plt.close(fig) |
|
|
| print(f"wrote 7 plots to {args.out} ({len(gb)} GB10 / {len(data) - len(gb)} Thor rows)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
