File size: 4,622 Bytes
3332e80 6b81445 3332e80 6b81445 3332e80 6b81445 3332e80 6b81445 3332e80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | #!/usr/bin/env python3
"""Appendix training-dynamics figures (4B / 8B), one single-figure PDF+PNG per
panel — composed into a 2x2 layout in LaTeX via \\subfigure. Follows doc/vis.md:
restricted academic palette, sans-serif, large fonts, tight crop, B&W-safe
(solid/dashed + circle/square markers), "Ours" (GroundFlow) gets the prominent
coral colour consistently.
Outputs (to doc/figures/ and figures/):
figA_loss_4b.{pdf,png} figA_loss_8b.{pdf,png}
figA_gn_4b.{pdf,png} figA_gn_8b.{pdf,png}
"""
import json, math, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
M = "/mnt/bn/leonworkspace/terry/model"
TRAINER = {
"4B": {
"SD-RPN": f"{M}/qwen3vl-4b-roi-K24T3-185k-ddp-verify/trainer_state.json",
"GroundFlow": f"{M}/qwen3vl-4b-roi-K24T3-185k-selfdistill-run7b-sdpa-iwa/trainer_state.json",
},
"8B": {
"SD-RPN": f"{M}/qwen3vl-8b-roi-K24T3-185k-sdrpn/trainer_state.json",
"GroundFlow": f"{M}/qwen3vl-8b-roi-K24T3-185k-a0.667-t1.5-iwa-g1.0/trainer_state.json",
},
}
# doc/vis.md palette: GroundFlow (=Ours) = coral #f57c6e (prominent, consistent).
# Per request, only the GroundFlow curve is plotted (SD-RPN's loss is not shown).
STYLE = {
"GroundFlow": dict(color="#f57c6e", ls="-", marker="o", label="GroundFlow"),
}
ORDER = ["GroundFlow"]
OUTDIRS = ["/opt/tiger/thothvl_pretrain/doc/figures", "/opt/tiger/thothvl_pretrain/figures"]
def load(path):
lh = json.load(open(path))["log_history"]
rows = [(d["step"], d["loss"], d.get("grad_norm")) for d in lh if "loss" in d]
step = np.array([r[0] for r in rows], float)
loss = np.array([r[1] for r in rows], float)
gn = np.array([(r[2] if r[2] is not None else np.nan) for r in rows], float)
fin = [d for d in lh if "train_loss" in d]
return step, loss, gn, (fin[-1]["train_loss"] if fin else float(loss[-1]))
def ema(x, alpha=0.06):
y = np.empty_like(x); m = next(v for v in x if math.isfinite(v))
for i, v in enumerate(x):
if math.isfinite(v):
m = alpha * v + (1 - alpha) * m
y[i] = m
return y
plt.rcParams.update({
"font.family": "sans-serif",
"font.sans-serif": ["Helvetica", "Arial", "DejaVu Sans"],
"mathtext.fontset": "dejavusans",
"font.size": 11, "axes.titlesize": 13, "axes.labelsize": 12,
"xtick.labelsize": 10.5, "ytick.labelsize": 10.5, "legend.fontsize": 10,
"axes.linewidth": 0.9, "xtick.direction": "in", "ytick.direction": "in",
"xtick.major.size": 3.2, "ytick.major.size": 3.2,
"legend.frameon": False, "axes.spines.top": False, "axes.spines.right": False,
"savefig.bbox": "tight", "savefig.pad_inches": 0.02,
})
def make_panel(model, kind, fname):
"""kind in {'loss','gn'}"""
fig, ax = plt.subplots(figsize=(4.0, 3.0))
for name in ORDER:
st, ls, gn, fl = load(TRAINER[model][name])
y = ls if kind == "loss" else gn
s = STYLE[name]
idx = np.linspace(0, len(st) - 1, 8).round().astype(int) # 8 markers, B&W aid
ax.plot(st, y, color=s["color"], lw=0.6, alpha=0.16) # raw, faint
ax.plot(st, ema(y), color=s["color"], ls=s["ls"], lw=2.0, # EMA
marker=s["marker"], markevery=list(idx), ms=4.5, mew=0,
label=s["label"], zorder=3 if name == "GroundFlow" else 2)
if kind == "loss":
yi = ema(y)[-1]
ax.annotate(f"{fl:.2f}" if fl >= 0.1 else f"{fl:.3f}",
xy=(st[-1], yi), xytext=(5, 0), textcoords="offset points",
va="center", fontsize=10, color=s["color"], fontweight="bold")
if kind == "loss":
ax.set_ylabel("training loss")
ax.set_ylim(bottom=0) # linear: clean single-curve shape
else:
ax.set_yscale("log"); ax.set_ylabel("gradient norm") # spans ~2 decades
ax.set_xlabel("optimizer step")
ax.set_title(f"Qwen3-VL-{model}")
ax.set_xlim(-30, 1480)
ax.xaxis.set_major_locator(mticker.MultipleLocator(500))
ax.grid(True, which="both", ls=":", lw=0.5, alpha=0.35)
ax.legend(loc="upper right", handlelength=2.4, borderaxespad=0.4)
fig.tight_layout()
for d in OUTDIRS:
os.makedirs(d, exist_ok=True)
for ext in ("pdf", "png"):
fig.savefig(f"{d}/{fname}.{ext}", dpi=300)
plt.close(fig)
print(f" wrote {fname}.{{pdf,png}}")
for model in ("4B", "8B"):
make_panel(model, "loss", f"figA_loss_{model.lower()}")
make_panel(model, "gn", f"figA_gn_{model.lower()}")
print("done.")
|