Add files using upload-large-folder tool
bcc8d35 verified 15 days ago
24.3 kB
#!/usr/bin/env python
"""OLMoE-1B-7B-0125 parity on the REAL checkpoint: align vs transformers, on CPU.
    .venv/bin/python scripts/olmoe_parity.py run [--dtype fp32|bf16] [--ckpt DIR] [--out DIR]
    .venv/bin/python scripts/olmoe_parity.py compare --a DIR/align --b DIR/hf --dtype fp32
Why this exists. The tiny-model oracle tier (tests/test_olmoe_parity.py) pins the
architecture at ~1e-6 on random 3-layer weights. It cannot see the things only the
released weights can expose: the HUB key layout at 64 experts x 16 layers (a transposed
gate/up stack is invisible when the weights are random and symmetric), a `config.json`
whose rope is spelled the transformers-4.x way, fp32-stored tensors, a 50304-row
embedding whose last rows are padding, and depth: a 1e-7 per-layer error is invisible at
3 layers and visible at 16. "Align is an OLMoE repo" is a claim about THIS checkpoint.
THE GATE (pinned in `evaluate`, threshold constants below):
    fp32 compute: max |logit_align - logit_hf| < 1e-3 AND top-1 agreement 100%
                  AND every layer's router top-8 SET agrees on >= 99.9% of tokens.
    bf16 compute: measured and reported, not gated (see BF16 note on `evaluate`).
Why fp32 and not bf16. The released weights are stored fp32, so fp32 compute is the same
function, not a rounded copy of it: any residual is accumulation order, which at 16
layers and logit magnitudes of ~10-25 was measured at 1.4e-5 (sequential `index_add_` combine)
and 3.0e-5 (fixed-order combine, a different summation order from transformers'). A bf16 gate (top-1 100% with
max-abs < 0.1) would be ~100x looser, would hide a real porting error (a mis-scaled norm, RoPE theta
off by a factor typically lands at 1e-2..10), and was in fact UNSATISFIABLE by two correct bf16
implementations (measured max-abs 1.3 against |logit| up to 25: the bf16 ulp there is 0.125). bf16 is
still RUN and recorded (`--dtype bf16`) because it is the dtype training uses; its disagreement is the
noise floor the fp32 gate must not be confused with.
Memory discipline (host: 186GB, no swap, shared; /apps/CLAUDE.md is binding):
  * fp32 weights are 27.7GB per copy (the checkpoint is fp32 on disk). align and
    transformers are run in SEPARATE processes, one after the other, and compare saved
    tensors, so the peak is one copy plus the loader's transient (align stacks 64
    per-expert tensors per layer: ~2x its weights at the peak), never two models.
  * Each worker runs the host preflight and takes the heavy-job flock before importing
    torch; the orchestrator itself holds nothing, so the workers cannot deadlock on it.
  * Workers see no GPU (CUDA_VISIBLE_DEVICES=""): this is a CPU measurement and it must not
    land on a card another job is using.
  * Everything is written under --out (default .tmp/olmoe/parity), never /tmp.
Import order is load-bearing (`apply_resource_guards` sets env torch reads at import), so
torch is imported inside functions, after the guards. The legacy-venv half is a source
string run by `python -c`, like tests/test_olmoe_parity.py: `transformers` is an oracle, not
an import of this file (tests/test_dependency_firewall.py scans it).
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import resource
import subprocess
import sys
import time

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python"
HF_HOME = "/apps/development/rlhf/hf_cache"
CKPT_GLOB = os.path.join(HF_HOME, "hub", "models--allenai--OLMoE-1B-7B-0125", "snapshots", "*")
DEFAULT_OUT = os.path.join(REPO, ".tmp", "olmoe", "parity")

#: Three fixed prompts of different character (prose, code, numbers) and length. They are
#: tokenized by the HF tokenizer in the legacy venv; align receives ids, never text, because
#: the align tokenizer is Mistral-only (ROADMAP P2-9).
PROMPTS: tuple[str, ...] = (
    "The capital of France is",
    "def fibonacci(n):\n    if n < 2:\n        return n\n    return fibonacci(n - 1) + fibonacci(n - 2)\n",
    "In 1969, Apollo 11 landed on the Moon after a journey of 384,400 kilometres; "
    "Neil Armstrong said that it was one small step for man.",
)

# ---------------------------------------------------------------------------
# The gate (pure functions over numpy arrays: no torch, unit-testable)
# ---------------------------------------------------------------------------

#: fp32: the same function up to accumulation order. See the module docstring.
FP32_MAX_ABS_LOGIT_DIFF = 1e-3
#: Fraction of tokens whose top-8 expert SET must agree, per layer, fp32. Not 1.0: a
#: router logit tie within ~1e-6 between the 8th and 9th expert may legitimately flip
#: under a different summation order. A wrong router (QK-norm, transposed gate) misses
#: by whole percents, far below this.
FP32_ROUTER_SET_AGREEMENT = 0.999
#: bf16 is MEASURED, NOT GATED. The alternative gate (top-1 100% and max-abs < 0.1) is unsatisfiable
#: by two CORRECT bf16 implementations: bf16 logits of magnitude 16-32 have an ulp of 0.125, and 16 layers
#: of different accumulation order put the two engines ~10 ulps apart (measured 1.3 on this checkpoint,
#: MEASUREMENTS.md). A tolerance chosen to pass would be a number fitted to the answer.
TOP_K = 8


def topk_sets(router_logits, k: int = TOP_K):
    """[T, E] router logits -> [T, k] expert ids, each row sorted ascending so two rows
    are equal iff they pick the same SET (order within top-k is irrelevant to the forward)."""
    import numpy as np

    idx = np.argpartition(-router_logits, k - 1, axis=-1)[..., :k]
    return np.sort(idx, axis=-1)


def compare_tensors(a: dict, b: dict, n_prompts: int, n_layers: int, k: int = TOP_K) -> dict:
    """Metrics between two engines' saved tensors (dicts of numpy arrays).
    Keys: `logits_{p}` [T, V], `router_logits_{p}_{l}` [T, E]. Returns plain python
    numbers; `evaluate` turns them into a verdict.
    """
    import numpy as np

    per_prompt = []
    worst = 0.0
    n_tok = n_agree = 0
    for p in range(n_prompts):
        la, lb = a[f"logits_{p}"].astype(np.float64), b[f"logits_{p}"].astype(np.float64)
        if la.shape != lb.shape:
            raise ValueError(f"logits_{p}: shape {la.shape} vs {lb.shape}")
        d = np.abs(la - lb)
        top_a, top_b = la.argmax(-1), lb.argmax(-1)
        agree = int((top_a == top_b).sum())
        per_prompt.append(
            {
                "tokens": int(la.shape[0]),
                "max_abs": float(d.max()),
                "mean_abs": float(d.mean()),
                "abs_max_logit": float(np.abs(lb).max()),
                "top1_agree": agree,
            }
        )
        worst = max(worst, float(d.max()))
        n_tok += la.shape[0]
        n_agree += agree

    layers = []
    for l in range(n_layers):
        match = total = 0
        rl_max = 0.0
        min_overlap = k
        for p in range(n_prompts):
            ra, rb = a[f"router_logits_{p}_{l}"], b[f"router_logits_{p}_{l}"]
            rl_max = max(rl_max, float(np.abs(ra.astype(np.float64) - rb.astype(np.float64)).max()))
            sa, sb = topk_sets(ra, k), topk_sets(rb, k)
            eq = (sa == sb).all(-1)
            match += int(eq.sum())
            total += int(eq.size)
            for i in np.nonzero(~eq)[0]:
                min_overlap = min(min_overlap, len(set(sa[i].tolist()) & set(sb[i].tolist())))
        layers.append(
            {
                "layer": l,
                "tokens": total,
                "set_agreement": match / total,
                "min_overlap": int(min_overlap),
                "router_logit_max_abs": rl_max,
            }
        )
    return {
        "max_abs_logit_diff": worst,
        "top1_agreement": n_agree / n_tok,
        "tokens": n_tok,
        "per_prompt": per_prompt,
        "router_layers": layers,
        "min_router_set_agreement": min(x["set_agreement"] for x in layers),
    }


def evaluate(metrics: dict, dtype: str) -> dict:
    """The verdict. Only fp32 is a gate; bf16 returns `gated=False, passed=None` with the same
    measurements. Anything else raises (a typo must not fall through to a default tolerance)."""
    if dtype == "bf16":
        return {"dtype": dtype, "gated": False, "passed": None, "thresholds": None, "checks": {},
                "note": "bf16 is recorded as the noise floor between two correct implementations, not gated"}
    if dtype != "fp32":
        raise ValueError(f"dtype must be 'fp32' or 'bf16', got {dtype!r}")
    checks = {
        "max_abs_logit_diff": metrics["max_abs_logit_diff"] < FP32_MAX_ABS_LOGIT_DIFF,
        "top1_agreement": metrics["top1_agreement"] == 1.0,
        "router_set_agreement": metrics["min_router_set_agreement"] >= FP32_ROUTER_SET_AGREEMENT,
    }
    return {
        "dtype": dtype, "gated": True,
        "thresholds": {"max_abs_logit_diff": FP32_MAX_ABS_LOGIT_DIFF, "top1_agreement": 1.0,
                       "router_set_agreement": FP32_ROUTER_SET_AGREEMENT},
        "checks": checks,
        "passed": all(checks.values()),
    }


def _load_dir(d: str) -> tuple[dict, dict]:
    from safetensors.numpy import load_file

    with open(os.path.join(d, "meta.json")) as f:
        meta = json.load(f)
    return load_file(os.path.join(d, "tensors.safetensors")), meta


def compare_dirs(a_dir: str, b_dir: str, dtype: str, allow_mutant: bool = False) -> dict:
    a, ma = _load_dir(a_dir)
    if "MUTANT" in ma["engine"] and not allow_mutant:
        raise ValueError(f"{a_dir} holds a MUTANT run ({ma['engine']}); refusing to issue a parity verdict on it")
    b, mb = _load_dir(b_dir)
    if ma["input_ids"] != mb["input_ids"]:
        raise ValueError("the two engines were fed different token ids; the comparison is void")
    if ma["dtype"] != dtype or mb["dtype"] != dtype:
        raise ValueError(f"dtype mismatch: gate={dtype}, a={ma['dtype']}, b={mb['dtype']}")
    n_layers = ma["n_layers"]
    metrics = compare_tensors(a, b, len(ma["input_ids"]), n_layers)
    verdict = evaluate(metrics, dtype)
    return {
        "metrics": metrics,
        "verdict": verdict,
        "engines": {"a": {k: v for k, v in ma.items() if k != "input_ids"}, "b": {k: v for k, v in mb.items() if k != "input_ids"}},
        "input_ids": ma["input_ids"],
    }


# ---------------------------------------------------------------------------
# Workers
# ---------------------------------------------------------------------------


def default_ckpt() -> str:
    hits = sorted(glob.glob(CKPT_GLOB))
    if not hits:
        raise FileNotFoundError(f"no OLMoE-1B-7B-0125 snapshot under {CKPT_GLOB}")
    return hits[-1]


def checkpoint_present(ckpt: str | None = None) -> bool:
    try:
        c = ckpt or default_ckpt()
    except FileNotFoundError:
        return False
    return os.path.isfile(os.path.join(c, "model.safetensors.index.json")) and len(
        glob.glob(os.path.join(c, "model-*-of-*.safetensors"))
    ) == 6


def _peak_rss_gb() -> float:
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2  # KB -> GiB


#: Known OLMoE porting defects, each injected into the align side ONLY so the gate can be
#: shown to bite on the real checkpoint (not only on the random tiny model of the oracle
#: tier). Name -> what is broken. A parity gate that cannot fail is decoration.
MUTANTS: dict[str, str] = {
    "per_head_qknorm": "q/k RMSNorm per head instead of over the full projection (Qwen3/Gemma3 style)",
    "swap_gate_up": "expert gate and up rows stacked in the wrong order",
    "topk_then_softmax": "Mixtral routing: top-k on logits, softmax over only the kept experts",
    "rope_theta_1e6": "Mistral's RoPE base (1e6) instead of OLMoE's 1e4",
}


def _apply_mutant(name: str, cfg):
    """Monkeypatch the named defect into the already-imported align modules; returns `cfg`
    (possibly replaced). Raises on an unknown name: a typo must not silently run the control."""
    import torch

    from align.model import moe
    from align.model import olmoe as om
    from align.model import olmoe_checkpoint as oc

    if name == "per_head_qknorm":

        def per_head(norm, proj):
            hd = cfg.head_dim
            shp = proj.shape
            w = norm.weight.view(-1, hd)
            v = proj.view(*shp[:-1], -1, hd).float()
            v = v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + norm.eps)
            return (w * v.to(proj.dtype)).view(shp)

        om._qk_norm = per_head
    elif name == "swap_gate_up":
        oc._stack_gate_up = lambda g, u: torch.cat([u, g], dim=0)
    elif name == "topk_then_softmax":

        def mixtral(x, weight, top_k, norm_topk_prob=False):
            logits = moe._router_logits(x, weight)
            probs = torch.softmax(logits, dim=-1)
            vals, idx = torch.topk(logits, top_k, dim=-1)
            return moe.RouterOutput(weights=torch.softmax(vals, dim=-1), indices=idx, logits=logits, probs=probs)

        moe.route_top_k = mixtral
    elif name == "rope_theta_1e6":
        cfg = cfg.replace(rope_theta=1e6)
    else:
        raise ValueError(f"unknown mutant {name!r}; known: {sorted(MUTANTS)}")
    return cfg


def worker_align(ckpt: str, ids_path: str, out_dir: str, dtype: str, mutant: str | None = None) -> None:
    """Runs in the align venv. Guards first, torch after."""
    from align.utils.guard import apply_resource_guards

    apply_resource_guards(cpu_threads=8, heavy=True, label=f"olmoe-parity-align-{dtype}")
    import torch

    from align.model import olmoe as om
    from align.model import olmoe_checkpoint as oc
    from align.model.checkpoint import load_hf_safetensors

    t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype]
    with open(ids_path) as f:
        ids = json.load(f)["input_ids"]
    t0 = time.perf_counter()
    cfg = oc.read_olmoe_config(ckpt)
    if mutant:
        cfg = _apply_mutant(mutant, cfg)
    # Loaded straight into the compute dtype shard by shard (load_hf_safetensors casts per
    # shard), so the transient is the target dtype's size, not fp32 + target.
    sd = oc.hf_to_align(load_hf_safetensors(ckpt, dtype=t_dtype), cfg)
    model = om.OlmoeForCausalLM.from_state_dict(cfg, sd, moe_backend="loop", attn_implementation="sdpa").eval()
    del sd
    t_load = time.perf_counter() - t0

    tensors: dict = {}
    t1 = time.perf_counter()
    with torch.no_grad():
        for p, row in enumerate(ids):
            x = torch.tensor([row], dtype=torch.long)
            logits, stats = model(input_ids=x, return_router_stats=True, keep_router_logits=True)
            tensors[f"logits_{p}"] = logits[0].float().numpy()
            for l, s in enumerate(stats):
                tensors[f"router_logits_{p}_{l}"] = s.logits.float().numpy()
    t_fwd = time.perf_counter() - t1
    _write(out_dir, tensors, {
        "engine": "align" + (f"+MUTANT:{mutant}" if mutant else ""), "dtype": dtype, "n_layers": cfg.n_layers, "input_ids": ids,
        "torch": torch.__version__, "attn": model.attn_implementation, "moe_backend": "loop",
        "load_s": round(t_load, 2), "forward_s": round(t_fwd, 2), "peak_rss_gb": round(_peak_rss_gb(), 2),
    })


def _write(out_dir: str, tensors: dict, meta: dict) -> None:
    import numpy as np
    from safetensors.numpy import save_file

    os.makedirs(out_dir, exist_ok=True)
    save_file({k: np.ascontiguousarray(v) for k, v in tensors.items()}, os.path.join(out_dir, "tensors.safetensors"))
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)


# Runs INSIDE the legacy venv (transformers 5.x). argv: mode ckpt ids_json out_dir dtype repo
_LEGACY_SCRIPT = r'''
import json, os, resource, sys, time
mode, ckpt, ids_path, out_dir, dtype, repo = sys.argv[1:7]
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
if mode == "tokenize":
    from transformers import AutoTokenizer
    prompts = json.loads(ids_path)
    tok = AutoTokenizer.from_pretrained(ckpt)
    ids = [tok(p)["input_ids"] for p in prompts]
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "ids.json"), "w") as f:
        json.dump({"prompts": prompts, "input_ids": ids, "tokenizer": type(tok).__name__,
                   "vocab_size": len(tok)}, f, indent=2)
    print(json.dumps({"n": [len(x) for x in ids]}))
    sys.exit(0)
sys.path.insert(0, repo)
from align.utils.guard import apply_resource_guards
apply_resource_guards(cpu_threads=8, heavy=True, label="olmoe-parity-hf-" + dtype)
import numpy as np, torch, transformers
from safetensors.numpy import save_file
from transformers import OlmoeForCausalLM
t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype]
ids = json.load(open(ids_path))["input_ids"]
t0 = time.perf_counter()
m = OlmoeForCausalLM.from_pretrained(ckpt, dtype=t_dtype, attn_implementation="eager").eval()
t_load = time.perf_counter() - t0
tensors = {}
t1 = time.perf_counter()
with torch.no_grad():
    for p, row in enumerate(ids):
        o = m(input_ids=torch.tensor([row]), output_router_logits=True)
        tensors[f"logits_{p}"] = o.logits[0].float().numpy()
        for l, rl in enumerate(o.router_logits):
            tensors[f"router_logits_{p}_{l}"] = rl.float().numpy()
t_fwd = time.perf_counter() - t1
os.makedirs(out_dir, exist_ok=True)
save_file({k: np.ascontiguousarray(v) for k, v in tensors.items()}, os.path.join(out_dir, "tensors.safetensors"))
meta = {"engine": "transformers", "dtype": dtype, "n_layers": m.config.num_hidden_layers, "input_ids": ids,
        "transformers": transformers.__version__, "torch": torch.__version__, "attn": "eager",
        "load_s": round(t_load, 2), "forward_s": round(t_fwd, 2),
        "peak_rss_gb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2, 2)}
json.dump(meta, open(os.path.join(out_dir, "meta.json"), "w"), indent=2)
print(json.dumps({"ok": True}))
'''


def _worker_env() -> dict:
    return {
        "PATH": os.environ.get("PATH", ""),
        "HOME": os.environ.get("HOME", ""),
        "HF_HOME": HF_HOME,
        "HF_HUB_CACHE": os.path.join(HF_HOME, "hub"),
        "HF_HUB_DISABLE_XET": "1",
        "TMPDIR": os.path.join(REPO, ".tmp"),
        "CUDA_VISIBLE_DEVICES": "",  # a CPU measurement must not touch a card someone else is using
        "OMP_NUM_THREADS": "8",
        "TOKENIZERS_PARALLELISM": "false",
        "PYTHONPATH": REPO,
    }


def _run(cmd: list[str], timeout: float) -> str:
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=_worker_env(), cwd=os.path.join(REPO, ".tmp"))
    if p.returncode != 0:
        raise RuntimeError(f"{cmd[:3]} exited {p.returncode}:\n{p.stderr[-3000:]}")
    return p.stdout


def run(ckpt: str, out: str, dtype: str, timeout: float = 3600.0) -> dict:
    """Tokenize (legacy) -> align worker -> HF worker -> compare -> report.json."""
    # Absolute: the workers run with cwd=.tmp, and a relative --out silently lands under .tmp/.tmp.
    out, ckpt = os.path.abspath(out), os.path.abspath(ckpt)
    os.makedirs(out, exist_ok=True)
    ids_dir = os.path.join(out, "ids")
    _run([LEGACY_PY, "-c", _LEGACY_SCRIPT, "tokenize", ckpt, json.dumps(list(PROMPTS)), ids_dir, dtype, REPO], 300)
    ids_path = os.path.join(ids_dir, "ids.json")
    a_dir, b_dir = os.path.join(out, f"align_{dtype}"), os.path.join(out, f"hf_{dtype}")
    # SEQUENTIAL by construction: two fp32 copies are 55GB. Each worker frees its memory by exiting.
    _run([sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, "--ids", ids_path, "--out", a_dir, "--dtype", dtype], timeout)
    _run([LEGACY_PY, "-c", _LEGACY_SCRIPT, "hf", ckpt, ids_path, b_dir, dtype, REPO], timeout)
    rep = compare_dirs(a_dir, b_dir, dtype)
    rep["checkpoint"] = ckpt
    with open(os.path.join(out, f"report_{dtype}.json"), "w") as f:
        json.dump(rep, f, indent=2)
    return rep


def run_mutants(ckpt: str, out: str, names: list[str] | None = None, timeout: float = 3600.0) -> dict:
    """Re-run ONLY the align side per mutant (fp32) and compare to the already-saved
    transformers tensors of `run`. Returns {name: max_abs_logit_diff, ...} plus the verdict
    that each mutant FAILS the fp32 gate."""
    out, ckpt = os.path.abspath(out), os.path.abspath(ckpt)
    ids_path = os.path.join(out, "ids", "ids.json")
    b_dir = os.path.join(out, "hf_fp32")
    if not (os.path.isfile(ids_path) and os.path.isfile(os.path.join(b_dir, "meta.json"))):
        raise FileNotFoundError("run `olmoe_parity.py run --dtype fp32` first: the mutants compare to its transformers tensors")
    res = {}
    for name in names or sorted(MUTANTS):
        a_dir = os.path.join(out, f"align_fp32_mut_{name}")
        _run([sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, "--ids", ids_path,
              "--out", a_dir, "--dtype", "fp32", "--mutant", name], timeout)
        rep = compare_dirs(a_dir, b_dir, "fp32", allow_mutant=True)
        res[name] = {
            "what": MUTANTS[name],
            "max_abs_logit_diff": rep["metrics"]["max_abs_logit_diff"],
            "top1_agreement": rep["metrics"]["top1_agreement"],
            "min_router_set_agreement": rep["metrics"]["min_router_set_agreement"],
            "gate_passed": rep["verdict"]["passed"],
        }
    with open(os.path.join(out, "mutants_fp32.json"), "w") as f:
        json.dump(res, f, indent=2)
    return res


def _summary(rep: dict) -> str:
    m, v = rep["metrics"], rep["verdict"]
    lines = [
        f"dtype={v['dtype']} gated={v['gated']} passed={v['passed']} checks={v['checks']}",
        f"max_abs_logit_diff={m['max_abs_logit_diff']:.3e} (gate {'< %g' % v['thresholds']['max_abs_logit_diff'] if v['gated'] else 'none: bf16 informational'}) "
        f"top1={m['top1_agreement']:.4f} tokens={m['tokens']} min_router_set_agreement={m['min_router_set_agreement']:.5f}",
    ]
    for p in m["per_prompt"]:
        lines.append(f"  prompt: tokens={p['tokens']} max_abs={p['max_abs']:.3e} mean_abs={p['mean_abs']:.3e} |logit|max={p['abs_max_logit']:.2f} top1={p['top1_agree']}/{p['tokens']}")
    for x in m["router_layers"]:
        lines.append(f"  layer {x['layer']:2d}: set_agreement={x['set_agreement']:.5f} min_overlap={x['min_overlap']} router_logit_max_abs={x['router_logit_max_abs']:.3e}")
    return "\n".join(lines)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run")
    r.add_argument("--ckpt")
    r.add_argument("--out", default=DEFAULT_OUT)
    r.add_argument("--dtype", choices=("fp32", "bf16"), default="fp32")
    r.add_argument("--timeout", type=float, default=3600.0)
    w = sub.add_parser("worker-align")
    w.add_argument("--ckpt", required=True)
    w.add_argument("--ids", required=True)
    w.add_argument("--out", required=True)
    w.add_argument("--dtype", choices=("fp32", "bf16"), required=True)
    w.add_argument("--mutant", choices=sorted(MUTANTS))
    mu = sub.add_parser("mutants")
    mu.add_argument("--ckpt")
    mu.add_argument("--out", default=DEFAULT_OUT)
    mu.add_argument("--only", nargs="*", choices=sorted(MUTANTS))
    c = sub.add_parser("compare")
    c.add_argument("--a", required=True)
    c.add_argument("--b", required=True)
    c.add_argument("--dtype", choices=("fp32", "bf16"), required=True)
    a = ap.parse_args(argv)
    if a.cmd == "worker-align":
        worker_align(a.ckpt, a.ids, a.out, a.dtype, a.mutant)
        return 0
    if a.cmd == "mutants":
        res = run_mutants(a.ckpt or default_ckpt(), a.out, a.only)
        for k, v in res.items():
            print(f"{k:20s} max_abs={v['max_abs_logit_diff']:.3e} top1={v['top1_agreement']:.3f} router>={v['min_router_set_agreement']:.3f} gate_passed={v['gate_passed']}  [{v['what']}]")
        return 1 if any(v["gate_passed"] for v in res.values()) else 0
    if a.cmd == "compare":
        rep = compare_dirs(a.a, a.b, a.dtype)
    else:
        rep = run(a.ckpt or default_ckpt(), a.out, a.dtype, a.timeout)
    print(_summary(rep))
    return 1 if rep["verdict"]["passed"] is False else 0


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