#!/usr/bin/env python3 """Find when a training run (CSV log) reaches a target loss, after smoothing. Works for fixed-batch runs and for dynamic-batch runs (batch size changes mid-run, e.g. train_390M_dynamicb.py). Two cases, both reported: - target already crossed within the logged data -> interpolated step/tokens. - target not yet crossed -> fit L(T) = L0 + A*T^-alpha (Eq.14) to the post-warmup smoothed trajectory and solve for T algebraically, flagging how far past the observed range the extrapolation reaches. Everything runs on a TOKEN axis, not steps: once the batch size changes, a step no longer stands for a fixed amount of data, so L(S) kinks at every switch and a predicted step can't be compared across runs. For a fixed batch T = B*S, so L0 + A*S^-alpha and L0 + A'*T^-alpha are the same curve (same L0 and alpha, A' = A*B^alpha) -- fixed-batch results are unchanged. Batch size per step: the CSV's `global_batch` column if present (the dynamic trainer logs it), else --batch-size, else parsed from the filename (_1Mbs). Tokens consumed before each step: the CSV's `tokens` column if present, else accumulated from the per-step batch (= step * B for a fixed batch). """ 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 EMA_TOKEN_WINDOW = 100_000_000 # smoothing window in TOKENS, not steps -- see # batchsize_scaling.py for why: a fixed step # count covers a different amount of data per # run depending on batch size, which biases # the cross-batch-size comparison. The same # holds across the stages of one dynamic run. TOKEN_UNIT = 1e9 # the fit's x-axis is tokens / TOKEN_UNIT, which # keeps A near O(1) so curve_fit's p0 is sane def power_law(t, l0, a, alpha): return l0 + a * np.power(t, -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 parse_batch_schedule(pairs): """TOKENS:BATCH pairs, as train_390M_dynamicb.py's --batch-schedule takes them -> [(start_tokens, batch), ...] sorted by start.""" try: return sorted((float(t), float(b)) for t, b in (s.split(":") for s in pairs)) except ValueError: sys.exit(f"--batch-schedule takes TOKENS:BATCH pairs, got {' '.join(pairs)}") def load_run(path, loss_col, batch_size_arg): """-> (df with per-step `batch` and `tokens` columns, where the batch came from).""" df = pd.read_csv(path) if loss_col not in df.columns: sys.exit(f"'{loss_col}' not in columns: {list(df.columns)}") # A resumed run appends to its CSV from the checkpoint step, so steps the # interrupted attempt logged past that checkpoint appear twice -- keep the # resumed (later) rows. df = df.drop_duplicates("step", keep="last").sort_values("step").reset_index(drop=True) if "global_batch" in df.columns: df["batch"] = df["global_batch"].astype(float) source = "CSV global_batch column" else: b = batch_size_arg or parse_batch_size_from_filename(path) if b is None: sys.exit("no global_batch column and could not infer batch size " "from filename -- pass --batch-size") df["batch"] = float(b) source = "--batch-size" if batch_size_arg else "filename" if "tokens" in df.columns: df["tokens"] = df["tokens"].astype(float) else: # Tokens consumed BEFORE each step (E = B*S for a fixed batch); a gap in # the logged steps is filled with the batch of the row before it. steps = df["step"].to_numpy(float) batch = df["batch"].to_numpy() tokens = np.empty(len(df)) tokens[0] = steps[0] * batch[0] tokens[1:] = tokens[0] + np.cumsum(batch[:-1] * np.diff(steps)) df["tokens"] = tokens df = df[df[loss_col].notna()].reset_index(drop=True) return df, source def batch_stages(df): """Runs of rows with a constant batch -> [(first_row, last_row, batch), ...].""" b = df["batch"].to_numpy() starts = np.flatnonzero(np.r_[True, b[1:] != b[:-1]]) ends = np.r_[starts[1:] - 1, len(b) - 1] return [(int(s), int(e), float(b[s])) for s, e in zip(starts, ends)] def token_window_ema(loss, batch, token_window): """EMA whose window is fixed in TOKENS: each step's weight uses the span token_window / (that step's batch), so a batch switch mid-run doesn't change how much data the smoothing averages over. For a fixed batch B this is exactly pandas' .ewm(span=token_window / B, adjust=False).""" span = token_window / batch if (span < 1).any(): print(f"[WARN] EMA token window ({token_window:,.0f}) < batch size " f"({batch.max():,.0f}) on some steps -- window is under one step " "there; clamping span to 1 (effectively no smoothing).") weight = 2.0 / (np.maximum(span, 1.0) + 1.0) out = np.empty(len(loss)) ema = loss[0] for i, (x, w) in enumerate(zip(loss, weight)): ema = w * x + (1 - w) * ema out[i] = ema return out def detect_warmup_end(df, stages): # Searched within the first batch stage only: an arm that scales the LR # with the batch (train_390_dynamic_sqrtlr.py) reaches its overall LR max # only in its last stage, long after warmup. first = df.iloc[: stages[0][1] + 1] lr_max = first["lr"].max() past_warmup = first["step"][first["lr"] >= 0.99 * lr_max] return int(past_warmup.iloc[0]) if len(past_warmup) else 0 def interpolated_crossing(df, smoothed, target): """-> (step, tokens) where the smoothed loss first reaches target, linearly interpolated between the two logged rows around it; None if it never does.""" below = np.flatnonzero(smoothed <= target) if not len(below): return None steps, tokens = df["step"].to_numpy(float), df["tokens"].to_numpy() i = below[0] if i == 0: return steps[0], tokens[0] l0, l1 = smoothed[i - 1], smoothed[i] frac = 1.0 if l0 == l1 else (l0 - target) / (l0 - l1) return (steps[i - 1] + frac * (steps[i] - steps[i - 1]), tokens[i - 1] + frac * (tokens[i] - tokens[i - 1])) def fit_and_solve(tokens, smoothed, target): """-> (predicted tokens, (l0, a, alpha)), or (None, reason). A is in units of TOKEN_UNIT tokens.""" x = tokens / TOKEN_UNIT p0 = [smoothed.min() * 0.9, max(smoothed[0] - smoothed.min(), 1e-3), 0.3] try: (l0, a, alpha), _ = curve_fit( power_law, x, 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}" t_pred = (a / (target - l0)) ** (1.0 / alpha) * TOKEN_UNIT return t_pred, (l0, a, alpha) def tokens_to_step(t, df, schedule=None): """Step at which `t` tokens have been consumed. Inside the logged range it is interpolated from the log. Past it, steps are counted forward from the last logged row at its batch -- or, when --batch-schedule is given, at each later stage's batch once its start is reached (the trainer may round a stage's batch to a micro-batch multiple, so steps there are approximate).""" steps, tokens = df["step"].to_numpy(float), df["tokens"].to_numpy() if t <= tokens[-1]: return float(np.interp(t, tokens, steps)) step, tok, batch = steps[-1], tokens[-1], float(df["batch"].iloc[-1]) for start, next_batch in (schedule or []): if start <= tok: continue n = np.ceil((start - tok) / batch) # steps still run at `batch` if t <= tok + n * batch: break step, tok, batch = step + n, tok + n * batch, next_batch return float(step + (t - tok) / batch) def stability_sweep(df, smoothed, fit_start, target, sweep_step, min_window, schedule): 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)).to_numpy() if mask.sum() < 10: continue tokens = df["tokens"].to_numpy()[mask] t_pred, info = fit_and_solve(tokens, smoothed[mask], target) s_pred = tokens_to_step(t_pred, df, schedule) if t_pred is not None else None rows.append((max_step, tokens[-1], t_pred, 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} {'t_pred':>18} {'ratio':>8} {'L0':>8} {'alpha':>8} note") for max_step, t_max, t_pred, s_pred, info in rows: if t_pred is None: print(f"{max_step:>10} {'--':>14} {'--':>18} {'--':>8} {'--':>8} {'--':>8} {info}") continue l0, a, alpha = info ratio = t_pred / t_max flag = " [beyond warn factor]" if ratio > warn_factor else "" print(f"{max_step:>10} {s_pred:>14,.0f} {t_pred:>18,.0f} {ratio:>7.2f}x " f"{l0:>8.4f} {alpha:>8.4f}{flag}") # Judged on predicted tokens; for a fixed batch that is the same verdict as on steps. valid = [t for _, _, t, _, _ in rows if t is not None] if len(valid) < 3: print("verdict: NOT ENOUGH DATA to judge convergence yet (need >=3 valid windows)") return t_prev2, t_prev1, t_last = valid[-3:] rel_changes = [abs(t_prev1 - t_prev2) / t_prev2, abs(t_last - t_prev1) / t_prev1] if all(c <= convergence_tol for c in rel_changes): print(f"verdict: CONVERGED -- last two window-to-window 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 window-to-window changes " f"({rel_changes[0]:.1%}, {rel_changes[1]:.1%}) exceed tolerance {convergence_tol:.0%}; " f"keep training and re-check") def print_stages(df, stages): print(f"{'stage':>7} {'steps':>17} {'tokens (before step)':>33} {'batch':>11}") for k, (r0, r1, b) in enumerate(stages, 1): s0, s1 = df["step"].iloc[r0], df["step"].iloc[r1] t0, t1 = df["tokens"].iloc[r0], df["tokens"].iloc[r1] print(f"{k:>7} {s0:>8} - {s1:<6} {t0:>15,.0f} - {t1:<15,.0f} {b:>11,.0f}") def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) 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-token-window", type=float, default=EMA_TOKEN_WINDOW, help=f"EMA smoothing window in TOKENS, not steps (default: " f"{EMA_TOKEN_WINDOW:,.0f}) -- converted to a step-based " "span using each step's own batch size, so different " "batch sizes (across runs, or across the stages of a " "dynamic run) get an equivalent amount of data smoothed, " "not an equivalent step count") p.add_argument("--batch-size", type=float, default=None, help="tokens per step for a fixed-batch run; parsed from the " "filename (e.g. _1Mbs) if omitted. Ignored when the CSV " "has a global_batch column (dynamic-batch runs)") p.add_argument("--batch-schedule", nargs="+", default=None, metavar="TOKENS:BATCH", help="the run's batch schedule, same format as " "train_390M_dynamicb.py --batch-schedule (e.g. " "'0:589824 2e9:786432 8.3e9:819200'). Only used to " "convert an extrapolated token count back to a step " "when it lands in a stage the log hasn't reached yet; " "without it, the last logged batch is assumed to continue") p.add_argument("--fit-start-step", type=int, default=None, help="default: auto-detected end of LR warmup") p.add_argument("--fit-last-stage", action="store_true", help="start the fit no earlier than the first step of the " "last batch stage in the log, so the extrapolation " "reflects the batch size the run is continuing at") 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 target 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, batch_source = load_run(args.csv, args.loss_col, args.batch_size) schedule = parse_batch_schedule(args.batch_schedule) if args.batch_schedule else None stages = batch_stages(df) dynamic = len(stages) > 1 fit_start = args.fit_start_step if fit_start is None: fit_start = detect_warmup_end(df, stages) if args.fit_last_stage: fit_start = max(fit_start, int(df["step"].iloc[stages[-1][0]])) batch = df["batch"].to_numpy() smoothed = token_window_ema(df[args.loss_col].to_numpy(float), batch, args.ema_token_window) post = (df["step"] >= fit_start).to_numpy() if post.sum() < 10: sys.exit(f"only {post.sum()} logged rows at or after fit start step {fit_start}") tokens_post = df["tokens"].to_numpy()[post] s_max_observed = df["step"].max() t_max_observed = df["tokens"].max() spans = args.ema_token_window / np.array([b for _, _, b in stages]) if dynamic: batch_desc = (f"batch: DYNAMIC, {len(stages)} stages " f"({stages[0][2]:,.0f} -> {stages[-1][2]:,.0f} tokens/step, " f"from {batch_source})") span_desc = f"~{spans.max():.1f}-{spans.min():.1f} steps across stages" else: batch_desc = f"batch size: {stages[0][2]:,.0f} tokens/step (from {batch_source})" span_desc = f"~{spans[0]:.1f} steps at this batch size" print(f"loss col: {args.loss_col} | ema token window: {args.ema_token_window:,.0f} " f"({span_desc}) | {batch_desc} | fit start step: {fit_start}") if dynamic: print_stages(df, stages) print("-" * 100) for target in args.target: cross = interpolated_crossing(df, smoothed, target) if cross is not None: s_cross, t_cross = cross print(f"target={target:<8} REACHED (interpolated): " f"step={s_cross:,.1f} tokens={t_cross:,.0f}") continue if args.stability_sweep: rows = stability_sweep(df, smoothed, fit_start, target, args.sweep_step, args.sweep_min_window, schedule) print_sweep(target, rows, args.convergence_tol, args.extrapolation_warn_factor) continue t_pred, info = fit_and_solve(tokens_post, smoothed[post], target) if t_pred is None: print(f"target={target:<8} NOT REACHED, fit says: {info}") continue l0, a, alpha = info s_pred = tokens_to_step(t_pred, df, schedule) ratio = t_pred / t_max_observed flag = " ** EXTRAPOLATION BEYOND WARN FACTOR **" if ratio > args.extrapolation_warn_factor else "" step_note = "" if dynamic and schedule is None: step_note = f" [step assumes batch stays {stages[-1][2]:,.0f}; pass --batch-schedule]" print(f"target={target:<8} EXTRAPOLATED (fit L0={l0:.4f} A={a:.4f} alpha={alpha:.4f}, " f"T in {TOKEN_UNIT:.0e} tokens): step={s_pred:,.1f} tokens={t_pred:,.0f} " f"({ratio:.2f}x past observed T={t_max_observed:,.0f}, " f"S={s_max_observed:,.0f}){flag}{step_note}") if __name__ == "__main__": main()