code/scripts/qwen3_parity.py · viai957/align at main

Add files using upload-large-folder tool
bcc8d35 verified 16 days ago
23.4 kB
#!/usr/bin/env python
"""Goedel-Prover-V2-8B (Qwen3ForCausalLM) parity on the REAL checkpoint: align vs transformers, CPU.
    .venv/bin/python scripts/qwen3_parity.py run [--dtype fp32|bf16] [--ckpt DIR] [--out DIR]
    .venv/bin/python scripts/qwen3_parity.py mutants [--only NAME ...]      # needs `run --dtype fp32` first
    .venv/bin/python scripts/qwen3_parity.py compare --a DIR/align --b DIR/hf --dtype fp32
Why this exists. The tiny-model oracle tier (tests/test_qwen3_parity.py) pins the architecture
at ~1e-6 on random 3-layer weights. It cannot see what only the released weights expose: the
real key layout and shard index, a config.json whose rope is spelled the transformers-4.x way
(`rope_theta` top-level, `rope_scaling: null`), a 151936-row embedding, and DEPTH -- a 1e-7
per-layer error is invisible at 3 layers and visible at 36. "Align runs Qwen3" is a claim about
THIS checkpoint.
THE GATE (pinned in `evaluate`):
    fp32 compute: max |logit_align - logit_hf| < 1e-3 AND top-1 agreement 100%.
    bf16 compute: measured and reported, NOT gated. Two correct bf16 implementations differ by
    several ulps of a logit whose magnitude is 10-30 (ulp 0.0625-0.125); a bf16 tolerance chosen
    to pass is a number fitted to the answer. fp32 is the same function up to accumulation order.
Mutants. The four porting defects of the tiny oracle tier that apply to a dense Qwen3 are
re-run against the real checkpoint, in ONE process (the model is loaded once, the defect is
installed, the prompts are forwarded, the defect is removed). Each must FAIL the fp32 gate:
if one passes, the gate has no teeth on this checkpoint.
Memory discipline (host: 186 GB, no swap, shared; /apps/CLAUDE.md is binding): fp32 weights are
32.8 GB per copy. align and transformers run in SEPARATE processes, one after the other, and
compare saved tensors; each takes the heavy-job flock and the host preflight before importing
torch; the orchestrator holds nothing, so the workers cannot deadlock on it. Workers see no GPU.
Everything is written under --out (default .tmp/p2/p2a/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` (`transformers` is an oracle, not an import of this file; the dependency firewall
scans it).
"""

from __future__ import annotations

import argparse
import glob
import hashlib
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--Goedel-LM--Goedel-Prover-V2-8B", "snapshots", "*")
DEFAULT_OUT = os.path.join(REPO, ".tmp", "p2", "p2a", "parity")

#: Four prompts of different character, tokenized by the HF tokenizer in the legacy venv (align
#: receives ids). The third is a Lean 4 / Mathlib prompt in the Goedel-Prover-V2 chat framing, the
#: distribution the model is actually used on.
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",
    "<|im_start|>user\nComplete the following Lean 4 code:\n\n```lean4\nimport Mathlib\nimport Aesop\n\n"
    "set_option maxHeartbeats 400000\n\nopen BigOperators Real Nat Topology Rat\n\n"
    "theorem mathd_algebra_10 : abs ((120 : ℝ) / 100 * 30 - 130 / 100 * 20) = 10 := by\n  sorry\n```\n"
    "<|im_end|>\n<|im_start|>assistant\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.",
)

FP32_MAX_ABS_LOGIT_DIFF = 1e-3

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


def compare_tensors(a: dict, b: dict, n_prompts: int) -> dict:
    """Metrics between two engines' saved logits (dicts of numpy arrays `logits_{p}` [T, V])."""
    import numpy as np

    per_prompt, worst, n_tok, n_agree = [], 0.0, 0, 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)
        agree = int((la.argmax(-1) == lb.argmax(-1)).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
    return {"max_abs_logit_diff": worst, "top1_agreement": n_agree / n_tok, "tokens": n_tok, "per_prompt": per_prompt}


def evaluate(metrics: dict, dtype: str) -> dict:
    """The verdict. Only fp32 is a gate; bf16 returns `gated=False, passed=None`. Anything else raises."""
    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,
    }
    return {"dtype": dtype, "gated": True,
            "thresholds": {"max_abs_logit_diff": FP32_MAX_ABS_LOGIT_DIFF, "top1_agreement": 1.0},
            "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']}")
    metrics = compare_tensors(a, b, len(ma["input_ids"]))
    return {"metrics": metrics, "verdict": evaluate(metrics, dtype),
            "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 Goedel-Prover-V2-8B snapshot under {CKPT_GLOB}")
    return hits[-1]


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


