|
|
| |
| """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() |
|
|
|
|
|
|
