| |
| """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 |
| |
| |
| |
| |
| |
| TOKEN_UNIT = 1e9 |
| |
|
|
|
|
| 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)}") |
|
|
| |
| |
| |
| 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: |
| |
| |
| 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): |
| lr_max = df["lr"].max() |
| past_warmup = df["step"][df["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) |
| 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}") |
|
|
| |
| 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) |
| 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() |
|
|
