Buckets:
| """Data, training and RL figures for the PDF report. | |
| Run with the report venv (matplotlib, numpy), after analysis/fig_train_compute.py (project venv): | |
| CUDA_VISIBLE_DEVICES="" <scratchpad>/reportvenv/bin/python analysis/fig_train_plot.py [slug ...] | |
| Inputs: the dashboard's assembled data (dashboard/build_data.py output; path in DASH), runs/v2/results/{lr_sweep, | |
| speed_bench}.json, runs/v2_2b/results/lr_sweep.json, results/report_figs/train_figs_data.json. | |
| Outputs: results/report_figs/<slug>.png (200 dpi) and results/report_figs/manifest_train.json. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import sys | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| R = "/workspace/decision-reasoner" | |
| DASH = "/tmp/claude-0/-workspace/8842c99f-448c-4cba-80ab-3e9c26b892df/scratchpad/decision_reasoner_dashboard.html.json" | |
| OUT = os.path.join(R, "results/report_figs") | |
| BLUE, G1, G2, G3, AMBER, TEAL = "#2563eb", "#9ca3af", "#6b7280", "#374151", "#d97706", "#059669" | |
| RED, GREEN, GRID, INK = "#dc2626", "#16a34a", "#e5e7eb", "#111827" | |
| FULL, HALF = 6.5, 3.2 | |
| plt.rcParams.update({ | |
| "font.family": "DejaVu Sans", "font.size": 9, "axes.titlesize": 9, "axes.titleweight": "normal", | |
| "axes.labelsize": 9, "xtick.labelsize": 8, "ytick.labelsize": 8, "legend.fontsize": 8, | |
| "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": G2, "axes.labelcolor": G3, | |
| "xtick.color": G3, "ytick.color": G3, "text.color": INK, "figure.facecolor": "white", "axes.facecolor": "white", | |
| "savefig.facecolor": "white", "axes.titlelocation": "left", "axes.titlepad": 5, | |
| }) | |
| MANIFEST = [] | |
| def jload(p): | |
| return json.load(open(p if os.path.isabs(p) else os.path.join(R, p))) | |
| def grid(ax, axis="y"): | |
| ax.grid(axis=axis, color=GRID, lw=0.7) | |
| ax.set_axisbelow(True) | |
| def title(fig, text, y=0.995): | |
| fig.suptitle(text, x=0.01, y=y, ha="left", va="top", fontsize=10.5, weight="bold", color=INK) | |
| def save(fig, slug, ttl, caption, key): | |
| """Left-align the figure title with the leftmost drawn element (tick labels included), then save tightly.""" | |
| st = fig._suptitle | |
| if st is not None: | |
| st.set_visible(False) | |
| fig.canvas.draw() | |
| bb = fig.get_tightbbox(fig.canvas.get_renderer()) | |
| st.set_x(bb.x0 / fig.get_figwidth()) | |
| st.set_visible(True) | |
| path = os.path.join(OUT, f"{slug}.png") | |
| fig.savefig(path, dpi=200, bbox_inches="tight", pad_inches=0.06) | |
| plt.close(fig) | |
| MANIFEST.append({"slug": slug, "title": ttl, "caption": caption, "key_numbers": key}) | |
| print("wrote", path) | |
| def smooth(x, y, w, width): | |
| """Example-weighted moving mean of per-step values over a window of `width` on the x axis (centred).""" | |
| x, y, w = np.asarray(x, float), np.asarray(y, float), np.asarray(w, float) | |
| ok = ~np.isnan(y) & (w > 0) | |
| out = np.full(len(x), np.nan) | |
| for i, xi in enumerate(x): | |
| m = ok & (np.abs(x - xi) <= width / 2) | |
| if m.any(): | |
| out[i] = (y[m] * w[m]).sum() / w[m].sum() | |
| return out | |
| def steps(run): | |
| st = DASHD["training"][run]["steps"] | |
| f = lambda k: np.array([np.nan if s.get(k) is None else s[k] for s in st], float) | |
| return {"tok": np.cumsum(f("step_tokens")), "rn": f("readout_examples"), "cn": f("reasoner_examples"), | |
| "ro": f("readout_loss"), "ch": f("chain_loss"), "ent": f("readout_target_entropy"), "meta": DASHD["training"][run]["meta"]} | |
| def wmean(y, w, lo, hi): | |
| """Example-weighted mean of a per-step value over the steps in [lo, hi) of the run (fractions).""" | |
| n = len(y) | |
| y, w = y[int(n * lo):int(n * hi)], w[int(n * lo):int(n * hi)] | |
| ok = ~np.isnan(y) & (w > 0) | |
| return float((y[ok] * w[ok]).sum() / w[ok].sum()) | |
| # ------------------------------------------------------------------ 1. data composition | |
| def data_composition(): | |
| d = TF["pool"] | |
| fig, (a, b) = plt.subplots(2, 1, figsize=(FULL, 3.25), gridspec_kw={"height_ratios": [2, 3], "hspace": 0.75}) | |
| qt = [("choice", "choice", G3, "white"), ("noul", "yes/no", G2, "white"), ("score", "score", G1, INK)] | |
| rows = [("Procedural generators\n(18 families x 1,000)", d["proc"]["by_qtype"]), | |
| ("tasksource\n(717 source datasets)", {k: d["ts_hard"]["by_qtype"].get(k, 0) + d["ts_soft"]["by_qtype"].get(k, 0) for k in ("choice", "noul", "score")})] | |
| for i, (lab, c) in enumerate(rows[::-1]): | |
| left = 0 | |
| for k, name, col, tc in qt: | |
| v = c.get(k, 0) | |
| a.barh(i, v, left=left, height=0.62, color=col, edgecolor="white", linewidth=1.2) | |
| a.text(left + v / 2, i, f"{name} {v:,}" if i == len(rows) - 1 else f"{v:,}", ha="center", va="center", fontsize=7.4, color=tc) | |
| left += v | |
| a.text(left + 250, i, f"{left:,}", va="center", ha="left", fontsize=8, color=G3) | |
| a.set_yticks(range(len(rows))) | |
| a.set_yticklabels([r[0] for r in rows[::-1]], fontsize=8) | |
| a.set_xlim(0, 27500) | |
| a.set_xlabel("readout training questions") | |
| a.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda v, _: f"{v / 1000:.0f}K")) | |
| grid(a, "x") | |
| a.tick_params(axis="y", length=0) | |
| a.set_title(f"Readout questions: {TF['pool_total']:,} (v1: 2,996); 7,320 targets mix in votes from the chains", | |
| fontsize=8.6, color=G3) | |
| ch = TF["chains"] | |
| s = lambda key, f: sum(ch[k][f] for k in ch if k.startswith(key)) | |
| crow = [("Qwen3.5-4B (own), procedural", s("4b_v1:proc", "trained") + s("4b_v2:proc", "trained"), | |
| s("4b_v1:proc", "trained_questions") + s("4b_v2:proc", "trained_questions"), BLUE), | |
| ("Qwen3.5-4B (own), tasksource", s("4b_v1:ts", "trained") + s("4b_v2:ts", "trained"), | |
| s("4b_v1:ts", "trained_questions") + s("4b_v2:ts", "trained_questions"), BLUE), | |
| ("Qwen3.5-35B-A3B teacher, tasksource", ch["teacher:ts"]["trained"], ch["teacher:ts"]["trained_questions"], AMBER)] | |
| for i, (lab, n, q, col) in enumerate(crow[::-1]): | |
| b.barh(i, n, height=0.62, color=col) | |
| b.text(n + 70, i, f"{n:,} chains on {q:,} questions", va="center", ha="left", fontsize=8, color=G3) | |
| b.set_yticks(range(len(crow))) | |
| b.set_yticklabels([r[0] for r in crow[::-1]], fontsize=8) | |
| b.set_xlim(0, 9000) | |
| b.xaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda v, _: f"{v / 1000:.0f}K")) | |
| b.set_xlabel("training chains (shortest correct of 4 samples; 2 when 3+ were correct)") | |
| grid(b, "x") | |
| b.tick_params(axis="y", length=0) | |
| tot = sum(r[1] for r in crow) | |
| b.set_title(f"Reasoning chains: {tot:,} on {sum(r[2] for r in crow):,} questions; the teacher wrote {ch['teacher:ts']['trained'] / tot:.1%}", | |
| fontsize=8.6, color=G3) | |
| title(fig, "V2 trained the readout on 14x v1's questions; chains were 92% the 4B's own", y=1.04) | |
| v1q = s("4b_v1", "trained") | |
| save(fig, "data_composition", "V2 training data", | |
| "V2's readout pool by source and question type, and the reasoning chains it was trained on by author. Sources: " | |
| "runs/v2/data/train_pool.jsonl and gen.jsonl, with chains selected by train.py's V2 rule (reproduces the run's " | |
| "10,974 chains on 5,779 questions exactly); computed by analysis/fig_train_compute.py.", | |
| {"readout_questions": TF["pool_total"], "procedural": d["proc"]["n"], "tasksource": d["ts_hard"]["n"] + d["ts_soft"]["n"], | |
| "tasksource_sources": 717, "qtype_all": {k: d["proc"]["by_qtype"][k] + d["ts_hard"]["by_qtype"][k] + d["ts_soft"]["by_qtype"][k] for k in ("choice", "noul", "score")}, | |
| "readout_targets_with_votes": 7320, "chains": tot, "chain_questions": sum(r[2] for r in crow), | |
| "chains_4b_procedural": crow[0][1], "chains_4b_tasksource": crow[1][1], "chains_teacher": crow[2][1], | |
| "chains_from_v1_generation": v1q, "teacher_share": round(ch["teacher:ts"]["trained"] / tot, 4)}) | |
| # ------------------------------------------------------------------ 2. LR sweep | |
| def lr_sweep(): | |
| s4, s2 = jload("runs/v2/results/lr_sweep.json"), jload("runs/v2_2b/results/lr_sweep.json") | |
| rates = ["1e-4", "2e-4", "3e-4"] | |
| x = np.arange(3) | |
| fig, (a, b) = plt.subplots(1, 2, figsize=(FULL, 2.55), gridspec_kw={"wspace": 0.42}) | |
| for s, col, lab, z in [(s2, G1, "MiniCPM5-2B", 2), (s4, BLUE, "Qwen3.5-4B", 3)]: | |
| nll = [s[r]["dev_nll"] for r in rates] | |
| acc = np.array([s[r]["dev_acc"] for r in rates]) * 100 | |
| ci = 1.96 * np.sqrt(acc / 100 * (1 - acc / 100) / 500) * 100 | |
| a.plot(x, nll, "-o", color=col, lw=2, ms=6, zorder=z) | |
| a.plot([1], [nll[1]], "o", ms=11, mfc="none", mec=col, mew=1.4, zorder=z) | |
| a.text(2.14, nll[2], lab, va="center", fontsize=8, color=col) | |
| a.text(1, nll[1] - 0.022, "picked", ha="center", va="top", fontsize=7.4, color=col) | |
| b.errorbar(x + (0.04 if col == BLUE else -0.04), acc, yerr=ci, fmt="-o", color=col, lw=2, ms=6, capsize=0, elinewidth=1, zorder=z) | |
| b.text(2.2, acc[2], lab, va="center", fontsize=8, color=col) | |
| for ax in (a, b): | |
| ax.set_xticks(x) | |
| ax.set_xticklabels(rates) | |
| ax.set_xlim(-0.25, 3.05) | |
| ax.set_xlabel("learning rate") | |
| grid(ax) | |
| a.set_title("Dev NLL, lower is better") | |
| a.set_ylabel("NLL (nats)") | |
| a.set_ylim(0.78, 1.1) | |
| b.set_title("Dev accuracy, 95% CI (n = 500)") | |
| b.set_ylabel("accuracy (%)") | |
| title(fig, "2e-4 had the lowest dev NLL for both base models", y=1.06) | |
| save(fig, "lr_sweep", "Learning-rate sweep", | |
| "Readout-only LoRA runs on 1,500 training questions (17 steps), scored on 500 held-out dev questions; the rule fixed " | |
| "beforehand took the lowest rate within 0.01 NLL of the best, and both models picked 2e-4 (the 2B's accuracy peaked at " | |
| "3e-4). Source: runs/v2/results/lr_sweep.json, runs/v2_2b/results/lr_sweep.json.", | |
| {"qwen4b": {r: {"acc": s4[r]["dev_acc"], "nll": round(s4[r]["dev_nll"], 3), "ece15": round(s4[r]["dev_ece15"], 3)} for r in rates}, | |
| "minicpm2b": {r: {"acc": s2[r]["dev_acc"], "nll": round(s2[r]["dev_nll"], 3), "ece15": round(s2[r]["dev_ece15"], 3)} for r in rates}, | |
| "picked": "2e-4", "dev_n": 500, "train_questions": 1500}) | |
| # ------------------------------------------------------------------ 3. speed benchmark | |
| def speed(): | |
| c = jload("runs/v2/results/speed_bench.json")["configs"] | |
| rows = [("A", "A checkpoint every layer,\n 16K micro-batches (old)"), ("B", "B A + Liger kernels"), | |
| ("C", "C no checkpointing up to 8K,\n 8K micro-batches (used)"), ("D", "D C + Liger kernels")] | |
| col = {"A": G3, "B": G1, "C": BLUE, "D": G1} | |
| lim = c["A"]["loss_err_vs_fp32"] * 1.5 * 100 | |
| gib = lambda g: g * 1e9 / 2 ** 30 | |
| fig, axs = plt.subplots(1, 3, figsize=(FULL, 2.7), sharey=True, gridspec_kw={"wspace": 0.16, "width_ratios": [1.15, 1, 1]}) | |
| y = np.arange(len(rows))[::-1] | |
| for (k, _), yi in zip(rows, y): | |
| r = c[k] | |
| axs[0].barh(yi, r["tok_per_s"] / 1000, color=col[k], height=0.6) | |
| axs[0].text(r["tok_per_s"] / 1000 + 0.08, yi, f"{r['tok_per_s'] / c['A']['tok_per_s']:.2f}x", va="center", fontsize=8, | |
| color=INK if k == "C" else G3, weight="bold" if k == "C" else "normal") | |
| e = r["loss_err_vs_fp32"] * 100 | |
| axs[1].barh(yi, e, color=col[k], height=0.6) | |
| ok = e <= lim | |
| axs[1].text(1.62, yi, "pass" if ok else "fail", va="center", fontsize=8, color=GREEN if ok else RED) | |
| axs[2].barh(yi, gib(r["peak_gb"]), color=col[k], height=0.6) | |
| axs[2].text(gib(r["peak_gb"]) - 1.5 if r["peak_gb"] > 50 else gib(r["peak_gb"]) + 1.5, yi, f"{gib(r['peak_gb']):.1f}", va="center", | |
| ha="right" if r["peak_gb"] > 50 else "left", fontsize=8, color="white" if r["peak_gb"] > 50 else G3) | |
| axs[0].set_yticks(y) | |
| axs[0].set_yticklabels([r[1] for r in rows], fontsize=7.6) | |
| axs[0].tick_params(axis="y", length=0) | |
| axs[0].set_xlim(0, 6.3) | |
| axs[0].set_xlabel("K tokens/s") | |
| axs[0].set_title("Throughput (x = vs A)") | |
| axs[1].axvline(lim, color=G2, lw=1, ls="--") | |
| axs[1].text(lim + 0.04, 3.55, "limit: 1.5x A", fontsize=7.2, color=G2, va="center") | |
| axs[1].set_xlim(0, 2.0) | |
| axs[1].set_xlabel("loss error vs fp32 (%)") | |
| axs[1].set_title("Numerics check") | |
| axs[2].axvline(gib(72), color=G2, lw=1, ls="--") | |
| axs[2].text(gib(72) - 1.5, 3.55, "72 GB limit", fontsize=7.2, color=G2, va="center", ha="right") | |
| axs[2].axvline(80, color=G3, lw=1) | |
| axs[2].text(82, -0.45, "80 GiB\ncard", fontsize=7.2, color=G3, ha="left", va="center") | |
| axs[2].set_xlim(0, 104) | |
| axs[2].set_xticks([0, 25, 50, 75]) | |
| axs[2].set_xlabel("GiB") | |
| axs[2].set_title("Peak memory") | |
| for ax in axs: | |
| grid(ax, "x") | |
| ax.set_ylim(-0.95, 3.85) | |
| ax.tick_params(axis="y", length=0) | |
| title(fig, "Skipping recomputation for short micro-batches trained 1.48x faster at equal numerics", y=1.05) | |
| save(fig, "speed", "Training-speed benchmark", | |
| "Each setup trained the same 260K tokens of V2 data on Qwen3.5-4B; numerics were judged against an fp32 reference " | |
| "(loss error at most 1.5x the old setup's). Liger kernels failed the loss check; C passed but peaked above the 72 GB " | |
| "guard set beforehand and was used on the user's call (the V2 joint run peaked at 71.5 GiB, no out-of-memory redos). " | |
| "Source: runs/v2/results/speed_bench.json.", | |
| {k: {"tok_per_s": round(c[k]["tok_per_s"]), "speedup_vs_A": round(c[k]["tok_per_s"] / c["A"]["tok_per_s"], 3), | |
| "loss_err_pct": round(c[k]["loss_err_vs_fp32"] * 100, 3), "grad_err": round(c[k]["grad_err_vs_fp32"], 3), | |
| "peak_GiB": round(gib(c[k]["peak_gb"]), 1), "equivalent": c[k]["equivalent"]} for k in "ABCD"} | |
| | {"loss_err_limit_pct": round(lim, 3), "used": "C", "v2_joint_run_tok_per_s": 4947, "v2_joint_run_peak_GiB": 71.5}) | |
| # ------------------------------------------------------------------ 4. V2 training curves | |
| def binned(pos, y, w, width): | |
| """Example-weighted means of per-step values in consecutive bins of `width` on the x axis; returns (bin centre, mean).""" | |
| pos, y, w = np.asarray(pos, float), np.asarray(y, float), np.asarray(w, float) | |
| ok = ~np.isnan(y) & (w > 0) | |
| idx = (pos // width).astype(int) | |
| xs, ys = [], [] | |
| for b in range(idx.max() + 1): | |
| m = ok & (idx == b) | |
| if m.any(): | |
| xs.append(min((b + 0.5) * width, pos[m].max())) | |
| ys.append((y[m] * w[m]).sum() / w[m].sum()) | |
| return np.array(xs), np.array(ys) | |
| def v2_training(): | |
| v2, v1 = steps("v2_joint"), steps("joint") | |
| W = 2.0e6 | |
| fig, (a, b) = plt.subplots(1, 2, figsize=(FULL, 2.75), gridspec_kw={"wspace": 0.3}) | |
| key = {"bin_tokens_M": W / 1e6} | |
| curves = {} | |
| for s, col, lab, z in [(v1, G1, "v1", 2), (v2, BLUE, "V2", 3)]: | |
| mid = s["tok"] - np.diff(np.concatenate([[0], s["tok"]])) / 2 # token position of each step's middle | |
| xr, ro = binned(mid, s["ro"], s["rn"], W) | |
| xc, ch = binned(mid, s["ch"], s["cn"], W) | |
| a.plot(xr / 1e6, ro, "-o", color=col, lw=1.8, ms=3.5, zorder=z) | |
| b.plot(xc / 1e6, ch, "-o", color=col, lw=1.8, ms=3.5, zorder=z) | |
| curves[lab] = (xr, ro, xc, ch) | |
| key[lab] = {"tokens_M": round(float(s["tok"][-1]) / 1e6, 2), "readout_examples": int(np.nansum(s["rn"])), "chains": int(np.nansum(s["cn"])), | |
| "readout_loss_bins": [round(float(v), 3) for v in ro], "chain_loss_bins": [round(float(v), 3) for v in ch], | |
| "readout_target_entropy_mean": round(float(np.nansum(s["ent"] * s["rn"]) / np.nansum(s["rn"][~np.isnan(s["ent"])])), 3)} | |
| xr1, ro1, xc1, ch1 = curves["v1"] | |
| xr2, ro2, xc2, ch2 = curves["V2"] | |
| a.text(1, 0.67, "v1, whole run", color=G2, fontsize=8, va="top") | |
| a.text(xr2[-1] / 1e6, ro2[-1] - 0.07, "V2", color=BLUE, fontsize=8, ha="center", va="top", weight="bold") | |
| b.text(xc1[-1] / 1e6 + 1.2, ch1[-1], "v1, whole run", color=G2, fontsize=8, va="center") | |
| b.text(xc2[-1] / 1e6, ch2[-1] + 0.015, "V2", color=BLUE, fontsize=8, ha="center", va="bottom", weight="bold") | |
| a.set_ylim(0, 1.0) | |
| b.set_ylim(0, 0.3) | |
| for ax in (a, b): | |
| ax.set_xlim(0, 50) | |
| ax.set_xlabel("training tokens (M)") | |
| grid(ax) | |
| a.set_title("Readout loss (soft cross-entropy, nats)") | |
| b.set_title("Chain loss (mean token cross-entropy, nats)") | |
| title(fig, "V2's readout loss kept falling across 48M tokens; chain loss stayed flat", y=1.06) | |
| save(fig, "v2_training", "V2 joint training", | |
| "Training losses of the V2 joint run and v1's joint run on one token axis, as example-weighted means over 2M-token bins " | |
| "(each run is a single pass, so every step scores examples it has not seen). v1's readout targets were all mixed with its " | |
| f"own votes (mean target entropy {key['v1']['readout_target_entropy_mean']:.2f} nats vs {key['V2']['readout_target_entropy_mean']:.2f} for V2), " | |
| "so its readout loss has a higher floor; the flat chain loss " | |
| "scores imitation of closed chains and could not show the loops V2 later produced (5.1% of eval chains). " | |
| "Source: train_log.jsonl of results/adapters/joint and runs/v2/results/adapters/v2_joint (via the dashboard data).", | |
| key) | |
| # ------------------------------------------------------------------ 5. RL dynamics | |
| def rl_dynamics(): | |
| rr = TF["rl_rounds"] | |
| x = np.array([r["step"] - 1 for r in rr]) # policy step the round's chains were sampled from | |
| ph_a = x < 40 | |
| fig, axs = plt.subplots(2, 2, figsize=(FULL, 4.5), sharex=True, gridspec_kw={"hspace": 0.42, "wspace": 0.27}) | |
| panels = [ | |
| (axs[0, 0], "Chains right (%)", [("right", 100, BLUE, "all"), ("right_among_finished", 100, G1, "finished")], (0, 100)), | |
| (axs[0, 1], "No answer by the 8,192-token cap (%)", [("unfinished", 100, BLUE, None)], (0, 35)), | |
| (axs[1, 0], "Median chain length (K tokens)", [("tokens_median_all", 1e-3, G1, "all"), ("tokens_median_finished", 1e-3, BLUE, "finished")], (0, 6)), | |
| (axs[1, 1], "\"Wait\" per finished chain", [("waits_per_finished", 1, BLUE, None)], (0, 12)), | |
| ] | |
| key = {} | |
| for ax, ttl, series, ylim in panels: | |
| for k, mul, col, lab in series: | |
| v = np.array([r[k] for r in rr]) * mul | |
| ax.plot(x, v, "-o", color=col, lw=1.6, ms=3.6, zorder=3 if col == BLUE else 2) | |
| if col == BLUE: | |
| for m, xs in [(ph_a, (0, 36)), (~ph_a, (40, 56))]: | |
| mu = v[m].mean() | |
| ax.plot(xs, [mu, mu], color=BLUE, lw=1, ls=(0, (4, 2)), alpha=0.75, zorder=2) | |
| key[k] = {"first_round": round(float(v[0]), 3), "last_round": round(float(v[-1]), 3), | |
| "mean_lr_1e-5_rounds": round(float(v[ph_a].mean()), 3), "mean_lr_2.5e-5_rounds": round(float(v[~ph_a].mean()), 3)} | |
| if lab: | |
| dy = {"all": 1, "finished": -1}[lab] * (0.3 if ylim[1] < 10 else 0) if ax is axs[1, 0] else 0 | |
| ax.text(57.8, v[-1] + dy, lab, color=col, fontsize=7.4, va="center") | |
| ax.axvline(40, color=G2, lw=0.9, ls="--", zorder=1) | |
| ax.set_ylim(*ylim) | |
| ax.set_title(ttl) | |
| grid(ax) | |
| ax.set_xlim(-2, 69) | |
| ax.set_xticks([0, 10, 20, 30, 40, 50, 60]) | |
| axs[0, 0].text(40.8, 6, "lr 1e-5 | 2.5e-5", fontsize=7.4, color=G2, va="bottom") | |
| # phase-mean labels on the decisive panels | |
| w = key["waits_per_finished"] | |
| axs[1, 1].text(18, w["mean_lr_1e-5_rounds"] + 0.5, f"mean {w['mean_lr_1e-5_rounds']:.1f}", color=BLUE, fontsize=7.4, ha="center", va="bottom") | |
| axs[1, 1].text(49, w["mean_lr_2.5e-5_rounds"] + 0.55, f"mean {w['mean_lr_2.5e-5_rounds']:.1f}", color=BLUE, fontsize=7.4, ha="center", va="bottom") | |
| u = key["unfinished"] | |
| axs[0, 1].text(18, 31, f"mean {u['mean_lr_1e-5_rounds']:.0f}%", color=BLUE, fontsize=7.4, ha="center") | |
| axs[0, 1].text(50, 31, f"mean {u['mean_lr_2.5e-5_rounds']:.0f}%", color=BLUE, fontsize=7.4, ha="center") | |
| fig.supxlabel("optimizer step (each point: one round of 16 questions x 8 chains, sampled at that step)", fontsize=8.5, color=G3, y=0.0) | |
| title(fig, "RL cut \"Wait\" per finished chain from 10 to 5; training-time accuracy stayed flat", y=1.0) | |
| # difficulty-adjusted trends per round (share right / unfinished ~ V2 readout p_gold + round), as in the notebook | |
| q = TF["rl_questions"] | |
| X = np.array([[1, z["p_gold"], (z["step"] - 1) // 4] for z in q], float) | |
| adj = {} | |
| for nm, y in [("right", np.array([z["right_of_8"] / 8 for z in q])), ("unfinished", np.array([z["unfinished"] / 8 for z in q]))]: | |
| bta = np.linalg.lstsq(X, y, rcond=None)[0] | |
| e = y - X @ bta | |
| se = np.sqrt(np.diag(e @ e / (len(y) - 3) * np.linalg.inv(X.T @ X))) | |
| adj[nm] = {"per_round_pts": round(float(bta[2]) * 100, 2), "se_pts": round(float(se[2]) * 100, 2)} | |
| key["difficulty_adjusted_trend_per_round"] = adj | |
| key["rounds"] = len(rr) | |
| key["questions"] = len(q) | |
| save(fig, "rl_dynamics", "RL training dynamics", | |
| "GRPO on V2's thinking mode, one point per generation round (16 questions not used before x 8 chains, feeding 4 optimizer steps); " | |
| "dashed lines are the means of the 1e-5 and 2.5e-5 rounds. Rounds draw different questions, so single points mix " | |
| f"difficulty with learning: adjusted for the V2 readout's p(gold), share right moved {adj['right']['per_round_pts']:+.1f} " | |
| f"points per round (se {adj['right']['se_pts']:.1f}) and no-answer {adj['unfinished']['per_round_pts']:+.1f} (se " | |
| f"{adj['unfinished']['se_pts']:.1f}). Source: runs/v2/checkpoints/rl_grpo*/completions (tokens counted with the model's tokenizer).", | |
| key) | |
| # ------------------------------------------------------------------ 6. 2B training curves | |
| def twob_training(): | |
| j, ro, v2 = steps("v2b_joint"), steps("v2b_ro"), steps("v2_joint") | |
| fig, (a, b) = plt.subplots(1, 2, figsize=(FULL, 2.8), gridspec_kw={"wspace": 0.28, "width_ratios": [1.45, 1]}) | |
| runs = [(v2, G1, "4B, V2 recipe", 1), (ro, AMBER, "2B, readout only", 2), (j, BLUE, "2B, V2 recipe", 3)] | |
| WR, WC = 2500, 1000 | |
| key = {"bin_readout_questions": WR, "bin_chains": WC} | |
| ends = {} | |
| for s, col, lab, z in runs: | |
| rn, cn = np.nan_to_num(s["rn"]), np.nan_to_num(s["cn"]) | |
| xr, yr = binned(np.cumsum(rn) - rn / 2, s["ro"], s["rn"], WR) | |
| a.plot(xr / 1000, yr, "-o", color=col, lw=1.8, ms=3.2, zorder=z) | |
| ends[lab] = (xr[-1] / 1000, yr[-1]) | |
| k = {"readout_loss_bins": [round(float(v), 3) for v in yr], "dev_acc": round(s["meta"]["dev_acc"], 3), | |
| "dev_nll": round(s["meta"]["dev_nll"], 3), "dev_ece15": round(s["meta"]["dev_ece15"], 3), "steps": s["meta"]["steps_done"], | |
| "tokens_M": round(s["meta"]["train_tokens"] / 1e6, 1)} | |
| if s["meta"]["mode"] == "joint": | |
| xc, yc = binned(np.cumsum(cn) - cn / 2, s["ch"], s["cn"], WC) | |
| b.plot(xc / 1000, yc, "-o", color=col, lw=1.8, ms=3.2, zorder=z) | |
| k["chain_loss_bins"] = [round(float(v), 3) for v in yc] | |
| b.text(xc[-1] / 1000 + 0.4, yc[-1], lab, color=col, fontsize=7.6, va="center", weight="bold" if col == BLUE else "normal") | |
| key[lab] = k | |
| off = {"4B, V2 recipe": -0.03, "2B, readout only": -0.012, "2B, V2 recipe": 0.035} | |
| for _, col, lab, _ in runs: | |
| x, y = ends[lab] | |
| a.text(x + 1.2, y + off[lab], f"{lab}: dev {key[lab]['dev_acc'] * 100:.1f}%", color=col, fontsize=7.4, va="center", | |
| weight="bold" if col == BLUE else "normal") | |
| a.set_xlim(0, 79) | |
| a.set_xticks([0, 10, 20, 30, 40]) | |
| a.set_ylim(0, 1.4) | |
| a.set_xlabel("readout questions seen (K)") | |
| a.set_title("Readout loss (soft cross-entropy, nats)") | |
| b.set_xlim(0, 17) | |
| b.set_xticks([0, 2, 4, 6, 8, 10]) | |
| b.set_ylim(0, 0.5) | |
| b.set_xlabel("training chains seen (K)") | |
| b.set_title("Chain loss (token cross-entropy, nats)") | |
| for ax in (a, b): | |
| grid(ax) | |
| title(fig, "The 2B fit the 4B's chains poorly, and they did not help its readout", y=1.06) | |
| save(fig, "twob_training", "2B training", | |
| "MiniCPM5-2B trained on V2's data with chains (V2 recipe) and without (readout only), with the 4B's V2 run for reference: " | |
| "example-weighted means over bins of 2,500 readout questions or 1,000 chains (an examples axis, since the tokenizers " | |
| "differ). Dev = readout accuracy on 999 held-out questions. Source: train_log.jsonl / train_meta.json of " | |
| "runs/v2_2b/results/adapters/{v2b_joint,v2b_ro} and runs/v2/results/adapters/v2_joint (via the dashboard data).", | |
| key) | |
| # ------------------------------------------------------------------ 7. extra: chain selection | |
| def chain_selection(): | |
| ch = TF["chains"] | |
| groups = [("4B own, procedural", ["4b_v1:proc", "4b_v2:proc"], BLUE), ("4B own, tasksource", ["4b_v1:ts", "4b_v2:ts"], BLUE), | |
| ("Teacher (35B-A3B), tasksource", ["teacher:ts"], AMBER)] | |
| fig, (a, b) = plt.subplots(1, 2, figsize=(FULL, 2.45), sharey=True, gridspec_kw={"wspace": 0.12, "width_ratios": [1.35, 1]}) | |
| key = {} | |
| for i, (lab, ks, col) in enumerate(groups[::-1]): | |
| samp = np.concatenate([ch[k]["sample_tokens"] for k in ks]) | |
| trn = np.concatenate([ch[k]["trained_tokens"] for k in ks]) | |
| sw = np.average([ch[k]["sample_waits_mean"] for k in ks], weights=[ch[k]["samples"] for k in ks]) | |
| tw = np.concatenate([ch[k]["trained_waits"] for k in ks]).mean() | |
| for vals, c, dy, nm in [(samp, G1, 0.17, "all samples"), (trn, col, -0.17, "trained")]: | |
| p10, p50, p90 = np.percentile(vals, [10, 50, 90]) | |
| a.plot([p10, p90], [i + dy, i + dy], color=c, lw=2.2, solid_capstyle="round") | |
| a.plot([p50], [i + dy], "o", color=c, ms=7, mec="white", mew=1) | |
| a.text(p90 * 1.12, i + dy, f"{p50:,.0f}", fontsize=7.4, va="center", color=G3) | |
| if i == 2: | |
| a.text(165, i + dy, nm, fontsize=7.4, va="center", ha="left", color=c) | |
| b.barh(i + 0.17, sw, height=0.3, color=G1) | |
| b.barh(i - 0.17, tw, height=0.3, color=col) | |
| b.text(sw + 0.15, i + 0.17, f"{sw:.1f}", fontsize=7.4, va="center", color=G3) | |
| b.text(tw + 0.15, i - 0.17, f"{tw:.1f}", fontsize=7.4, va="center", color=G3) | |
| key[lab] = {"sampled_median_tokens": float(np.median(samp)), "trained_median_tokens": float(np.median(trn)), | |
| "sampled_p90": round(float(np.percentile(samp, 90))), "trained_p90": round(float(np.percentile(trn, 90))), | |
| "waits_per_closed_sample": round(float(sw), 2), "waits_per_trained_chain": round(float(tw), 2), | |
| "samples": int(len(samp)), "trained": int(len(trn))} | |
| a.set_xscale("log") | |
| a.set_xlim(150, 40000) | |
| a.set_xticks([200, 500, 1000, 2000, 5000, 10000, 30000]) | |
| a.set_xticklabels(["200", "500", "1K", "2K", "5K", "10K", "30K"]) | |
| a.set_xlabel("tokens (log scale)") | |
| a.set_yticks(range(3)) | |
| a.set_yticklabels([g[0] for g in groups[::-1]], fontsize=8) | |
| a.tick_params(axis="y", length=0) | |
| a.set_title("Chain length (median, 10th-90th pct.)") | |
| a.minorticks_off() | |
| b.set_xlim(0, 10.5) | |
| b.set_xlabel("mean per chain") | |
| b.set_title("\"Wait\" per chain") | |
| b.tick_params(axis="y", length=0) | |
| for ax in (a, b): | |
| grid(ax, "x") | |
| title(fig, "Selection trimmed tasksource chains most: half the tokens, a quarter of the \"Wait\"s", y=1.07) | |
| save(fig, "chain_selection", "Chain selection", | |
| "Chain length and \"Wait\" count of all sampled chains (grey; \"Wait\" over chains that closed) vs the chains V2 " | |
| "trained on (shortest correct of 4, second-shortest too when 3+ were correct; invented look-ups and repetition filtered), " | |
| "by author. After selection the teacher's tasksource chains are only 12% shorter than the 4B's own. " | |
| "Source: runs/v2/data/gen.jsonl via analysis/fig_train_compute.py.", | |
| key) | |
| # ------------------------------------------------------------------ 8. extra: RL signal per question | |
| def rl_signal(): | |
| q = TF["rl_questions"] | |
| kinds = [("fresh_reasoning", "New, reasoning-heavy"), ("fresh_other", "New, other sources"), | |
| ("in_domain", "From V2's training pool")] | |
| fig, axs = plt.subplots(1, 3, figsize=(FULL, 2.5), sharey=True, gridspec_kw={"wspace": 0.1}) | |
| key = {} | |
| for ax, (k, lab) in zip(axs, kinds): | |
| r = np.array([z["right_of_8"] for z in q if z["kind"] == k]) | |
| h = np.bincount(r, minlength=9) / len(r) * 100 | |
| cols = [G1] + [BLUE] * 7 + [G1] | |
| ax.bar(range(9), h, color=cols, width=0.78) | |
| mixed = ((r > 0) & (r < 8)).mean() * 100 | |
| ax.set_title(f"{lab}\n(n = {len(r)})", fontsize=8.2) | |
| ax.text(4, 52, f"{mixed:.0f}% split", color=BLUE, ha="center", fontsize=8, weight="bold") | |
| ax.text(0, h[0] + 1.2, f"{h[0]:.0f}%", ha="center", fontsize=7.2, color=G3) | |
| ax.text(8, h[8] + 1.2, f"{h[8]:.0f}%", ha="center", fontsize=7.2, color=G3) | |
| ax.set_xticks([0, 2, 4, 6, 8]) | |
| ax.set_xlabel("chains right, of 8") | |
| ax.set_ylim(0, 60) | |
| grid(ax) | |
| key[k] = {"n": int(len(r)), "split_pct": round(mixed, 1), "all_wrong_pct": round(h[0], 1), "all_right_pct": round(h[8], 1), | |
| "mean_right": round(float(r.mean() / 8), 3), "mean_v2_readout_p_gold": round(float(np.mean([z["p_gold"] for z in q if z["kind"] == k])), 3)} | |
| r = np.array([z["right_of_8"] for z in q]) | |
| key["all"] = {"n": int(len(r)), "split_pct": round(((r > 0) & (r < 8)).mean() * 100, 1), "all_wrong_pct": round((r == 0).mean() * 100, 1), | |
| "all_right_pct": round((r == 8).mean() * 100, 1)} | |
| axs[0].set_ylabel("questions (%)") | |
| title(fig, f"Only {key['all']['split_pct']:.0f}% of RL questions split right/wrong across their 8 chains", y=1.13) | |
| save(fig, "rl_signal", "RL questions by outcome", | |
| "For each of the 240 RL questions (new = tasksource train-split questions never trained on), how many of its 8 chains " | |
| "were right (reward > 0; grey = all wrong or all right, where GRPO gets no " | |
| "right-vs-wrong signal, only the length and no-answer penalties). New questions were screened only by the fast readout's " | |
| "p(gold) (0.15-0.85), which did not predict thinking-mode difficulty: 46% of the reasoning-heavy ones were 8 of 8. " | |
| "Source: runs/v2/checkpoints/rl_grpo*/completions and runs/v2/data/rl_prompts.jsonl.", | |
| key) | |
| FIGS = {"data_composition": data_composition, "lr_sweep": lr_sweep, "speed": speed, "v2_training": v2_training, | |
| "rl_dynamics": rl_dynamics, "twob_training": twob_training, "chain_selection": chain_selection, "rl_signal": rl_signal} | |
| if __name__ == "__main__": | |
| DASHD = jload(DASH) | |
| TF = jload("results/report_figs/train_figs_data.json") | |
| want = sys.argv[1:] or list(FIGS) | |
| for name in want: | |
| FIGS[name]() | |
| mp = os.path.join(OUT, "manifest_train.json") | |
| old = [m for m in (json.load(open(mp)) if os.path.exists(mp) else []) if m["slug"] not in {m["slug"] for m in MANIFEST}] | |
| order = {k: i for i, k in enumerate(FIGS)} | |
| allm = sorted(old + MANIFEST, key=lambda m: order.get(m["slug"], 99)) | |
| with open(mp, "w") as f: | |
| json.dump(allm, f, indent=1) | |
| print("manifest", mp, len(allm)) | |
Xet Storage Details
- Size:
- 31.3 kB
- Xet hash:
- 9aeafa13458f62daa1b0ccd26255cd39049332e561cbb335617439ac48568454
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.