#!/usr/bin/env python
"""Merge a LoRA SFT checkpoint into the base, export HF Qwen3, and prove vLLM reads what align wrote.

    CUDA_VISIBLE_DEVICES=0,1 .venv/bin/python scripts/qwen3_export_check.py run \
        --ckpt RUN/checkpoints/micro_N --out-dir .tmp/p2/p2a/export [--keep]

The claim under test: *the file align writes is the model align trained.* Five independent checks,
each of which fails if the export is wrong (the process exit status gates ALL of them, `export_ok`):

  1. MERGE PATH vs TRAINER FILE. align merges the adapter (fp32 B@A * alpha/r, one bf16 cast) into
     the pristine snapshot weights from the adapter stored in the checkpoint's optimizer payload,
     and writes it with `save_hf_qwen3`. The trainer's own checkpoint weight file was built by a
     different route (FSDP gather -> `export_merged_state_dict`). The two must be bit-identical.
  2. vLLM LOADS IT. A separate venv (`.venv-rollout`, vLLM 0.26) loads the exported directory
     unmodified -- sharded safetensors, config.json (rope_theta under both spellings), tokenizer --
     and greedy-generates 64 tokens on 2 prompts, returning the top-20 logprobs at every position.
  3. FIRST-TOKEN LOGPROBS. align's own forward on the merged weights (bf16, GPU 1) scores the same
     prompt; the first-token top-1 id must agree and the logprob difference is reported.
  4. TEACHER-FORCED AGREEMENT over all 64 generated tokens: align scores prompt+generation in one
     pass; per-position argmax agreement with what vLLM sampled, and max |dlogprob| of the sampled
     token. bf16 TOLERANCE (documented, not fitted): two correct bf16 engines with different
     attention/GEMM kernels and accumulation order differ by a few ulps of a logit of magnitude
     10-30 (ulp 0.0625-0.125) amplified through 36 layers; the gate is |dlogprob| <= 0.25 on the
     first token and argmax agreement >= 95% over the 64 positions (a near-tie legitimately flips
     argmax; a wrong merge, a missing QK-norm or a wrong rope_theta disagrees on most positions and
     by whole nats). The thresholds are constants below, asserted by `evaluate`, and exercised on
     injected mutants in tests/test_qwen3_scripts.py.
  5. ADAPTER vs MERGED. The single bf16 cast of W+delta rounds away any delta below half an ulp of
     W (a dead zone, not noise), so at small LoRA movement the merged file can be a different
     policy from base+adapter -- and checks 1-4 all score the MERGED file against itself. Here the
     SAME teacher-forced tokens are scored by align with the adapter applied live (base bf16 +
     fp32 delta) and by the merged file; the gates are the first-token / argmax-agreement
     constants of check 4 (reused, not refitted). The retained-delta cosine (merged-W vs the
     intended delta) is recorded alongside. Also gated: the real snapshot's hf<->align mapping
     round trip, and that the adapter MOVED every adapted matrix (a zero adapter exports the base).

The export is DELETED afterwards unless --keep (16.4 GB on a 94%-full disk).
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import time

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO)
ROLLOUT_PY = os.path.join(REPO, ".venv-rollout", "bin", "python")
SNAPSHOT = os.environ.get("P2_EXPORT_BASE_SNAPSHOT") or ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/"
            "dfd02e6271a58375dfbf3ece0175277cf6b6a89a")
WORK = os.path.join(REPO, ".tmp", "p2", "p2a", "export_check")

N_NEW = 64
TOPK = 20
FIRST_TOKEN_MAX_ABS_DLOGPROB = 0.25
MIN_ARGMAX_AGREEMENT = 0.95

PROMPTS = (
    "<|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",
    "<|im_start|>user\nWhat is 17 * 23? Answer briefly.<|im_end|>\n<|im_start|>assistant\n",
)


# ---------------------------------------------------------------------------
# merge (pure tensor math; unit-tested)
# ---------------------------------------------------------------------------


def merge_lora_state(base: dict, lora: dict, scaling: float) -> dict:
    """`base` (align-keyed, `<path>.weight`) with `scaling * B @ A` folded in for every adapter path
    in `lora` (`<path>.A` / `<path>.B`, as `align.ppo.lora.lora_state_dict` writes them).

    The delta is accumulated in fp32 and cast ONCE to the weight's dtype (the same arithmetic as
    `export_merged_state_dict`, so the two routes can be compared bit for bit). Every adapter path
    must resolve to a base weight and must have both halves: a half-present adapter or a path that
    matches nothing would otherwise "export" an unmodified base.
    """
    import torch

    a_keys = {k[:-2] for k in lora if k.endswith(".A")}
    b_keys = {k[:-2] for k in lora if k.endswith(".B")}
    stray = sorted(k for k in lora if not (k.endswith(".A") or k.endswith(".B")))
    if stray:
        raise KeyError(f"adapter state has keys that are neither .A nor .B: {stray[:4]}")
    if a_keys != b_keys:
        raise KeyError(f"adapter halves do not pair: A-only={sorted(a_keys - b_keys)[:3]} B-only={sorted(b_keys - a_keys)[:3]}")
    if not a_keys:
        raise KeyError("adapter state is empty; refusing to export an unmodified base as a merge")
    out = dict(base)
    for p in sorted(a_keys):
        wk = f"{p}.weight"
        if wk not in base:
            raise KeyError(f"adapter path {p!r} has no base weight {wk!r}")
        w, A, B = base[wk], lora[f"{p}.A"], lora[f"{p}.B"]
        if A.shape[0] != B.shape[1] or (B.shape[0], A.shape[1]) != tuple(w.shape):
            raise ValueError(f"{p}: A{tuple(A.shape)} B{tuple(B.shape)} do not compose to {tuple(w.shape)}")
        out[wk] = (w.float() + (B.float() @ A.float()) * scaling).to(w.dtype)
    return out


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


def compare_generation(vllm: dict, align: dict) -> dict:
    """Per-prompt metrics between vLLM's generation record and align's teacher-forced scoring.

    `vllm[i]`: {"token_ids": [64 ints], "chosen_logprobs": [64 floats], "first_top": {id: logprob}}.
    `align[i]`: {"argmax": [64 ints], "chosen_logprobs": [64 floats], "first_logprobs": {id: logprob}}
    (align's log-softmax at the same 64 positions, and at position 0 for vLLM's top-20 ids).
    """
    per = []
    for v, a in zip(vllm, align):
        n = len(v["token_ids"])
        if not (len(a["argmax"]) == len(a["chosen_logprobs"]) == n):
            raise ValueError("generation records have different lengths")
        agree = sum(int(x == y) for x, y in zip(v["token_ids"], a["argmax"]))
        dchosen = [abs(x - y) for x, y in zip(v["chosen_logprobs"], a["chosen_logprobs"])]
        v_top1 = max(v["first_top"], key=v["first_top"].get)
        a_first = a["first_logprobs"]
        per.append({
            "tokens": n,
            "argmax_agreement": agree / n,
            "max_abs_dlogprob_chosen": max(dchosen),
            "mean_abs_dlogprob_chosen": sum(dchosen) / n,
            "first_token_top1_same": int(v_top1) == int(max(a_first, key=a_first.get)),
            "first_token_abs_dlogprob": abs(v["first_top"][v_top1] - a_first[v_top1]),
            "first_token_max_abs_dlogprob_top20": max(abs(v["first_top"][k] - a_first[k]) for k in v["first_top"]),
        })
    return {"per_prompt": per}


def evaluate(metrics: dict) -> dict:
    checks = {}
    for i, p in enumerate(metrics["per_prompt"]):
        checks[f"prompt{i}.first_token_top1_same"] = bool(p["first_token_top1_same"])
        checks[f"prompt{i}.first_token_dlogprob"] = p["first_token_abs_dlogprob"] <= FIRST_TOKEN_MAX_ABS_DLOGPROB
        checks[f"prompt{i}.argmax_agreement"] = p["argmax_agreement"] >= MIN_ARGMAX_AGREEMENT
    return {"checks": checks, "passed": bool(checks) and all(checks.values()),
            "thresholds": {"first_token_max_abs_dlogprob": FIRST_TOKEN_MAX_ABS_DLOGPROB,
                           "min_argmax_agreement": MIN_ARGMAX_AGREEMENT}}


def merge_fidelity_stats(base: dict, merged: dict, lora: dict, scaling: float) -> dict:
    """How much of the intended delta `scaling * B @ A` survived the cast: aggregate cosine and
    relative error of (merged - base) against the delta, and the fraction of nonzero delta
    elements that rounded to exactly zero (the bf16 dead zone). Same arithmetic as
    `align.ppo.lora.export_merged_state_dict`'s report."""
    import math

    import torch

    dot = ne = nd = 0.0
    nz_total = dropped = layers = 0
    for p in sorted({k[:-2] for k in lora if k.endswith(".A")}):
        wk = f"{p}.weight"
        d = ((lora[f"{p}.B"].float() @ lora[f"{p}.A"].float()) * scaling).flatten()
        e = (merged[wk].float() - base[wk].float()).flatten()
        dot, ne, nd = dot + float(torch.dot(e, d)), ne + float(torch.dot(e, e)), nd + float(torch.dot(d, d))
        nz = d != 0
        nz_total += int(nz.sum())
        dropped += int((nz & (e == 0)).sum())
        layers += 1
    if nd == 0.0:
        return {"layers": layers, "cosine": 1.0, "rel_error": 0.0, "dropped_fraction": 0.0, "delta_norm": 0.0}
    return {"layers": layers, "cosine": dot / math.sqrt(max(ne * nd, 1e-300)),
            "rel_error": math.sqrt(max(ne - 2.0 * dot + nd, 0.0) / nd),
            "dropped_fraction": dropped / nz_total if nz_total else 0.0, "delta_norm": math.sqrt(nd)}


