model-3/scripts/patch_fastinfer.py

"""Speed up reranker inference. Zero effect on the scores - only the schedule changes.

MEASURED PROBLEM: GPU 1 oscillates 0% <-> 100% with CPU pinned at ~1.5 cores.
Tokenisation runs inline between GPU forward passes, so the card idles roughly half
the time: 988 pairs/s against a 4,033/s benchmark.

FIXES, in order of expected value:
  1. prefetch tokenisation on THREADS, 3 batches ahead, so encoding overlaps the GPU.
     Threads not processes: lexical_topk still forks every chunk, and fork + a warmed
     fast tokenizer is exactly the deadlock that cost us an hour earlier. Each worker
     thread gets its OWN tokenizer instance, so no object is shared across threads.
  2. bigger eval batch (1024 -> 2048): we use 21 GB of a 98 GB card, and fewer, larger
     launches amortise Python-loop and kernel-launch overhead.
  3. torch.inference_mode() instead of no_grad() - skips version counter bookkeeping.
  4. SDPA attention when available.
  5. non_blocking host->device copies.

Ordering is preserved exactly: batches are consumed in the same sorted order, and the
resume checkpoint still records a prefix of that order.
"""
import io, re, sys

P = "src/cross_encoder.py"
s = io.open(P, encoding="utf-8").read()

if "_prefetch" in s:
    print("= already patched")
    sys.exit(0)

start = s.index("@torch.no_grad()\ndef predict_cross_encoder(")
new_fn = '''@torch.inference_mode()
def predict_cross_encoder(a, b, model_dir, cfg, cache=None):
    """P(match) per pair. If `cache` is given, scoring is resumable: the finished
    prefix is reloaded and only the remainder is computed.

    Tokenisation is prefetched on threads so it overlaps GPU compute - see the module
    docstring. Scores and their order are identical to the serial version.
    """
    import threading
    from collections import deque
    from concurrent.futures import ThreadPoolExecutor

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    n = len(a)
    fp = _texts_fp(a, b)
    done = 0
    out = np.empty(n, np.float32)

    part_p = done_p = None
    if cache:
        os.makedirs(os.path.dirname(cache) or ".", exist_ok=True)
        part_p, done_p = cache + ".part.npy", cache + ".done.json"
        if os.path.exists(done_p) and os.path.exists(part_p):
            try:
                st = json.load(open(done_p))
                if st.get("fp") == fp and st.get("n") == n:
                    prev = np.load(part_p)
                    done = int(st.get("done_sorted", st.get("done", 0)))
                    out[:len(prev)] = prev[:len(prev)]
                    print(f"[ce] resuming scoring: {done:,}/{n:,} already done")
                else:
                    print("[ce] prediction cache mismatch - rescoring")
            except Exception as e:                            # noqa: BLE001
                print(f"[ce] unreadable prediction cache ({e}) - rescoring")
    if done >= n:
        return out

    try:
        model = AutoModelForSequenceClassification.from_pretrained(
            model_dir, attn_implementation="sdpa").to(dev).eval()
    except Exception:                                          # noqa: BLE001
        model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(dev).eval()
    ngpu = torch.cuda.device_count()
    use_dp = ngpu > 1 and cfg.get("ce_data_parallel", False)
    net = torch.nn.DataParallel(model) if use_dp else model
    amp = _amp_dtype()
    bs = int(cfg["ce_eval_batch"])
    ck_rows = int(cfg.get("ce_pred_ckpt_rows", 2_000_000))
    max_len = cfg["ce_max_len"]

    # one tokenizer per worker thread: HF fast tokenizers are not documented as safe
    # to share across threads, and a per-thread instance costs nothing.
    _tl = threading.local()

    def _tok():
        t = getattr(_tl, "tok", None)
        if t is None:
            t = AutoTokenizer.from_pretrained(model_dir)
            _tl.tok = t
        return t

    def _encode(idx):
        tk = _tok()
        enc = tk([a[i] for i in idx], [b[i] for i in idx], truncation="longest_first",
                 max_length=max_len, padding=True, return_tensors="pt")
        return idx, enc

    order = np.argsort([len(a[i]) + len(b[i]) for i in range(n)], kind="stable")
    buf = np.empty(n, np.float32)
    if done:
        buf[:] = out[:]

    batches = [order[s:s + bs] for s in range(done, len(order), bs)]
    n_prefetch = int(cfg.get("ce_prefetch", 3))
    pos, scored, since = done, 0, 0
    t0 = time.time()

    with ThreadPoolExecutor(max_workers=n_prefetch) as ex:
        it = iter(batches)
        futs = deque()
        for _ in range(n_prefetch):
            try:
                futs.append(ex.submit(_encode, next(it)))
            except StopIteration:
                break
        while futs:
            idx, enc = futs.popleft().result()
            try:
                futs.append(ex.submit(_encode, next(it)))
            except StopIteration:
                pass
            enc = {k: v.to(dev, non_blocking=True) for k, v in enc.items()}
            with torch.autocast("cuda", dtype=amp, enabled=amp is not None):
                lg = net(**enc).logits.squeeze(-1).float()
            buf[idx] = torch.sigmoid(lg).cpu().numpy()
            pos += len(idx); scored += len(idx); since += len(idx)
            if scored % (bs * 1000) < bs:
                rate = scored / max(time.time() - t0, 1e-9)
                print(f"[ce] scored {pos:,}/{n:,} ({rate:,.0f}/s, "
                      f"{time.time()-t0:.0f}s)", flush=True)
            if cache and since >= ck_rows:
                np.save(part_p + ".tmp.npy", buf)
                os.replace(part_p + ".tmp.npy", part_p)
                _atomic_json(done_p, {"fp": fp, "n": n, "done_sorted": int(pos)})
                print(f"[ce] scoring checkpoint: {pos:,}/{n:,} rows persisted "
                      f"({100*pos/max(n,1):.1f}%)", flush=True)
                since = 0

    if cache:
        np.save(part_p + ".tmp.npy", buf)
        os.replace(part_p + ".tmp.npy", part_p)
        _atomic_json(done_p, {"fp": fp, "n": n, "done_sorted": n})
    return buf
'''
s = s[:start] + new_fn
io.open(P, "w", encoding="utf-8").write(s)
print("+ cross_encoder.py: threaded tokenisation prefetch, sdpa, inference_mode")

C = "src/config.py"
c = io.open(C, encoding="utf-8").read()
old = '    "ce_eval_batch": 1024,'
new = ('    "ce_eval_batch": 2048,   # 21 GB used of 98 GB; larger batches amortise launch overhead\n'
       '    "ce_prefetch": 3,        # tokenisation batches queued ahead of the GPU')
if '"ce_prefetch"' in c:
    print("= config already patched")
else:
    assert old in c, "ce_eval_batch anchor missing"
    io.open(C, "w", encoding="utf-8").write(c.replace(old, new, 1))
    print("+ config.py: ce_eval_batch 2048, ce_prefetch 3")