code/scripts/m0_spike.py

#!/usr/bin/env python3
"""M0 spike (B200 edition): module table, verbalizer table, equivalence gate on real rows,
fwd-only and LoRA fwd+bwd throughput at several micro-batch sizes, peak memory.

python3 scripts/m0_spike.py --model /root/models/Qwen3.5-9B --out reports/m0_qwen35_9b.md
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
import pandas as pd
import torch

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))

from jev_judge.data import JevDataset, KindBatchSampler, collate, load_split  # noqa: E402
from jev_judge.losses import judge_loss  # noqa: E402
from jev_judge.model import JevJudge, LoraSpec, masked_probs  # noqa: E402


def gb(x: int) -> float:
    return x / 1024**3


def bench(judge: JevJudge, ds: JevDataset, micro_batch: int, steps: int, train: bool, pad: int, lambda_rps: float = 0.5):
    sampler = KindBatchSampler(ds.kind_ids, ds.lengths, micro_batch, seed=1)
    it = iter(sampler)
    batches = [collate([ds[i] for i in next(it)], pad) for _ in range(steps + 2)]
    dev = judge.device
    params = [p for p in judge.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(params, lr=1e-4) if train else None
    torch.cuda.reset_peak_memory_stats()
    step_tok, step_dt = [], []
    for i, b in enumerate(batches):
        ids = b["input_ids"].to(dev); am = b["attention_mask"].to(dev)
        torch.cuda.synchronize(); t0 = time.time()
        if train:
            z, m = judge(ids, am, b["lengths"].to(dev), b["kind_ids"].to(dev), b["n_options"].to(dev))
            loss, _ = judge_loss(z, b["target"].to(dev), m, b["kind_ids"].to(dev), b["weight"].to(dev), lambda_rps)
            loss.backward()
            opt.step(); opt.zero_grad(set_to_none=True)
        else:
            with torch.no_grad():
                judge(ids, am, b["lengths"].to(dev), b["kind_ids"].to(dev), b["n_options"].to(dev))
        torch.cuda.synchronize(); dt = time.time() - t0
        if i >= 2:
            step_tok.append(int(b["lengths"].sum())); step_dt.append(dt)
        print(f"    [{'train' if train else 'fwd'} mb={micro_batch}] step {i}: T={ids.shape[1]} {dt*1000:.0f} ms {int(b['lengths'].sum())/dt:.0f} tok/s", flush=True)
    tok, el = sum(step_tok), sum(step_dt)
    padded = sum(int(b["input_ids"].numel()) for b in batches[2:])
    rates = np.array(step_tok) / np.array(step_dt)
    return {"micro_batch": micro_batch, "tok_per_s": tok / el, "median_tok_per_s": float(np.median(rates)),
            "padded_tok_per_s": padded / el, "samples_per_s": micro_batch * steps / el,
            "pad_efficiency": tok / padded, "peak_mem_gb": gb(torch.cuda.max_memory_allocated())}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="/root/models/Qwen3.5-9B")
    ap.add_argument("--data", default="data")
    ap.add_argument("--out", required=True)
    ap.add_argument("--fwd-batches", default="16,64,128")
    ap.add_argument("--train-batches", default="8,16,32")
    ap.add_argument("--steps", type=int, default=12)
    ap.add_argument("--gate-rows", type=int, default=96)
    ap.add_argument("--skip-train", action="store_true")
    args = ap.parse_args()

    t0 = time.time()
    judge = JevJudge.from_base(args.model, keep_lm_head=True)
    load_s = time.time() - t0
    tok = judge.tokenizer
    pad = tok.pad_token_id
    n_params = sum(p.numel() for p in judge.lm.parameters())
    md = [f"# M0 spike — `{args.model}` on {torch.cuda.get_device_name()}", "",
          f"load time {load_s:.0f}s · text params {n_params/1e9:.2f}B · hidden {judge.hidden_size} · weights {gb(torch.cuda.memory_allocated()):.1f} GB (incl. lm_head)", ""]
    md += ["## Verbalizer table (bare, line-start single token)", "", judge.verbalizers.as_markdown(), ""]
    leaves = judge.linear_leaf_table()
    md += ["## Linear leaves inside decoder layers (LoRA candidates)", "", "| leaf | count |", "|---|---|",
           *[f"| {k} | {v} |" for k, v in leaves.items()], ""]
    print("linear leaves:", leaves, flush=True)

    # equivalence gate on real validation rows (fp32 recompute space)
    val = load_split(args.data, "validation")
    val = pd.concat([g.sample(args.gate_rows // 3, random_state=0) for _, g in val.groupby("kind")]).reset_index(drop=True)
    ds = JevDataset(val, tok, 1024)
    diffs = []
    for s in range(0, len(ds), 16):
        b = collate([ds[i] for i in range(s, min(s + 16, len(ds)))], pad)
        dev = judge.device
        ids, am, L = b["input_ids"].to(dev), b["attention_mask"].to(dev), b["lengths"].to(dev)
        with torch.no_grad():
            z, m = judge(ids, am, L, b["kind_ids"].to(dev), b["n_options"].to(dev))
            zr = judge.restricted_reference(ids, am, L)
        diffs.append((masked_probs(z, m) - masked_probs(zr, m)).abs().max().item())
    gate = max(diffs)
    md += ["## Equivalence gate (step-0 head vs restricted decoding, fp32)", "",
           f"rows: {len(ds)} · max |Δp| = **{gate:.3e}** · gate 1e-5 → {'**PASS**' if gate < 1e-5 else '**FAIL**'}", ""]
    print(f"equivalence gate max|Δp| = {gate:.3e}", flush=True)
    judge.drop_lm_head()

    # throughput
    train_df = load_split(args.data, "train").sample(20000, random_state=0).reset_index(drop=True)
    tds = JevDataset(train_df, tok, 1024)
    md += ["## Throughput (real length distribution, kind-stratified batches, fla kernels)", "",
           "| mode | micro_batch | real tok/s | median tok/s | padded tok/s | samples/s | pad eff. | peak mem GB |", "|---|---|---|---|---|---|---|---|"]
    fwd_rows = []
    for mb in [int(x) for x in args.fwd_batches.split(",")]:
        r = bench(judge, tds, mb, args.steps, train=False, pad=pad)
        fwd_rows.append(r)
        md.append(f"| fwd-only | {mb} | {r['tok_per_s']:.0f} | {r['median_tok_per_s']:.0f} | {r['padded_tok_per_s']:.0f} | {r['samples_per_s']:.1f} | {r['pad_efficiency']:.2f} | {r['peak_mem_gb']:.1f} |")
        print("fwd", r, flush=True)
    train_rows = []
    if not args.skip_train:
        targets = judge.attach_lora(LoraSpec())
        judge.set_stage_grads("s2")
        n_lora = sum(p.numel() for n, p in judge.lm.named_parameters() if "lora_" in n)
        md += [f"", f"LoRA r=16 attached to {targets} → {n_lora/1e6:.1f}M adapter params", ""]
        md += ["| mode | micro_batch | real tok/s | median tok/s | padded tok/s | samples/s | pad eff. | peak mem GB |", "|---|---|---|---|---|---|---|---|"]
        for mb in [int(x) for x in args.train_batches.split(",")]:
            try:
                r = bench(judge, tds, mb, args.steps, train=True, pad=pad)
            except torch.cuda.OutOfMemoryError:
                md.append(f"| S2 fwd+bwd | {mb} | OOM | | | | | |"); torch.cuda.empty_cache(); continue
            train_rows.append(r)
            md.append(f"| S2 fwd+bwd | {mb} | {r['tok_per_s']:.0f} | {r['median_tok_per_s']:.0f} | {r['padded_tok_per_s']:.0f} | {r['samples_per_s']:.1f} | {r['pad_efficiency']:.2f} | {r['peak_mem_gb']:.1f} |")
            print("train", r, flush=True)
        if train_rows:
            best = max(train_rows, key=lambda r: r["tok_per_s"])
            epoch_tok = 84.6e6
            md += ["", f"**Projection** (best S2 config micro_batch={best['micro_batch']}, {best['tok_per_s']:.0f} real tok/s): "
                   f"1 epoch of train ({epoch_tok/1e6:.1f}M tok) ≈ **{epoch_tok/best['tok_per_s']/3600:.2f} h**; 2 epochs ≈ {2*epoch_tok/best['tok_per_s']/3600:.2f} h; "
                   f"10% scan (2 ep) ≈ {0.2*epoch_tok/best['tok_per_s']/60:.0f} min", ""]
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        f.write("\n".join(md))
    with open(os.path.splitext(args.out)[0] + ".json", "w") as f:
        json.dump({"gate_max_diff": gate, "leaves": leaves, "fwd": fwd_rows, "train": train_rows, "load_s": load_s}, f, indent=2)
    print("report ->", args.out)


if __name__ == "__main__":
    main()