def compare_adapter_vs_merged(adapter_lps: list, merged_lps: list) -> dict:
    """Per-prompt agreement between align's log-softmax at the generated positions with the adapter
    applied live (`adapter_lps[i]`, [T, V]) and on the merged bf16 file (`merged_lps[i]`, [T, V])."""
    import torch

    per = []
    for a, m in zip(adapter_lps, merged_lps):
        if a.shape != m.shape or a.dim() != 2:
            raise ValueError(f"log-prob records differ in shape: {tuple(a.shape)} vs {tuple(m.shape)}")
        a, m = a.float(), m.float()
        a_top, m_top = a.argmax(-1), m.argmax(-1)
        t0 = int(a_top[0])
        per.append({
            "tokens": a.shape[0],
            "argmax_agreement": float((a_top == m_top).float().mean()),
            "first_token_abs_dlogprob": abs(float(a[0, t0] - m[0, t0])),
            "max_abs_dlogprob_top1": float((a.gather(1, a_top[:, None]) - m.gather(1, a_top[:, None])).abs().max()),
            "kl_adapter_to_merged_mean": float((a.exp() * (a - m)).sum(-1).mean()),
        })
    return {"per_prompt": per}


def evaluate_adapter_vs_merged(metrics: dict) -> dict:
    checks = {}
    for i, p in enumerate(metrics["per_prompt"]):
        checks[f"prompt{i}.first_token_dlogprob"] = p["first_token_abs_dlogprob"] <= FIRST_TOKEN_MAX_ABS_DLOGPROB
        checks[f"prompt{i}.argmax_agreement"] = p["argmax_agreement"] >= MIN_ARGMAX_AGREEMENT
    return {"checks": checks, "passed": bool(checks) and all(checks.values()),
            "thresholds": {"first_token_max_abs_dlogprob": FIRST_TOKEN_MAX_ABS_DLOGPROB,
                           "min_argmax_agreement": MIN_ARGMAX_AGREEMENT}}