#: Porting defects injected into the align side ONLY, so the gate is shown to bite on the real
#: checkpoint. Name -> what is broken.
MUTANTS: dict[str, str] = {
    "drop_qk_norm": "per-head q/k RMSNorm skipped entirely",
    "fullwidth_qk_norm": "q/k RMSNorm over the whole projection width (OLMoE style) instead of per head",
    "wrong_rope_theta": "RoPE base 1e4 instead of Qwen3's 1e6",
    "swap_gate_up": "SwiGLU gate and up operands swapped: silu(up) * gate",
}


def _install_mutant(name: str, model):
    """Install the named defect; return an `undo()` closure. Raises on an unknown name (a typo must
    not silently run the control)."""
    import torch
    import torch.nn.functional as F

    from align.model import qwen3 as q3

    if name == "drop_qk_norm":
        orig = q3._head_norm
        q3._head_norm = lambda norm, x: x
        return lambda: setattr(q3, "_head_norm", orig)
    if name == "fullwidth_qk_norm":
        orig = q3._head_norm

        def full(norm, x):
            shp = x.shape
            flat = x.reshape(*shp[:-2], -1).float()
            flat = flat * torch.rsqrt(flat.pow(2).mean(-1, keepdim=True) + norm.eps)
            return (norm.weight.repeat(shp[-2]) * flat.to(x.dtype)).view(shp)

        q3._head_norm = full
        return lambda: setattr(q3, "_head_norm", orig)
    if name == "wrong_rope_theta":
        rope = model.model.rope
        saved = rope.inv_freq
        rope.theta = 1e4
        rope.inv_freq = rope._compute()
        def undo():
            rope.theta = 1e6
            rope.inv_freq = saved
        return undo
    if name == "swap_gate_up":
        orig = q3._swiglu
        q3._swiglu = lambda g, u: F.silu(u) * g
        return lambda: setattr(q3, "_swiglu", orig)
    raise ValueError(f"unknown mutant {name!r}; known: {sorted(MUTANTS)}")


def worker_align(ckpt: str, ids_path: str, out_dir: str, dtype: str, mutants: list[str] | None = None) -> None:
    """Runs in the align venv. Guards first, torch after. With `mutants`, the control is NOT
    written; each mutant gets `out_dir/mut_<name>/`."""
    from align.utils.guard import apply_resource_guards

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

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

    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, sd = qc.load_hf_qwen3(ckpt, dtype=t_dtype)  # per-shard cast: transient is the target dtype's size
    model = q3.Qwen3ForCausalLM.from_state_dict(cfg, sd, attn_implementation="sdpa", device="cpu").eval()
    del sd
    t_load = time.perf_counter() - t0

    def forward_all() -> dict:
        tensors = {}
        with torch.no_grad():
            for p, row in enumerate(ids):
                tensors[f"logits_{p}"] = model(input_ids=torch.tensor([row], dtype=torch.long))[0].float().numpy()
        return tensors

    base_meta = {"dtype": dtype, "n_layers": cfg.n_layers, "input_ids": ids, "torch": torch.__version__,
                 "attn": model.attn_implementation, "load_s": round(t_load, 2)}
    for name in (mutants or [None]):
        undo = _install_mutant(name, model) if name else None
        t1 = time.perf_counter()
        try:
            tensors = forward_all()
        finally:
            if undo:
                undo()
        _write(os.path.join(out_dir, f"mut_{name}") if name else out_dir, tensors,
               {**base_meta, "engine": "align" + (f"+MUTANT:{name}" if name else ""),
                "forward_s": round(time.perf_counter() - t1, 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, add_special_tokens=False)["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="qwen3-parity-hf-" + dtype)
import numpy as np, torch, transformers
from safetensors.numpy import save_file
from transformers import Qwen3ForCausalLM
t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype]
ids = json.load(open(ids_path))["input_ids"]
t0 = time.perf_counter()
m = Qwen3ForCausalLM.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):
        tensors[f"logits_{p}"] = m(input_ids=torch.tensor([row])).logits[0].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", "p2", "p2a"),
        "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:
    cwd = os.path.join(REPO, ".tmp", "p2", "p2a")
    os.makedirs(cwd, exist_ok=True)
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=_worker_env(), cwd=cwd)
    if p.returncode != 0:
        raise RuntimeError(f"{cmd[:3]} exited {p.returncode}:\n{p.stderr[-3000:]}")
    return p.stdout


def _ckpt_fingerprint(ckpt: str) -> str:
    """Identity of the checkpoint a worker consumed: path + (name, size, mtime) of every weight /
    index / config file. Cheap (no 16 GB read) and changes on any re-download or edit."""
    h = hashlib.sha256(os.path.realpath(ckpt).encode())
    for name in sorted(os.listdir(ckpt)):
        if name.endswith((".safetensors", ".json")):
            st = os.stat(os.path.join(ckpt, name))  # follows the hub snapshot's blob symlinks
            h.update(f"{name}:{st.st_size}:{st.st_mtime_ns}".encode())
    return h.hexdigest()


