check.py

#!/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):
    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)   # 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)
    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()