def export_ok(report: dict) -> bool:
    """The process verdict. EVERY advertised check must have run and passed: a check that did not
    run (key absent) is a failure, not a pass -- a zero adapter, a lossy hf<->align mapping or a
    merge that dropped the delta must not exit 0."""
    rel_min = (report.get("merge_relative_weight_change") or {}).get("min")
    return bool(
        (report.get("verdict") or {}).get("passed") is True
        and (report.get("merge_vs_trainer_file") or {}).get("bitwise_mismatches") == 0
        and (report.get("snapshot_mapping_roundtrip") or {}).get("bitwise_equal") is True
        and (report.get("merge_changed_tensors_vs_snapshot") or 0) > 0
        and isinstance(rel_min, (int, float)) and rel_min > 0
        and ((report.get("adapter_vs_merged") or {}).get("verdict") or {}).get("passed") is True
    )


# ---------------------------------------------------------------------------
# vLLM side: runs in .venv-rollout (source string: not an import of this file)
# ---------------------------------------------------------------------------

_VLLM_SCRIPT = r'''
import json, os, sys, time
model_dir, prompts_json, out_path, n_new, topk = sys.argv[1:6]
n_new, topk = int(n_new), int(topk)
os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
tok = AutoTokenizer.from_pretrained(model_dir)
prompts = json.loads(prompts_json)
ids = [tok(p, add_special_tokens=False)["input_ids"] for p in prompts]
t0 = time.time()
llm = LLM(model=model_dir, tokenizer=model_dir, dtype="bfloat16", max_model_len=2048,
          gpu_memory_utilization=0.85, enforce_eager=True, seed=0)
t_load = time.time() - t0
sp = SamplingParams(temperature=0.0, max_tokens=n_new, logprobs=topk, ignore_eos=True)
outs = llm.generate([{"prompt_token_ids": x} for x in ids], sp)
res = []
for x, o in zip(ids, outs):
    c = o.outputs[0]
    lps = c.logprobs
    res.append({
        "prompt_ids": x,
        "token_ids": list(c.token_ids),
        "text": c.text,
        "chosen_logprobs": [float(lps[t][c.token_ids[t]].logprob) for t in range(len(c.token_ids))],
        "first_top": {str(k): float(v.logprob) for k, v in lps[0].items()},
    })
json.dump({"results": res, "load_s": t_load, "model": model_dir}, open(out_path, "w"))
print(json.dumps({"ok": True, "n": len(res), "load_s": round(t_load, 1)}), flush=True)
# vLLM 0.26's EngineCore subprocess can hang at interpreter shutdown (measured: the result file was
# written at 00:44 and the process was still alive 20 minutes later holding 19 GiB of GPU 0). The
# result is on disk and flushed; leave without running the atexit/engine teardown.
sys.stdout.flush(); sys.stderr.flush()
os._exit(0)
'''


