#!/usr/bin/env python3 """Find when a training run (CSV log) reaches a target loss, after smoothing. Two cases, both reported: - target already crossed within the logged data -> interpolated step/tokens. - target not yet crossed -> fit L(S) = L0 + A*S^-alpha (Eq.14) to the post-warmup smoothed trajectory and solve for S algebraically, flagging how far past the observed range the extrapolation reaches. """ import argparse import re import sys import numpy as np import pandas as pd from scipy.optimize import curve_fit EXTRAPOLATION_WARN_FACTOR = 3.0 CONVERGENCE_TOL = 0.05 SWEEP_STEP = 300 SWEEP_MIN_WINDOW = 500 def power_law(s, l0, a, alpha): return l0 + a * np.power(s, -alpha) def parse_batch_size_from_filename(path): m = re.search(r"_([\d.]+)([MK])bs", path, re.IGNORECASE) if not m: return None value, unit = float(m.group(1)), m.group(2).upper() return value * (1_000_000 if unit == "M" else 1_000) def detect_warmup_end(df): lr_max = df["lr"].max() past_warmup = df.index[df["lr"] >= 0.99 * lr_max] return int(past_warmup[0]) if len(past_warmup) else 0 def interpolated_crossing(steps, smoothed, target): below = smoothed <= target if not below.any(): return None i = below.idxmax() if i == smoothed.index[0]: return float(steps[i]) i_prev = i - 1 s0, s1 = steps[i_prev], steps[i] l0, l1 = smoothed[i_prev], smoothed[i] if l0 == l1: return float(s1) frac = (l0 - target) / (l0 - l1) return float(s0 + frac * (s1 - s0)) def fit_and_solve(steps, smoothed, target): p0 = [smoothed.min() * 0.9, smoothed.iloc[0] - smoothed.min(), 0.3] try: (l0, a, alpha), _ = curve_fit( power_law, steps, smoothed, p0=p0, bounds=([0, 1e-6, 1e-3], [smoothed.min(), np.inf, 10]), maxfev=20000, ) except RuntimeError as e: return None, f"fit failed to converge: {e}" if target <= l0: return None, f"unreachable under this fit -- asymptote L0={l0:.4f} >= target {target}" s_pred = (a / (target - l0)) ** (1.0 / alpha) return s_pred, (l0, a, alpha) def stability_sweep(df, smoothed_full, fit_start, target, sweep_step, min_window): s_max_total = int(df["step"].max()) checkpoints = list(range(fit_start + min_window, s_max_total, sweep_step)) if not checkpoints or checkpoints[-1] != s_max_total: checkpoints.append(s_max_total) rows = [] for max_step in checkpoints: mask = (df["step"] >= fit_start) & (df["step"] <= max_step) if mask.sum() < 10: continue steps = df["step"][mask] smoothed = smoothed_full.loc[steps.index] s_pred, info = fit_and_solve(steps, smoothed, target) rows.append((max_step, s_pred, info)) return rows def print_sweep(target, rows, convergence_tol, warn_factor): print(f"\ntarget={target} -- stability sweep") print(f"{'max_step':>10} {'s_pred':>14} {'ratio':>8} {'L0':>8} {'alpha':>8} note") for max_step, s_pred, info in rows: if s_pred is None: print(f"{max_step:>10} {'--':>14} {'--':>8} {'--':>8} {'--':>8} {info}") continue l0, a, alpha = info ratio = s_pred / max_step flag = " [beyond warn factor]" if ratio > warn_factor else "" print(f"{max_step:>10} {s_pred:>14,.0f} {ratio:>7.2f}x {l0:>8.4f} {alpha:>8.4f}{flag}") valid = [(m, s) for m, s, _ in rows if s is not None] if len(valid) < 3: print("verdict: NOT ENOUGH DATA to judge convergence yet (need >=3 valid windows)") return (_, s_prev2), (_, s_prev1), (_, s_last) = valid[-3:] rel_changes = [abs(s_prev1 - s_prev2) / s_prev2, abs(s_last - s_prev1) / s_prev1] if all(c <= convergence_tol for c in rel_changes): print(f"verdict: CONVERGED -- last two step-to-step changes " f"({rel_changes[0]:.1%}, {rel_changes[1]:.1%}) both within tolerance {convergence_tol:.0%}") else: print(f"verdict: NOT YET CONVERGED -- last two step-to-step changes " f"({rel_changes[0]:.1%}, {rel_changes[1]:.1%}) exceed tolerance {convergence_tol:.0%}; " f"keep training and re-check") def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("csv") p.add_argument("--target", type=float, nargs="+", required=True) p.add_argument("--loss-col", default="train_loss") p.add_argument("--ema-span", type=float, default=50, help="EMA span in steps for smoothing (default: 50)") p.add_argument("--batch-size", type=float, default=None, help="tokens per step; parsed from filename (e.g. _1Mbs) if omitted") p.add_argument("--fit-start-step", type=int, default=None, help="default: auto-detected end of LR warmup") p.add_argument("--extrapolation-warn-factor", type=float, default=EXTRAPOLATION_WARN_FACTOR) p.add_argument("--stability-sweep", action="store_true", help="instead of a single prediction, refit on growing windows " "of data and report whether the extrapolated step has " "converged yet") p.add_argument("--sweep-step", type=int, default=SWEEP_STEP, help="steps between successive sweep windows (default: 300)") p.add_argument("--sweep-min-window", type=int, default=SWEEP_MIN_WINDOW, help="post-warmup steps before the first sweep window (default: 500)") p.add_argument("--convergence-tol", type=float, default=CONVERGENCE_TOL, help="max relative change between successive windows to call " "it converged (default: 0.05 = 5%%)") args = p.parse_args() df = pd.read_csv(args.csv) if args.loss_col not in df.columns: sys.exit(f"'{args.loss_col}' not in columns: {list(df.columns)}") batch_size = args.batch_size or parse_batch_size_from_filename(args.csv) if batch_size is None: sys.exit("could not infer batch size from filename -- pass --batch-size") fit_start = args.fit_start_step if fit_start is None: fit_start = detect_warmup_end(df) smoothed_full = df[args.loss_col].ewm(span=args.ema_span, adjust=False).mean() post_warmup = df[df["step"] >= fit_start] smoothed_post = smoothed_full.loc[post_warmup.index] steps_post = post_warmup["step"] s_max_observed = steps_post.max() print(f"loss col: {args.loss_col} | ema span: {args.ema_span} steps | " f"batch size: {batch_size:,.0f} tokens/step | fit start step: {fit_start}") print("-" * 100) for target in args.target: s_cross = interpolated_crossing(df["step"], smoothed_full, target) if s_cross is not None: tokens = s_cross * batch_size print(f"target={target:<8} REACHED (interpolated): " f"step={s_cross:,.1f} tokens={tokens:,.0f}") continue if args.stability_sweep: rows = stability_sweep(df, smoothed_full, fit_start, target, args.sweep_step, args.sweep_min_window) print_sweep(target, rows, args.convergence_tol, args.extrapolation_warn_factor) continue s_pred, info = fit_and_solve(steps_post, smoothed_post, target) if s_pred is None: print(f"target={target:<8} NOT REACHED, fit says: {info}") continue l0, a, alpha = info ratio = s_pred / s_max_observed tokens = s_pred * batch_size flag = " ** EXTRAPOLATION BEYOND WARN FACTOR **" if ratio > args.extrapolation_warn_factor else "" print(f"target={target:<8} EXTRAPOLATED (fit L0={l0:.4f} A={a:.4f} alpha={alpha:.4f}): " f"step={s_pred:,.1f} tokens={tokens:,.0f} " f"({ratio:.2f}x past observed S={s_max_observed:,.0f}){flag}") if __name__ == "__main__": main()