Upload folder using huggingface_hub

390M_dynamic_sqrtlr/blumodel_muon_dynamic_step00001000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6c95346bd3c7098ca6a5f56850ec5036cbbdaa117fe9c5a58825d8226f3dc9e
+size 4205808528

390M_dynamic_sqrtlr/blumodel_muon_dynamic_step00002000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69431dad4d781999264aabd41cc9385dfd8578376be65c4461455a2f0396d6af
+size 4205808528

390M_dynamic_sqrtlr/blumodel_muon_dynamic_step00003543.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46b54aa4edca016657c6086e1276921ac1ad950810a043b10ac69a6d0e4f4e40
+size 4205808528

390M_dynamic_sqrtlr/blumodel_muon_dynamic_step00004000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:914cd97d61da301cbe82fc261eefb9069c016173009fc69426ba7fbe2cdd11aa
+size 4205808592

check_target_loss.py

@@ -0,0 +1,363 @@
+#!/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()