def _align_code_sha() -> str:
    """Hash of the align-side code a parity verdict describes: align/model/*.py plus this script."""
    h = hashlib.sha256()
    paths = sorted(glob.glob(os.path.join(REPO, "align", "model", "*.py"))) + [os.path.abspath(__file__)]
    for path in paths:
        with open(path, "rb") as f:
            h.update(os.path.relpath(path, REPO).encode() + b"\0" + f.read())
    return h.hexdigest()


def _provenance(side: str, ckpt: str, ids_path: str, dtype: str) -> dict:
    """What a worker's saved tensors depend on. `meta.json` carries none of this, so `run` stamps
    it beside the tensors (`provenance.json`) and refuses to reuse them on any mismatch."""
    with open(ids_path, "rb") as f:
        ids_sha = hashlib.sha256(f.read()).hexdigest()
    prov = {"side": side, "dtype": dtype, "ckpt": os.path.realpath(ckpt),
            "ckpt_fingerprint": _ckpt_fingerprint(ckpt), "ids_sha256": ids_sha}
    if side == "align":
        prov["align_code_sha256"] = _align_code_sha()
    return prov


def _read_provenance(out_dir: str) -> dict | None:
    try:
        with open(os.path.join(out_dir, "provenance.json")) as f:
            return json.load(f)
    except (OSError, ValueError):
        return None


def _reusable(out_dir: str, prov: dict) -> bool:
    return os.path.isfile(os.path.join(out_dir, "meta.json")) and _read_provenance(out_dir) == prov


def _run_worker(out_dir: str, prov: dict, cmd: list[str], timeout: float) -> None:
    # Invalidate FIRST: a worker that dies mid-way must not leave the previous run's meta.json
    # standing beside a stamp for a run that never finished.
    for name in ("meta.json", "provenance.json"):
        if os.path.isfile(os.path.join(out_dir, name)):
            os.remove(os.path.join(out_dir, name))
    _run(cmd, timeout)
    with open(os.path.join(out_dir, "provenance.json"), "w") as f:
        json.dump(prov, f, indent=2)


def run(ckpt: str, out: str, dtype: str, timeout: float = 3600.0, fresh: bool = False) -> dict:
    """Tokenize (legacy) -> align worker -> HF worker -> compare -> report.json."""
    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 66 GB. Each worker frees its memory by exiting.
    # A worker whose output is already on disk is not re-run: the heavy-job flock is shared with
    # other agents, a worker can lose the race for it AFTER its sibling finished, and re-doing
    # 33 GB of CPU work to retry the one that failed is the wrong default. Reuse is gated on the
    # PROVENANCE stamp (checkpoint identity, align code hash, token ids): tensors that another
    # checkpoint or older code produced must never carry this run's verdict. --fresh redoes both.
    a_prov = _provenance("align", ckpt, ids_path, dtype)
    b_prov = _provenance("hf", ckpt, ids_path, dtype)
    if fresh or not _reusable(a_dir, a_prov):
        _run_worker(a_dir, a_prov, [sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt,
                                    "--ids", ids_path, "--out", a_dir, "--dtype", dtype], timeout)
    if fresh or not _reusable(b_dir, b_prov):
        _run_worker(b_dir, b_prov, [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, all mutants in one process (fp32), against the saved
    transformers tensors of `run`."""
    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 `qwen3_parity.py run --dtype fp32` first: the mutants compare to its transformers tensors")
    if _read_provenance(b_dir) != _provenance("hf", ckpt, ids_path, "fp32"):
        raise RuntimeError(f"{b_dir} was not produced from {ckpt} (missing or stale provenance.json): "
                           "re-run `qwen3_parity.py run --dtype fp32` before the mutants")
    names = names or sorted(MUTANTS)
    a_root = os.path.join(out, "align_fp32_mutants")
    _run([sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, "--ids", ids_path,
          "--out", a_root, "--dtype", "fp32", "--mutants", *names], timeout)
    res = {}
    for name in names:
        rep = compare_dirs(os.path.join(a_root, f"mut_{name}"), 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"], "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']}",
    ]
    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} "
                     f"|logit|max={p['abs_max_logit']:.2f} top1={p['top1_agree']}/{p['tokens']}")
    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)
    r.add_argument("--fresh", action="store_true", help="re-run workers whose output already exists")
    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("--mutants", nargs="*", 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.mutants)
        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} 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, a.fresh)
    print(_summary(rep))
    return 1 if rep["verdict"]["passed"] is False else 0


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