Blame for code/scripts/qwen3_sft_probe.py

#!/usr/bin/env python
"""Qwen3-8B LoRA SFT memory / throughput / resume probe on 4x A10G (P2-A).

    .venv/bin/python scripts/qwen3_sft_probe.py prepare --lens 4096 8192 16384 [--prompt-len 1024]
    torchrun --standalone --nproc_per_node=4 scripts/qwen3_sft_probe.py train CONFIG [dotted.overrides=...]
    .venv/bin/python scripts/qwen3_sft_probe.py report RUN_DIR [--accum A --world W --seq-len L]
    .venv/bin/python scripts/qwen3_sft_probe.py batches RUN_DIR

What it measures, and why it is not just `scripts/sft.py`. `train` runs the SAME functions
`scripts/sft.py` runs (`build_model`, `run_gates`, `train`: arch dispatch, LoRA wiring, the
Trainer, the checkpoint format), with two additions the production entry deliberately does not
carry: a `Tracker.log` wrapper that stamps every history record with `rank`, wall time and the
CUDA allocator's peak, and a heavy-job guard identical to sft.py's. The numbers it reports are
therefore the numbers a real launch would see.

DATA is REAL TOKENS, not random ids: the formal statements and informal prefixes of the Goedel
reference repo's miniF2F + MOBench files (604 problems), tokenized by the checkpoint's own
tokenizer in the legacy venv, concatenated into one stream, and cut into rows of exactly `len`
tokens at distinct offsets. Loss values on these rows are meaningless as a quality signal (rows
are arbitrary windows); compute, memory and the sharding/resume mechanics are exactly those of
real rows of that length. Every micro-batch is one row of the target length, so every micro-step
costs the same, and tokens/s is `len / micro_step_time`.

Import order is load-bearing: this module imports only the stdlib at top level.
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import subprocess
import sys

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO)

LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python"
GOEDEL_REPO = "/apps/development/rlhf/reference-repos/Goedel-Prover-V2"
SNAPSHOT = ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
            "dfd02e6271a58375dfbf3ece0175277cf6b6a89a")
OUT_ROOT = os.path.join(REPO, ".tmp", "p2", "p2a")
STREAM = os.path.join(OUT_ROOT, "probe_stream.npy")

_TOKENIZE = r'''
import json, sys
import numpy as np
from transformers import AutoTokenizer
ckpt, files, out = sys.argv[1], sys.argv[2].split(","), sys.argv[3]
tok = AutoTokenizer.from_pretrained(ckpt)
ids = []
for fn in files:
    for line in open(fn):
        d = json.loads(line)
        text = (d.get("informal_prefix") or "") + (d.get("formal_statement") or "") + "\n"
        ids.extend(tok(text, add_special_tokens=False)["input_ids"] + [tok.eos_token_id])
np.save(out, np.asarray(ids, dtype=np.int32))
print(json.dumps({"tokens": len(ids), "vocab": len(tok)}))
'''


def cut_rows(stream, length: int, n_rows: int, seed: int = 0):
    """`n_rows` windows of exactly `length` tokens at distinct, seeded offsets (wrapping the stream
    when it is shorter than needed). Pure numpy; unit-tested."""
    import numpy as np

    if length < 2:
        raise ValueError("length must be >= 2")
    n = len(stream)
    rng = np.random.default_rng(seed)
    offs = rng.choice(n, size=n_rows, replace=n_rows > n)
    doubled = np.concatenate([stream] * (length // n + 2))
    return [doubled[o:o + length] for o in offs]


def make_tiny_snapshot(out_dir: str, vocab: int = 2048, layers: int = 4, hidden: int = 256, seed: int = 0) -> dict:
    """A random-weight, structurally-faithful Qwen3 snapshot (GQA 8:4, head_dim 32, untied head) with
    the Goedel tokenizer files, for wiring smokes on the 4-GPU path in seconds instead of minutes.
    NOT a model of anything: it exists so FSDP2 + LoRA + checkpoint + resume bugs surface before the
    8B load. Returns the config overrides a TOML needs to agree with it."""
    import torch

    from align.model import qwen3 as q3
    from align.model.qwen3_checkpoint import save_hf_qwen3

    cfg = q3.Qwen3Config(name="qwen3-tiny", vocab_size=vocab, hidden_size=hidden, intermediate_size=2 * hidden,
                         n_layers=layers, n_heads=8, n_kv_heads=4, head_dim=hidden // 8, max_position_embeddings=4096)
    torch.manual_seed(seed)
    m = q3.Qwen3ForCausalLM(cfg, attn_implementation="sdpa", device="cpu")
    snap = os.path.join(out_dir, "snapshots", "tiny")
    save_hf_qwen3(m, snap, dtype=torch.bfloat16, tokenizer_dir=SNAPSHOT)
    return {"paths.model_snapshot": os.path.join(out_dir, "snapshots"), "model.n_params": q3.count_params(cfg),
            "model.n_layers": layers, "model.vocab_size": vocab, "model.hidden_size": hidden}


def prepare(lens: list[int], rows: int, prompt_len: int, out_root: str, vocab_mod: int | None = None) -> dict:
    import numpy as np

    from align.data.store import ShardWriter

    os.makedirs(out_root, exist_ok=True)
    if not os.path.isfile(STREAM):
        env = {"PATH": os.environ.get("PATH", ""), "HOME": os.environ.get("HOME", ""),
               "HF_HOME": "/apps/development/rlhf/hf_cache", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1",
               "TOKENIZERS_PARALLELISM": "false", "TMPDIR": OUT_ROOT}
        files = ",".join(os.path.join(GOEDEL_REPO, "dataset", f) for f in ("minif2f.jsonl", "MOBench.jsonl"))
        p = subprocess.run([LEGACY_PY, "-c", _TOKENIZE, SNAPSHOT, files, STREAM], capture_output=True, text=True,
                           timeout=600, env=env, cwd=OUT_ROOT)
        if p.returncode != 0:
            raise RuntimeError(p.stderr[-2000:])
    stream = np.load(STREAM)
    if vocab_mod:
        stream = stream % vocab_mod  # tiny-vocab smokes only: ids must index the tiny embedding
    out = {"stream_tokens": int(len(stream))}
    for L in lens:
        d = os.path.join(out_root, f"L{L}", "sft", "train")
        if prompt_len >= L:
            raise ValueError(f"prompt_len {prompt_len} >= row length {L}")
        with ShardWriter(d, "sft-train-0000") as w:
            for i, row in enumerate(cut_rows(stream, L, rows, seed=L)):
                w.add_row(row.tolist(), problem_id=f"probe-{L}-{i}", subject="lean", boundary_positions=[prompt_len],
                          prompt_len=prompt_len)
        out[f"L{L}"] = {"dir": d, "rows": rows}
    return out


# ---------------------------------------------------------------------------
# analysis (pure; unit-tested)
# ---------------------------------------------------------------------------


def read_jsonl(path: str) -> list[dict]:
    if not os.path.exists(path):
        return []
    lines = [l for l in open(path).read().split("\n") if l.strip()]
    out = []
    for i, l in enumerate(lines):
        try:
            out.append(json.loads(l))
        except json.JSONDecodeError:
            if i != len(lines) - 1:
                raise
    return out


def rank_histories(run_dir: str, world: int) -> dict[int, list[dict]]:
    """rank -> history records (rank 0 in `run_dir`, rank r in `run_dir/rank{r}`)."""
    return {r: read_jsonl(os.path.join(run_dir if r == 0 else os.path.join(run_dir, f"rank{r}"), "history.jsonl"))
            for r in range(world)}


def rank_disjointness(hist: dict[int, list[dict]]) -> dict:
    """Do the ranks see DIFFERENT rows? `batch` is the list of dataset row indices of a micro-step.
    Returns the per-rank row sets and whether they are pairwise disjoint, plus per-micro-step
    disjointness (the contract: at micro-step m no two ranks share a row)."""
    seen = {r: [tuple(rec["batch"]) for rec in recs] for r, recs in hist.items()}
    sets = {r: {i for b in bs for i in b} for r, bs in seen.items()}
    ranks = sorted(sets)
    pairwise = all(not (sets[a] & sets[b]) for i, a in enumerate(ranks) for b in ranks[i + 1:])
    n_micro = min(len(v) for v in seen.values())
    per_micro = all(len({i for r in ranks for i in seen[r][m]}) == sum(len(seen[r][m]) for r in ranks) for m in range(n_micro))
    return {"rows_per_rank": {r: len(s) for r, s in sets.items()}, "pairwise_disjoint": pairwise,
            "per_micro_step_disjoint": per_micro, "micro_steps_compared": n_micro,
            "first_batches": {r: [list(b) for b in seen[r][:4]] for r in ranks}}


def step_stats(recs: list[dict], seq_len: int, accum: int, world: int, drop_first: int = 2,
               save_every: int | None = None) -> dict:
    """Median micro-step time from consecutive `t_wall` stamps of ONE rank (skipping the first
    `drop_first` intervals: allocator warmup, first-use kernels, and any interval containing a
    checkpoint save), peak memory over the run, and tokens/s = world * seq_len / micro_step_time
    (every rank processes one row of `seq_len` tokens per micro-step)."""
    recs = sorted(recs, key=lambda r: r["micro_step"])
    dts = [b["t_wall"] - a["t_wall"] for a, b in zip(recs, recs[1:])
           if b["micro_step"] == a["micro_step"] + 1 and a["micro_step"] >= drop_first
           and not (save_every and a["micro_step"] % save_every == 0)]
    if not dts:
        raise ValueError("no usable step intervals")
    med = statistics.median(dts)
    return {
        "n_intervals": len(dts), "micro_median_s": med, "micro_min_s": min(dts), "micro_max_s": max(dts),
        "optimizer_step_s": med * accum, "tokens_per_s_total": world * seq_len / med,
        "tokens_per_s_per_gpu": seq_len / med,
        "alloc_peak_gb": max(r["mem/alloc_peak_gb"] for r in recs),
        "reserved_peak_gb": max(r["mem/reserved_peak_gb"] for r in recs),
        "loss_first": recs[0]["loss"], "loss_last": recs[-1]["loss"],
        "grad_norm_last": next((r["grad_norm"] for r in reversed(recs) if r["grad_norm"] == r["grad_norm"]), float("nan")),
    }


def dedupe_by_micro(recs: list[dict]) -> dict[int, dict]:
    """micro_step -> record, the LAST occurrence winning (a relaunch may re-log the boundary it resumed at)."""
    return {int(r["micro_step"]): r for r in recs}


def compare_resume(ref: dict[int, list[dict]], sup: dict[int, list[dict]], stop_micro: int,
                   floor_mult: float = 2.0, floor_pad: float = 5e-3) -> dict:
    """Does a SIGTERM'd-and-resumed run continue the uninterrupted one?

    Three separate questions, per rank, because a loss curve alone cannot answer the first two:

      * DATA: the resumed segment (micro > stop) saw EXACTLY the batches the reference saw at the
        same micro-steps (the stream offset round-trips), and nothing was replayed or skipped
        (micro-steps are contiguous across the stop).
      * CONTINUITY: |dloss| of the resumed segment vs the reference stays within
        `floor_mult * (pre-stop |dloss|) + floor_pad`. The pre-stop segment is the SAME computation
        in both runs, so its |dloss| measures pure run-to-run nondeterminism (NCCL reduction order,
        SDPA backward, 8-bit rounding); the resumed segment must not be worse than that floor by
        more than the factor. A vacuous comparison (one side empty) raises.
      * PROVENANCE is checked separately and exactly by the caller (the relaunch log names the
        checkpoint it loaded; this function does not claim it).
    """
    out = {}
    for r in sorted(ref):
        a, b = dedupe_by_micro(ref[r]), dedupe_by_micro(sup[r])
        common = sorted(set(a) & set(b))
        pre = [abs(a[m]["loss"] - b[m]["loss"]) for m in common if m <= stop_micro]
        post = [abs(a[m]["loss"] - b[m]["loss"]) for m in common if m > stop_micro]
        if not pre or not post:
            raise ValueError(f"rank {r}: need micro-steps on both sides of the stop ({stop_micro}); pre={len(pre)} post={len(post)}")
        batches_equal = all(a[m]["batch"] == b[m]["batch"] for m in common)
        contiguous = sorted(b) == list(range(min(b), max(b) + 1))
        tol = floor_mult * max(pre) + floor_pad
        out[r] = {"n_pre": len(pre), "n_post": len(post), "batches_equal_all_common": batches_equal,
                  "micro_steps_contiguous": contiguous, "pre_stop_max_dloss": max(pre), "resumed_max_dloss": max(post),
                  "tolerance": tol, "continuity_passed": max(post) <= tol,
                  "last_micro_ref": max(a), "last_micro_resumed": max(b)}
    out["passed"] = all(v["batches_equal_all_common"] and v["micro_steps_contiguous"] and v["continuity_passed"]
                        and v["last_micro_ref"] == v["last_micro_resumed"] for k, v in out.items())
    return out


def adapter_diff(lora_a: dict, lora_b: dict) -> dict:
    """Relative difference of two adapter states (`lora_state_dict` format, `<path>.A` / `<path>.B`).

    The strongest resume check a loss curve cannot give: the FINAL adapter of an interrupted +
    resumed run against an uninterrupted run of the same schedule. `rel_B` is
    ||B_a - B_b|| / ||B_b||, global over all layers. B starts at EXACTLY zero, so ||B|| is the whole
    distance trained; a resume that restarts the adapter from the bare base (B = 0) after k of n
    optimizer steps lands near rel_B ~ (1 - k/n)-ish, orders of magnitude above run-to-run
    nondeterminism (flash-attention backward atomics) which stays at the 1e-3 level.
    """
    import torch

    if set(lora_a) != set(lora_b):
        raise KeyError(f"adapter key sets differ: {sorted(set(lora_a) ^ set(lora_b))[:4]}")
    num = {"A": 0.0, "B": 0.0}
    den = {"A": 0.0, "B": 0.0}
    for k in lora_a:
        kind = k.rsplit(".", 1)[1]
        a, b = lora_a[k].double(), lora_b[k].double()
        if a.shape != b.shape:
            raise ValueError(f"{k}: {tuple(a.shape)} vs {tuple(b.shape)}")
        num[kind] += float((a - b).pow(2).sum())
        den[kind] += float(b.pow(2).sum())
    if den["B"] == 0.0:
        raise ValueError("reference adapter has B == 0: nothing was trained, the comparison is vacuous")
    return {"rel_A": (num["A"] / den["A"]) ** 0.5, "rel_B": (num["B"] / den["B"]) ** 0.5,
            "norm_B": den["B"] ** 0.5, "n_tensors": len(lora_a)}


def adapter_diff_main(a) -> int:
    import torch

    ld = lambda d: torch.load(os.path.join(d, "optimizer_rank0.pt"), weights_only=True, map_location="cpu")["lora"]
    res = adapter_diff(ld(a.ckpt_a), ld(a.ckpt_b))
    print(json.dumps(res, indent=2))
    return 0


# ---------------------------------------------------------------------------
# train (torchrun target)
# ---------------------------------------------------------------------------


def tree_equal(a, b, path: str = "") -> list[str]:
    """Paths at which two nested state structures (dict / list / tensor / scalar) are NOT bitwise
    equal. Tensors compare by dtype, shape and `torch.equal`; everything else by `==`. Empty list =
    identical. Pure; unit-tested."""
    import torch

    if isinstance(a, dict) and isinstance(b, dict):
        out = [f"{path}/<keys>" for _ in [0] if set(a) != set(b)]
        for k in sorted(set(a) & set(b), key=str):
            out += tree_equal(a[k], b[k], f"{path}/{k}")
        return out
    if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)):
        if len(a) != len(b):
            return [f"{path}/<len {len(a)} vs {len(b)}>"]
        return [d for i, (x, y) in enumerate(zip(a, b)) for d in tree_equal(x, y, f"{path}[{i}]")]
    if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
        same = a.dtype == b.dtype and a.shape == b.shape and torch.equal(a.cpu(), b.cpu())
        return [] if same else [path]
    return [] if a == b else [path]


def _verify_resume_only(trainer):
    """Stand-in for `Trainer.train` under --verify-resume-only: after `resume_from_latest()` the live
    sharded state (the gathered adapter, the 8-bit moments, the step counters) must equal what the
    checkpoint holds, BITWISE. No step is taken, so there is no GPU nondeterminism to excuse a
    difference -- this isolates the scatter path (`distribute_tensor` / local-shard copies) from
    training noise, which the adapter-diff of two uninterrupted runs (rel_B up to ~0.4) cannot do."""
    import torch

    from align.utils.dist import rank

    R = rank()
    ck = trainer.latest_checkpoint()
    if ck is None:
        raise RuntimeError("--verify-resume-only needs an existing checkpoint to resume from")
    saved = torch.load(os.path.join(ck, f"optimizer_rank{R}.pt"), weights_only=True, map_location="cpu")
    live = trainer.train_module.optim_state_dict()  # collective: every rank gathers the adapter
    diffs = tree_equal(saved, live)
    meta = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu")
    counters = {"micro_step": (meta["micro_step"], trainer.micro_step),
                "optimizer_step": (meta["optimizer_step"], trainer.optimizer_step)}
    bad_counters = {k: v for k, v in counters.items() if v[0] != v[1]}
    res = {"rank": R, "checkpoint": ck, "state_tensors_compared": len(tree_leaves(saved)),
           "unequal_paths": diffs[:10], "n_unequal": len(diffs), "counters": counters, "bad_counters": bad_counters,
           "passed": not diffs and not bad_counters}
    with open(os.path.join(trainer.run_dir, f"verify_resume_rank{R}.json"), "w") as f:
        json.dump(res, f, indent=2, default=str)
    print(json.dumps(res, default=str), flush=True)
    trainer.stop_reason = "verify_only"
    trainer.verify_passed = bool(res["passed"])
    return {"optimizer_step": trainer.optimizer_step, "verify_only": True, **{"passed": res["passed"]}}


def final_exit_code(trainer, verify_only: bool, any_rank=None) -> int:
    """Process exit status. Under --verify-resume-only a FAILED (or never-recorded) bitwise verdict is
    nonzero on EVERY rank (the verdict is MAX-reduced first), so a wrapper / CI cannot read success
    from a run whose per-rank verdict only went to verify_resume_rank{R}.json and stdout."""
    if verify_only:
        if any_rank is None:
            from align.utils.dist import any_rank
        if any_rank(not getattr(trainer, "verify_passed", False)):
            return 1
    return trainer.stop_exit_code or 0


def tree_leaves(x) -> list:
    import torch

    if isinstance(x, dict):
        return [l for v in x.values() for l in tree_leaves(v)]
    if isinstance(x, (list, tuple)):
        return [l for v in x for l in tree_leaves(v)]
    return [x] if isinstance(x, torch.Tensor) else []


def train_main(argv: list[str]) -> int:
    import importlib.util

    # scripts/sft.py by path: the same build_model / run_gates / train the production entry runs.
    spec = importlib.util.spec_from_file_location("sft_script", os.path.join(REPO, "scripts", "sft.py"))
    sft_script = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(sft_script)

    from align.config import load_config
    from align.train.stage_gates import GateLedger, eval_source_hash, parse_args
    from align.utils.guard import apply_resource_guards

    verify_only = "--verify-resume-only" in argv
    argv = [x for x in argv if x != "--verify-resume-only"]
    args = parse_args(argv)
    cfg = load_config(args.config, args.overrides)
    apply_resource_guards(cpu_threads=8, heavy=True, label="p2a-qwen3-sft-probe", rss_cfg=cfg)
    import logging
    import time

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    from align.utils.dist import init_distributed

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    init_distributed(run_dir=run_dir)
    import torch

    from align.train.stage_gates import assert_distinct_devices
    from align.utils.dist import barrier, rank

    assert_distinct_devices()
    R = rank()
    os.makedirs(run_dir, exist_ok=True)
    with open(os.path.join(run_dir, f"rank{R}.pid"), "w") as f:
        f.write(str(os.getpid()))

    from align.utils import tracking

    orig_log = tracking.Tracker.log

    def log(self, step, metrics=None, **kw):
        torch.cuda.synchronize()
        m = {**(dict(metrics) if metrics else {}), **kw}
        m.update({"rank": R, "t_wall": time.time(), "mem/alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30,
                  "mem/reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30})
        orig_log(self, step, m)

    tracking.Tracker.log = log

    ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())
    banner = sft_script.run_gates(cfg, ledger)
    from align.registry import render_banner

    if R == 0:
        print(render_banner(banner, title=f"align SFT probe :: {cfg.run_name}"), flush=True)
    model = sft_script.build_model(cfg)
    torch.cuda.synchronize()
    post_wrap = torch.cuda.memory_allocated() / 2**30
    # The load phase held a full 16.4 GB bf16 replica per rank before FSDP sharded it; the caching
    # allocator keeps that block reserved. Without this, `reserved_peak` reports the LOAD, not the
    # training step (measured: 18.875 GiB at both seq 4096 and 8192).
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    if verify_only:
        from align.train import trainer as trainer_mod

        trainer_mod.Trainer.train = _verify_resume_only
    trainer = sft_script.train(cfg, ledger, model, banner)
    barrier()
    rec = {"rank": R, "post_wrap_alloc_gb": post_wrap, "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30,
           "reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30, "stop_reason": trainer.stop_reason,
           "micro_step": trainer.micro_step, "optimizer_step": trainer.optimizer_step,
           "load_stats": getattr(model, "_stage_load_stats", {})}
    with open(os.path.join(run_dir, f"final_rank{R}.json"), "w") as f:
        json.dump(rec, f, indent=2)
    return final_exit_code(trainer, verify_only)


def report_main(a) -> int:
    run = os.path.join(a.run_dir, "sft") if os.path.isdir(os.path.join(a.run_dir, "sft")) else a.run_dir
    hist = rank_histories(run, a.world)
    out = {"disjoint": rank_disjointness(hist), "per_rank": {r: step_stats(h, a.seq_len, a.accum, a.world, save_every=a.save_every) for r, h in hist.items()}}
    print(json.dumps(out, indent=2))
    return 0


def compare_main(a) -> int:
    pick = lambda d: os.path.join(d, "sft") if os.path.isdir(os.path.join(d, "sft")) else d
    res = compare_resume(rank_histories(pick(a.ref), a.world), rank_histories(pick(a.sup), a.world), a.stop_micro)
    print(json.dumps(res, indent=2))
    return 0 if res["passed"] else 1


def batches_main(a) -> int:
    run = os.path.join(a.run_dir, "sft") if os.path.isdir(os.path.join(a.run_dir, "sft")) else a.run_dir
    hist = rank_histories(run, a.world)
    for r, recs in hist.items():
        print(f"rank {r}: " + " ".join(f"m{rec['micro_step']}:{rec['batch']}" for rec in recs[:12]))
    print(json.dumps(rank_disjointness(hist), indent=2))
    return 0


def main(argv=None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    if argv and argv[0] == "train":  # config + dotted overrides are parsed by the repo's own parser
        return train_main(argv[1:])
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("prepare")
    p.add_argument("--lens", type=int, nargs="+", required=True)
    p.add_argument("--rows", type=int, default=256)
    p.add_argument("--prompt-len", type=int, default=1024)
    p.add_argument("--out-root", default=os.path.join(OUT_ROOT, "shards"))
    p.add_argument("--vocab-mod", type=int, default=None)
    t = sub.add_parser("make-tiny")
    t.add_argument("--out-dir", default=os.path.join(OUT_ROOT, "tiny_model"))
    for name in ("report", "batches"):
        r = sub.add_parser(name)
        r.add_argument("run_dir")
        r.add_argument("--world", type=int, default=4)
        r.add_argument("--accum", type=int, default=1)
        r.add_argument("--seq-len", type=int, default=4096)
        r.add_argument("--save-every", type=int, default=None)
    d = sub.add_parser("adapter-diff")
    d.add_argument("ckpt_a")
    d.add_argument("ckpt_b")
    c = sub.add_parser("compare")
    c.add_argument("ref")
    c.add_argument("sup")
    c.add_argument("--stop-micro", type=int, required=True)
    c.add_argument("--world", type=int, default=4)
    a = ap.parse_args(argv)
    if a.cmd == "adapter-diff":
        return adapter_diff_main(a)
    if a.cmd == "compare":
        return compare_main(a)
    if a.cmd == "make-tiny":
        print(json.dumps(make_tiny_snapshot(a.out_dir), indent=2))
        return 0
    if a.cmd == "prepare":
        print(json.dumps(prepare(a.lens, a.rows, a.prompt_len, a.out_root, a.vocab_mod), indent=2))
        return 0
    return report_main(a) if a.cmd == "report" else batches_main(a)


if __name__ == "__main__":
    raise SystemExit(main())