Spaces:
Running
Running
File size: 2,315 Bytes
dbd3aaf | 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 | #!/usr/bin/env python3
"""Plot GRPO saturation evidence for the Blog: frac_reward_zero_std + grad_norm vs step.
Reads the aggressive warm-start single-task run logged under docs/hf_runs/.
"""
from __future__ import annotations
import csv
from pathlib import Path
import matplotlib.pyplot as plt
ROOT = Path(__file__).resolve().parents[2]
DEFAULT_CSV = (
ROOT / "docs/hf_runs/continue_warm_start/continue_deep_seed_5823_lr5e-6/training_metrics.csv"
)
OUT = ROOT / "docs" / "plots" / "grpo_aggressive_lr_saturation.svg"
def main() -> None:
csv_path = DEFAULT_CSV
if not csv_path.exists():
raise SystemExit(f"Missing {csv_path}")
steps: list[int] = []
frac: list[float] = []
grad: list[float] = []
with csv_path.open() as f:
for row in csv.DictReader(f):
steps.append(int(row["step"]))
fz = row.get("frac_reward_zero_std") or ""
frac.append(float(fz) if fz not in ("", "nan") else float("nan"))
g = row.get("grad_norm") or ""
grad.append(float(g) if g not in ("", "nan") else float("nan"))
fig, ax1 = plt.subplots(figsize=(9, 4.2))
ax1.set_xlabel("Training step", fontsize=11)
ax1.set_ylabel("Fraction of prompt groups with zero reward variance", fontsize=10, color="#b45309")
(ln1,) = ax1.plot(steps, frac, color="#b45309", linewidth=2, marker="o", markersize=3, label="frac_reward_zero_std")
ax1.set_ylim(-0.05, 1.05)
ax1.grid(True, alpha=0.25)
ax2 = ax1.twinx()
ax2.set_ylabel("Gradient norm (policy update size)", fontsize=10, color="#1d4ed8")
(ln2,) = ax2.plot(steps, grad, color="#1d4ed8", linewidth=1.8, linestyle="--", marker="s", markersize=2, label="grad_norm")
ax2.set_ylim(-0.02, max(0.5, max((g for g in grad if g == g), default=0.1) * 1.2))
fig.suptitle(
"Aggressive warm-start run (continue_deep_seed_5823_lr5e-6): GRPO loses signal when all completions tie",
fontsize=11,
fontweight="bold",
y=1.02,
)
fig.legend([ln1, ln2], [ln1.get_label(), ln2.get_label()], loc="upper right", fontsize=9)
fig.tight_layout()
OUT.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(OUT, format="svg", bbox_inches="tight")
plt.close(fig)
print("wrote", OUT)
if __name__ == "__main__":
main()
|