def _first_visible_gpu(env=None) -> str:
    """The child's CUDA_VISIBLE_DEVICES is relative to PHYSICAL indices, so a hard-coded "0" put
    vLLM (0.85 of the card) on physical GPU 0 even when the operator pinned this check to 2,3.
    Use the parent's FIRST visible entry ("0" when the parent is unpinned)."""
    v = (os.environ if env is None else env).get("CUDA_VISIBLE_DEVICES", "")
    first = v.split(",")[0].strip()
    return first or "0"


def _rollout_env(gpu: str) -> dict:
    scrub = ("RANK", "LOCAL_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "TORCHELASTIC")
    env = {k: v for k, v in os.environ.items() if not k.startswith(scrub)}
    env.update({
        "CUDA_VISIBLE_DEVICES": gpu, "VLLM_USE_FLASHINFER_SAMPLER": "0",
        "TMPDIR": os.path.join(REPO, ".tmp", "p2", "p2a"), "XDG_CACHE_HOME": os.path.join(REPO, ".cache"),
        "HF_HOME": "/apps/development/rlhf/hf_cache", "HF_HUB_CACHE": "/apps/development/rlhf/hf_cache/hub",
        "HF_HUB_DISABLE_XET": "1", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1",
        "TRITON_CACHE_DIR": os.path.join(REPO, ".cache", "triton"),
        "TORCHINDUCTOR_CACHE_DIR": os.path.join(REPO, ".cache", "torchinductor"),
        "PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false",
    })
    return env


# ---------------------------------------------------------------------------
# the run
# ---------------------------------------------------------------------------


def run_vllm_child(cmd, result_path, *, env, cwd, log_path, grace_s: float = 30.0, timeout_s: float = 3000.0):
    """Run the vLLM child WITHOUT inheriting pipes, and stop its process group once it has written
    its result.

    DEFECT THIS PREVENTS (measured in the P2 smoke loop): vLLM 0.26's EngineCore can outlive its
    parent at interpreter exit (re-parented to init, holding the parent's stdout/stderr pipes). With
    `subprocess.run(capture_output=True)` the parent then blocks in communicate() until the ORPHAN
    closes the pipe -- 12+ minutes, the result file already written, GPU 0 held at 19 GB, the heavy
    lock held. Here the child logs to a file (nothing for an orphan to hold open), runs in its own
    session so the whole group can be signalled by its exact pgid, and is terminated `grace_s` after
    the result file appears. Returns (returncode or None if it had to be terminated, log tail).
    """
    import signal

    with open(log_path, "w") as log:
        proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, env=env, cwd=cwd, start_new_session=True)
        t0 = time.time()
        done_at = None
        rc = None
        while True:
            rc = proc.poll()
            if rc is not None:
                break
            if os.path.isfile(result_path) and done_at is None:
                done_at = time.time()
            if (done_at is not None and time.time() - done_at > grace_s) or time.time() - t0 > timeout_s:
                break
            time.sleep(1.0)
        # signal the group even when the leader exited: orphaned workers keep the pgid
        for sig in (signal.SIGTERM, signal.SIGKILL):
            try:
                os.killpg(proc.pid, sig)
            except ProcessLookupError:
                break
            time.sleep(5.0)
        if rc is None:
            try:
                rc = proc.wait(timeout=10)
            except subprocess.TimeoutExpired:
                rc = None
    with open(log_path) as f:
        return rc, f.read()[-3000:]


