"""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")