def _lora_scaling(ckpt: str, r: int | None, alpha: float | None) -> float:
    if r is None or alpha is None:
        raise SystemExit("--r and --alpha are required: the checkpoint stores A/B but not alpha/r, and guessing the scaling would export a wrong merge silently")
    return alpha / r


def run(a) -> dict:
    from align.utils.guard import apply_resource_guards

    apply_resource_guards(cpu_threads=8, heavy=True, label="p2a-qwen3-export-check")
    import torch

    from align.model import qwen3 as q3
    from align.model.checkpoint import load_hf_safetensors
    from align.model.qwen3_checkpoint import align_to_hf, hf_to_align, read_qwen3_config, save_hf_qwen3

    work = os.path.abspath(a.work)
    out_dir = os.path.abspath(a.out_dir)
    os.makedirs(work, exist_ok=True)
    report: dict = {"ckpt": a.ckpt, "out_dir": out_dir}
    qcfg = read_qwen3_config(SNAPSHOT)
    scaling = _lora_scaling(a.ckpt, a.r, a.alpha)

    # -- 1. merge path vs the trainer's file ---------------------------------------------------
    t0 = time.time()
    raw = load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16)
    base = hf_to_align(raw, qcfg)
    # real-checkpoint mapping round trip: hf -> align -> hf must return every tensor bit for bit
    # (any tensor dropped, renamed or transposed by the table fails here, on the real shapes).
    back = align_to_hf(base, qcfg)
    report["snapshot_mapping_roundtrip"] = {"tensors": len(raw), "bitwise_equal": all(torch.equal(back[k], raw[k]) for k in raw)
                                            and set(back) == set(raw)}
    del raw, back
    payload = torch.load(os.path.join(a.ckpt, "optimizer_rank0.pt"), weights_only=True, map_location="cpu")
    merged = merge_lora_state(base, payload["lora"], scaling)
    report["merge_delta_retention"] = merge_fidelity_stats(base, merged, payload["lora"], scaling)
    del base
    save_hf_qwen3(merged, out_dir, cfg=qcfg, dtype=torch.bfloat16, tokenizer_dir=SNAPSHOT)
    report["merge_s"] = round(time.time() - t0, 1)
    ckpt_state = hf_to_align(load_hf_safetensors(a.ckpt, dtype=None), qcfg)
    mism = [k for k, v in merged.items() if not torch.equal(v, ckpt_state[k])]
    n_changed = sum(1 for p in {k[:-2] for k in payload["lora"] if k.endswith(".A")})
    report["merge_vs_trainer_file"] = {"tensors": len(merged), "bitwise_mismatches": len(mism), "first_mismatches": mism[:3],
                                       "adapted_layers": n_changed}
    del ckpt_state
    # the adapter must have MOVED the weights: a zero adapter exports the base and "passes" check 1
    base2 = hf_to_align(load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16), qcfg)
    moved = sum(1 for k, v in merged.items() if not torch.equal(v, base2[k]))
    report["merge_changed_tensors_vs_snapshot"] = moved
    # How far the adapter moved the weights: ||W' - W||_F / ||W||_F per adapted matrix (the number
    # an LR choice is judged by; bf16 rounding of W' contributes ~2^-9 ~ 0.2% and is included).
    # Over the ADAPTED matrices only (every one must have moved: `export_ok` gates the minimum).
    adapted = {f"{k[:-2]}.weight" for k in payload["lora"] if k.endswith(".A")}
    rel = [((merged[k].float() - base2[k].float()).norm() / base2[k].float().norm()).item() for k in sorted(adapted)]
    report["merge_relative_weight_change"] = {"n": len(rel), "mean": sum(rel) / len(rel), "max": max(rel), "min": min(rel)}
    del base2, merged

    # -- 2. vLLM generates from the export -----------------------------------------------------
    vout = os.path.join(work, "vllm_out.json")
    env = _rollout_env(_first_visible_gpu())
    if os.path.isfile(vout):
        os.remove(vout)  # a stale result from an earlier run must never be read as this run's
    rc, tail = run_vllm_child([ROLLOUT_PY, "-c", _VLLM_SCRIPT, out_dir, json.dumps(list(PROMPTS)), vout, str(N_NEW), str(TOPK)],
                              vout, env=env, cwd=work, log_path=os.path.join(work, "vllm.log"))
    report["vllm_returncode"] = rc
    if not os.path.isfile(vout):
        report["vllm_stderr_tail"] = tail
        return report
    vres = json.load(open(vout))
    report["vllm_load_s"] = vres["load_s"]
    report["vllm_text"] = [r["text"][:200] for r in vres["results"]]

    # -- 3/4. align scores the same tokens on the merged weights (GPU 1) ------------------------
    dev = torch.device("cuda", 1 if torch.cuda.device_count() > 1 else 0)
    cfg2, sd = (qcfg, hf_to_align(load_hf_safetensors(out_dir, dtype=torch.bfloat16), qcfg))
    model = q3.Qwen3ForCausalLM.from_state_dict(cfg2, sd, attn_implementation="sdpa", device="cuda").to(dev).eval()
    del sd
    arec, mlps = [], []
    with torch.no_grad():
        for r in vres["results"]:
            ids = r["prompt_ids"] + r["token_ids"]
            logits = model(input_ids=torch.tensor([ids], device=dev))[0].float()
            lp = torch.log_softmax(logits, dim=-1)
            n0 = len(r["prompt_ids"])
            pos = lp[n0 - 1: n0 - 1 + len(r["token_ids"])]  # position t predicts generated token t
            chosen = pos.gather(1, torch.tensor(r["token_ids"], device=dev)[:, None])[:, 0]
            first = {k: float(pos[0, int(k)]) for k in r["first_top"]}
            arec.append({"argmax": pos.argmax(-1).tolist(), "chosen_logprobs": chosen.tolist(), "first_logprobs": first})
            mlps.append(pos.cpu())
    metrics = compare_generation(vres["results"], arec)
    report["metrics"] = metrics
    report["verdict"] = evaluate(metrics)

    # -- 5. the same tokens through base + live adapter (fp32 delta) vs the merged file ---------
    from align.ppo.lora import LoRAConfig, inject_lora, load_lora_state_dict

    del model
    torch.cuda.empty_cache()
    base_sd = hf_to_align(load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16), qcfg)
    amodel = q3.Qwen3ForCausalLM.from_state_dict(qcfg, base_sd, attn_implementation="sdpa", device="cuda").to(dev).eval()
    del base_sd
    inject_lora(amodel, LoRAConfig(r=a.r, alpha=a.alpha), adapter="sft", generator=torch.Generator().manual_seed(0))
    load_lora_state_dict(amodel, "sft", payload["lora"], strict=True)
    alps = []
    with torch.no_grad():
        for r in vres["results"]:
            ids = r["prompt_ids"] + r["token_ids"]
            lp = torch.log_softmax(amodel(input_ids=torch.tensor([ids], device=dev))[0].float(), dim=-1)
            n0 = len(r["prompt_ids"])
            alps.append(lp[n0 - 1: n0 - 1 + len(r["token_ids"])].cpu())
    m5 = compare_adapter_vs_merged(alps, mlps)
    report["adapter_vs_merged"] = {"metrics": m5, "verdict": evaluate_adapter_vs_merged(m5)}
    return report


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", required=True, help="a trainer checkpoint dir (micro_N) with optimizer_rank0.pt")
    r.add_argument("--out-dir", default=os.path.join(WORK, "export"))
    r.add_argument("--work", default=WORK)
    r.add_argument("--r", type=int, default=None)
    r.add_argument("--alpha", type=float, default=None)
    r.add_argument("--keep", action="store_true")
    a = ap.parse_args(argv)
    rep = {}
    try:
        rep = run(a)
    finally:
        if not a.keep and os.path.isdir(a.out_dir):
            shutil.rmtree(a.out_dir, ignore_errors=True)
            rep["export_deleted"] = True
        os.makedirs(a.work, exist_ok=True)
        with open(os.path.join(a.work, "report.json"), "w") as f:
            json.dump(rep, f, indent=2, default=str)
    print(json.dumps(rep, indent=2, default=str))
    return 0 if export_ok(rep) else 1


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